DeepHealth/tests/test_calibration_metrics.py

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()