Derive censoring logprob from minimum, maximum, rounding and casting operations

pymc/logprob/censoring.py

@@ -38,45 +38,110 @@
 import numpy as np
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Apply
+from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.scalar.basic import Ceil, Clip, Floor, RoundHalfToEven
+from pytensor.scalar.basic import (
+    Ceil,
+    Clip,
+    Floor,
+    Maximum,
+    RoundHalfAwayFromZero,
+    RoundHalfToEven,
+    Trunc,
+)
 from pytensor.scalar.basic import clip as scalar_clip
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.math import ceil, clip, floor, round_half_to_even
+from pytensor.tensor.math import (
+    ceil,
+    clip,
+    floor,
+    maximum,
+    minimum,
+    round_half_away_from_zero,
+    round_half_to_even,
+    trunc,
+)
+from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.variable import TensorConstant
 
-from pymc.logprob.abstract import MeasurableElemwise, _logccdf_helper, _logcdf, _logprob
+from pymc.logprob.abstract import (
+    MeasurableElemwise,
+    _icdf,
+    _icdf_helper,
+    _logccdf_helper,
+    _logcdf,
+    _logcdf_helper,
+    _logprob,
+)
 from pymc.logprob.rewriting import measurable_ir_rewrites_db
-from pymc.logprob.utils import CheckParameterValue, filter_measurable_variables
+from pymc.logprob.utils import (
+    CheckParameterValue,
+    check_potential_measurability,
+    filter_measurable_variables,
+)
 
 
 class MeasurableClip(MeasurableElemwise):
-    """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
-
-    valid_scalar_types = (Clip,)
+    """A placeholder used to specify a log-likelihood for a clipped RV sub-graph.
 
