HalfCauchy, Gamma, Weibull and LogNormal moments

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

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)

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():