Refactor Student T Distribution (#4694)

* feat: adapt student t
* Change default parameterization in terms of sigma and tweak tests

Co-authored-by: Ricardo <[email protected]>

Author: themrzmaster

pymc3/distributions/continuous.py

@@ -1793,6 +1793,21 @@ def logcdf(value, mu, sigma):
         )
 
 
+class StudentTRV(RandomVariable):
+    name = "studentt"
+    ndim_supp = 0
+    ndims_params = [0, 0, 0]
+    dtype = "floatX"
+    _print_name = ("StudentT", "\\operatorname{StudentT}")
+
+    @classmethod
+    def rng_fn(cls, rng, nu, mu, sigma, size=None):
+        return stats.t.rvs(nu, mu, sigma, size=size, random_state=rng)
+
+
+studentt = StudentTRV()
+
+
 class StudentT(Continuous):
     r"""
     Student's T log-likelihood.
@@ -1856,45 +1871,22 @@ class StudentT(Continuous):
         with pm.Model():
             x = pm.StudentT('x', nu=15, mu=0, lam=1/23)
     """
+    rv_op = studentt
 
-    def __init__(self, nu, mu=0, lam=None, sigma=None, sd=None, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    @classmethod
+    def dist(cls, nu, mu=0, lam=None, sigma=None, sd=None, *args, **kwargs):
         if sd is not None:
             sigma = sd
-        self.nu = nu = at.as_tensor_variable(floatX(nu))
+        nu = at.as_tensor_variable(floatX(nu))
         lam, sigma = get_tau_sigma(tau=lam, sigma=sigma)
-        self.lam = lam = at.as_tensor_variable(lam)
-        self.sigma = self.sd = sigma = at.as_tensor_variable(sigma)
-        self.mean = self.median = self.mode = self.mu = mu = at.as_tensor_variable(mu)
-
-        self.variance = at.switch((nu > 2) * 1, (1 / self.lam) * (nu / (nu - 2)), np.inf)
+        sigma = at.as_tensor_variable(sigma)
 
-        assert_negative_support(lam, "lam (sigma)", "StudentT")
+        assert_negative_support(sigma, "sigma (lam)", "StudentT")
         assert_negative_support(nu, "nu", "StudentT")
 
-    def random(self, point=None, size=None):
-        """
-        Draw random values from StudentT 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).
+        return super().dist([nu, mu, sigma], **kwargs)
 
-        Returns
-        -------
-        array
-        """
-        # nu, mu, lam = draw_values([self.nu, self.mu, self.lam], point=point, size=size)
-        # return generate_samples(
-        #     stats.t.rvs, nu, loc=mu, scale=lam ** -0.5, dist_shape=self.shape, size=size
-        # )
-
-    def logp(self, value):
+    def logp(value, nu, mu, sigma):
         """
         Calculate log-probability of StudentT distribution at specified value.
 
@@ -1908,11 +1900,7 @@ def logp(self, value):
         -------
         TensorVariable
         """
-        nu = self.nu
-        mu = self.mu
-        lam = self.lam
-        sigma = self.sigma
-
+        lam, sigma = get_tau_sigma(sigma=sigma)
         return bound(
             gammaln((nu + 1.0) / 2.0)
             + 0.5 * at.log(lam / (nu * np.pi))
@@ -1923,10 +1911,7 @@ def logp(self, value):
             sigma > 0,
         )
 
-    def _distr_parameters_for_repr(self):
-        return ["nu", "mu", "lam"]
-
-    def logcdf(self, value):
+    def logcdf(value, nu, mu, sigma):
         """
         Compute the log of the cumulative distribution function for Student's T distribution
         at the specified value.
@@ -1946,10 +1931,8 @@ def logcdf(self, value):
                 f"StudentT.logcdf expects a scalar value but received a {np.ndim(value)}-dimensional object."
             )
 
-        nu = self.nu
-        mu = self.mu
-        sigma = self.sigma
-        lam = self.lam
+        lam, sigma = get_tau_sigma(sigma=sigma)
+
         t = (value - mu) / sigma
         sqrt_t2_nu = at.sqrt(t ** 2 + nu)
         x = (t + sqrt_t2_nu) / (2.0 * sqrt_t2_nu)

pymc3/tests/test_distributions.py

@@ -1301,20 +1301,33 @@ def test_lognormal(self):
             n_samples=5,  # Just testing alternative parametrization
         )
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_t(self):
         self.check_logp(
             StudentT,
             R,
             {"nu": Rplus, "mu": R, "lam": Rplus},
             lambda value, nu, mu, lam: sp.t.logpdf(value, nu, mu, lam ** -0.5),
         )
+        self.check_logp(
+            StudentT,
+            R,
+            {"nu": Rplus, "mu": R, "sigma": Rplus},
+            lambda value, nu, mu, sigma: sp.t.logpdf(value, nu, mu, sigma),
+            n_samples=5,  # Just testing alternative parametrization
+        )
         self.check_logcdf(
             StudentT,
             R,
             {"nu": Rplus, "mu": R, "lam": Rplus},
             lambda value, nu, mu, lam: sp.t.logcdf(value, nu, mu, lam ** -0.5),
-            n_samples=10,
+            n_samples=10,  # relies on slow incomplete beta
+        )
+        self.check_logcdf(
+            StudentT,
+            R,
+            {"nu": Rplus, "mu": R, "sigma": Rplus},
+            lambda value, nu, mu, sigma: sp.t.logcdf(value, nu, mu, sigma),
+            n_samples=5,  # Just testing alternative parametrization
         )
 
     def test_cauchy(self):

pymc3/tests/test_distributions_random.py

@@ -283,12 +283,6 @@ class TestAsymmetricLaplace(BaseTestCases.BaseTestCase):
     params = {"kappa": 1.0, "b": 1.0, "mu": 0.0}
 
 
-@pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-class TestStudentT(BaseTestCases.BaseTestCase):
-    distribution = pm.StudentT
-    params = {"nu": 5.0, "mu": 0.0, "lam": 1.0}
-
-
 @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
 class TestChiSquared(BaseTestCases.BaseTestCase):
     distribution = pm.ChiSquared
@@ -481,6 +475,29 @@ class TestGumbel(BaseTestDistribution):
     ]
 
 
+class TestStudentT(BaseTestDistribution):
+    pymc_dist = pm.StudentT
+    pymc_dist_params = {"nu": 5.0, "mu": -1.0, "sigma": 2.0}
+    expected_rv_op_params = {"nu": 5.0, "mu": -1.0, "sigma": 2.0}
+    reference_dist_params = {"df": 5.0, "loc": -1.0, "scale": 2.0}
+    reference_dist = seeded_scipy_distribution_builder("t")
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_pymc_draws_match_reference",
+        "check_rv_size",
+    ]
+
+
+class TestStudentTLam(BaseTestDistribution):
+    pymc_dist = pm.StudentT
+    lam, sigma = get_tau_sigma(tau=2.0)
+    pymc_dist_params = {"nu": 5.0, "mu": -1.0, "lam": lam}
+    expected_rv_op_params = {"nu": 5.0, "mu": -1.0, "lam": sigma}
+    reference_dist_params = {"df": 5.0, "loc": -1.0, "scale": sigma}
+    reference_dist = seeded_scipy_distribution_builder("t")
+    tests_to_run = ["check_pymc_params_match_rv_op"]
+
+
 class TestNormal(BaseTestDistribution):
     pymc_dist = pm.Normal
     pymc_dist_params = {"mu": 5.0, "sigma": 10.0}
@@ -1134,13 +1151,6 @@ def ref_rand(size, kappa, b, mu):
 
         pymc3_random(pm.AsymmetricLaplace, {"b": Rplus, "kappa": Rplus, "mu": R}, ref_rand=ref_rand)
 
-    @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-    def test_student_t(self):
-        def ref_rand(size, nu, mu, lam):
-            return st.t.rvs(nu, mu, lam ** -0.5, size=size)
-
-        pymc3_random(pm.StudentT, {"nu": Rplus, "mu": R, "lam": Rplus}, ref_rand=ref_rand)
-
     @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_ex_gaussian(self):
         def ref_rand(size, mu, sigma, nu):

pymc3/tests/test_posteriors.py

@@ -76,7 +76,6 @@ class TestNUTSBetaBinomial(sf.NutsFixture, sf.BetaBinomialFixture):
     min_n_eff = 400
 
 
-@pytest.mark.xfail(reason="StudentT not refactored for v4")
 class TestNUTSStudentT(sf.NutsFixture, sf.StudentTFixture):
     n_samples = 10000
     tune = 1000
@@ -98,7 +97,7 @@ class TestNUTSNormalLong(sf.NutsFixture, sf.NormalFixture):
     atol = 0.001
 
 
-@pytest.mark.xfail(reason="StudentT not refactored for v4")
+@pytest.mark.xfail(reason="LKJCholeskyCov not refactored for v4")
 class TestNUTSLKJCholeskyCov(sf.NutsFixture, sf.LKJCholeskyCovFixture):
     n_samples = 2000
     tune = 1000