+    A bound that is the clipped variable itself marks that side as unbounded, which is
+    what the logprob methods read it back as.
+    """
 
-measurable_clip = MeasurableClip(scalar_clip)
+    valid_scalar_types = (Clip,)
 
 
 @node_rewriter(tracks=[clip])
 def find_measurable_clips(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     # TODO: Canonicalize x[x>ub] = ub -> clip(x, x, ub)
 
-    if not filter_measurable_variables(node.inputs):
-        return None
-
+    lower_bound: Variable | None
+    upper_bound: Variable | None
     base_var, lower_bound, upper_bound = node.inputs
 
-    # Replace bounds by `+-inf` if `y = clip(x, x, ?)` or `y=clip(x, ?, x)`
-    # This is used in `clip_logprob` to generate a more succinct logprob graph
-    # for one-sided clipped random variables
-    lower_bound = lower_bound if (lower_bound is not base_var) else pt.constant(-np.inf)
-    upper_bound = upper_bound if (upper_bound is not base_var) else pt.constant(np.inf)
+    # Clipping a constant by measurable bounds is an order statistic, not censoring
+    if not filter_measurable_variables([base_var]):
+        return None
 
-    clipped_rv = measurable_clip.make_node(base_var, lower_bound, upper_bound).outputs[0]
+    # A bound that is the clipped variable itself declares one-sided clipping
+    # (`y = clip(x, x, ub)`), as do literal +-inf constants
+    if lower_bound is base_var or (
+        isinstance(lower_bound, TensorConstant) and np.all(np.isneginf(lower_bound.value))
+    ):
+        lower_bound = None
+    if upper_bound is base_var or (
+        isinstance(upper_bound, TensorConstant) and np.all(np.isinf(upper_bound.value))
+    ):
+        upper_bound = None
+
+    # Fuse clips of clips into a single clip, so that e.g. two-sided censoring written
+    # as maximum(minimum(x, ub), lb) becomes one node. Mass pooled at an inner bound
+    # either stays there when it lies inside the outer interval, or is re-pooled at
+    # the outer bound, so the bounds combine with maximum/minimum (constant bounds
+    # fold; crossed bounds are caught by the CheckParameterValue in the logprob)
+    if isinstance(base_var.owner.op, MeasurableClip):
+        inner_base, inner_lower, inner_upper = base_var.owner.inputs
+        if inner_lower is not inner_base:
+            lower_bound = (
+                inner_lower if lower_bound is None else pt.maximum(lower_bound, inner_lower)
+            )
+        if inner_upper is not inner_base:
+            upper_bound = (
+                inner_upper if upper_bound is None else pt.minimum(upper_bound, inner_upper)
+            )
+        base_var = inner_base
+
+    # The variable itself is used as the unbounded sentinel because, unlike +-inf
+    # constants, it does not upcast discrete variables
+    if lower_bound is None:
+        lower_bound = base_var
+    if upper_bound is None:
+        upper_bound = base_var
+
+    clipped_rv = (
+        MeasurableClip(scalar_clip).make_node(base_var, lower_bound, upper_bound).outputs[0]
+    )
     return [clipped_rv]
 
 
@@ -88,6 +153,40 @@ def find_measurable_clips(fgraph: FunctionGraph, node: Apply) -> list[TensorVari
 )
 
 
+@node_rewriter(tracks=[maximum, minimum])
+def measurable_max_min_to_clip(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
+    """Convert one-sided censoring maximum(x, c) and minimum(x, c) to clip form.
+
+    The unbounded side uses the measurable variable itself, which `find_measurable_clips`
+    understands as one-sided clipping and which preserves the dtype of discrete
+    variables. Two-sided censoring maximum(minimum(x, ub), lb) is fused into a single
+    two-sided clip by `find_measurable_clips`.
+    """
+    measurable_inputs = filter_measurable_variables(node.inputs)
+    if len(measurable_inputs) != 1:
+        return None
+
+    [measurable_input] = measurable_inputs
+    [other_input] = [inp for inp in node.inputs if inp is not measurable_input]
+
+    # If both inputs are potentially measurable this is an order statistic, not censoring
+    if check_potential_measurability([other_input]):  # type: ignore[list-item]
+        return None
+
+    if isinstance(node.op.scalar_op, Maximum):
+        return [pt.clip(measurable_input, other_input, measurable_input)]
+    else:
+        return [pt.clip(measurable_input, measurable_input, other_input)]
+
+
+measurable_ir_rewrites_db.register(
+    "measurable_max_min_to_clip",
+    measurable_max_min_to_clip,
+    "basic",
+    "censoring",
+)
+
+
 @_logprob.register(MeasurableClip)
 def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
     r"""Logprob of a clipped censored distribution.
@@ -115,10 +214,7 @@ def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
         logprob.name = f"{base_rv_op}_logprob"
         logcdf.name = f"{base_rv_op}_logcdf"
 
-    is_lower_bounded, is_upper_bounded = False, False
-    if not (isinstance(upper_bound, TensorConstant) and np.all(np.isinf(upper_bound.value))):
-        is_upper_bounded = True
-
+    if upper_bound is not base_rv:
         logccdf = _logccdf_helper(base_rv, value, **kwargs)
 
         # For right clipped discrete RVs, we need to add an extra term
@@ -131,34 +227,87 @@ def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
             logccdf,
             pt.switch(pt.gt(value, upper_bound), -np.inf, logprob),
         )
-    if not (isinstance(lower_bound, TensorConstant) and np.all(np.isneginf(lower_bound.value))):
-        is_lower_bounded = True
+    if lower_bound is not base_rv:
         logprob = pt.switch(
             pt.eq(value, lower_bound),
             logcdf,
             pt.switch(pt.lt(value, lower_bound), -np.inf, logprob),
         )
 
-    if is_lower_bounded and is_upper_bounded:
+    if lower_bound is not base_rv and upper_bound is not base_rv:
         logprob = CheckParameterValue("lower_bound <= upper_bound")(
             logprob, pt.all(pt.le(lower_bound, upper_bound))
         )
 
     return logprob
 
 
