Obtain step information from dims and observed

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 import warnings
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import aesara
 import aesara.tensor as at
@@ -31,13 +31,20 @@
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.utils import normalize_size_param
 
-from pymc.aesaraf import change_rv_size, floatX, intX
+from pymc.aesaraf import change_rv_size, convert_observed_data, floatX, intX
 from pymc.distributions import distribution, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import SymbolicDistribution, _moment, moment
 from pymc.distributions.logprob import ignore_logprob, logp
-from pymc.distributions.shape_utils import Shape, rv_size_is_none, to_tuple
+from pymc.distributions.shape_utils import (
+    Dims,
+    Shape,
+    convert_dims,
+    rv_size_is_none,
+    to_tuple,
+)
+from pymc.model import modelcontext
 from pymc.util import check_dist_not_registered
 
 __all__ = [
@@ -50,51 +57,61 @@
 ]
 
 
-def get_steps_from_shape(
+def get_steps(
     steps: Optional[Union[int, np.ndarray, TensorVariable]],
-    shape: Optional[Shape],
+    *,
+    shape: Optional[Shape] = None,
+    dims: Optional[Dims] = None,
+    observed: Optional[Any] = None,
     step_shape_offset: int = 0,
 ):
-    """Extract number of steps from shape information
+    """Extract number of steps from shape / dims / observed information
 
     Parameters
     ----------
     steps:
         User specified steps for timeseries distribution
     shape:
         User specified shape for timeseries distribution
+    dims:
+        User specified dims for timeseries distribution
+    observed:
+        User specified observed data from timeseries distribution
     step_shape_offset:
         Difference between last shape dimension and number of steps in timeseries
         distribution, defaults to 0
 
-    Raises
-    ------
-    ValueError
-        If neither shape nor steps are provided
-
     Returns
     -------
     steps
         Steps, if specified directly by user, or inferred from the last dimension of
-        shape. When both steps and shape are provided, a symbolic Assert is added
-        to make sure they are consistent.
+        shape / dims / observed. When two sources of step information are provided,
+        a symbolic Assert is added to ensure they are consistent.
     """
-    steps_from_shape = None
+    inferred_steps = None
     if shape is not None:
         shape = to_tuple(shape)
         if shape[-1] is not ...:
-            steps_from_shape = shape[-1] - step_shape_offset
-    if steps is None:
-        if steps_from_shape is not None:
-            steps = steps_from_shape
-        else:
-            raise ValueError("Must specify steps or shape parameter")
-    elif steps_from_shape is not None:
-        # Assert that steps and shape are consistent
-        steps = Assert(msg="Steps do not match last shape dimension")(
-            steps, at.eq(steps, steps_from_shape)
+            inferred_steps = shape[-1] - step_shape_offset
+
+    if inferred_steps is None and dims is not None:
+        dims = convert_dims(dims)
+        if dims[-1] is not ...:
+            model = modelcontext(None)
+            inferred_steps = model.dim_lengths[dims[-1]] - step_shape_offset
+
+    if inferred_steps is None and observed is not None:
+        observed = convert_observed_data(observed)
+        inferred_steps = observed.shape[-1] - step_shape_offset
+
+    if inferred_steps is None:
+        inferred_steps = steps
+    # If there are two sources of information for the steps, assert they are consistent
+    elif steps is not None:
+        inferred_steps = Assert(msg="Steps do not match last shape dimension")(
+            inferred_steps, at.eq(inferred_steps, steps)
         )
-    return steps
+    return inferred_steps
 
 
 class GaussianRandomWalkRV(RandomVariable):
@@ -212,26 +229,38 @@ class GaussianRandomWalk(distribution.Continuous):
 
         .. warning:: init will be cloned, rendering them independent of the ones passed as input.
 
-    steps : int
-        Number of steps in Gaussian Random Walks (steps > 0).
+    steps : int, optional
+        Number of steps in Gaussian Random Walk (steps > 0). Only needed if size is
+        used to specify distribution
     """
 
     rv_op = gaussianrandomwalk
 
-    def __new__(cls, name, mu=0.0, sigma=1.0, init=None, steps=None, **kwargs):
-        if init is not None:
-            check_dist_not_registered(init)
-        return super().__new__(cls, name, mu, sigma, init, steps, **kwargs)
+    def __new__(cls, *args, steps=None, **kwargs):
+        steps = get_steps(
+            steps=steps,
+            shape=None,  # Shape will be checked in `cls.dist`
+            dims=kwargs.get("dims", None),
+            observed=kwargs.get("observed", None),
+            step_shape_offset=1,
+        )
+        return super().__new__(cls, *args, steps=steps, **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=None, steps=None, size=None, **kwargs
     ) -> at.TensorVariable:
 
         mu = at.as_tensor_variable(floatX(mu))
         sigma = at.as_tensor_variable(floatX(sigma))
 
-        steps = get_steps_from_shape(steps, kwargs.get("shape", None), step_shape_offset=1)
+        steps = get_steps(
+            steps=steps,
+            shape=kwargs.get("shape", None),
+            step_shape_offset=1,
+        )
+        if steps is None:
+            raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps))
 
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
@@ -245,6 +274,7 @@ def dist(
                 and init.owner.op.ndim_supp == 0
             ):
                 raise TypeError("init must be a univariate distribution variable")
+            check_dist_not_registered(init)
 
         # Ignores logprob of init var because that's accounted for in the logp method
         init = ignore_logprob(init)
