From f36c433866099bbdfae6be8ae4a84d88d35dba8f Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Mon, 8 Nov 2021 21:21:00 +0100 Subject: [PATCH 1/2] Add DensityDist moment --- pymc/distributions/distribution.py | 11 ++++++++--- pymc/tests/test_distributions_moments.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 37b5483297..0eb9808322 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -24,6 +24,7 @@ import aesara from aeppl.logprob import _logcdf, _logprob +from aesara import tensor as at from aesara.tensor.basic import as_tensor_variable from aesara.tensor.random.op import RandomVariable from aesara.tensor.random.var import RandomStateSharedVariable @@ -472,9 +473,9 @@ def __new__( as the first argument ``rv``. ``size`` is the random variable's size implied by the ``dims``, ``size`` and parameters supplied to the distribution. Finally, ``rv_inputs`` is the sequence of the distribution parameters, in the same order - as they were supplied when the DensityDist was created. If ``None``, a - ``NotImplemented`` error will be raised when trying to draw random samples from - the distribution's prior or posterior predictive. + as they were supplied when the DensityDist was created. If ``None``, a default + ``get_moment`` function will be assigned that will always return 0, or an array + of zeros. ndim_supp : int The number of dimensions in the support of the distribution. Defaults to assuming a scalar distribution, i.e. ``ndim_supp = 0``. @@ -614,3 +615,7 @@ def func(*args, **kwargs): raise NotImplementedError(message) return func + + +def default_get_moment(rv, size, *rv_inputs): + return at.zeros(size, dtype=rv.dtype) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 5af9900550..8425a97fbd 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from aesara import tensor as at from scipy import special from pymc.distributions import ( @@ -13,6 +14,7 @@ Cauchy, ChiSquared, Constant, + DensityDist, Dirichlet, DiscreteUniform, ExGaussian, @@ -898,3 +900,20 @@ def test_rice_moment(nu, sigma, size, expected): with Model() as model: Rice("x", nu=nu, sigma=sigma, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "get_moment, size, expected", + [ + (None, None, 0.0), + (None, 5, np.zeros(5)), + ("custom_moment", None, 5), + ("custom_moment", (2, 5), np.full((2, 5), 5)), + ], +) +def test_density_dist_moment(get_moment, size, expected): + if get_moment == "custom_moment": + get_moment = lambda rv, size, *rv_inputs: 5 * at.ones(size, dtype=rv.dtype) + with Model() as model: + DensityDist("x", get_moment=get_moment, size=size) + assert_moment_is_expected(model, expected) From 45e33bcfbff2094e2ef597cf9ad7ba0db5f2a6dd Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Tue, 9 Nov 2021 12:04:58 +0100 Subject: [PATCH 2/2] Specialize get_moment for multivariate density dists --- pymc/distributions/distribution.py | 24 ++++++-- pymc/tests/test_distributions_moments.py | 71 +++++++++++++++++++++++- pymc/tests/test_moment.py | 38 ------------- 3 files changed, 88 insertions(+), 45 deletions(-) delete mode 100644 pymc/tests/test_moment.py diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 0eb9808322..fd849be49e 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -551,12 +551,17 @@ def random(mu, rng=None, size=None): if logcdf is None: logcdf = default_not_implemented(name, "logcdf") + if get_moment is None: + get_moment = functools.partial( + default_get_moment, + rv_name=name, + has_fallback=random is not None, + ndim_supp=ndim_supp, + ) + if random is None: random = default_not_implemented(name, "random") - if get_moment is None: - get_moment = default_not_implemented(name, "get_moment") - rv_op = type( f"DensityDist_{name}", (DensityDistRV,), @@ -617,5 +622,14 @@ def func(*args, **kwargs): return func -def default_get_moment(rv, size, *rv_inputs): - return at.zeros(size, dtype=rv.dtype) +def default_get_moment(rv, size, *rv_inputs, rv_name=None, has_fallback=False, ndim_supp=0): + if ndim_supp == 0: + return at.zeros(size, dtype=rv.dtype) + elif has_fallback: + return at.zeros_like(rv) + else: + raise TypeError( + "Cannot safely infer the size of a multivariate random variable's moment. " + f"Please provide a get_moment function when instantiating the {rv_name} " + "random variable." + ) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 8425a97fbd..fd37eaac8b 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1,9 +1,12 @@ +import aesara import numpy as np import pytest from aesara import tensor as at from scipy import special +import pymc as pm + from pymc.distributions import ( AsymmetricLaplace, Bernoulli, @@ -49,8 +52,9 @@ ZeroInflatedBinomial, ZeroInflatedPoisson, ) +from pymc.distributions.distribution import get_moment from pymc.distributions.multivariate import MvNormal -from pymc.distributions.shape_utils import rv_size_is_none +from pymc.distributions.shape_utils import rv_size_is_none, to_tuple from pymc.initial_point import make_initial_point_fn from pymc.model import Model @@ -911,9 +915,72 @@ def test_rice_moment(nu, sigma, size, expected): ("custom_moment", (2, 5), np.full((2, 5), 5)), ], ) -def test_density_dist_moment(get_moment, size, expected): +def test_density_dist_default_moment_univariate(get_moment, size, expected): if get_moment == "custom_moment": get_moment = lambda rv, size, *rv_inputs: 5 * at.ones(size, dtype=rv.dtype) with Model() as model: DensityDist("x", get_moment=get_moment, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str) +def test_density_dist_custom_moment_univariate(size): + def moment(rv, size, mu): + return (at.ones(size) * mu).astype(rv.dtype) + + mu_val = np.array(np.random.normal(loc=2, scale=1)).astype(aesara.config.floatX) + with pm.Model(): + mu = pm.Normal("mu") + a = pm.DensityDist("a", mu, get_moment=moment, size=size) + evaled_moment = get_moment(a).eval({mu: mu_val}) + assert evaled_moment.shape == to_tuple(size) + assert np.all(evaled_moment == mu_val) + + +@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str) +def test_density_dist_custom_moment_multivariate(size): + def moment(rv, size, mu): + return (at.ones(size)[..., None] * mu).astype(rv.dtype) + + mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX) + with pm.Model(): + mu = pm.Normal("mu", size=5) + a = pm.DensityDist("a", mu, get_moment=moment, ndims_params=[1], ndim_supp=1, size=size) + evaled_moment = get_moment(a).eval({mu: mu_val}) + assert evaled_moment.shape == to_tuple(size) + (5,) + assert np.all(evaled_moment == mu_val) + + +@pytest.mark.parametrize( + "with_random, size", + [ + (True, ()), + (True, (2,)), + (True, (3, 2)), + (False, ()), + (False, (2,)), + ], +) +def test_density_dist_default_moment_multivariate(with_random, size): + def _random(mu, rng=None, size=None): + return rng.normal(mu, scale=1, size=to_tuple(size) + mu.shape) + + if with_random: + random = _random + else: + random = None + + mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX) + with pm.Model(): + mu = pm.Normal("mu", size=5) + a = pm.DensityDist("a", mu, random=random, ndims_params=[1], ndim_supp=1, size=size) + if with_random: + evaled_moment = get_moment(a).eval({mu: mu_val}) + assert evaled_moment.shape == to_tuple(size) + (5,) + assert np.all(evaled_moment == 0) + else: + with pytest.raises( + TypeError, + match="Cannot safely infer the size of a multivariate random variable's moment.", + ): + evaled_moment = get_moment(a).eval({mu: mu_val}) diff --git a/pymc/tests/test_moment.py b/pymc/tests/test_moment.py deleted file mode 100644 index 17ab2423a8..0000000000 --- a/pymc/tests/test_moment.py +++ /dev/null @@ -1,38 +0,0 @@ -import aesara -import numpy as np -import pytest - -from aesara import tensor as at - -import pymc as pm - -from pymc.distributions.distribution import get_moment -from pymc.distributions.shape_utils import to_tuple - - -@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str) -def test_density_dist_moment_scalar(size): - def moment(rv, size, mu): - return (at.ones(size) * mu).astype(rv.dtype) - - mu_val = np.array(np.random.normal(loc=2, scale=1)).astype(aesara.config.floatX) - with pm.Model(): - mu = pm.Normal("mu") - a = pm.DensityDist("a", mu, get_moment=moment, size=size) - evaled_moment = get_moment(a).eval({mu: mu_val}) - assert evaled_moment.shape == to_tuple(size) - assert np.all(evaled_moment == mu_val) - - -@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str) -def test_density_dist_moment_multivariate(size): - def moment(rv, size, mu): - return (at.ones(size)[..., None] * mu).astype(rv.dtype) - - mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX) - with pm.Model(): - mu = pm.Normal("mu", size=5) - a = pm.DensityDist("a", mu, get_moment=moment, ndims_params=[1], ndim_supp=1, size=size) - evaled_moment = get_moment(a).eval({mu: mu_val}) - assert evaled_moment.shape == to_tuple(size) + (5,) - assert np.all(evaled_moment == mu_val)