144 lines
5.0 KiB
Python
144 lines
5.0 KiB
Python
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()
|