Rename GaussianRandomWalk init argument to init_dist and fix some type hints

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -125,22 +125,24 @@ class GaussianRandomWalkRV(RandomVariable):
     dtype = "floatX"
     _print_name = ("GaussianRandomWalk", "\\operatorname{GaussianRandomWalk}")
 
-    def make_node(self, rng, size, dtype, mu, sigma, init, steps):
+    def make_node(self, rng, size, dtype, mu, sigma, init_dist, steps):
         steps = at.as_tensor_variable(steps)
         if not steps.ndim == 0 or not steps.dtype.startswith("int"):
             raise ValueError("steps must be an integer scalar (ndim=0).")
 
         mu = at.as_tensor_variable(mu)
         sigma = at.as_tensor_variable(sigma)
-        init = at.as_tensor_variable(init)
+        init_dist = at.as_tensor_variable(init_dist)
 
         # Resize init distribution
         size = normalize_size_param(size)
         # If not explicit, size is determined by the shapes of mu, sigma, and init
-        init_size = size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init)
-        init = change_rv_size(init, init_size)
+        init_dist_size = (
+            size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init_dist)
+        )
+        init_dist = change_rv_size(init_dist, init_dist_size)
 
-        return super().make_node(rng, size, dtype, mu, sigma, init, steps)
+        return super().make_node(rng, size, dtype, mu, sigma, init_dist, steps)
 
     def _supp_shape_from_params(self, dist_params, reop_param_idx=0, param_shapes=None):
         steps = dist_params[3]
@@ -153,7 +155,7 @@ def rng_fn(
         rng: np.random.RandomState,
         mu: Union[np.ndarray, float],
         sigma: Union[np.ndarray, float],
-        init: float,
+        init_dist: Union[np.ndarray, float],
         steps: int,
         size: Tuple[int],
     ) -> np.ndarray:
@@ -170,16 +172,16 @@ def rng_fn(
         ----------
         rng: np.random.RandomState
            Numpy random number generator
-        mu: array_like
+        mu: array_like of float
            Random walk mean
-        sigma: np.ndarray
+        sigma: array_like of float
             Standard deviation of innovation (sigma > 0)
-        init: float
+        init_dist: array_like of float
             Initialization value for GaussianRandomWalk
         steps: int
             Length of random walk, must be greater than 1. Returned array will be of size+1 to
             account as first value is initial value
-        size: int
+        size: tuple of int
             The number of Random Walk time series generated
 
         Returns
@@ -196,7 +198,7 @@ def rng_fn(
             bcast_shape = np.broadcast_shapes(
                 np.asarray(mu).shape,
                 np.asarray(sigma).shape,
-                np.asarray(init).shape,
+                np.asarray(init_dist).shape,
             )
             dist_shape = (*bcast_shape, int(steps))
 
@@ -207,7 +209,7 @@ def rng_fn(
         # Add one dimension to the right, so that mu and sigma broadcast safely along
         # the steps dimension
         innovations = rng.normal(loc=mu[..., None], scale=sigma[..., None], size=dist_shape)
-        grw = np.concatenate([init[..., None], innovations], axis=-1)
+        grw = np.concatenate([init_dist[..., None], innovations], axis=-1)
         return np.cumsum(grw, axis=-1)
 
 
@@ -223,7 +225,7 @@ class GaussianRandomWalk(distribution.Continuous):
         innovation drift, defaults to 0.0
     sigma : tensor_like of float, optional
         sigma > 0, innovation standard deviation, defaults to 1.0
-    init : unnamed distribution
+    init_dist : unnamed distribution
         Univariate distribution of the initial value, created with the `.dist()` API.
         Defaults to a unit Normal.
 
@@ -248,7 +250,7 @@ def __new__(cls, *args, steps=None, **kwargs):
 
     @classmethod
     def dist(
-        cls, mu=0.0, sigma=1.0, *, init=None, steps=None, size=None, **kwargs
+        cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, size=None, **kwargs
     ) -> at.TensorVariable:
 
         mu = at.as_tensor_variable(floatX(mu))
@@ -263,27 +265,34 @@ def dist(
             raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps))
 
+        if "init" in kwargs:
+            warnings.warn(
+                "init parameter is now called init_dist. Using init will raise an error in a future release.",
+                FutureWarning,
+            )
+            init_dist = kwargs.pop("init")
+
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
-        if init is None:
-            init = Normal.dist(0, 1)
+        if init_dist is None:
+            init_dist = Normal.dist(0, 1)
         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
+                isinstance(init_dist, at.TensorVariable)
+                and init_dist.owner is not None
+                and isinstance(init_dist.owner.op, RandomVariable)
+                and init_dist.owner.op.ndim_supp == 0
             ):
                 raise TypeError("init must be a univariate distribution variable")
-            check_dist_not_registered(init)
+            check_dist_not_registered(init_dist)
 
         # Ignores logprob of init var because that's accounted for in the logp method
-        init = ignore_logprob(init)
+        init_dist = ignore_logprob(init_dist)
 
-        return super().dist([mu, sigma, init, steps], size=size, **kwargs)
+        return super().dist([mu, sigma, init_dist, steps], size=size, **kwargs)
 
-    def moment(rv, size, mu, sigma, init, steps):
+    def moment(rv, size, mu, sigma, init_dist, steps):
         grw_moment = at.zeros_like(rv)
-        grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init))
+        grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init_dist))
         # Add one dimension to the right, so that mu broadcasts safely along the steps
         # dimension
         grw_moment = at.set_subtensor(grw_moment[..., 1:], mu[..., None])
