Add random method to GaussianRandomWalk

Author: ColCarroll

pymc3/distributions/timeseries.py

@@ -1,10 +1,12 @@
+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 . import multivariate
 from . import distribution
+from .shape_utils import to_tuple
 
 
 __all__ = [
@@ -166,17 +168,22 @@ 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())
     """
@@ -187,9 +194,14 @@ def __init__(self, tau=None, init=Flat.dist(), sigma=None, mu=0.,
         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)
+        if sigma.ndim > 0:
+            sigma = sigma[:-1]
+        self.sigma = self.sd = sigma
+        self.mu = tt.as_tensor_variable(mu)
+        if self.mu.ndim > 0:
+            self.mu = self.mu[:-1]
         self.init = init
         self.mean = tt.as_tensor_variable(0.)
 
@@ -206,15 +218,41 @@ 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:]
 
-        x_im1 = x[:-1]
-        x_i = x[1:]
+            sigma = self.sigma
+            mu = self.mu
+
+            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)
+
+    def random(self, point=None, size=None):
+        """Draw random values from GaussianRandomWalk.
 
-        innov_like = Normal.dist(mu=x_im1 + mu, sigma=sigma).logp(x_i)
-        return init.logp(x[0]) + tt.sum(innov_like)
+        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).
+
+        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)
+
+    def _random(self, sigma, mu, size):
+        """Implement a Gaussian random walk as a cumulative sum of normals."""
+        rv = stats.norm(mu, sigma)
+        return rv.rvs(size).cumsum(axis=0)
 
     def _repr_latex_(self, name=None, dist=None):
         if dist is None:

pymc3/tests/test_distributions_random.py

@@ -251,6 +251,14 @@ def test_different_shapes_and_sample_sizes(self, shape):
                 assert e == a
 
 
+class TestGaussianRandomWalk(BaseTestCases.BaseTestCase):
+    distribution = pm.GaussianRandomWalk
+    params = {'mu': 1., 'sigma': 1.}
+
+    @pytest.mark.xfail(reason="Supporting this makes a nasty API")
+    def test_broadcast_shape(self):
+        super().test_broadcast_shape()
+
 class TestNormal(BaseTestCases.BaseTestCase):
     distribution = pm.Normal
     params = {'mu': 0., 'tau': 1.}
@@ -1006,7 +1014,7 @@ def test_density_dist_without_random_not_sampleable(self):
             normal_dist = pm.Normal.dist(mu, 1)
             pm.DensityDist('density_dist', normal_dist.logp, observed=np.random.randn(100))
             trace = pm.sample(100)
-    
+
         samples = 500
         with pytest.raises(ValueError):
             pm.sample_posterior_predictive(trace, samples=samples, model=model, size=100)