diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index f43ea7b0da..d3bfe644c4 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -1698,6 +1698,12 @@ def dist(cls, mu=0, sigma=None, tau=None, sd=None, *args, **kwargs): return super().dist([mu, sigma], *args, **kwargs) + def get_moment(rv, size, mu, sigma): + mean = at.exp(mu + 0.5 * sigma ** 2) + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logcdf(value, mu, sigma): """ Compute the log of the cumulative distribution function for LogNormal distribution @@ -2100,6 +2106,12 @@ def dist(cls, beta, *args, **kwargs): assert_negative_support(beta, "beta", "HalfCauchy") return super().dist([0.0, beta], **kwargs) + def get_moment(rv, size, loc, beta): + mean = beta + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logcdf(value, loc, beta): """ Compute the log of the cumulative distribution function for HalfCauchy distribution @@ -2214,6 +2226,13 @@ def get_alpha_beta(cls, alpha=None, beta=None, mu=None, sigma=None): return alpha, beta + def get_moment(rv, size, alpha, inv_beta): + # The Aesara `GammaRV` `Op` inverts the `beta` parameter itself + mean = alpha * inv_beta + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logcdf(value, alpha, inv_beta): """ Compute the log of the cumulative distribution function for Gamma distribution @@ -2495,6 +2514,12 @@ def dist(cls, alpha, beta, *args, **kwargs): return super().dist([alpha, beta], *args, **kwargs) + def get_moment(rv, size, alpha, beta): + mean = beta * at.gamma(1 + 1 / alpha) + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logcdf(value, alpha, beta): r""" Compute the log of the cumulative distribution function for Weibull distribution diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index d0c2d63324..36fe591e69 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1,15 +1,21 @@ import numpy as np import pytest +from scipy import special + from pymc import Bernoulli, Flat, HalfFlat, Normal, TruncatedNormal, Uniform from pymc.distributions import ( Beta, Cauchy, Exponential, + Gamma, + HalfCauchy, HalfNormal, Kumaraswamy, Laplace, + LogNormal, StudentT, + Weibull, ) from pymc.distributions.shape_utils import rv_size_is_none from pymc.initial_point import make_initial_point_fn @@ -241,3 +247,82 @@ def test_kumaraswamy_moment(a, b, size, expected): with Model() as model: Kumaraswamy("x", a=a, b=b, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "mu, sigma, size, expected", + [ + (0, 1, None, np.exp(0.5)), + (0, 1, 5, np.full(5, np.exp(0.5))), + (np.arange(5), 1, None, np.exp(np.arange(5) + 0.5)), + ( + np.arange(5), + np.arange(1, 6), + (2, 5), + np.full((2, 5), np.exp(np.arange(5) + 0.5 * np.arange(1, 6) ** 2)), + ), + ], +) +def test_lognormal_moment(mu, sigma, size, expected): + with Model() as model: + LogNormal("x", mu=mu, sigma=sigma, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "beta, size, expected", + [ + (1, None, 1), + (1, 5, np.ones(5)), + (np.arange(5), None, np.arange(5)), + ( + np.arange(5), + (2, 5), + np.full((2, 5), np.arange(5)), + ), + ], +) +def test_halfcauchy_moment(beta, size, expected): + with Model() as model: + HalfCauchy("x", beta=beta, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "alpha, beta, size, expected", + [ + (1, 1, None, 1), + (1, 1, 5, np.full(5, 1)), + (np.arange(1, 6), 1, None, np.arange(1, 6)), + ( + np.arange(1, 6), + 2 * np.arange(1, 6), + (2, 5), + np.full((2, 5), 0.5), + ), + ], +) +def test_gamma_moment(alpha, beta, size, expected): + with Model() as model: + Gamma("x", alpha=alpha, beta=beta, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "alpha, beta, size, expected", + [ + (1, 1, None, 1), + (1, 1, 5, np.full(5, 1)), + (np.arange(1, 6), 1, None, special.gamma(1 + 1 / np.arange(1, 6))), + ( + np.arange(1, 6), + np.arange(2, 7), + (2, 5), + np.full((2, 5), np.arange(2, 7) * special.gamma(1 + 1 / np.arange(1, 6))), + ), + ], +) +def test_weibull_moment(alpha, beta, size, expected): + with Model() as model: + Weibull("x", alpha=alpha, beta=beta, size=size) + assert_moment_is_expected(model, expected) diff --git a/pymc/tests/test_sampling.py b/pymc/tests/test_sampling.py index aaa6634a74..4bf6b8283d 100644 --- a/pymc/tests/test_sampling.py +++ b/pymc/tests/test_sampling.py @@ -1074,7 +1074,7 @@ def test_density_dist(self): obs = np.random.normal(-1, 0.1, size=10) with pm.Model(): mu = pm.Normal("mu", 0, 1) - sd = pm.Gamma("sd", 1, 2) + sd = pm.HalfNormal("sd", 1e-6) a = pm.DensityDist( "a", mu, @@ -1084,7 +1084,7 @@ def test_density_dist(self): ) prior = pm.sample_prior_predictive(return_inferencedata=False) - npt.assert_almost_equal(prior["a"].mean(), 0, decimal=1) + npt.assert_almost_equal((prior["a"] - prior["mu"][..., None]).mean(), 0, decimal=3) def test_shape_edgecase(self): with pm.Model():