Make MvStudentT distribution v4 compatible

Author: Sayam753

pymc3/distributions/multivariate.py

@@ -27,6 +27,7 @@
 from aesara.tensor import gammaln
 from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace
 from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal
+from aesara.tensor.random.op import RandomVariable, default_shape_from_params
 from aesara.tensor.random.utils import broadcast_params
 from aesara.tensor.slinalg import (
     Cholesky,
@@ -248,6 +249,66 @@ def _distr_parameters_for_repr(self):
         return ["mu", "cov"]
 
 
+def safe_multivariate_t(nu, mu, cov, size=None, rng=None):
+    res = np.atleast_1d(
+        stats.multivariate_t(loc=mu, shape=cov, df=nu, allow_singular=True).rvs(
+            size=size, random_state=rng
+        )
+    )
+
+    if size is not None:
+        res = res.reshape(list(size) + [-1])
+
+    return res
+
+
+class MvStudentTRV(RandomVariable):
+    name = "multivariate_studentt"
+    ndim_supp = 1
+    ndims_params = [0, 1, 2]
+    dtype = "floatX"
+    _print_name = ("MvStudentT", "\\operatorname{MvStudentT}")
+
+    def __call__(self, nu, mu=None, cov=None, size=None, **kwargs):
+
+        dtype = aesara.config.floatX if self.dtype == "floatX" else self.dtype
+
+        if mu is None:
+            mu = np.array([0.0], dtype=dtype)
+        if cov is None:
+            cov = np.array([[1.0]], dtype=dtype)
+        return super().__call__(nu, mu, cov, size=size, **kwargs)
+
+    def _shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
+        return default_shape_from_params(self.ndim_supp, dist_params, rep_param_idx, param_shapes)
+
+    @classmethod
+    def rng_fn(cls, rng, nu, mu, cov, size):
+
+        if mu.ndim > 1 or cov.ndim > 2:
+            # Neither SciPy nor NumPy implement parameter broadcasting for
+            # multivariate normals (or many other multivariate distributions),
+            # so we have implement a quick and dirty one here
+            mu, cov = broadcast_params([mu, cov], cls.ndims_params[1:])
+            size = tuple(size or ())
+
+            if size:
+                mu = np.broadcast_to(mu, size + mu.shape)
+                cov = np.broadcast_to(cov, size + cov.shape)
+
+            res = np.empty(mu.shape)
+            for idx in np.ndindex(mu.shape[:-1]):
+                m = mu[idx]
+                c = cov[idx]
+                res[idx] = safe_multivariate_t(nu, m, c, rng=rng)
+            return res
+        else:
+            return safe_multivariate_t(nu, mu, cov, size=size, rng=rng)
+
+
+mv_studentt = MvStudentTRV()
+
+
 class MvStudentT(Continuous):
     r"""
     Multivariate Student-T log-likelihood.
@@ -288,55 +349,20 @@ class MvStudentT(Continuous):
     lower: bool, default=True
         Whether the cholesky fatcor is given as a lower triangular matrix.
     """
+    rv_op = mv_studentt
 
-    def __init__(
-        self, nu, Sigma=None, mu=None, cov=None, tau=None, chol=None, lower=True, *args, **kwargs
-    ):
+    @classmethod
+    def dist(cls, nu, Sigma=None, mu=None, cov=None, tau=None, chol=None, lower=True, **kwargs):
         if Sigma is not None:
             if cov is not None:
                 raise ValueError("Specify only one of cov and Sigma")
             cov = Sigma
-        super().__init__(mu=mu, cov=cov, tau=tau, chol=chol, lower=lower, *args, **kwargs)
-        self.nu = nu = at.as_tensor_variable(nu)
-        self.mean = self.median = self.mode = self.mu = self.mu
-
-    def random(self, point=None, size=None):
-        """
-        Draw random values from Multivariate Student's T distribution.
-
-        Parameters
-        ----------
-        point: dict, optional
-            Dict of variable values on which random values are to be
-            conditioned (uses default point if not specified).
-        size: int, optional
-            Desired size of random sample (returns one sample if not
-            specified).
-
-        Returns
-        -------
-        array
-        """
-        # with _DrawValuesContext():
-        #     nu, mu = draw_values([self.nu, self.mu], point=point, size=size)
-        #     if self._cov_type == "cov":
-        #         (cov,) = draw_values([self.cov], point=point, size=size)
-        #         dist = MvNormal.dist(mu=np.zeros_like(mu), cov=cov, shape=self.shape)
-        #     elif self._cov_type == "tau":
-        #         (tau,) = draw_values([self.tau], point=point, size=size)
-        #         dist = MvNormal.dist(mu=np.zeros_like(mu), tau=tau, shape=self.shape)
-        #     else:
-        #         (chol,) = draw_values([self.chol_cov], point=point, size=size)
-        #         dist = MvNormal.dist(mu=np.zeros_like(mu), chol=chol, shape=self.shape)
-        #
-        #     samples = dist.random(point, size)
-        #
-        # chi2_samples = np.random.chisquare(nu, size)
-        # # Add distribution shape to chi2 samples
-        # chi2_samples = chi2_samples.reshape(chi2_samples.shape + (1,) * len(self.shape))
-        # return (samples / np.sqrt(chi2_samples / nu)) + mu
+        nu = at.as_tensor_variable(nu)
+        mu = at.as_tensor_variable(mu)
+        cov = quaddist_matrix(cov, chol, tau, lower)
+        return super().dist([nu, mu, cov], **kwargs)
 
-    def logp(value, nu, cov):
+    def logp(value, nu, mu, cov):
         """
         Calculate log-probability of Multivariate Student's T distribution
         at specified value.
@@ -350,15 +376,15 @@ def logp(value, nu, cov):
         -------
         TensorVariable
         """
-        quaddist, logdet, ok = quaddist_parse(value, nu, cov)
+        quaddist, logdet, ok = quaddist_parse(value, mu, cov)
         k = floatX(value.shape[-1])
 
         norm = gammaln((nu + k) / 2.0) - gammaln(nu / 2.0) - 0.5 * k * floatX(np.log(nu * np.pi))
         inner = -(nu + k) / 2.0 * at.log1p(quaddist / nu)
         return bound(norm + inner - logdet, ok)
 
     def _distr_parameters_for_repr(self):
-        return ["mu", "nu", "cov"]
+        return ["nu", "mu", "cov"]
 
 
 class Dirichlet(Continuous):

pymc3/tests/test_distributions.py

@@ -2019,7 +2019,6 @@ def test_kroneckernormal(self, n, m, sigma):
         )
 
     @pytest.mark.parametrize("n", [1, 2])
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_mvt(self, n):
         self.check_logp(
             MvStudentT,
@@ -2032,6 +2031,7 @@ def test_mvt(self, n):
             RealMatrix(2, n),
             {"nu": Rplus, "Sigma": PdMatrix(n), "mu": Vector(R, n)},
             mvt_logpdf,
+            extra_args={"size": 2},
         )
 
     @pytest.mark.parametrize("n", [2, 3, 4])

pymc3/tests/test_distributions_random.py

@@ -1402,7 +1402,6 @@ def ref_rand_evd(size, mu, evds, sigma):
                 model_args=evd_args,
             )
 
-    @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_mv_t(self):
         def ref_rand(size, nu, Sigma, mu):
             normal = st.multivariate_normal.rvs(cov=Sigma, size=size)