@@ -340,6 +370,9 @@ class AR(SymbolicDistribution):
     ar_order: int, optional
         Order of the AR process. Inferred from length of the last dimension of rho, if
         possible. ar_order = rho.shape[-1] if constant else rho.shape[-1] - 1
+    steps : int, optional
+        Number of steps in AR process (steps > 0). Only needed if size is used to
+        specify distribution
 
     Notes
     -----
@@ -360,6 +393,15 @@ class AR(SymbolicDistribution):
 
     """
 
+    def __new__(cls, *args, steps=None, **kwargs):
+        steps = get_steps(
+            steps=steps,
+            shape=None,  # Shape will be checked in `cls.dist`
+            dims=kwargs.get("dims", None),
+            observed=kwargs.get("observed", None),
+        )
+        return super().__new__(cls, *args, steps=steps, **kwargs)
+
     @classmethod
     def dist(
         cls,
@@ -384,7 +426,9 @@ def dist(
             )
             init_dist = kwargs["init"]
 
-        steps = get_steps_from_shape(steps, kwargs.get("shape", None))
+        steps = get_steps(steps=steps, shape=kwargs.get("shape", None))
+        if steps is None:
+            raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps), ndim=0)
 
         if ar_order is None:

pymc/tests/test_distributions_timeseries.py

@@ -16,21 +16,84 @@
 import pytest
 import scipy.stats
 
+from aesara.tensor import TensorVariable
+
 import pymc as pm
 
 from pymc.aesaraf import floatX
 from pymc.distributions.continuous import Flat, HalfNormal, Normal
 from pymc.distributions.discrete import Constant
 from pymc.distributions.logprob import logp
 from pymc.distributions.multivariate import Dirichlet
-from pymc.distributions.timeseries import AR, GARCH11, EulerMaruyama, GaussianRandomWalk
+from pymc.distributions.timeseries import (
+    AR,
+    GARCH11,
+    EulerMaruyama,
+    GaussianRandomWalk,
+    get_steps,
+)
 from pymc.model import Model
 from pymc.sampling import draw, sample, sample_posterior_predictive
 from pymc.tests.helpers import select_by_precision
 from pymc.tests.test_distributions_moments import assert_moment_is_expected
 from pymc.tests.test_distributions_random import BaseTestDistributionRandom
 
 
+@pytest.mark.parametrize(
+    "steps, shape, step_shape_offset, expected_steps, consistent",
+    [
+        (10, None, 0, 10, True),
+        (10, None, 1, 10, True),
+        (None, (10,), 0, 10, True),
+        (None, (10,), 1, 9, True),
+        (None, (10, 5), 0, 5, True),
+        (None, (10, ...), 0, None, True),
+        (None, None, 0, None, True),
+        (10, (10,), 0, 10, True),
+        (10, (11,), 1, 10, True),
+        (10, (5, ...), 1, 10, True),
+        (10, (5, 5), 0, 5, False),
+        (10, (5, 10), 1, 9, False),
+    ],
+)
+@pytest.mark.parametrize("info_source", ("shape", "dims", "observed"))
+def test_get_steps(info_source, steps, shape, step_shape_offset, expected_steps, consistent):
+    if info_source == "shape":
+        inferred_steps = get_steps(steps=steps, shape=shape, step_shape_offset=step_shape_offset)
+
+    elif info_source == "dims":
+        if shape is None:
+            dims = None
+            coords = {}
+        else:
+            dims = tuple(str(i) if shape is not ... else ... for i, shape in enumerate(shape))
+            coords = {str(i): range(shape) for i, shape in enumerate(shape) if shape is not ...}
+        with Model(coords=coords):
+            inferred_steps = get_steps(steps=steps, dims=dims, step_shape_offset=step_shape_offset)
+
+    elif info_source == "observed":
+        if shape is None:
+            observed = None
+        else:
+            if ... in shape:
+                # There is no equivalent to implied dims in observed
+                return
+            observed = np.zeros(shape)
+        inferred_steps = get_steps(
+            steps=steps, observed=observed, step_shape_offset=step_shape_offset
+        )
+
+    if not isinstance(inferred_steps, TensorVariable):
+        assert inferred_steps == expected_steps
+    else:
+        if consistent:
+            assert inferred_steps.eval() == expected_steps
+        else:
+            assert inferred_steps.owner.inputs[0].eval() == expected_steps
+            with pytest.raises(AssertionError, match="Steps do not match"):
+                inferred_steps.eval()
+
+
 class TestGaussianRandomWalk:
     class TestGaussianRandomWalkRandom(BaseTestDistributionRandom):
         # Override default size for test class
@@ -127,6 +190,18 @@ def test_inconsistent_steps_and_shape(self):
         with pytest.raises(AssertionError, match="Steps do not match last shape dimension"):
             x = GaussianRandomWalk.dist(steps=12, shape=45)
 
+    def test_inferred_steps_from_dims(self):
+        with pm.Model(coords={"batch": range(5), "steps": range(20)}):
+            x = GaussianRandomWalk("x", dims=("batch", "steps"))
+        steps = x.owner.inputs[-1]
+        assert steps.eval() == 19
+
+    def test_inferred_steps_from_observed(self):
+        with pm.Model():
+            x = GaussianRandomWalk("x", observed=np.zeros(10))
+        steps = x.owner.inputs[-1]
+        assert steps.eval() == 9
+
     @pytest.mark.parametrize(
         "init",
         [