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, token, death_idx, selected_rho in ( ("exponential", 4, 9, None), ("weibull", 4, 9, rho), ("mixed", 9, 9, rho), ("mixed", 4, 9, None), ): actual = _risk_probability_matrix( logits=logits, rho=selected_rho, horizons=horizons, dist_mode=dist_mode, token=token, death_idx=death_idx, ) expected = np.vstack( [ _score_to_probability( logits, selected_rho, score_mode="risk", horizon=float(horizon), dist_mode=dist_mode, token=token, death_idx=death_idx, ) 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), "death_index": 9, "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()