Random walk, random method

RELEASE-NOTES.md

@@ -12,6 +12,8 @@
 - Sampling from variational approximation now allows for alternative trace backends [#3550].
 - Infix `@` operator now works with random variables and deterministics [#3619](https://github.com/pymc-devs/pymc3/pull/3619).
 - [ArviZ](https://arviz-devs.github.io/arviz/) is now a requirement, and handles plotting, diagnostics, and statistical checks.
+- Can use GaussianRandomWalk in sample_prior_predictive and sample_prior_predictive [#3682](https://github.com/pymc-devs/pymc3/pull/3682)
+- Now 11 years of S&P returns in data set[#3682](https://github.com/pymc-devs/pymc3/pull/3682)
 
 ### Maintenance
 - Moved math operations out of `Rice`, `TruncatedNormal`, `Triangular` and `ZeroInflatedNegativeBinomial` `random` methods. Math operations on values returned by `draw_values` might not broadcast well, and all the `size` aware broadcasting is left to `generate_samples`. Fixes [#3481](https://github.com/pymc-devs/pymc3/issues/3481) and [#3508](https://github.com/pymc-devs/pymc3/issues/3508)

pymc3/distributions/timeseries.py

@@ -1,8 +1,10 @@
+from scipy import stats
 import theano.tensor as tt
 from theano import scan
 
 from pymc3.util import get_variable_name
 from .continuous import get_tau_sigma, Normal, Flat
+from .shape_utils import to_tuple
 from . import multivariate
 from . import distribution
 
@@ -166,33 +168,50 @@ def logp(self, value):
 
 
 class GaussianRandomWalk(distribution.Continuous):
-    R"""
-    Random Walk with Normal innovations
+    R"""Random Walk with Normal innovations
 
     Parameters
     ----------
     mu: tensor
         innovation drift, defaults to 0.0
+        For vector valued mu, first dimension must match shape of the random walk, and
+        the first element will be discarded (since there is no innovation in the first timestep)
     sigma : tensor
         sigma > 0, innovation standard deviation (only required if tau is not specified)
+        For vector valued sigma, first dimension must match shape of the random walk, and
+        the first element will be discarded (since there is no innovation in the first timestep)
     tau : tensor
         tau > 0, innovation precision (only required if sigma is not specified)
+        For vector valued tau, first dimension must match shape of the random walk, and
+        the first element will be discarded (since there is no innovation in the first timestep)
     init : distribution
         distribution for initial value (Defaults to Flat())
     """
 
     def __init__(self, tau=None, init=Flat.dist(), sigma=None, mu=0.,
                  sd=None, *args, **kwargs):
+        kwargs.setdefault('shape', 1)
         super().__init__(*args, **kwargs)
+        if sum(self.shape) == 0:
+            raise TypeError("GaussianRandomWalk must be supplied a non-zero shape argument!")
         if sd is not None:
             sigma = sd
         tau, sigma = get_tau_sigma(tau=tau, sigma=sigma)
-        self.tau = tau = tt.as_tensor_variable(tau)
-        self.sigma = self.sd = sigma = tt.as_tensor_variable(sigma)
-        self.mu = mu = tt.as_tensor_variable(mu)
+        self.tau = tt.as_tensor_variable(tau)
+        sigma = tt.as_tensor_variable(sigma)
+        self.sigma = self.sd = sigma
+        self.mu = tt.as_tensor_variable(mu)
         self.init = init
         self.mean = tt.as_tensor_variable(0.)
 
+    def _mu_and_sigma(self, mu, sigma):
+        """Helper to get mu and sigma if they are high dimensional."""
+        if sigma.ndim > 0:
+            sigma = sigma[1:]
+        if mu.ndim > 0:
+            mu = mu[1:]
+        return mu, sigma
+
     def logp(self, x):
         """
         Calculate log-probability of Gaussian Random Walk distribution at specified value.
@@ -206,15 +225,45 @@ def logp(self, x):
         -------
         TensorVariable
         """
-        sigma = self.sigma
-        mu = self.mu
-        init = self.init
+        if x.ndim > 0:
+            x_im1 = x[:-1]
+            x_i = x[1:]
+            mu, sigma = self._mu_and_sigma(self.mu, self.sigma)
+            innov_like = Normal.dist(mu=x_im1 + mu, sigma=sigma).logp(x_i)
+            return self.init.logp(x[0]) + tt.sum(innov_like)
+        return self.init.logp(x)
 
-        x_im1 = x[:-1]
-        x_i = x[1:]
+    def random(self, point=None, size=None):
+        """Draw random values from GaussianRandomWalk.
+
+        Parameters
+        ----------
+        point : dict, optional
+            Dict of variable values on which random values are to be
+            conditioned (uses default point if not specified).
+        size : int, optional
+            Desired size of random sample (returns one sample if not
+            specified).
 
-        innov_like = Normal.dist(mu=x_im1 + mu, sigma=sigma).logp(x_i)
-        return init.logp(x[0]) + tt.sum(innov_like)
+        Returns
+        -------
+        array
+        """
+        sigma, mu = distribution.draw_values([self.sigma, self.mu], point=point, size=size)
+        return distribution.generate_samples(self._random, sigma=sigma, mu=mu, size=size,
+                                             dist_shape=self.shape,
+                                             not_broadcast_kwargs={"sample_shape": to_tuple(size)})
+
+    def _random(self, sigma, mu, size, sample_shape):
+        """Implement a Gaussian random walk as a cumulative sum of normals."""
+        if size[len(sample_shape)] == sample_shape:
+            axis = len(sample_shape)
+        else:
+            axis = 0
+        rv = stats.norm(mu, sigma)
+        data = rv.rvs(size).cumsum(axis=axis)
+        data = data - data[0]  # TODO: this should be a draw from `init`, if available
+        return data
 
     def _repr_latex_(self, name=None, dist=None):
         if dist is None: