458 lines
16 KiB
Python
458 lines
16 KiB
Python
import math
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
|
|
from evaluate_calibration import (
|
|
_censoring_km,
|
|
_evaluate_calibration_token,
|
|
_risk_probability_matrix,
|
|
aggregate_metric_rows,
|
|
compute_ipcw_cell,
|
|
compute_ipcw_horizons,
|
|
evaluate_landmark_calibration,
|
|
fit_weighted_logistic_calibration,
|
|
fit_weighted_logistic_calibration_batch,
|
|
)
|
|
from evaluate_auc_v2 import _score_to_probability
|
|
|
|
|
|
class IPCWCalibrationMetricTests(unittest.TestCase):
|
|
@staticmethod
|
|
def _naive_censoring_km(observed_times, censor_events):
|
|
observed_times = np.asarray(observed_times, dtype=np.float64)
|
|
censor_events = np.asarray(censor_events, dtype=bool)
|
|
event_times = np.unique(observed_times[censor_events])
|
|
survival = 1.0
|
|
survival_after = []
|
|
for time_value in event_times:
|
|
at_risk = np.sum(observed_times >= time_value)
|
|
censored = np.sum(
|
|
censor_events & (observed_times == time_value)
|
|
)
|
|
survival *= 1.0 - float(censored) / float(at_risk)
|
|
survival_after.append(survival)
|
|
return event_times, np.asarray(survival_after)
|
|
|
|
def test_sorted_censoring_km_matches_naive_reference(self):
|
|
rng = np.random.RandomState(12)
|
|
observed_times = rng.randint(1, 20, size=500).astype(np.float64)
|
|
censor_events = rng.uniform(size=500) < 0.4
|
|
|
|
expected_times, expected_survival = self._naive_censoring_km(
|
|
observed_times,
|
|
censor_events,
|
|
)
|
|
actual_times, actual_survival = _censoring_km(
|
|
observed_times,
|
|
censor_events,
|
|
)
|
|
|
|
np.testing.assert_array_equal(actual_times, expected_times)
|
|
np.testing.assert_allclose(
|
|
actual_survival,
|
|
expected_survival,
|
|
rtol=1e-14,
|
|
atol=1e-14,
|
|
)
|
|
|
|
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_batched_calibration_fits_multiple_horizons(self):
|
|
probabilities = np.vstack(
|
|
[
|
|
np.repeat([0.1, 0.3, 0.7, 0.9], 100),
|
|
np.repeat([0.2, 0.4, 0.6, 0.8], 100),
|
|
]
|
|
)
|
|
outcomes = np.vstack(
|
|
[
|
|
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)],
|
|
]
|
|
),
|
|
np.concatenate(
|
|
[
|
|
np.r_[np.ones(20), np.zeros(80)],
|
|
np.r_[np.ones(40), np.zeros(60)],
|
|
np.r_[np.ones(60), np.zeros(40)],
|
|
np.r_[np.ones(80), np.zeros(20)],
|
|
]
|
|
),
|
|
]
|
|
)
|
|
|
|
calibration_in_large, intercept, slope = (
|
|
fit_weighted_logistic_calibration_batch(
|
|
probabilities,
|
|
outcomes,
|
|
np.ones_like(probabilities),
|
|
)
|
|
)
|
|
|
|
np.testing.assert_allclose(calibration_in_large, 0.0, atol=1e-7)
|
|
np.testing.assert_allclose(intercept, 0.0, atol=1e-7)
|
|
np.testing.assert_allclose(slope, 1.0, atol=1e-7)
|
|
|
|
def test_all_horizons_reuse_one_censoring_km(self):
|
|
probabilities = np.asarray(
|
|
[
|
|
[0.05, 0.10, 0.15, 0.20, 0.25],
|
|
[0.10, 0.20, 0.30, 0.40, 0.50],
|
|
[0.20, 0.35, 0.50, 0.65, 0.80],
|
|
]
|
|
)
|
|
event_times = np.asarray([0.5, 1.5, 4.0, np.inf, np.inf])
|
|
censor_times = np.asarray([5.0, 5.0, 5.0, 2.5, 5.0])
|
|
|
|
with patch(
|
|
"evaluate_calibration._censoring_km",
|
|
wraps=_censoring_km,
|
|
) as km:
|
|
results = compute_ipcw_horizons(
|
|
probabilities=probabilities,
|
|
event_times=event_times,
|
|
censor_times=censor_times,
|
|
horizons=np.asarray([1.0, 2.0, 5.0]),
|
|
min_cases=1,
|
|
min_controls=1,
|
|
max_ipcw_weight=0.0,
|
|
)
|
|
|
|
self.assertEqual(km.call_count, 1)
|
|
self.assertEqual(len(results), 3)
|
|
self.assertTrue(all(result is not None for result in results))
|
|
self.assertEqual([result[0]["n_events"] for result in results], [1, 2, 3])
|
|
|
|
def test_batched_risk_probabilities_match_scalar_reference(self):
|
|
logits = np.asarray([-2.0, -0.5, 0.2, 1.5], dtype=np.float32)
|
|
rho = np.asarray([0.8, 1.0, 1.2, 1.5], dtype=np.float32)
|
|
horizons = np.asarray([0.1, 1.0, 5.0], dtype=np.float32)
|
|
|
|
for dist_mode, selected_rho in (
|
|
("exponential", None),
|
|
("weibull", rho),
|
|
):
|
|
actual = _risk_probability_matrix(
|
|
logits=logits,
|
|
rho=selected_rho,
|
|
horizons=horizons,
|
|
dist_mode=dist_mode,
|
|
)
|
|
expected = np.vstack(
|
|
[
|
|
_score_to_probability(
|
|
logits,
|
|
selected_rho,
|
|
score_mode="risk",
|
|
horizon=float(horizon),
|
|
dist_mode=dist_mode,
|
|
)
|
|
for horizon in horizons
|
|
]
|
|
)
|
|
np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-7)
|
|
|
|
def test_per_disease_worker_is_thread_safe(self):
|
|
logits_chunk = np.asarray(
|
|
[
|
|
[-2.0, -1.5],
|
|
[-1.0, -0.5],
|
|
[0.0, 0.5],
|
|
[0.5, 1.0],
|
|
[1.0, 1.5],
|
|
[1.5, 2.0],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
common = {
|
|
"logits_chunk": logits_chunk,
|
|
"rho_chunk": None,
|
|
"strata": [("Female", 50.0, np.arange(6, dtype=np.int64))],
|
|
"row_patient_id": np.arange(6, dtype=np.int32),
|
|
"row_followup_end": np.full(6, 65.0, dtype=np.float32),
|
|
"row_death_time": np.full(6, np.inf, dtype=np.float32),
|
|
"first_occurrence_by_token": {
|
|
4: (
|
|
np.asarray([0, 1], dtype=np.int32),
|
|
np.asarray([50.5, 52.0], dtype=np.float32),
|
|
),
|
|
5: (
|
|
np.asarray([2, 3], dtype=np.int32),
|
|
np.asarray([50.7, 53.0], dtype=np.float32),
|
|
),
|
|
},
|
|
"patient_count": 6,
|
|
"death_tokens": set(),
|
|
"label_id_to_code": {4: "D4", 5: "D5"},
|
|
"dist_mode": "exponential",
|
|
"horizons": np.asarray([1.0, 5.0], dtype=np.float32),
|
|
"min_cases": 1,
|
|
"min_controls": 1,
|
|
"max_ipcw_weight": 0.0,
|
|
"exclude_death_competing": True,
|
|
"probability_bins": np.asarray([0.0, 0.5, 1.0]),
|
|
}
|
|
tasks = [
|
|
{"column_index": 0, "token": 4, **common},
|
|
{"column_index": 1, "token": 5, **common},
|
|
]
|
|
serial = [
|
|
_evaluate_calibration_token(**task)
|
|
for task in tasks
|
|
]
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
parallel = list(
|
|
executor.map(
|
|
lambda task: _evaluate_calibration_token(**task),
|
|
tasks,
|
|
)
|
|
)
|
|
|
|
for (serial_rows, serial_curve), (
|
|
parallel_rows,
|
|
parallel_curve,
|
|
) in zip(serial, parallel):
|
|
pd.testing.assert_frame_equal(
|
|
pd.DataFrame(serial_rows),
|
|
pd.DataFrame(parallel_rows),
|
|
)
|
|
self.assertEqual(set(serial_curve), set(parallel_curve))
|
|
for key in serial_curve:
|
|
self.assertEqual(
|
|
set(serial_curve[key]),
|
|
set(parallel_curve[key]),
|
|
)
|
|
np.testing.assert_allclose(
|
|
list(serial_curve[key].values()),
|
|
list(parallel_curve[key].values()),
|
|
equal_nan=True,
|
|
)
|
|
|
|
def test_landmark_evaluation_parallel_matches_serial(self):
|
|
logits_chunk = np.asarray(
|
|
[
|
|
[-2.0, -1.5],
|
|
[-1.0, -0.5],
|
|
[0.0, 0.5],
|
|
[0.5, 1.0],
|
|
[1.0, 1.5],
|
|
[1.5, 2.0],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
row_arrays = {
|
|
"patient_id": np.arange(6, dtype=np.int32),
|
|
"sex": np.zeros(6, dtype=np.int8),
|
|
"landmark_age": np.full(6, 50.0, dtype=np.float32),
|
|
"followup_end_time": np.full(6, 65.0, dtype=np.float32),
|
|
"death_time": np.full(6, np.inf, dtype=np.float32),
|
|
}
|
|
landmark_dataset = SimpleNamespace(
|
|
subset_indices=np.arange(6, dtype=np.int64),
|
|
death_token_ids=[],
|
|
first_occurrence_by_token={
|
|
4: (
|
|
np.asarray([0, 1], dtype=np.int32),
|
|
np.asarray([50.5, 52.0], dtype=np.float32),
|
|
),
|
|
5: (
|
|
np.asarray([2, 3], dtype=np.int32),
|
|
np.asarray([50.7, 53.0], dtype=np.float32),
|
|
),
|
|
},
|
|
dataset=SimpleNamespace(
|
|
label_id_to_code={4: "D4", 5: "D5"}
|
|
),
|
|
)
|
|
|
|
class FakeModel:
|
|
death_idx = 9
|
|
vocab_size = 10
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
def to(self, _device):
|
|
return self
|
|
|
|
common = {
|
|
"model": FakeModel(),
|
|
"loader": [],
|
|
"landmark_dataset": landmark_dataset,
|
|
"disease_ids": [4, 5],
|
|
"dist_mode": "exponential",
|
|
"horizons": np.asarray([1.0, 5.0], dtype=np.float32),
|
|
"device": torch.device("cpu"),
|
|
"use_amp": False,
|
|
"hidden_cache_dtype": "float16",
|
|
"logit_batch_size": 8,
|
|
"disease_chunk_size": 2,
|
|
"min_cases": 1,
|
|
"min_controls": 1,
|
|
"max_ipcw_weight": 0.0,
|
|
"exclude_death_competing": True,
|
|
"probability_bins": np.asarray([0.0, 0.5, 1.0]),
|
|
}
|
|
with (
|
|
patch(
|
|
"evaluate_calibration.infer_landmark_hidden",
|
|
return_value=(
|
|
np.zeros((6, 4), dtype=np.float16),
|
|
row_arrays,
|
|
),
|
|
),
|
|
patch(
|
|
"evaluate_calibration.project_distribution_chunk",
|
|
return_value=(logits_chunk, None),
|
|
),
|
|
):
|
|
serial_metrics, serial_curve = evaluate_landmark_calibration(
|
|
**common,
|
|
num_workers_calibration=1,
|
|
)
|
|
parallel_metrics, parallel_curve = evaluate_landmark_calibration(
|
|
**common,
|
|
num_workers_calibration=2,
|
|
)
|
|
|
|
pd.testing.assert_frame_equal(serial_metrics, parallel_metrics)
|
|
pd.testing.assert_frame_equal(serial_curve, parallel_curve)
|
|
|
|
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()
|