Add logcdf for CensoredRV

pymc/distributions/censored.py

@@ -24,6 +24,7 @@
 )
 from pymc.distributions.shape_utils import _change_dist_size, change_dist_size
 from pymc.util import check_dist_not_registered
+from pymc.logprob.abstract import _logcdf
 
 
 class CensoredRV(SymbolicRandomVariable):
@@ -148,3 +149,35 @@ def moment_censored(op, rv, dist, lower, upper):
     )
     moment = pt.full_like(dist, moment)
     return moment
+
+@_logcdf.register(CensoredRV)
+def censored_logcdf(op, value, *inputs, **kwargs):
+    *rv_inputs, lower, upper, rng = inputs
+    rv_inputs = [rng, *rv_inputs]
+
+    base_rv_op = op.base_rv_op
+    logcdf_cens = _logcdf(base_rv_op, value, *rv_inputs, **kwargs)
+
+    lower_logcdf = _logcdf(base_rv_op, lower_value, *rv_inputs, **kwargs)
+    upper_logcdf = _logcdf(base_rv_op, upper, *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 function is left censored, set logcdf to -np,inf
+    if is_lower_bounded:
+        logcdf_trunc = pt.switch(value < lower, -np.inf, logcdf_cens)
+
+    #If function is right censored, set logcdf to 0
+    if is_upper_bounded:
+        logcdf_trunc = pt.switch(value <= upper, logcdf_cens, 0.0)
+
+    #If in domain, set logcdf as if uncensored
+    if is_lower_bounded and is_upper_bounded:
+        logcdf_trunc = check_parameters(
+            logcdf_cens,
+            pt.le(lower, upper),
+            msg="lower_bound <= upper_bound",
+        )
+
+    return logcdf_cens

tests/distributions/test_censored.py

@@ -18,6 +18,7 @@
 import pymc as pm
 
 from pymc.distributions.shape_utils import change_dist_size
+from pymc.logprob.basic import logcdf
 
 
 class TestCensored:
@@ -110,3 +111,45 @@ def test_dist_broadcasted_by_lower_upper(self):
             pm.Normal.dist(size=(3, 4, 2)), lower=np.zeros((2,)), upper=np.zeros((4, 2))
         )
         assert tuple(x.owner.inputs[0].shape.eval()) == (3, 4, 2)
+
+    @pytest.mark.parametrize("lower, upper", [(-1, np.inf), (-1, 1.5), (-np.inf, 1.5)])
+    @pytest.mark.parametrize("op_type", ["icdf", "rejection"])
+    def test_censoring_continuous_logcdf(op_type, lower, upper):
+        loc = 0.15
+        scale = 10
+        op = icdf_normal if op_type == "icdf" else rejection_normal
+
+        x = op(loc, scale, name="x")
+        xt = pm.Censored.dist(x, lower=lower, upper=upper)
+        assert isinstance(xt.owner.op, pm.CensoredRV)
+
+        xt_vv = xt.clone()
+        xt_logcdf_fn = pytensor.function([xt_vv], logcdf(xt, xt_vv))
+
+        for bound in (lower, upper):
+            if np.isinf(bound):
+                return
+            for offset in (-1, 0, 1):
+                test_xt_v = bound + offset
+                assert xt_logcdf_fn(test_xt_v) is not None
+
+
+    @pytest.mark.parametrize("lower, upper", [(2, np.inf), (2, 5), (-np.inf, 5)])
+    @pytest.mark.parametrize("op_type", ["icdf", "rejection"])
+    def test_censoring_discrete_logcdf(op_type, lower, upper):
+        p = 0.7
+        op = icdf_geometric if op_type == "icdf" else rejection_geometric
+
+        x = op(p, name="x")
+        xt = pm.Censored.dist(x, lower=lower, upper=upper)
+        assert isinstance(xt.owner.op, pm.CensoredRV)
+
+        xt_vv = xt.clone()
+        xt_logcdf_fn = pytensor.function([xt_vv], logcdf(xt, xt_vv))
+
+        for bound in (lower, upper):
+            if np.isinf(bound):
+                continue
+            for offset in (-1, 0, 1):
+                test_xt_v = bound + offset
+                assert xt_logcdf_fn(test_xt_v) is not None