Add private outcome-weighted calibration bin primitive

src/rtichoke/calibration/_weights.py

@@ -0,0 +1,95 @@
+"""Private outcome-weighted calibration preparation helpers."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+
+from .calibration import _make_deciles_dat_binary
+
+
+def _prepare_calibration_bins(
+    probs: np.ndarray,
+    reals: np.ndarray,
+    outcome_weights: np.ndarray | None = None,
+    n_bins: int = 10,
+) -> pl.DataFrame:
+    """Prepare vector-level calibration bins with optional outcome weights.
+
+    Prediction-bin membership and displayed mean predicted risk are always
+    defined from the target-population predictions. ``outcome_weights`` affect
+    only the observed outcome estimate inside each bin. This separation is
+    intentional for later counterfactual calibration, where treatment/IPW
+    weights identify the observed risk under an intervention but should not
+    redefine the target population's prediction distribution.
+
+    With no weights, this delegates to the established factual calibration
+    helper so existing semantics remain unchanged.
+    """
+    p = np.asarray(probs).ravel()
+    y = np.asarray(reals).ravel()
+
+    if p.shape[0] != y.shape[0]:
+        raise ValueError("probs and reals must have the same length.")
+    if not isinstance(n_bins, int) or isinstance(n_bins, bool) or n_bins < 1:
+        raise ValueError("n_bins must be a positive integer.")
+
+    if outcome_weights is None:
+        return _make_deciles_dat_binary({"model": p}, y, n_bins=n_bins).with_columns(
+            pl.col("n").cast(pl.Float64).alias("outcome_weight_sum")
+        )
+
+    weights = np.asarray(outcome_weights, dtype=float).ravel()
+    if weights.shape[0] != y.shape[0]:
+        raise ValueError("outcome_weights must have the same length as reals.")
+    if not np.all(np.isfinite(weights)) or np.any(weights < 0):
+        raise ValueError(
+            "outcome_weights must contain only finite, non-negative values."
+        )
+
+    df = pl.DataFrame(
+        {
+            "reference_group": ["model"] * p.shape[0],
+            "model": ["model"] * p.shape[0],
+            "prob": p.astype(float, copy=False),
+            "real": y.astype(float, copy=False),
+            "outcome_weight": weights,
+        }
+    ).with_columns(
+        ((pl.col("prob").rank("ordinal") - 1) * n_bins // pl.len() + 1).alias(
+            "decile"
+        )
+    )
+
+    bins = (
+        df.group_by(["reference_group", "model", "decile"])
+        .agg(
+            pl.len().alias("n"),
+            pl.mean("prob").alias("x"),
+            pl.sum("real").alias("n_reals"),
+            pl.sum("outcome_weight").alias("outcome_weight_sum"),
+            (pl.col("outcome_weight") * pl.col("real"))
+            .sum()
+            .alias("weighted_sum_reals"),
+        )
+        .sort(["reference_group", "model", "decile"])
+    )
+
+    if bins.filter(pl.col("outcome_weight_sum") <= 0).height:
+        raise ValueError(
+            "Every calibration bin must have positive total outcome_weights."
+        )
+
+    return bins.with_columns(
+        (pl.col("weighted_sum_reals") / pl.col("outcome_weight_sum")).alias("y")
+    ).select(
+        "reference_group",
+        "model",
+        "decile",
+        "n",
+        "x",
+        "y",
+        "n_reals",
+        "outcome_weight_sum",
+        "weighted_sum_reals",
+    )

tests/test_calibration_weights.py

@@ -0,0 +1,102 @@
+import numpy as np
+import pytest
+
+from rtichoke.calibration._weights import _prepare_calibration_bins
+from rtichoke.calibration.calibration import _make_deciles_dat_binary
+
+
+def test_unweighted_private_calibration_bins_preserve_factual_semantics():
+    probs = np.linspace(0.05, 0.95, 20)
+    reals = np.tile(np.array([0, 1]), 10)
+
+    expected = _make_deciles_dat_binary({"model": probs}, reals)
+    actual = _prepare_calibration_bins(probs, reals)
+
+    assert actual.select(expected.columns).equals(expected)
+    assert actual.get_column("outcome_weight_sum").to_list() == [
+        float(value) for value in expected.get_column("n").to_list()
+    ]
+
+
+def test_outcome_weights_change_observed_calibration_not_bins_or_predicted_means():
+    probs = np.array([0.10, 0.20, 0.30, 0.40, 0.60, 0.70, 0.80, 0.90])
+    reals = np.array([0, 1, 0, 1, 1, 0, 1, 0])
+    weights = np.array([1, 3, 1, 1, 2, 1, 4, 1], dtype=float)
+
+    weighted = _prepare_calibration_bins(
+        probs, reals, outcome_weights=weights, n_bins=2
+    )
+    unweighted = _make_deciles_dat_binary({"model": probs}, reals, n_bins=2)
+
+    assert weighted.get_column("decile").to_list() == unweighted.get_column(
+        "decile"
+    ).to_list()
+    np.testing.assert_allclose(
+        weighted.get_column("x").to_numpy(),
+        unweighted.get_column("x").to_numpy(),
+    )
+    np.testing.assert_allclose(weighted.get_column("y").to_numpy(), [4 / 6, 6 / 8])
+    np.testing.assert_allclose(
+        weighted.get_column("outcome_weight_sum").to_numpy(), [6, 8]
+    )
+    np.testing.assert_allclose(
+        weighted.get_column("weighted_sum_reals").to_numpy(), [4, 6]
+    )
+
+
+def test_all_one_outcome_weights_reproduce_unweighted_bin_estimates():
+    probs = np.linspace(0.05, 0.95, 20)
+    reals = np.tile(np.array([0, 1]), 10)
+
+    unweighted = _prepare_calibration_bins(probs, reals)
+    weighted = _prepare_calibration_bins(
+        probs, reals, outcome_weights=np.ones(reals.shape[0])
+    )
+
+    np.testing.assert_allclose(
+        weighted.get_column("x").to_numpy(), unweighted.get_column("x").to_numpy()
+    )
+    np.testing.assert_allclose(
+        weighted.get_column("y").to_numpy(), unweighted.get_column("y").to_numpy()
+    )
+    assert weighted.get_column("n_reals").to_list() == unweighted.get_column(
+        "n_reals"
+    ).to_list()
+    assert weighted.get_column("n").to_list() == unweighted.get_column("n").to_list()
+
+
+@pytest.mark.parametrize(
+    "weights, message",
+    [
+        (np.array([1.0, 1.0]), "same length"),
+        (np.array([1.0, -1.0, 1.0, 1.0]), "finite, non-negative"),
+        (np.array([1.0, np.inf, 1.0, 1.0]), "finite, non-negative"),
+    ],
+)
+def test_private_weighted_calibration_bins_validate_weights(weights, message):
+    probs = np.array([0.1, 0.2, 0.8, 0.9])
+    reals = np.array([0, 1, 1, 0])
+
+    with pytest.raises(ValueError, match=message):
+        _prepare_calibration_bins(probs, reals, outcome_weights=weights, n_bins=2)
+
+
+def test_private_weighted_calibration_bins_require_positive_weight_in_each_bin():
+    probs = np.array([0.1, 0.2, 0.8, 0.9])
+    reals = np.array([0, 1, 1, 0])
+
+    with pytest.raises(ValueError, match="positive total"):
+        _prepare_calibration_bins(
+            probs,
+            reals,
+            outcome_weights=np.array([1.0, 1.0, 0.0, 0.0]),
+            n_bins=2,
+        )
+
+
+def test_private_weighted_calibration_bins_validate_bin_count():
+    probs = np.array([0.1, 0.2])
+    reals = np.array([0, 1])
+
+    with pytest.raises(ValueError, match="positive integer"):
+        _prepare_calibration_bins(probs, reals, n_bins=0)