From 135fb37e4ccfea2d8329f6e971183471a023be1d Mon Sep 17 00:00:00 2001 From: Ricardo Date: Thu, 5 May 2022 11:22:28 +0200 Subject: [PATCH 1/3] Rename `pandas_to_array` to `convert_observed_data` --- pymc/aesaraf.py | 9 +++------ pymc/data.py | 6 +++--- pymc/distributions/shape_utils.py | 4 ++-- pymc/model.py | 6 +++--- pymc/tests/test_aesaraf.py | 10 +++++----- 5 files changed, 16 insertions(+), 19 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index e2de6d9fe2..0767bee29d 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -81,16 +81,13 @@ "set_at_rng", "at_rng", "take_along_axis", - "pandas_to_array", + "convert_observed_data", ] -def pandas_to_array(data): - """Convert a pandas object to a NumPy array. +def convert_observed_data(data): + """Convert user provided dataset to accepted formats.""" - XXX: When `data` is a generator, this will return an Aesara tensor! - - """ if hasattr(data, "to_numpy") and hasattr(data, "isnull"): # typically, but not limited to pandas objects vals = data.to_numpy() diff --git a/pymc/data.py b/pymc/data.py index a240712982..0f1a54d09c 100644 --- a/pymc/data.py +++ b/pymc/data.py @@ -33,7 +33,7 @@ import pymc as pm -from pymc.aesaraf import pandas_to_array +from pymc.aesaraf import convert_observed_data __all__ = [ "get_data", @@ -636,9 +636,9 @@ def Data( ) name = model.name_for(name) - # `pandas_to_array` takes care of parameter `value` and + # `convert_observed_data` takes care of parameter `value` and # transforms it to something digestible for Aesara. - arr = pandas_to_array(value) + arr = convert_observed_data(value) if mutable is None: major, minor = (int(v) for v in pm.__version__.split(".")[:2]) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 8110b3d280..23c9237199 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -26,7 +26,7 @@ from aesara.tensor.var import TensorVariable from typing_extensions import TypeAlias -from pymc.aesaraf import pandas_to_array +from pymc.aesaraf import convert_observed_data __all__ = [ "to_tuple", @@ -558,7 +558,7 @@ def resize_from_observed( Observations as numpy array or `Variable`. """ if not hasattr(observed, "shape"): - observed = pandas_to_array(observed) + observed = convert_observed_data(observed) ndim_resize = observed.ndim - ndim_implied resize_shape = tuple(observed.shape[d] for d in range(ndim_resize)) return resize_shape, observed diff --git a/pymc/model.py b/pymc/model.py index d13091be69..47ecabc732 100644 --- a/pymc/model.py +++ b/pymc/model.py @@ -50,10 +50,10 @@ from pymc.aesaraf import ( compile_pymc, + convert_observed_data, gradient, hessian, inputvars, - pandas_to_array, rvs_to_value_vars, ) from pymc.blocking import DictToArrayBijection, RaveledVars @@ -1158,7 +1158,7 @@ def set_data( if isinstance(values, list): values = np.array(values) - values = pandas_to_array(values) + values = convert_observed_data(values) dims = self.RV_dims.get(name, None) or () coords = coords or {} @@ -1290,7 +1290,7 @@ def make_obs_var( """ name = rv_var.name - data = pandas_to_array(data).astype(rv_var.dtype) + data = convert_observed_data(data).astype(rv_var.dtype) if data.ndim != rv_var.ndim: raise ShapeError( diff --git a/pymc/tests/test_aesaraf.py b/pymc/tests/test_aesaraf.py index 709f21c6ef..6b032ff2c6 100644 --- a/pymc/tests/test_aesaraf.py +++ b/pymc/tests/test_aesaraf.py @@ -38,8 +38,8 @@ _conversion_map, change_rv_size, compile_pymc, + convert_observed_data, extract_obs_data, - pandas_to_array, rvs_to_value_vars, take_along_axis, walk_model, @@ -413,9 +413,9 @@ def test_extract_obs_data(): @pytest.mark.parametrize("input_dtype", ["int32", "int64", "float32", "float64"]) -def test_pandas_to_array(input_dtype): +def test_convert_observed_data(input_dtype): """ - Ensure that pandas_to_array returns the dense array, masked array, + Ensure that convert_observed_data returns the dense array, masked array, graph variable, TensorVariable, or sparse matrix as appropriate. """ pd = pytest.importorskip("pandas") @@ -437,7 +437,7 @@ def test_pandas_to_array(input_dtype): square_generator = (np.array([i**2], dtype=int) for i in range(100)) # Alias the function to be tested - func = pandas_to_array + func = convert_observed_data ##### # Perform the various tests @@ -496,7 +496,7 @@ def test_pandas_to_array(input_dtype): def test_pandas_to_array_pandas_index(): pd = pytest.importorskip("pandas") data = pd.Index([1, 2, 3]) - result = pandas_to_array(data) + result = convert_observed_data(data) expected = np.array([1, 2, 3]) np.testing.assert_array_equal(result, expected) From 614bb0698514d38d4d8ee599d8845d3e68d64d13 Mon Sep 17 00:00:00 2001 From: Ricardo Date: Thu, 5 May 2022 11:31:54 +0200 Subject: [PATCH 2/3] Obtain step information from dims and observed --- pymc/distributions/timeseries.py | 112 ++++++++++++++------ pymc/tests/test_distributions_timeseries.py | 77 +++++++++++++- 2 files changed, 154 insertions(+), 35 deletions(-) diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 7649d0069d..67ff11fcd2 100644 --- a/pymc/distributions/timeseries.py +++ b/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,12 +57,15 @@ ] -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 ---------- @@ -63,38 +73,45 @@ def get_steps_from_shape( 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: diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py index d2fe7275d2..d6df46d61f 100644 --- a/pymc/tests/test_distributions_timeseries.py +++ b/pymc/tests/test_distributions_timeseries.py @@ -16,6 +16,8 @@ import pytest import scipy.stats +from aesara.tensor import TensorVariable + import pymc as pm from pymc.aesaraf import floatX @@ -23,7 +25,13 @@ 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 @@ -31,6 +39,61 @@ 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", [ From 504da82af6023e278e8fd4f5c2baae0e75c6f407 Mon Sep 17 00:00:00 2001 From: Ricardo Date: Thu, 5 May 2022 15:09:06 +0200 Subject: [PATCH 3/3] Make AR steps extend shape beyond initial_dist This is consistent with the meaning of steps in the GaussianRandomWalk and translates directly to the number of scan steps taken --- pymc/distributions/timeseries.py | 72 +++++++++++++-------- pymc/tests/test_distributions_timeseries.py | 6 +- 2 files changed, 48 insertions(+), 30 deletions(-) diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 67ff11fcd2..d57d5b170f 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -393,14 +393,19 @@ class AR(SymbolicDistribution): """ - def __new__(cls, *args, steps=None, **kwargs): + def __new__(cls, name, rho, *args, steps=None, constant=False, ar_order=None, **kwargs): + rhos = at.atleast_1d(at.as_tensor_variable(floatX(rho))) + ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order) 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=ar_order, + ) + return super().__new__( + cls, name, rhos, *args, steps=steps, constant=constant, ar_order=ar_order, **kwargs ) - return super().__new__(cls, *args, steps=steps, **kwargs) @classmethod def dist( @@ -426,34 +431,12 @@ def dist( ) init_dist = kwargs["init"] - steps = get_steps(steps=steps, shape=kwargs.get("shape", None)) + 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) 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: - # If ar_order is not specified we do constant folding on the shape of rhos - # to retrieve it. For example, this will detect that - # Normal(size=(5, 3)).shape[-1] == 3, which is not known by Aesara before. - shape_fg = FunctionGraph( - outputs=[rhos.shape[-1]], - features=[ShapeFeature()], - clone=True, - ) - (folded_shape,) = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs - folded_shape = getattr(folded_shape, "data", None) - if folded_shape is None: - raise ValueError( - "Could not infer ar_order from last dimension of rho. Pass it " - "explictily or make sure rho have a static shape" - ) - ar_order = int(folded_shape) - int(constant) - if ar_order < 1: - raise ValueError( - "Inferred ar_order is smaller than 1. Increase the last dimension " - "of rho or remove constant_term" - ) - if init_dist is not None: if not isinstance(init_dist, TensorVariable) or not isinstance( init_dist.owner.op, RandomVariable @@ -477,6 +460,41 @@ def dist( return super().dist([rhos, sigma, init_dist, steps, ar_order, constant], **kwargs) + @classmethod + def _get_ar_order(cls, rhos: TensorVariable, ar_order: Optional[int], constant: bool) -> int: + """Compute ar_order given inputs + + If ar_order is not specified we do constant folding on the shape of rhos + to retrieve it. For example, this will detect that + Normal(size=(5, 3)).shape[-1] == 3, which is not known by Aesara before. + + Raises + ------ + ValueError + If inferred ar_order cannot be inferred from rhos or if it is less than 1 + """ + if ar_order is None: + shape_fg = FunctionGraph( + outputs=[rhos.shape[-1]], + features=[ShapeFeature()], + clone=True, + ) + (folded_shape,) = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs + folded_shape = getattr(folded_shape, "data", None) + if folded_shape is None: + raise ValueError( + "Could not infer ar_order from last dimension of rho. Pass it " + "explictily or make sure rho have a static shape" + ) + ar_order = int(folded_shape) - int(constant) + if ar_order < 1: + raise ValueError( + "Inferred ar_order is smaller than 1. Increase the last dimension " + "of rho or remove constant_term" + ) + + return ar_order + @classmethod def num_rngs(cls, *args, **kwargs): return 2 @@ -540,7 +558,7 @@ def step(*args): fn=step, outputs_info=[{"initial": init_.T, "taps": range(-ar_order, 0)}], non_sequences=[rhos_bcast_.T[::-1], sigma_.T, noise_rng], - n_steps=at.max((0, steps_ - ar_order)), + n_steps=steps_, strict=True, ) (noise_next_rng,) = tuple(innov_updates_.values()) diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py index d6df46d61f..eadafc2d63 100644 --- a/pymc/tests/test_distributions_timeseries.py +++ b/pymc/tests/test_distributions_timeseries.py @@ -363,7 +363,7 @@ def test_batched_sigma(self): beta_tp.set_value(np.zeros((ar_order,))) # Should always be close to zero sigma_tp = np.full(batch_size, [0.01, 0.1, 1, 10, 100]) y_eval = t0["y"].eval({t0["sigma"]: sigma_tp}) - assert y_eval.shape == (*batch_size, steps) + assert y_eval.shape == (*batch_size, steps + ar_order) assert np.allclose(y_eval.std(axis=(0, 2)), [0.01, 0.1, 1, 10, 100], rtol=0.1) def test_batched_init_dist(self): @@ -389,7 +389,7 @@ def test_batched_init_dist(self): init_dist = t0["y"].owner.inputs[2] init_dist_tp = np.full((batch_size, ar_order), (np.arange(batch_size) * 100)[:, None]) y_eval = t0["y"].eval({init_dist: init_dist_tp}) - assert y_eval.shape == (batch_size, steps) + assert y_eval.shape == (batch_size, steps + ar_order) assert np.allclose( y_eval[:, -10:].mean(-1), np.arange(batch_size) * 100, rtol=0.1, atol=0.5 ) @@ -429,7 +429,7 @@ def test_multivariate_init_dist(self): def test_moment(self, size, expected): with Model() as model: init_dist = Constant.dist([[1.0, 2.0], [3.0, 4.0]]) - AR("x", rho=[0, 0], init_dist=init_dist, steps=7, size=size) + AR("x", rho=[0, 0], init_dist=init_dist, steps=5, size=size) assert_moment_is_expected(model, expected, check_finite_logp=False)