diff --git a/src/rtichoke/calibration/calibration.py b/src/rtichoke/calibration/calibration.py index c809ed6d..0d102d45 100644 --- a/src/rtichoke/calibration/calibration.py +++ b/src/rtichoke/calibration/calibration.py @@ -106,19 +106,25 @@ def create_calibration_curve_times( competing events as censored, which calibration does not support. """ - unsupported = [ - heuristics + real_values = reals.values() if isinstance(reals, dict) else [reals] + has_censoring = any(np.any(np.asarray(values) == 0) for values in real_values) + + unsupported_adjusted_censoring = has_censoring and any( + heuristics.get("censoring_heuristic") == "adjusted" for heuristics in heuristics_sets - if heuristics.get("censoring_heuristic") == "adjusted" - or heuristics.get("competing_heuristic") == "adjusted_as_censored" - ] - if unsupported: + ) + unsupported_competing_as_censored = any( + heuristics.get("competing_heuristic") == "adjusted_as_censored" + for heuristics in heuristics_sets + ) + + if unsupported_adjusted_censoring or unsupported_competing_as_censored: raise ValueError( "Unsupported calibration heuristics: " "create_calibration_curve_times() does not support " - "censoring_heuristic='adjusted' or " - "competing_heuristic='adjusted_as_censored'. " - "Use a supported heuristic combination such as " + "censoring_heuristic='adjusted' when censored observations are present, " + "or competing_heuristic='adjusted_as_censored'. " + "When censoring is present, use a supported heuristic combination such as " "censoring_heuristic='excluded' with " "competing_heuristic='adjusted_as_negative'." ) @@ -1091,6 +1097,7 @@ def _create_calibration_curve_list_times( """ # Part 1: Prepare initial dataframe from inputs initial_df = _build_initial_df_for_times(probs, reals, times) + has_censoring = initial_df.filter(pl.col("real") == 0).height > 0 # Part 2: Iterate and generate calibration data for each horizon/heuristic all_deciles = [] @@ -1105,7 +1112,7 @@ def _create_calibration_curve_list_times( competing_heuristic = heuristics["competing_heuristic"] if ( - censoring_heuristic == "adjusted" + (censoring_heuristic == "adjusted" and has_censoring) or competing_heuristic == "adjusted_as_censored" ): continue diff --git a/tests/test_calibration_times.py b/tests/test_calibration_times.py index 6fbc86ed..5be7d7e3 100644 --- a/tests/test_calibration_times.py +++ b/tests/test_calibration_times.py @@ -64,24 +64,58 @@ def test_create_calibration_curve_times_unequal_size_populations(): "entry_point", [create_calibration_curve_times_top_level, create_calibration_curve_times_direct], ) +def test_create_calibration_curve_times_allows_adjusted_without_censoring(entry_point): + probs = {"model_1": np.array([0.1, 0.2, 0.3, 0.4])} + reals = np.array([1, 1, 1, 1]) + times = np.array([1.0, 2.0, 3.0, 4.0]) + + fig = entry_point( + probs, + reals, + times, + fixed_time_horizons=[2.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + ) + + assert fig is not None + + @pytest.mark.parametrize( - "heuristics", - [ - { - "censoring_heuristic": "adjusted", - "competing_heuristic": "adjusted_as_negative", - }, - { - "censoring_heuristic": "excluded", - "competing_heuristic": "adjusted_as_censored", - }, - ], + "entry_point", + [create_calibration_curve_times_top_level, create_calibration_curve_times_direct], +) +def test_create_calibration_curve_times_rejects_adjusted_with_censoring(entry_point): + probs = {"model_1": np.array([0.1, 0.2, 0.3, 0.4])} + reals = np.array([0, 1, 1, 1]) + times = np.array([1.0, 2.0, 3.0, 4.0]) + + with pytest.raises(ValueError, match="Unsupported calibration heuristics"): + entry_point( + probs, + reals, + times, + fixed_time_horizons=[2.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + ) + + +@pytest.mark.parametrize( + "entry_point", + [create_calibration_curve_times_top_level, create_calibration_curve_times_direct], ) -def test_create_calibration_curve_times_rejects_unsupported_heuristics( - entry_point, heuristics -): +def test_create_calibration_curve_times_rejects_competing_as_censored(entry_point): probs = {"model_1": np.array([0.1, 0.2, 0.3, 0.4])} - reals = np.array([0, 1, 0, 1]) + reals = np.array([1, 1, 1, 1]) times = np.array([1.0, 2.0, 3.0, 4.0]) with pytest.raises(ValueError, match="Unsupported calibration heuristics"): @@ -90,5 +124,10 @@ def test_create_calibration_curve_times_rejects_unsupported_heuristics( reals, times, fixed_time_horizons=[2.0], - heuristics_sets=[heuristics], + heuristics_sets=[ + { + "censoring_heuristic": "excluded", + "competing_heuristic": "adjusted_as_censored", + } + ], )