feat-#6686-implementlogcdf-for-censoredRV

Author: asifzubair

pymc/distributions/censored.py

@@ -17,7 +17,9 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.utils import normalize_size_param
+from pytensor.tensor.variable import TensorConstant
 
+from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import (
     Distribution,
     SymbolicRandomVariable,
@@ -29,6 +31,7 @@
     implicit_size_from_params,
     rv_size_is_none,
 )
+from pymc.logprob.abstract import _logcdf
 from pymc.util import check_dist_not_registered
 
 
@@ -156,3 +159,30 @@ def support_point_censored(op, rv, dist, lower, upper):
     )
     support_point = pt.full_like(dist, support_point)
     return support_point
+
+
+@_logcdf.register(CensoredRV)
+def censored_logcdf(op, value, *inputs, **kwargs):
+    base_rv, lower, upper = inputs
+
+    base_rv_op = base_rv.owner.op
+    base_rv_inputs = base_rv.owner.inputs
+    logcdf_val = _logcdf(base_rv_op, value, *base_rv_inputs, **kwargs)
+
+    is_lower_bounded = not (isinstance(lower, TensorConstant) and np.all(np.isneginf(lower.value)))
+    is_upper_bounded = not (isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value)))
+
+    if is_lower_bounded:
+        logcdf_val = pt.switch(pt.lt(value, lower), -np.inf, logcdf_val)
+
+    if is_upper_bounded:
+        logcdf_val = pt.switch(pt.ge(value, upper), 0.0, logcdf_val)
+
+    if is_lower_bounded and is_upper_bounded:
+        logcdf_val = check_parameters(
+            logcdf_val,
+            pt.le(lower, upper),
+            msg="lower_bound <= upper_bound",
+        )
+
+    return logcdf_val

tests/distributions/test_censored.py

@@ -17,7 +17,7 @@
 
 import pymc as pm
 
-from pymc import logp
+from pymc import logcdf, logp
 from pymc.distributions.shape_utils import change_dist_size
 
 
@@ -126,3 +126,77 @@ def test_censored_categorical(self):
             logp(censored_cat, [-1, 0, 1, 2, 3, 4, 5]).exp().eval(),
             [0, 0, 0.3, 0.2, 0.5, 0, 0],
         )
+
+    def test_censored_logcdf_continuous(self):
+        norm = pm.Normal.dist(0, 1)
+        eval_points = np.array([-np.inf, -2, -1, 0, 1, 2, np.inf])
+
+        # No censoring
+        censored_norm = pm.Censored.dist(norm, lower=None, upper=None)
+        with pytest.warns(RuntimeWarning, match="divide by zero encountered in log"):
+            censored_eval = logcdf(censored_norm, eval_points).eval()
+        with pytest.warns(RuntimeWarning, match="divide by zero encountered in log"):
+            norm_eval = logcdf(norm, eval_points).eval()
+        np.testing.assert_allclose(censored_eval, norm_eval)
+
+        # Left censoring
+        censored_norm = pm.Censored.dist(norm, lower=-1, upper=None)
+        with pytest.warns(RuntimeWarning, match="divide by zero encountered in log"):
+            censored_eval = logcdf(censored_norm, eval_points).eval()
+        np.testing.assert_allclose(
+            censored_eval,
+            np.array([-np.inf, -np.inf, -1.84102167, -0.69314718, -0.17275377, -0.02301291, 0.0]),
+            rtol=1e-6,
+        )
+
+        # Right censoring
+        censored_norm = pm.Censored.dist(norm, lower=None, upper=1)
+        with pytest.warns(RuntimeWarning, match="divide by zero encountered in log"):
+            censored_eval = logcdf(censored_norm, eval_points).eval()
+        np.testing.assert_allclose(
+            censored_eval,
+            np.array([-np.inf, -3.78318435, -1.84102167, -0.69314718, 0, 0, 0.0]),
+            rtol=1e-6,
+        )
+
+        # Interval censoring
+        censored_norm = pm.Censored.dist(norm, lower=-1, upper=1)
+        with pytest.warns(RuntimeWarning, match="divide by zero encountered in log"):
+            censored_eval = logcdf(censored_norm, eval_points).eval()
+        np.testing.assert_allclose(
+            censored_eval,
+            np.array([-np.inf, -np.inf, -1.84102167, -0.69314718, 0, 0, 0.0]),
+            rtol=1e-6,
+        )
+
+    def test_censored_logcdf_discrete(self):
+        cat = pm.Categorical.dist([0.1, 0.2, 0.2, 0.3, 0.2])
+        eval_points = np.array([-1, 0, 1, 2, 3, 4, 5])
+
+        # No censoring
+        censored_cat = pm.Censored.dist(cat, lower=None, upper=None)
+        np.testing.assert_allclose(
+            logcdf(censored_cat, eval_points).eval(),
+            logcdf(cat, eval_points).eval(),
+        )
+
+        # Left censoring
+        censored_cat = pm.Censored.dist(cat, lower=1, upper=None)
+        np.testing.assert_allclose(
+            logcdf(censored_cat, eval_points).eval(),
+            np.array([-np.inf, -np.inf, -1.2039728, -0.69314718, -0.22314355, 0, 0]),
+        )
+
+        # Right censoring
+        censored_cat = pm.Censored.dist(cat, lower=None, upper=3)
+        np.testing.assert_allclose(
+            logcdf(censored_cat, eval_points).eval(),
+            np.array([-np.inf, -2.30258509, -1.2039728, -0.69314718, 0, 0, 0]),
+        )
+
+        # Interval censoring
+        censored_cat = pm.Censored.dist(cat, lower=1, upper=3)
+        np.testing.assert_allclose(
+            logcdf(censored_cat, eval_points).eval(),
+            np.array([-np.inf, -np.inf, -1.2039728, -0.69314718, 0, 0, 0]),
+        )