From dbeb80168093051b725abc50485bfc7839bcc0c8 Mon Sep 17 00:00:00 2001 From: Larry Dong Date: Thu, 26 May 2022 11:26:26 -0400 Subject: [PATCH] Use Aeppl for implementing GaussianRandomWalk Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> Co-authored-by: lucianopaz --- pymc/aesaraf.py | 52 ++++- pymc/distributions/timeseries.py | 211 ++++++-------------- pymc/tests/test_distributions_moments.py | 41 ++++ pymc/tests/test_distributions_timeseries.py | 85 ++++---- 4 files changed, 192 insertions(+), 197 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 7deed9fe83..1300593ca6 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -33,11 +33,12 @@ import scipy.sparse as sps from aeppl.abstract import MeasurableVariable -from aeppl.logprob import CheckParameterValue +from aeppl.logprob import CheckParameterValue, _logprob, logprob +from aeppl.tensor import MeasurableJoin from aesara import config, scalar from aesara.compile.mode import Mode, get_mode from aesara.gradient import grad -from aesara.graph import local_optimizer +from aesara.graph import local_optimizer, optimize_graph from aesara.graph.basic import ( Apply, Constant, @@ -52,6 +53,7 @@ from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream from aesara.scalar.basic import Cast from aesara.tensor.basic import _as_tensor_variable +from aesara.tensor.basic_opt import topo_constant_folding from aesara.tensor.elemwise import Elemwise from aesara.tensor.random.op import RandomVariable from aesara.tensor.random.var import ( @@ -875,6 +877,52 @@ def largest_common_dtype(tensors): return np.stack([np.ones((), dtype=dtype) for dtype in dtypes]).dtype +@_logprob.register(MeasurableJoin) +def logprob_join_constant_shapes(op, values, axis, *base_vars, **kwargs): + """Compute the log-likelihood graph for a `Join`. + + This overrides the implementation in Aeppl, to constant fold the shapes + of the base vars so that RandomVariables do not show up in the logp graph + """ + (value,) = values + + base_var_shapes = at.stack([base_var.shape[axis] for base_var in base_vars]) + + base_var_shapes = optimize_graph( + base_var_shapes, + custom_opt=topo_constant_folding, + ) + + split_values = at.split( + value, + splits_size=[base_var_shape for base_var_shape in base_var_shapes], + n_splits=len(base_vars), + axis=axis, + ) + + logps = [ + logprob(base_var, split_value) + for base_var, split_value in zip(base_vars, split_values) + ] + + if len(set(logp.ndim for logp in logps)) != 1: + raise ValueError( + "Joined logps have different number of dimensions, this can happen when " + "joining univariate and multivariate distributions", + ) + + base_vars_ndim_supp = split_values[0].ndim - logps[0].ndim + join_logprob = at.concatenate( + [ + at.atleast_1d(logprob(base_var, split_value)) + for base_var, split_value in zip(base_vars, split_values) + ], + axis=axis - base_vars_ndim_supp, + ) + + return join_logprob + + @local_optimizer(tracks=[CheckParameterValue]) def local_remove_check_parameter(fgraph, node): """Rewrite that removes Aeppl's CheckParameterValue diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 347ed90d0f..807365cf93 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -13,7 +13,7 @@ # limitations under the License. import warnings -from typing import Any, Optional, Tuple, Union +from typing import Any, Optional, Union import aesara import aesara.tensor as at @@ -28,22 +28,15 @@ from aesara.raise_op import Assert from aesara.tensor import TensorVariable from aesara.tensor.basic_opt import ShapeFeature, topo_constant_folding +from aesara.tensor.extra_ops import CumOp from aesara.tensor.random.op import RandomVariable -from aesara.tensor.random.utils import normalize_size_param 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 ( - Dims, - Shape, - convert_dims, - rv_size_is_none, - to_tuple, -) +from pymc.distributions.shape_utils import Dims, Shape, convert_dims, to_tuple from pymc.model import modelcontext from pymc.util import check_dist_not_registered @@ -114,110 +107,8 @@ def get_steps( return inferred_steps -class GaussianRandomWalkRV(RandomVariable): - """ - GaussianRandomWalk Random Variable - """ - - name = "GaussianRandomWalk" - ndim_supp = 1 - ndims_params = [0, 0, 0, 0] - dtype = "floatX" - _print_name = ("GaussianRandomWalk", "\\operatorname{GaussianRandomWalk}") - - 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_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_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_dist, steps) - - def _supp_shape_from_params(self, dist_params, reop_param_idx=0, param_shapes=None): - steps = dist_params[3] - - return (steps + 1,) - - @classmethod - def rng_fn( - cls, - rng: np.random.RandomState, - mu: Union[np.ndarray, float], - sigma: Union[np.ndarray, float], - init_dist: Union[np.ndarray, float], - steps: int, - size: Tuple[int], - ) -> np.ndarray: - """Gaussian Random Walk generator. - - The init value is drawn from the Normal distribution with the same sigma as the - innovations. - - Notes - ----- - Currently does not support custom init distribution - - Parameters - ---------- - rng: np.random.RandomState - Numpy random number generator - mu: array_like of float - Random walk mean - sigma: array_like of float - Standard deviation of innovation (sigma > 0) - 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: tuple of int - The number of Random Walk time series generated - - Returns - ------- - ndarray - """ - - if steps < 1: - raise ValueError("Steps must be greater than 0") - - # If size is None then the returned series should be (*implied_dims, 1+steps) - if size is None: - # broadcast parameters with each other to find implied dims - bcast_shape = np.broadcast_shapes( - np.asarray(mu).shape, - np.asarray(sigma).shape, - np.asarray(init_dist).shape, - ) - dist_shape = (*bcast_shape, int(steps)) - - # If size is None then the returned series should be (*size, 1+steps) - else: - dist_shape = (*size, int(steps)) - - # 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_dist[..., None], innovations], axis=-1) - return np.cumsum(grw, axis=-1) - - -gaussianrandomwalk = GaussianRandomWalkRV() - - -class GaussianRandomWalk(distribution.Continuous): - r"""Random Walk with Normal innovations. +class GaussianRandomWalk(SymbolicDistribution): + r"""Random Walk with Normal innovations Parameters ---------- @@ -236,8 +127,6 @@ class GaussianRandomWalk(distribution.Continuous): provided. """ - rv_op = gaussianrandomwalk - def __new__(cls, *args, steps=None, **kwargs): steps = get_steps( steps=steps, @@ -269,6 +158,8 @@ def dist(cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, **kwargs) -> at. FutureWarning, ) init_dist = kwargs.pop("init") + if not steps.ndim == 0: + raise ValueError("steps must be an integer scalar (ndim=0).") # If no scalar distribution is passed then initialize with a Normal of same mu and sigma if init_dist is None: @@ -293,39 +184,67 @@ def dist(cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, **kwargs) -> at. return super().dist([mu, sigma, init_dist, steps], **kwargs) - 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_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]) - return at.cumsum(grw_moment, axis=-1) - - def logp( - value: at.Variable, - mu: at.Variable, - sigma: 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_dist, value[..., 0]) - - # Make time series stationary around the mean value - stationary_series = value[..., 1:] - value[..., :-1] - # Add one dimension to the right, so that mu and sigma broadcast safely along - # the steps dimension - series_logp = logp(Normal.dist(mu[..., None], sigma[..., None]), stationary_series) - - return check_parameters( - init_logp + series_logp.sum(axis=-1), - steps > 0, - msg="steps > 0", + @classmethod + def ndim_supp(cls, *args): + return 1 + + @classmethod + def rv_op(cls, mu, sigma, init_dist, steps, size=None): + # If not explicit, size is determined by the shapes of mu, sigma, and init + if size is not None: + # we have all the information regarding size from users or .dist() + init_size = size + else: + # we infer size from parameters + init_size = at.broadcast_shape( + mu, + sigma, + init_dist, + ) + + # TODO: extend for multivariate init + init_dist = change_rv_size(init_dist, init_size) + innovation_dist = Normal.dist(mu[..., None], sigma[..., None], size=(*init_size, steps)) + rv_out = at.cumsum(at.concatenate([init_dist[..., None], innovation_dist], axis=-1), axis=-1) + + rv_out.tag.mu = mu + rv_out.tag.sigma = sigma + rv_out.tag.init_dist = init_dist + rv_out.tag.innovation_dist = innovation_dist + rv_out.tag.steps = steps + rv_out.tag.is_grw = True # for moment dispatching + + return rv_out + + @classmethod + def change_size(cls, rv, new_size, expand=False): + if expand: + new_size = at.concatenate([new_size, rv.shape[:-1]]) + + return cls.rv_op( + mu=rv.tag.mu, + sigma=rv.tag.sigma, + init_dist=rv.tag.init_dist, + steps=rv.tag.steps, + size=new_size, ) +@_moment.register(CumOp) +def moment_grw(op, rv, dist_params): + """ + This moment dispatch is currently only applicable for a GaussianRandomWalk. + TODO: Encapsulate GRW graph in an OpFromGraph so that we can dispatch + the moment directly on it + """ + if not getattr(rv.tag, "is_grw", False): + raise NotImplementedError("Moment not implemented for `CumOp`") + init_dist = rv.tag.init_dist + innovation_dist = rv.tag.innovation_dist + grw_moment = at.concatenate([moment(init_dist)[..., None], moment(innovation_dist)], axis=-1) + return at.cumsum(grw_moment, axis=-1) + + class AutoRegressiveRV(OpFromGraph): """A placeholder used to specify a log-likelihood for an AR sub-graph.""" diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 75c29849bc..11b865bf32 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -28,6 +28,7 @@ Exponential, Flat, Gamma, + GaussianRandomWalk, Geometric, Gumbel, HalfCauchy, @@ -1506,3 +1507,43 @@ def test_dirichlet_multinomial_moment(a, n, size, expected): with Model() as model: DirichletMultinomial("x", n=n, a=a, size=size) assert_moment_is_expected(model, expected) + + +even_numbers = np.insert(np.cumsum(np.tile(2, 10)), 0, 0) # for test_gaussianrandomwalk below + + +@pytest.mark.parametrize( + "mu, sigma, init_dist, steps, size, expected", + [ + (0, 3, StudentT.dist(5), 10, None, np.zeros(11)), + (Normal.dist(2, 3), Gamma.dist(1, 1), StudentT.dist(5), 10, None, even_numbers), + ( + Normal.dist(2, 3, size=(3, 5)), + Gamma.dist(1, 1), + StudentT.dist(5), + 10, + None, + np.broadcast_to(even_numbers, shape=(3, 5, 11)), + ), + ( + Normal.dist(2, 3, size=(3, 1)), + Gamma.dist(1, 1, size=(1, 5)), + StudentT.dist(5), + 10, + (3, 5), + np.broadcast_to(even_numbers, shape=(3, 5, 11)), + ), + ( + Normal.dist(2, 3), + Gamma.dist(1, 1), + StudentT.dist(5), + 10, + (3, 5), + np.broadcast_to(even_numbers, shape=(3, 5, 11)), + ), + ], +) +def test_gaussianrandomwalk(mu, sigma, init_dist, steps, size, expected): + with Model() as model: + GaussianRandomWalk("x", mu=mu, sigma=sigma, init_dist=init_dist, steps=steps, size=size) + assert_moment_is_expected(model, expected) diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py index 186c865451..fbcf7f4677 100644 --- a/pymc/tests/test_distributions_timeseries.py +++ b/pymc/tests/test_distributions_timeseries.py @@ -21,7 +21,7 @@ import pymc as pm from pymc.aesaraf import floatX -from pymc.distributions.continuous import Flat, HalfNormal, Normal +from pymc.distributions.continuous import Flat, HalfNormal, Normal, StudentT from pymc.distributions.discrete import DiracDelta from pymc.distributions.logprob import logp from pymc.distributions.multivariate import Dirichlet @@ -36,7 +36,6 @@ 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( @@ -95,37 +94,21 @@ def test_get_steps(info_source, steps, shape, step_shape_offset, expected_steps, class TestGaussianRandomWalk: - class TestGaussianRandomWalkRandom(BaseTestDistributionRandom): - # Override default size for test class - size = None - - pymc_dist = pm.GaussianRandomWalk - pymc_dist_params = {"mu": 1.0, "sigma": 2, "init_dist": pm.DiracDelta.dist(0), "steps": 4} - expected_rv_op_params = { - "mu": 1.0, - "sigma": 2, - "init_dist": pm.DiracDelta.dist(0), - "steps": 4, - } - - checks_to_run = [ - "check_pymc_params_match_rv_op", - "check_rv_inferred_size", - ] + pymc_dist = pm.GaussianRandomWalk - def check_rv_inferred_size(self): - steps = self.pymc_dist_params["steps"] - sizes_to_check = [None, (), 1, (1,)] - sizes_expected = [(steps + 1,), (steps + 1,), (1, steps + 1), (1, steps + 1)] + def check_rv_inferred_size(self): + steps = self.pymc_dist_params["steps"] + sizes_to_check = [None, (), 1, (1,)] + sizes_expected = [(steps + 1,), (steps + 1,), (1, steps + 1), (1, steps + 1)] - for size, expected in zip(sizes_to_check, sizes_expected): - pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size) - expected_symbolic = tuple(pymc_rv.shape.eval()) - assert expected_symbolic == expected + for size, expected in zip(sizes_to_check, sizes_expected): + pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size) + expected_symbolic = tuple(pymc_rv.shape.eval()) + assert expected_symbolic == expected - def test_steps_scalar_check(self): - with pytest.raises(ValueError, match="steps must be an integer scalar"): - self.pymc_dist.dist(steps=[1]) + def test_steps_scalar_check(self): + with pytest.raises(ValueError, match="steps must be an integer scalar"): + self.pymc_dist.dist(steps=[1]) def test_gaussianrandomwalk_inference(self): mu, sigma, steps = 2, 1, 1000 @@ -147,33 +130,37 @@ 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""" + + def get_init_dist(grw_dist): + return grw_dist.owner.inputs[0].owner.inputs[1].owner.inputs[0] + grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init) - assert tuple(grw.owner.inputs[-2].shape.eval()) == () + assert tuple(get_init_dist(grw).shape.eval()) == () grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init, size=(5,)) - assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,) + assert tuple(get_init_dist(grw).shape.eval()) == (5,) grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init, shape=2) - assert tuple(grw.owner.inputs[-2].shape.eval()) == () + assert tuple(get_init_dist(grw).shape.eval()) == () 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,) + assert tuple(get_init_dist(grw).shape.eval()) == (5,) grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init_dist=init) - assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,) + assert tuple(get_init_dist(grw).shape.eval()) == (2,) grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init_dist=init) - assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,) + assert tuple(get_init_dist(grw).shape.eval()) == (2,) 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) + assert tuple(get_init_dist(grw).shape.eval()) == (3, 2) def test_shape_ellipsis(self): grw = pm.GaussianRandomWalk.dist( 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,) + assert tuple(grw.owner.inputs[0].owner.inputs[1].owner.inputs[0].shape.eval()) == (3,) def test_gaussianrandomwalk_broadcasted_by_init_dist(self): grw = pm.GaussianRandomWalk.dist( @@ -185,7 +172,7 @@ def test_gaussianrandomwalk_broadcasted_by_init_dist(self): @pytest.mark.parametrize("shape", ((6,), (3, 6))) def test_inferred_steps_from_shape(self, shape): x = GaussianRandomWalk.dist(shape=shape) - steps = x.owner.inputs[-1] + steps = x.owner.inputs[0].owner.inputs[2].owner.inputs[1][-1] assert steps.eval() == 5 @pytest.mark.parametrize("shape", (None, (5, ...))) @@ -200,13 +187,13 @@ def test_inconsistent_steps_and_shape(self): 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] + steps = x.owner.inputs[0].owner.inputs[2].owner.inputs[1][-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] + steps = x.owner.inputs[0].owner.inputs[2].owner.inputs[1] assert steps.eval() == 9 @pytest.mark.parametrize( @@ -219,30 +206,30 @@ def test_inferred_steps_from_observed(self): def test_gaussian_random_walk_init_dist_logp(self, init): grw = pm.GaussianRandomWalk.dist(init_dist=init, steps=1) assert np.isclose( - pm.logp(grw, [0, 0]).eval(), + pm.logp(grw, [0, 0]).eval().sum(), pm.logp(init, 0).eval() + scipy.stats.norm.logpdf(0), ) @pytest.mark.parametrize( - "mu, sigma, init_dist, steps, size, expected", + "mu, sigma, init, 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))), + (0, 1, StudentT.dist(5), 10, None, np.zeros((11,))), + (1, 1, StudentT.dist(5), 10, (2,), np.full((2, 11), np.arange(11))), (1, 1, Normal.dist([0, 1]), 10, None, np.vstack((np.arange(11), np.arange(11) + 1))), - (0, [1, 1], Normal.dist(0), 10, None, np.zeros((2, 11))), + (0, [1, 1], StudentT.dist(5), 10, None, np.zeros((2, 11))), ( [1, -1], 1, - Normal.dist(0), + StudentT.dist(5), 10, (4, 2), np.full((4, 2, 11), np.vstack((np.arange(11), -np.arange(11)))), ), ], ) - def test_moment(self, mu, sigma, init_dist, steps, size, expected): + def test_moment(self, mu, sigma, init, steps, size, expected): with Model() as model: - GaussianRandomWalk("x", mu=mu, sigma=sigma, init_dist=init_dist, steps=steps, size=size) + GaussianRandomWalk("x", mu=mu, sigma=sigma, init_dist=init, steps=steps, size=size) assert_moment_is_expected(model, expected) def test_init_deprecated_arg(self):