Implement logprob rewrites for rounding operations

Author: Ricardo Vieira

aeppl/censoring.py

@@ -5,7 +5,7 @@
 from aesara.graph.basic import Node
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.rewriting.basic import node_rewriter
-from aesara.scalar.basic import Clip
+from aesara.scalar.basic import Ceil, Clip, Floor, RoundHalfToEven
 from aesara.scalar.basic import clip as scalar_clip
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.var import TensorConstant
@@ -15,7 +15,7 @@
     MeasurableVariable,
     assign_custom_measurable_outputs,
 )
-from aeppl.logprob import CheckParameterValue, _logcdf, _logprob
+from aeppl.logprob import CheckParameterValue, _logcdf, _logprob, logdiffexp
 from aeppl.rewriting import measurable_ir_rewrites_db
 
 
@@ -142,3 +142,113 @@ def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
         )
 
     return logprob
+
+
+class MeasurableRound(MeasurableElemwise):
+    """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
+
+    valid_scalar_types = (RoundHalfToEven, Floor, Ceil)
+
+
+@node_rewriter(tracks=[Elemwise])
+def find_measurable_roundings(
+    fgraph: FunctionGraph, node: Node
+) -> Optional[List[MeasurableRound]]:
+
+    rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    if isinstance(node.op, MeasurableRound):
+        return None  # pragma: no cover
+
+    if not (
+        isinstance(node.op, Elemwise)
+        and isinstance(node.op.scalar_op, MeasurableRound.valid_scalar_types)
+    ):
+        return None
+
+    (rounded_var,) = node.outputs
+    (base_var,) = node.inputs
+
+    if not (
+        base_var.owner
+        and isinstance(base_var.owner.op, MeasurableVariable)
+        and base_var not in rv_map_feature.rv_values
+        # Rounding only makes sense for continuous variables
+        and base_var.dtype.startswith("float")
+    ):
+        return None
+
+    # Make base_var unmeasurable
+    unmeasurable_base_var = assign_custom_measurable_outputs(base_var.owner)
+
+    rounded_op = MeasurableRound(node.op.scalar_op)
+    rounded_rv = rounded_op.make_node(unmeasurable_base_var).default_output()
+    rounded_rv.name = rounded_var.name
+    return [rounded_rv]
+
+
+measurable_ir_rewrites_db.register(
+    "find_measurable_roundings",
+    find_measurable_roundings,
+    0,
+    "basic",
+    "censoring",
+)
+
+
+@_logprob.register(MeasurableRound)
+def round_logprob(op, values, base_rv, **kwargs):
+    r"""Logprob of a rounded censored distribution
+
+    The probability of a distribution rounded to the nearest integer is given by
+    .. math::
+        \begin{cases}
+            \text{CDF}(x+\frac{1}{2}, dist) - \text{CDF}(x-\frac{1}{2}, dist) & \text{for } x \in \mathbb{Z}, \\
+            0 & \text{otherwise},
+        \end{cases}
+
+    The probability of a distribution rounded up is given by
+    .. math::
+        \begin{cases}
+            \text{CDF}(x, dist) - \text{CDF}(x-1, dist) & \text{for } x \in \mathbb{Z}, \\
+            0 & \text{otherwise},
+        \end{cases}
+
+    The probability of a distribution rounded down is given by
+    .. math::
+        \begin{cases}
+            \text{CDF}(x+1, dist) - \text{CDF}(x, dist) & \text{for } x \in \mathbb{Z}, \\
+            0 & \text{otherwise},
+        \end{cases}
+
+    """
+    (value,) = values
+
+    if isinstance(op.scalar_op, RoundHalfToEven):
+        value = at.round(value)
+        value_upper = value + 0.5
+        value_lower = value - 0.5
+    elif isinstance(op.scalar_op, Floor):
+        value = at.floor(value)
+        value_upper = value + 1.0
+        value_lower = value
+    elif isinstance(op.scalar_op, Ceil):
+        value = at.ceil(value)
+        value_upper = value
+        value_lower = value - 1.0
+    else:
+        raise TypeError(f"Unsupported scalar_op {op.scalar_op}")  # pragma: no cover
+
+    base_rv_op = base_rv.owner.op
+    base_rv_inputs = base_rv.owner.inputs
+
+    logcdf_upper = _logcdf(base_rv_op, value_upper, *base_rv_inputs, **kwargs)
+    logcdf_lower = _logcdf(base_rv_op, value_lower, *base_rv_inputs, **kwargs)
+
+    if base_rv_op.name:
+        logcdf_upper.name = f"{base_rv_op}_logcdf_upper"
+        logcdf_lower.name = f"{base_rv_op}_logcdf_lower"
+
+    return logdiffexp(logcdf_upper, logcdf_lower)

aeppl/logprob.py

@@ -40,6 +40,11 @@ def binomln(n, k):
     return at.gammaln(n + 1) - at.gammaln(k + 1) - at.gammaln(n - k + 1)
 
 
+def logdiffexp(a, b):
+    """log(exp(a) - exp(b))"""
+    return a + at.log1mexp(b - a)
+
+
 def xlogy0(m, x):
     # TODO: This should probably be a basic Aesara stabilization
     return at.switch(at.eq(x, 0), at.switch(at.eq(m, 0), 0.0, -np.inf), m * at.log(x))

tests/test_censoring.py

@@ -188,3 +188,33 @@ def test_clip_transform():
     )
 
     assert np.isclose(obs_logp, exp_logp)
+
+
+@pytest.mark.parametrize("rounding_op", (at.round, at.floor, at.ceil))
+def test_rounding(rounding_op):
+    loc = 1
+    scale = 2
+    test_value = np.arange(-3, 4)
+
+    x = at.random.normal(loc, scale, size=test_value.shape, name="x")
+    xr = rounding_op(x)
+    xr.name = "xr"
+
+    xr_vv = xr.clone()
+    logp = joint_logprob({xr: xr_vv}, sum=False)
+    assert logp is not None
+
+    x_sp = st.norm(loc, scale)
+    if rounding_op == at.round:
+        expected_logp = np.log(x_sp.cdf(test_value + 0.5) - x_sp.cdf(test_value - 0.5))
+    elif rounding_op == at.floor:
+        expected_logp = np.log(x_sp.cdf(test_value + 1.0) - x_sp.cdf(test_value))
+    elif rounding_op == at.ceil:
+        expected_logp = np.log(x_sp.cdf(test_value) - x_sp.cdf(test_value - 1.0))
+    else:
+        raise NotImplementedError()
+
+    assert np.allclose(
+        logp.eval({xr_vv: test_value}),
+        expected_logp,
+    )