Refer to specific censoring type in IR

Author: Ricardo Vieira

aeppl/censoring.py

@@ -15,22 +15,24 @@
 from aeppl.rewriting import measurable_ir_rewrites_db
 
 
-class CensoredRV(Elemwise):
-    """A placeholder used to specify a log-likelihood for a censored RV sub-graph."""
+class MeasurableClip(Elemwise):
+    """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
 
 
-MeasurableVariable.register(CensoredRV)
+MeasurableVariable.register(MeasurableClip)
 
 
 @node_rewriter(tracks=[Elemwise])
-def find_censored_rvs(fgraph: FunctionGraph, node: Node) -> Optional[List[CensoredRV]]:
+def find_measurable_clips(
+    fgraph: FunctionGraph, node: Node
+) -> Optional[List[MeasurableClip]]:
     # TODO: Canonicalize x[x>ub] = ub -> clip(x, x, ub)
 
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
     if rv_map_feature is None:
         return None  # pragma: no cover
 
-    if isinstance(node.op, CensoredRV):
+    if isinstance(node.op, MeasurableClip):
         return None  # pragma: no cover
 
     if not (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Clip)):
@@ -47,35 +49,48 @@ def find_censored_rvs(fgraph: FunctionGraph, node: Node) -> Optional[List[Censor
         return None
 
     # Replace bounds by `+-inf` if `y = clip(x, x, ?)` or `y=clip(x, ?, x)`
-    # This is used in `censor_logprob` to generate a more succint logprob graph
-    # for one-sided censored random variables
+    # This is used in `clip_logprob` to generate a more succint logprob graph
+    # for one-sided clipped random variables
     lower_bound = lower_bound if (lower_bound is not base_var) else at.constant(-np.inf)
     upper_bound = upper_bound if (upper_bound is not base_var) else at.constant(np.inf)
 
-    censored_op = CensoredRV(scalar_clip)
+    clipped_op = MeasurableClip(scalar_clip)
     # Make base_var unmeasurable
     unmeasurable_base_var = assign_custom_measurable_outputs(base_var.owner)
-    censored_rv_node = censored_op.make_node(
+    clipped_rv_node = clipped_op.make_node(
         unmeasurable_base_var, lower_bound, upper_bound
     )
-    censored_rv = censored_rv_node.outputs[0]
+    clipped_rv = clipped_rv_node.outputs[0]
 
-    censored_rv.name = clipped_var.name
+    clipped_rv.name = clipped_var.name
 
-    return [censored_rv]
+    return [clipped_rv]
 
 
 measurable_ir_rewrites_db.register(
-    "find_censored_rvs",
-    find_censored_rvs,
+    "find_measurable_clips",
+    find_measurable_clips,
     0,
     "basic",
     "censoring",
 )
 
 
-@_logprob.register(CensoredRV)
-def censor_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
+@_logprob.register(MeasurableClip)
+def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
+    r"""Logprob of a clipped censored distribution
+
+    The probability is given by
+    .. math::
+        \begin{cases}
+            0 & \text{for } x < lower, \\
+            \text{CDF}(lower, dist) & \text{for } x = lower, \\
+            \text{P}(x, dist) & \text{for } lower < x < upper, \\
+            1-\text{CDF}(upper, dist) & \text {for} x = upper, \\
+            0 & \text{for } x > upper,
+        \end{cases}
+
+    """
     (value,) = values
 
     base_rv_op = base_rv.owner.op
@@ -95,7 +110,7 @@ def censor_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
         is_upper_bounded = True
 
         logccdf = at.log1mexp(logcdf)
-        # For right censored discrete RVs, we need to add an extra term
+        # For right clipped discrete RVs, we need to add an extra term
         # corresponding to the pmf at the upper bound
         if base_rv.dtype.startswith("int"):
             logccdf = at.logaddexp(logccdf, logprob)

tests/test_censoring.py

@@ -11,7 +11,7 @@
 
 
 @aesara.config.change_flags(compute_test_value="raise")
-def test_continuous_rv_censoring():
+def test_continuous_rv_clip():
     x_rv = at.random.normal(0.5, 1)
     cens_x_rv = at.clip(x_rv, -2, 2)
 
@@ -32,7 +32,7 @@ def test_continuous_rv_censoring():
     assert np.isclose(logp_fn(0), ref_scipy.logpdf(0))
 
 
-def test_discrete_rv_censoring():
+def test_discrete_rv_clip():
     x_rv = at.random.poisson(2)
     cens_x_rv = at.clip(x_rv, 1, 4)
 
@@ -52,7 +52,7 @@ def test_discrete_rv_censoring():
     assert np.isclose(logp_fn(2), ref_scipy.logpmf(2))
 
 
-def test_one_sided_censoring():
+def test_one_sided_clip():
     x_rv = at.random.normal(0, 1)
     lb_cens_x_rv = at.clip(x_rv, -1, x_rv)
     ub_cens_x_rv = at.clip(x_rv, x_rv, 1)
@@ -74,7 +74,7 @@ def test_one_sided_censoring():
     np.testing.assert_almost_equal(logp_fn(1, -1), ref_scipy.logpdf(-1))
 
 
-def test_useless_censoring():
+def test_useless_clip():
     x_rv = at.random.normal(0.5, 1, size=3)
     cens_x_rv = at.clip(x_rv, x_rv, x_rv)
 
@@ -89,7 +89,7 @@ def test_useless_censoring():
     np.testing.assert_allclose(logp_fn([-2, 0, 2]), ref_scipy.logpdf([-2, 0, 2]))
 
 
-def test_random_censoring():
+def test_random_clip():
     lb_rv = at.random.normal(0, 1, size=2)
     x_rv = at.random.normal(0, 2)
     cens_x_rv = at.clip(x_rv, lb_rv, [1, 1])
@@ -105,7 +105,7 @@ def test_random_censoring():
     assert res[1] != -np.inf
 
 
-def test_broadcasted_censoring_constant():
+def test_broadcasted_clip_constant():
     lb_rv = at.random.uniform(0, 1)
     x_rv = at.random.normal(0, 2)
     cens_x_rv = at.clip(x_rv, lb_rv, [1, 1])
@@ -117,7 +117,7 @@ def test_broadcasted_censoring_constant():
     assert_no_rvs(logp)
 
 
-def test_broadcasted_censoring_random():
+def test_broadcasted_clip_random():
     lb_rv = at.random.normal(0, 1)
     x_rv = at.random.normal(0, 2, size=2)
     cens_x_rv = at.clip(x_rv, lb_rv, 1)
@@ -129,7 +129,7 @@ def test_broadcasted_censoring_random():
     assert_no_rvs(logp)
 
 
-def test_fail_base_and_censored_have_values():
+def test_fail_base_and_clip_have_values():
     """Test failure when both base_rv and clipped_rv are given value vars"""
     x_rv = at.random.normal(0, 1)
     cens_x_rv = at.clip(x_rv, x_rv, 1)
@@ -141,7 +141,7 @@ def test_fail_base_and_censored_have_values():
         factorized_joint_logprob({cens_x_rv: cens_x_vv, x_rv: x_vv})
 
 
-def test_fail_multiple_censored_single_base():
+def test_fail_multiple_clip_single_base():
     """Test failure when multiple clipped_rvs share a single base_rv"""
     base_rv = at.random.normal(0, 1)
     cens_rv1 = at.clip(base_rv, -1, 1)
@@ -172,7 +172,7 @@ def test_deterministic_clipping():
     )
 
 
-def test_censored_transform():
+def test_clip_transform():
     x_rv = at.random.normal(0.5, 1)
     cens_x_rv = at.clip(x_rv, 0, x_rv)
 

tests/test_composite_logprob.py

@@ -4,8 +4,8 @@
 import scipy.stats as st
 
 from aeppl import joint_logprob
+from aeppl.censoring import MeasurableClip
 from aeppl.rewriting import construct_ir_fgraph
-from aeppl.truncation import CensoredRV
 from tests.utils import assert_no_rvs
 
 
@@ -97,7 +97,7 @@ def test_unvalued_ir_reversion():
     assert memo[y_rv] in z_fgraph.preserve_rv_mappings.measurable_conversions
 
     measurable_y_rv = z_fgraph.preserve_rv_mappings.measurable_conversions[memo[y_rv]]
-    assert isinstance(measurable_y_rv.owner.op, CensoredRV)
+    assert isinstance(measurable_y_rv.owner.op, MeasurableClip)
 
     # `construct_ir_fgraph` should've reverted the un-valued measurable IR
     # change