@@ -293,13 +302,13 @@ def logp(
         value: at.Variable,
         mu: at.Variable,
         sigma: at.Variable,
-        init: at.Variable,
+        init_dist: at.Variable,
         steps: at.Variable,
     ) -> at.TensorVariable:
         """Calculate log-probability of Gaussian Random Walk distribution at specified value."""
 
         # Calculate initialization logp
-        init_logp = logp(init, value[..., 0])
+        init_logp = logp(init_dist, value[..., 0])
 
         # Make time series stationary around the mean value
         stationary_series = value[..., 1:] - value[..., :-1]
@@ -429,7 +438,7 @@ def dist(
                 "init parameter is now called init_dist. Using init will raise an error in a future release.",
                 FutureWarning,
             )
-            init_dist = kwargs["init"]
+            init_dist = kwargs.pop("init")
 
         ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
         steps = get_steps(steps=steps, shape=kwargs.get("shape", None), step_shape_offset=ar_order)

pymc/tests/test_distributions_timeseries.py

@@ -100,8 +100,8 @@ class TestGaussianRandomWalkRandom(BaseTestDistributionRandom):
         size = None
 
         pymc_dist = pm.GaussianRandomWalk
-        pymc_dist_params = {"mu": 1.0, "sigma": 2, "init": pm.Constant.dist(0), "steps": 4}
-        expected_rv_op_params = {"mu": 1.0, "sigma": 2, "init": pm.Constant.dist(0), "steps": 4}
+        pymc_dist_params = {"mu": 1.0, "sigma": 2, "init_dist": pm.Constant.dist(0), "steps": 4}
+        expected_rv_op_params = {"mu": 1.0, "sigma": 2, "init_dist": pm.Constant.dist(0), "steps": 4}
 
         checks_to_run = [
             "check_pymc_params_match_rv_op",
@@ -142,36 +142,38 @@ def test_gaussianrandomwalk_inference(self):
     @pytest.mark.parametrize("init", [None, pm.Normal.dist()])
     def test_gaussian_random_walk_init_dist_shape(self, init):
         """Test that init_dist is properly resized"""
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, size=(5,))
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init, size=(5,))
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,)
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=2)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init, shape=2)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=(5, 2))
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init, shape=(5, 2))
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,)
 
-        grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,)
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,)
 
-        grw = pm.GaussianRandomWalk.dist(mu=np.zeros((3, 1)), sigma=[1, 1], steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=np.zeros((3, 1)), sigma=[1, 1], steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (3, 2)
 
     def test_shape_ellipsis(self):
         grw = pm.GaussianRandomWalk.dist(
-            mu=0, sigma=1, steps=5, init=pm.Normal.dist(), shape=(3, ...)
+            mu=0, sigma=1, steps=5, init_dist=pm.Normal.dist(), shape=(3, ...)
         )
         assert tuple(grw.shape.eval()) == (3, 6)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (3,)
 
     def test_gaussianrandomwalk_broadcasted_by_init_dist(self):
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=4, init=pm.Normal.dist(size=(2, 3)))
+        grw = pm.GaussianRandomWalk.dist(
+            mu=0, sigma=1, steps=4, init_dist=pm.Normal.dist(size=(2, 3))
+        )
         assert tuple(grw.shape.eval()) == (2, 3, 5)
         assert grw.eval().shape == (2, 3, 5)
 
@@ -210,14 +212,14 @@ def test_inferred_steps_from_observed(self):
         ],
     )
     def test_gaussian_random_walk_init_dist_logp(self, init):
-        grw = pm.GaussianRandomWalk.dist(init=init, steps=1)
+        grw = pm.GaussianRandomWalk.dist(init_dist=init, steps=1)
         assert np.isclose(
             pm.logp(grw, [0, 0]).eval(),
             pm.logp(init, 0).eval() + scipy.stats.norm.logpdf(0),
         )
 
     @pytest.mark.parametrize(
-        "mu, sigma, init, steps, size, expected",
+        "mu, sigma, init_dist, steps, size, expected",
         [
             (0, 1, Normal.dist(1), 10, None, np.ones((11,))),
             (1, 1, Normal.dist(0), 10, (2,), np.full((2, 11), np.arange(11))),
@@ -233,11 +235,15 @@ def test_gaussian_random_walk_init_dist_logp(self, init):
             ),
         ],
     )
-    def test_moment(self, mu, sigma, init, steps, size, expected):
+    def test_moment(self, mu, sigma, init_dist, steps, size, expected):
         with Model() as model:
-            GaussianRandomWalk("x", mu=mu, sigma=sigma, init=init, steps=steps, size=size)
+            GaussianRandomWalk("x", mu=mu, sigma=sigma, init_dist=init_dist, steps=steps, size=size)
         assert_moment_is_expected(model, expected)
 
+    def test_init_deprecated_arg(self):
+        with pytest.warns(FutureWarning, match="init parameter is now called init_dist"):
+            pm.GaussianRandomWalk.dist(init=Normal.dist(), shape=(10,))
+
 
 class TestAR:
     def test_order1_logp(self):
@@ -434,6 +440,10 @@ def test_moment(self, size, expected):
             AR("x", rho=[0, 0], init_dist=init_dist, steps=5, size=size)
         assert_moment_is_expected(model, expected, check_finite_logp=False)
 
+    def test_init_deprecated_arg(self):
+        with pytest.warns(FutureWarning, match="init parameter is now called init_dist"):
+            pm.AR.dist(rho=[1, 2, 3], init=Normal.dist(), shape=(10,))
+
 
 @pytest.mark.xfail(reason="Timeseries not refactored", raises=NotImplementedError)
 def test_GARCH11():