Use automatic logprob for LogitNormal to support icdf

Co-authored-by: Ricardo Vieira <28983449+ricardov94@users.noreply.github.com>

Author: jnetzel1

pymc/distributions/continuous.py

@@ -53,6 +53,7 @@
 from pytensor.tensor.random.utils import normalize_size_param
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
+from pymc.distributions.custom import CustomDist
 from pymc.logprob.abstract import _logprob_helper
 from pymc.logprob.basic import TensorLike, icdf
 from pymc.pytensorf import normalize_rng_param
@@ -92,7 +93,7 @@ def polyagamma_cdf(*args, **kwargs):
 from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Continuous, SymbolicRandomVariable
 from pymc.distributions.shape_utils import implicit_size_from_params, rv_size_is_none
 from pymc.distributions.transforms import _default_transform
-from pymc.math import invlogit, logdiffexp, logit
+from pymc.math import invlogit, logdiffexp
 
 __all__ = [
     "AsymmetricLaplace",
@@ -3603,28 +3604,7 @@ def icdf(value, mu, s):
         )
 
 
-class LogitNormalRV(SymbolicRandomVariable):
-    name = "logit_normal"
-    extended_signature = "[rng],[size],(),()->[rng],()"
-    _print_name = ("LogitNormal", "\\operatorname{LogitNormal}")
-
-    @classmethod
-    def rv_op(cls, mu, sigma, *, size=None, rng=None):
-        mu = pt.as_tensor(mu)
-        sigma = pt.as_tensor(sigma)
-        rng = normalize_rng_param(rng)
-        size = normalize_size_param(size)
-
-        next_rng, normal_draws = normal(loc=mu, scale=sigma, size=size, rng=rng).owner.outputs
-        draws = pt.expit(normal_draws)
-
-        return cls(
-            inputs=[rng, size, mu, sigma],
-            outputs=[next_rng, draws],
-        )(rng, size, mu, sigma)
-
-
-class LogitNormal(UnitContinuous):
+class LogitNormal:
     r"""
     Logit-Normal distribution.
 
@@ -3672,37 +3652,26 @@ class LogitNormal(UnitContinuous):
         Defaults to 1.
     """
 
-    rv_type = LogitNormalRV
-    rv_op = LogitNormalRV.rv_op
+    @staticmethod
+    def logitnormal_dist(mu, sigma, size):
+        return invlogit(Normal.dist(mu=mu, sigma=sigma, size=size))
 
-    @classmethod
-    def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
+    def __new__(cls, name, mu=0, sigma=None, tau=None, **kwargs):
         _, sigma = get_tau_sigma(tau=tau, sigma=sigma)
-        return super().dist([mu, sigma], **kwargs)
-
-    def support_point(rv, size, mu, sigma):
-        median, _ = pt.broadcast_arrays(invlogit(mu), sigma)
-        if not rv_size_is_none(size):
-            median = pt.full(size, median)
-        return median
-
-    def logp(value, mu, sigma):
-        tau, _ = get_tau_sigma(sigma=sigma)
-
-        res = pt.switch(
-            pt.or_(pt.le(value, 0), pt.ge(value, 1)),
-            -np.inf,
-            (
-                -0.5 * tau * (logit(value) - mu) ** 2
-                + 0.5 * pt.log(tau / (2.0 * np.pi))
-                - pt.log(value * (1 - value))
-            ),
+        return CustomDist(
+            name,
+            mu,
+            sigma,
+            dist=cls.logitnormal_dist,
+            class_name="LogitNormal",
+            **kwargs,
         )
 
-        return check_parameters(
-            res,
-            tau > 0,
-            msg="tau > 0",
+    @classmethod
+    def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
+        _, sigma = get_tau_sigma(tau=tau, sigma=sigma)
+        return CustomDist.dist(
+            mu, sigma, dist=cls.logitnormal_dist, class_name="LogitNormal", **kwargs
         )
 
 

pymc/distributions/moments/means.py

@@ -59,7 +59,6 @@
     HalfFlatRV,
     HalfStudentTRV,
     KumaraswamyRV,
-    LogitNormalRV,
     MoyalRV,
     PolyaGammaRV,
     RiceRV,
@@ -290,11 +289,6 @@ def logistic_mean(op, rv, rng, size, mu, s):
     return maybe_resize(pt.broadcast_arrays(mu, s)[0], size)
 
 
-@_mean.register(LogitNormalRV)
-def logitnormal_mean(op, rv, rng, size, mu, sigma):
-    raise UndefinedMomentException("The mean of the LogitNormal distribution is undefined")
-
-
 @_mean.register(LogNormalRV)
 def lognormal_mean(op, rv, rng, size, mu, sigma):
     return maybe_resize(pt.exp(mu + 0.5 * sigma**2), size)

tests/distributions/moments/test_means.py

@@ -274,5 +274,5 @@ def test_mean_equal_expected(dist, dist_params, expected):
     ],
 )
 def test_no_mean(dist, dist_params):
-    with pytest.raises(UndefinedMomentException):
+    with pytest.raises((UndefinedMomentException, NotImplementedError)):
         mean(dist.dist(**dist_params))

tests/distributions/test_continuous.py

@@ -872,6 +872,12 @@ def test_logitnormal(self):
             ),
             decimal=select_by_precision(float64=6, float32=1),
         )
+        check_icdf(
+            pm.LogitNormal,
+            {"mu": R, "sigma": Rplus},
+            lambda q, mu, sigma: sp.expit(mu + sigma * st.norm.ppf(q)),
+            decimal=select_by_precision(float64=12, float32=5),
+        )
 
     @pytest.mark.skipif(
         condition=(pytensor.config.floatX == "float32"),