diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 0980356432..ed8b4469e8 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -86,7 +86,7 @@ def polyagamma_cdf(*args, **kwargs): ) from pymc.distributions.distribution import Continuous from pymc.distributions.shape_utils import rv_size_is_none -from pymc.math import logdiffexp, logit +from pymc.math import invlogit, logdiffexp, logit from pymc.util import UNSET __all__ = [ @@ -3101,6 +3101,12 @@ def dist(cls, lower=0, upper=1, c=0.5, *args, **kwargs): return super().dist([lower, c, upper], *args, **kwargs) + def get_moment(rv, size, lower, c, upper): + mean = (lower + upper + c) / 3 + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logcdf(value, lower, c, upper): """ Compute the log of the cumulative distribution function for Triangular distribution @@ -3198,6 +3204,12 @@ def dist( return super().dist([mu, beta], **kwargs) + def get_moment(rv, size, mu, beta): + mean = mu + beta * np.euler_gamma + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def _distr_parameters_for_repr(self): return ["mu", "beta"] @@ -3501,6 +3513,12 @@ def dist(cls, mu=0, sigma=None, tau=None, sd=None, **kwargs): return super().dist([mu, sigma], **kwargs) + def get_moment(rv, size, mu, sigma): + median, _ = at.broadcast_arrays(invlogit(mu), sigma) + if not rv_size_is_none(size): + median = at.full(size, median) + return median + def logp(value, mu, sigma): """ Calculate log-probability of LogitNormal distribution at specified value. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index bbe478f8ac..fa5144aac5 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -17,6 +17,7 @@ Flat, Gamma, Geometric, + Gumbel, HalfCauchy, HalfFlat, HalfNormal, @@ -25,12 +26,14 @@ Kumaraswamy, Laplace, Logistic, + LogitNormal, LogNormal, NegativeBinomial, Normal, Pareto, Poisson, StudentT, + Triangular, TruncatedNormal, Uniform, Wald, @@ -650,3 +653,62 @@ def test_dirichlet_moment(a, size, expected): with Model() as model: Dirichlet("x", a=a, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "mu, beta, size, expected", + [ + (0, 2, None, 2 * np.euler_gamma), + (1, np.arange(1, 4), None, 1 + np.arange(1, 4) * np.euler_gamma), + (np.arange(5), 2, None, np.arange(5) + 2 * np.euler_gamma), + (1, 2, 5, np.full(5, 1 + 2 * np.euler_gamma)), + ( + np.arange(5), + np.arange(1, 6), + (2, 5), + np.full((2, 5), np.arange(5) + np.arange(1, 6) * np.euler_gamma), + ), + ], +) +def test_gumbel_moment(mu, beta, size, expected): + with Model() as model: + Gumbel("x", mu=mu, beta=beta, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "c, lower, upper, size, expected", + [ + (1, 0, 5, None, 2), + (3, np.arange(-3, 6, 3), np.arange(3, 12, 3), None, np.array([1, 3, 5])), + (np.arange(-3, 6, 3), -3, 3, None, np.array([-1, 0, 1])), + (3, -3, 6, 5, np.full(5, 2)), + ( + np.arange(-3, 6, 3), + np.arange(-9, -2, 3), + np.arange(3, 10, 3), + (2, 3), + np.full((2, 3), np.array([-3, 0, 3])), + ), + ], +) +def test_triangular_moment(c, lower, upper, size, expected): + with Model() as model: + Triangular("x", c=c, lower=lower, upper=upper, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "mu, sigma, size, expected", + [ + (1, 2, None, special.expit(1)), + (0, np.arange(1, 5), None, special.expit(np.zeros(4))), + (np.arange(4), 1, None, special.expit(np.arange(4))), + (1, 5, 4, special.expit(np.ones(4))), + (np.arange(4), np.arange(1, 5), (2, 4), np.full((2, 4), special.expit(np.arange(4)))), + ], +) +def test_logitnormal_moment(mu, sigma, size, expected): + with Model() as model: + LogitNormal("x", mu=mu, sigma=sigma, size=size) + assert_moment_is_expected(model, expected)