+@_logcdf.register(MeasurableClip)
+def clip_logcdf(op, value, base_rv, lower_bound, upper_bound, **kwargs):
+    r"""Log-CDF of a clipped censored distribution.
+
+    .. math::
+        \begin{cases}
+            0 & \text{for } x < lower, \\
+            \text{CDF}(x, dist) & \text{for } lower <= x < upper, \\
+            1 & \text{for } x >= upper,
+        \end{cases}
+    """
+    logcdf = _logcdf_helper(base_rv, value)
+
+    if upper_bound is not base_rv:
+        logcdf = pt.switch(pt.ge(value, upper_bound), 0.0, logcdf)
+    if lower_bound is not base_rv:
+        logcdf = pt.switch(pt.lt(value, lower_bound), -np.inf, logcdf)
+
+    return logcdf
+
+
+@_icdf.register(MeasurableClip)
+def clip_icdf(op, value, base_rv, lower_bound, upper_bound, **kwargs):
+    # The point masses at the bounds absorb the respective tail quantiles
+    icdf = _icdf_helper(base_rv, value)
+
+    if lower_bound is not base_rv:
+        icdf = pt.maximum(icdf, lower_bound)
+    if upper_bound is not base_rv:
+        icdf = pt.minimum(icdf, upper_bound)
+
+    return icdf
+
+
 class MeasurableRound(MeasurableElemwise):
-    """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
+    """A placeholder used to specify a log-likelihood for a rounded RV sub-graph."""
 
-    valid_scalar_types = (RoundHalfToEven, Floor, Ceil)
+    valid_scalar_types = (RoundHalfToEven, RoundHalfAwayFromZero, Floor, Ceil, Trunc)
 
 
-@node_rewriter(tracks=[ceil, floor, round_half_to_even])
+@node_rewriter(tracks=[ceil, floor, round_half_to_even, round_half_away_from_zero, trunc])
 def find_measurable_roundings(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     if not filter_measurable_variables(node.inputs):
         return None
 
     [base_var] = node.inputs
+
+    if np.dtype(base_var.type.dtype).kind != "f" or isinstance(base_var.owner.op, MeasurableRound):
+        # The base already sits on the integers, either because its dtype guarantees it
+        # or because it is itself a rounding, so this one leaves it untouched save for
+        # the upcast the Ops apply to a discrete input. Reducing the rounding to that
+        # cast, rather than deriving the mass of a degenerate cell, leaves the base's
+        # own logprob, logcdf and icdf to apply. The cast folds away when the base is
+        # already a float, as it is for a rounded base.
+        return [pt.cast(base_var, node.outputs[0].type.dtype)]
+
+    # The reverse does not hold, so from here on the base must be shown to be continuous
+    # for the intervals below to be the right ones: an intermediate MeasurableVariable
+    # can be supported on the integers while carrying a float dtype, as `floor(x)` does.
+    # Only a RandomVariable states its own support, so anything else is declined rather
+    # than assumed continuous. Once MeasurableVariables carry the meta-information of
+    # the RV they encapsulate (https://github.com/pymc-devs/pymc/issues/6360), the
+    # support can be read off the base variable directly and this may be relaxed.
+    if not isinstance(base_var.owner.op, RandomVariable):
+        return None
+
     rounded_op = MeasurableRound(node.op.scalar_op)
     rounded_rv = rounded_op.make_node(base_var).default_output()
     rounded_rv.name = node.outputs[0].name
@@ -198,10 +347,23 @@ def round_logprob(op, values, base_rv, **kwargs):
             0 & \text{otherwise},
         \end{cases}
 
+    The probability of a distribution rounded towards zero is given by the
+    rounded-down probability for positive values and the rounded-up probability
+    for negative values, with both intervals pooled at zero.
+
     """
     (value,) = values
 
