Fix kronecker_normal signature, fix rice mean, add wishhart mean

Author: jessegrabowski

pymc/distributions/moments/means.py

@@ -78,9 +78,10 @@
     MatrixNormalRV,
     MvStudentTRV,
     StickBreakingWeightsRV,
+    WishartRV,
     _LKJCholeskyCovRV,
 )
-from pymc.distributions.shape_utils import rv_size_is_none, maybe_resize
+from pymc.distributions.shape_utils import maybe_resize, rv_size_is_none
 from pymc.exceptions import UndefinedMomentException
 
 __all__ = ["mean"]
@@ -244,7 +245,7 @@ def invgamma_mean(op, rv, rng, size, alpha, beta):
 
 
 @_mean.register(KroneckerNormalRV)
-def kronecker_normal_mean(op, rv, rng, size, mu, covs, chols, evds):
+def kronecker_normal_mean(op, rv, rng, size, mu, sigma, *covs):
     mean = mu
     if not rv_size_is_none(size):
         mean_size = pt.concatenate([size, mu.shape])
@@ -370,7 +371,9 @@ def polya_gamma_mean(op, rv, rng, size, h, z):
 
 
 @_mean.register(RiceRV)
-def rice_mean(op, rv, rng, size, nu, sigma):
+def rice_mean(op, rv, rng, size, b, sigma):
+    # b is the shape parameter, nu = b * sigma is the noncentrality parameter
+    nu = b * sigma
     nu_sigma_ratio = -(nu**2) / (2 * sigma**2)
     return maybe_resize(
         sigma
@@ -448,3 +451,12 @@ def wald_mean(op, rv, rng, size, mu, lam, alpha):
 @_mean.register(WeibullBetaRV)
 def weibull_mean(op, rv, rng, size, alpha, beta):
     return maybe_resize(beta * pt.gamma(1 + 1 / alpha), size)
+
+
+@_mean.register(WishartRV)
+def wishart_mean(op, rv, rng, size, nu, V):
+    mean = nu * V
+    if not rv_size_is_none(size):
+        mean_size = pt.concatenate([size, V.shape[-2:]])
+        mean = pt.full(mean_size, mean)
+    return mean

tests/distributions/moments/test_means.py

@@ -12,6 +12,8 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import warnings
+
 import numpy as np
 import pytest
 
@@ -53,6 +55,7 @@
     uniform,
     vonmises,
     weibull_min,
+    wishart,
 )
 
 from pymc import (
@@ -168,9 +171,7 @@
         [Normal, norm, {"mu": 2, "sigma": 2}, {"loc": 2, "scale": 2}],
         [Pareto, pareto, {"alpha": 5, "m": 2}, {"b": 5, "scale": 2}],
         [Poisson, poisson, {"mu": 20}, {"mu": 20}],
-        pytest.param(
-            Rice, rice, {"b": 2, "sigma": 2}, {"b": 2, "scale": 2}, marks=pytest.mark.xfail
-        ),  # Something is wrong with the Rice mean, maybe a Bessel function in pytensor?
+        [Rice, rice, {"b": 2, "sigma": 2}, {"b": 2, "scale": 2}],
         [SkewNormal, skewnorm, {"mu": 2, "sigma": 2, "alpha": 2}, {"loc": 2, "scale": 2, "a": 2}],
         [
             SkewStudentT,
@@ -262,6 +263,22 @@ def test_mean_equal_expected(dist, dist_params, expected):
     )
 
 
+def test_wishart_mean():
+    nu = 10
+    V = np.array([[2.0, 0.5], [0.5, 1.5]])
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        from pymc import Wishart
+
+        rv = Wishart.dist(nu=nu, V=V)
+
+    pymc_mean_val = mean(rv).eval()
+    scipy_mean = wishart(df=nu, scale=V).mean()
+
+    np.testing.assert_almost_equal(pymc_mean_val, scipy_mean)
+
+
 @pytest.mark.parametrize(
     ["dist", "dist_params"],
     [