Allow init to be a distribution again

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -20,7 +20,7 @@
 from aesara import scan
 from aesara.tensor.random.op import RandomVariable
 
-from pymc.aesaraf import floatX, intX
+from pymc.aesaraf import change_rv_size, floatX, intX
 from pymc.distributions import distribution, logprob, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.shape_utils import to_tuple
@@ -35,6 +35,8 @@
     "MvStudentTRandomWalk",
 ]
 
+from pymc.util import check_dist_not_registered
+
 
 class GaussianRandomWalkRV(RandomVariable):
     """
@@ -107,10 +109,10 @@ def rng_fn(
             init_size = (*size, 1)
             steps_size = (*size, steps)
 
-        init_val = rng.normal(init, sigma, size=init_size)
+        init = np.reshape(init, init_size)
         steps = rng.normal(loc=mu, scale=sigma, size=steps_size)
 
-        grw = np.concatenate([init_val, steps], axis=-1)
+        grw = np.concatenate([init, steps], axis=-1)
 
         return np.cumsum(grw, axis=-1)
 
@@ -132,8 +134,9 @@ class GaussianRandomWalk(distribution.Continuous):
         innovation drift, defaults to 0.0
     sigma: tensor_like of float, optional
         sigma > 0, innovation standard deviation, defaults to 0.0
-    init: tensor_like of float, optional
-        Mean value of initialization, defaults to 0.0
+    init: Scalar PyMC distribution
+        Scalar distribution of the initial value, created with the `.dist()` API. Defaults to
+        Normal with same `mu` and `sigma` as the GaussianRandomWalk
     steps: int
         Number of steps in Gaussian Random Walks
     size: int
@@ -142,14 +145,37 @@ class GaussianRandomWalk(distribution.Continuous):
 
     rv_op = gaussianrandomwalk
 
+    def __new__(cls, name, mu=0.0, sigma=1.0, init=None, steps: int = 1, **kwargs):
+        check_dist_not_registered(init)
+        return super().__new__(cls, name, mu, sigma, init, steps, **kwargs)
+
     @classmethod
-    def dist(cls, mu=0.0, sigma=1.0, *, steps: int, init=0.0, **kwargs) -> RandomVariable:
+    def dist(
+        cls, mu=0.0, sigma=1.0, init=None, steps: int = 1, size=None, **kwargs
+    ) -> RandomVariable:
+
+        mu = at.as_tensor_variable(floatX(mu))
+        sigma = at.as_tensor_variable(floatX(sigma))
+        steps = at.as_tensor_variable(intX(steps))
 
-        params = [at.as_tensor_variable(floatX(param)) for param in (mu, sigma, init)] + [
-            at.as_tensor_variable(intX(steps))
-        ]
+        if init is None:
+            init = Normal.dist(mu, sigma, size=size)
+        else:
+            if not (
+                isinstance(init, at.TensorVariable)
+                and init.owner is not None
+                and isinstance(init.owner.op, RandomVariable)
+                and init.owner.op.ndim_supp == 0
+            ):
+                raise TypeError("init must be a scalar distribution variable")
+            if size is not None or shape is not None:
+                init = change_rv_size(init, to_tuple(size or shape))
+            else:
+                # If not explicit, size is determined by the shape of mu and sigma
+                mu_ = at.broadcast_arrays(mu, sigma)[0]
+                init = change_rv_size(init, mu_.shape)
 
-        return super().dist(params, **kwargs)
+        return super().dist([mu, sigma, init, steps], size=size, **kwargs)
 
     def logp(
         value: at.Variable,
@@ -174,11 +200,13 @@ def logp(
         """
 
         # Calculate initialization logp
+        init_logp = logprob.logp(init, value[..., 0])
+
         # Make time series stationary around the mean value
-        stationary_series = at.diff(value)
+        stationary_series = at.diff(value, axis=-1)
         series_logp = logprob.logp(Normal.dist(mu, sigma), stationary_series)
 
-        return series_logp
+        return init_logp + series_logp.sum(axis=-1)
 
 
 class AR1(distribution.Continuous):

pymc/tests/test_distributions.py

@@ -2622,7 +2622,7 @@ def test_grw(self):
             pm.GaussianRandomWalk,
             R,
             {"mu": R, "sigma": Rplus, "steps": Nat},
-            lambda value, mu, sigma: sp.norm.logpdf(value, mu, sigma).cumsum(),
+            lambda value, mu, sigma: sp.norm.logpdf(value, mu, sigma).cumsum().sum(),
             decimal=select_by_precision(float64=6, float32=1),
         )
 

pymc/tests/test_distributions_timeseries.py

@@ -61,7 +61,7 @@ def test_grw_logp(self):
         vals = [0, 1, 2]
         mu = 1
         sigma = 1
-        init = 0
+        init = pm.Normal.dist(mu, sigma)
 
         with pm.Model():
             grw = GaussianRandomWalk("grw", mu, sigma, init=init, steps=2)
@@ -88,7 +88,7 @@ def test_grw_inference(self):
         with pm.Model():
             _mu = pm.Uniform("mu", -10, 10)
             _sigma = pm.Uniform("sigma", 0, 10)
-            grw = GaussianRandomWalk("grw", _mu, _sigma, init=0, steps=steps, observed=obs)
+            grw = GaussianRandomWalk("grw", _mu, _sigma, steps=steps, observed=obs)
 
             with pytest.raises(TypeError) as err:
                 trace = pm.sample()
@@ -113,6 +113,16 @@ def test_grw_shape(self, steps, size, expected):
         expected_symbolic = tuple(grw_dist.shape.eval())
         assert expected_symbolic == expected
 
+    @pytest.mark.parametrize("size", (None, (1, 2), (10, 2), (3, 100, 2)))
+    def test_init_automatically_resized(self, size):
+        x = GaussianRandomWalk.dist(mu=[0, 1], init=pm.Normal.dist(), size=size)
+        init = x.owner.inputs[-2]
+        assert init.eval().shape == size if size is not None else (2,)
+
+        x = GaussianRandomWalk.dist(mu=[0, 1], init=pm.Normal.dist(size=5), shape=size)
+        init = x.owner.inputs[-2]
+        assert init.eval().shape == size if size is not None else (2,)
+
 
 @pytest.mark.xfail(reason="Timeseries not refactored")
 def test_AR():