-    if isinstance(op.scalar_op, RoundHalfToEven):
+    if not value.type.dtype.startswith("float"):
+        # The Ops below snap the value onto the cell whose mass it asks for, and only
+        # accept float inputs. An integral value already sits on a cell, so the snapping
+        # is an identity and only the dtype needs adjusting (as when the rounding is
+        # followed by a cast to int).
+        value = pt.cast(value, "float64")
+
+    if isinstance(op.scalar_op, RoundHalfToEven | RoundHalfAwayFromZero):
+        # The tie-breaking rule only matters on a measure-zero set of the
+        # continuous base variable, so both variants share the same intervals
         value = pt.round(value)
         value_upper = value + 0.5
         value_lower = value - 0.5
@@ -213,6 +375,13 @@ def round_logprob(op, values, base_rv, **kwargs):
         value = pt.ceil(value)
         value_upper = value
         value_lower = value - 1.0
+    elif isinstance(op.scalar_op, Trunc):
+        # Truncation rounds towards zero: [x, x+1) for x >= 0, (x-1, x] for x < 0,
+        # and (-1, 1) for x == 0 (open/closed bounds are equivalent for the
+        # continuous base variables this rewrite applies to)
+        value = pt.trunc(value)
+        value_upper = value + (value >= 0)
+        value_lower = value - (value <= 0)
     else:
         raise TypeError(f"Unsupported scalar_op {op.scalar_op}")  # pragma: no cover
 

pymc/logprob/tensor.py

@@ -67,6 +67,7 @@
     _logprob_helper,
     promised_valued_rv,
 )
+from pymc.logprob.censoring import MeasurableRound
 from pymc.logprob.rewriting import (
     assume_valued_outputs,
     early_measurable_ir_rewrites_db,
@@ -487,16 +488,33 @@ def find_measurable_casts(fgraph, node) -> list[TensorVariable] | None:
 
     [out] = node.outputs
     # bool < integer < float; casting to a lower kind discretizes the base variable,
-    # which is not a measure-preserving identity (float -> int would need to be
-    # treated as a form of censoring)
+    # which is not a measure-preserving identity
     kind_order = {"b": 0, "u": 1, "i": 1, "f": 2}
-    in_kind = kind_order.get(np.dtype(base_var.type.dtype).kind)
-    out_kind = kind_order.get(np.dtype(out.type.dtype).kind)
+    in_dtype = np.dtype(base_var.type.dtype)
+    out_dtype = np.dtype(out.type.dtype)
+    in_kind = kind_order.get(in_dtype.kind)
+    out_kind = kind_order.get(out_dtype.kind)
     if in_kind is None or out_kind is None:
         # Kinds we can't order, such as complex
         return None
     if out_kind < in_kind:
-        return None
+        # A float -> signed int cast rounds towards zero, so it is a `trunc` composed
+        # with a cast that merely relabels the dtype. Introduce the `trunc` explicitly
+        # and let `find_measurable_roundings` claim it; the relabelling cast is then
+        # measure-preserving and claimed on a later pass.
+        # The other narrowing casts are not truncations: unsigned ints wrap around for
+        # negative values (-2.7 -> 254 for uint8), and bool tests `x != 0`, which
+        # collapses the support onto two points instead of partitioning it.
+        if in_dtype.kind != "f" or out_dtype.kind != "i":
+            return None
+        # Rewriting to the `trunc` the user could have written themselves leaves the
+        # judgement of whether the base may be truncated to `find_measurable_roundings`,
+        # which declines the bases whose support it cannot establish, rather than
+        # duplicating that reasoning here. Once the `trunc` is in place the cast only
+        # relabels the dtype and is claimed as measure-preserving below; skipping it for
+        # an already rounded base is what brings the rewrite to a fixpoint.
+        if not isinstance(base_var.owner.op, MeasurableRound):
+            return [pt.cast(pt.trunc(base_var), out_dtype.name)]
 
     if in_kind < kind_order["f"] and out_kind == kind_order["f"]:
         # Casting a discrete variable to float hides its discreteness from other