Add all-future calibration evaluation

2026-07-30 17:21:57 +08:00
parent c622ec50f7
commit e471de030d
3 changed files with 2039 additions and 0 deletions
--- a/tests/test_calibration_metrics.py
+++ b/tests/test_calibration_metrics.py
@@ -0,0 +1,143 @@
+import math
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from evaluate_calibration import (
+    aggregate_metric_rows,
+    compute_ipcw_cell,
+    fit_weighted_logistic_calibration,
+)
+
+
+class IPCWCalibrationMetricTests(unittest.TestCase):
+    def test_no_censoring_matches_binary_metrics(self):
+        result = compute_ipcw_cell(
+            probabilities=np.asarray([0.2, 0.8]),
+            event_times=np.asarray([np.inf, 0.5]),
+            censor_times=np.asarray([2.0, 2.0]),
+            horizon=1.0,
+            min_cases=1,
+            min_controls=1,
+            max_ipcw_weight=0.0,
+        )
+
+        self.assertIsNotNone(result)
+        row, arrays = result
+        self.assertEqual(row["n_events"], 1)
+        self.assertEqual(row["n_controls"], 1)
+        self.assertAlmostEqual(row["brier_ipcw"], 0.04)
+        self.assertAlmostEqual(row["nll_ipcw"], -math.log(0.8))
+        self.assertAlmostEqual(row["predicted_mean"], 0.5)
+        self.assertAlmostEqual(row["observed_rate_ipcw"], 0.5)
+        np.testing.assert_allclose(arrays["metric_weights"], [1.0, 1.0])
+
+    def test_censored_before_horizon_gets_zero_outcome_weight(self):
+        result = compute_ipcw_cell(
+            probabilities=np.asarray([0.8, 0.2, 0.4]),
+            event_times=np.asarray([0.5, np.inf, np.inf]),
+            censor_times=np.asarray([2.0, 2.0, 0.5]),
+            horizon=1.0,
+            min_cases=1,
+            min_controls=1,
+            max_ipcw_weight=0.0,
+        )
+
+        self.assertIsNotNone(result)
+        row, arrays = result
+        self.assertEqual(row["n_censored_before_horizon"], 1)
+        self.assertAlmostEqual(row["known_fraction"], 2.0 / 3.0)
+        np.testing.assert_allclose(
+            arrays["metric_weights"],
+            [1.0, 1.5, 0.0],
+        )
+        self.assertAlmostEqual(row["brier_ipcw"], 0.1 / 3.0)
+        self.assertAlmostEqual(
+            row["nll_ipcw"],
+            -2.5 * math.log(0.8) / 3.0,
+        )
+        self.assertAlmostEqual(row["observed_rate_ipcw"], 1.0 / 3.0)
+
+    def test_calibration_intercept_and_slope_recover_identity(self):
+        probabilities = np.repeat([0.1, 0.3, 0.7, 0.9], 100)
+        outcomes = np.concatenate(
+            [
+                np.r_[np.ones(10), np.zeros(90)],
+                np.r_[np.ones(30), np.zeros(70)],
+                np.r_[np.ones(70), np.zeros(30)],
+                np.r_[np.ones(90), np.zeros(10)],
+            ]
+        )
+        weights = np.ones_like(probabilities)
+
+        calibration_in_large, intercept, slope = (
+            fit_weighted_logistic_calibration(
+                probabilities,
+                outcomes,
+                weights,
+            )
+        )
+
+        self.assertAlmostEqual(calibration_in_large, 0.0, places=7)
+        self.assertAlmostEqual(intercept, 0.0, places=7)
+        self.assertAlmostEqual(slope, 1.0, places=7)
+
+    def test_metric_aggregation_uses_contribution_sums(self):
+        metrics = pd.DataFrame(
+            [
+                {
+                    "outcome": "Disease",
+                    "sex": "Female",
+                    "horizon": 5.0,
+                    "n_at_risk": 10,
+                    "n_events": 2,
+                    "n_controls": 7,
+                    "n_censored_before_horizon": 1,
+                    "prediction_sum": 2.0,
+                    "event_weight_sum": 2.0,
+                    "brier_ipcw_sum": 1.0,
+                    "nll_ipcw_sum": 3.0,
+                    "calibration_in_the_large": 0.1,
+                    "calibration_intercept": 0.2,
+                    "calibration_slope": 0.9,
+                    "ipcw_weight_max": 1.2,
+                    "ipcw_weights_clipped": 0,
+                },
+                {
+                    "outcome": "Disease",
+                    "sex": "Female",
+                    "horizon": 5.0,
+                    "n_at_risk": 10,
+                    "n_events": 3,
+                    "n_controls": 6,
+                    "n_censored_before_horizon": 1,
+                    "prediction_sum": 3.0,
+                    "event_weight_sum": 3.0,
+                    "brier_ipcw_sum": 2.0,
+                    "nll_ipcw_sum": 4.0,
+                    "calibration_in_the_large": -0.1,
+                    "calibration_intercept": -0.2,
+                    "calibration_slope": 1.1,
+                    "ipcw_weight_max": 1.4,
+                    "ipcw_weights_clipped": 1,
+                },
+            ]
+        )
+
+        aggregated = aggregate_metric_rows(
+            metrics,
+            group_columns=["outcome", "sex", "horizon"],
+        ).iloc[0]
+
+        self.assertEqual(aggregated["n_at_risk"], 20)
+        self.assertAlmostEqual(aggregated["predicted_mean"], 0.25)
+        self.assertAlmostEqual(aggregated["observed_rate_ipcw"], 0.25)
+        self.assertAlmostEqual(aggregated["brier_ipcw"], 0.15)
+        self.assertAlmostEqual(aggregated["nll_ipcw"], 0.35)
+        self.assertAlmostEqual(aggregated["calibration_slope_median"], 1.0)
+        self.assertEqual(aggregated["ipcw_weights_clipped"], 1)
+
+
+if __name__ == "__main__":
+    unittest.main()