Refactor LogitNormal

pymc3/distributions/continuous.py

@@ -3726,6 +3726,18 @@ def logcdf(value, mu, s):
             0 < s,
         )
 
+class LogitNormalRV(RandomVariable):
+    name = "logit_normal"
+    ndim_supp = 0
+    ndims_params = [0, 0]
+    dtype = "floatX"
+    _print_name = ("logitNormal", "\\operatorname{logitNormal}")
+
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, size=None):
+        return invlogit(stats.norm.rvs(loc=mu, scale=sigma, size=size, random_state=rng))
+
+logit_normal = LogitNormalRV()
 
 class LogitNormal(UnitContinuous):
     r"""
@@ -3771,44 +3783,22 @@ class LogitNormal(UnitContinuous):
     tau: float
         Scale parameter (tau > 0).
     """
+    rv_op = logit_normal
 
-    def __init__(self, mu=0, sigma=None, tau=None, sd=None, **kwargs):
+    @classmethod
+    def dist(cls, mu=0, sigma=None, tau=None, sd=None, **kwargs):
         if sd is not None:
             sigma = sd
-        self.mu = mu = at.as_tensor_variable(floatX(mu))
+        mu = at.as_tensor_variable(floatX(mu))
         tau, sigma = get_tau_sigma(tau=tau, sigma=sigma)
-        self.sigma = self.sd = at.as_tensor_variable(sigma)
-        self.tau = tau = at.as_tensor_variable(tau)
-
-        self.median = invlogit(mu)
+        sigma = sd = at.as_tensor_variable(sigma)
+        tau = at.as_tensor_variable(tau)
         assert_negative_support(sigma, "sigma", "LogitNormal")
         assert_negative_support(tau, "tau", "LogitNormal")
 
-        super().__init__(**kwargs)
-
-    def random(self, point=None, size=None):
-        """
-        Draw random values from LogitNormal 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
-        """
-        # mu, _, sigma = draw_values([self.mu, self.tau, self.sigma], point=point, size=size)
-        # return expit(
-        #     generate_samples(stats.norm.rvs, loc=mu, scale=sigma, dist_shape=self.shape, size=size)
-        # )
+        return super().dist([mu, sigma], **kwargs)
 
-    def logp(self, value):
+    def logp(value, mu, sigma):
         """
         Calculate log-probability of LogitNormal distribution at specified value.
 
@@ -3822,8 +3812,7 @@ def logp(self, value):
         -------
         TensorVariable
         """
-        mu = self.mu
-        tau = self.tau
+        tau, sigma = get_tau_sigma(sigma=sigma)
         return bound(
             -0.5 * tau * (logit(value) - mu) ** 2
             + 0.5 * at.log(tau / (2.0 * np.pi))
@@ -3833,9 +3822,8 @@ def logp(self, value):
             tau > 0,
         )
 
-    def _distr_parameters_for_repr(self):
-        return ["mu", "sigma"]
 
+
 
 class Interpolated(BoundedContinuous):
     r"""

pymc3/tests/test_distributions.py

@@ -2515,7 +2515,7 @@ def test_logistic(self):
             decimal=select_by_precision(float64=6, float32=1),
         )
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
+
     def test_logitnormal(self):
         self.check_logp(
             LogitNormal,

pymc3/tests/test_distributions_random.py

@@ -15,6 +15,7 @@
 import itertools
 
 from contextlib import ExitStack as does_not_raise
+from pymc3.math import invlogit
 from typing import Callable, List, Optional
 
 import aesara
@@ -301,11 +302,6 @@ class TestExGaussian(BaseTestCases.BaseTestCase):
     params = {"mu": 0.0, "sigma": 1.0, "nu": 1.0}
 
 
-@pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-class TestLogitNormal(BaseTestCases.BaseTestCase):
-    distribution = pm.LogitNormal
-    params = {"mu": 0.0, "sigma": 1.0}
-
 
 @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
 class TestZeroInflatedNegativeBinomial(BaseTestCases.BaseTestCase):
@@ -513,6 +509,22 @@ class TestNormal(BaseTestDistribution):
         "check_pymc_draws_match_reference",
         "check_rv_size",
     ]
+class TestLogitNormal(BaseTestDistribution):
+    def logit_normal_rng_fn(self, rng, size, loc, scale):
+        return expit(st.norm.rvs(loc=loc, scale=scale, size=size, random_state=rng))
+
+    pymc_dist = pm.LogitNormal
+    pymc_dist_params = {"mu": 5.0, "sigma": 10.0}
+    expected_rv_op_params = {"mu": 5.0, "sigma": 10.0}
+    reference_dist_params = {"loc": 5.0, "scale": 10.0}
+    reference_dist = lambda self: functools.partial(
+        self.logit_normal_rng_fn, rng=self.get_random_state()
+    )
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_pymc_draws_match_reference",
+        "check_rv_size",
+    ]
 
 
 class TestNormalTau(BaseTestDistribution):
@@ -1384,7 +1396,6 @@ def test_dirichlet_multinomial_dist_ShapeError(self, n, a, shape, expectation):
         with expectation:
             m.random()
 
-    @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_logitnormal(self):
         def ref_rand(size, mu, sigma):
             return expit(st.norm.rvs(loc=mu, scale=sigma, size=size))