diff --git a/src/rtichoke/calibration/_weights.py b/src/rtichoke/calibration/_weights.py new file mode 100644 index 00000000..7d870fa0 --- /dev/null +++ b/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", + ) diff --git a/tests/test_calibration_weights.py b/tests/test_calibration_weights.py new file mode 100644 index 00000000..74e3710d --- /dev/null +++ b/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)