diff --git a/src/rtichoke/calibration/calibration.py b/src/rtichoke/calibration/calibration.py index 0d102d45..44f430dc 100644 --- a/src/rtichoke/calibration/calibration.py +++ b/src/rtichoke/calibration/calibration.py @@ -10,6 +10,7 @@ from plotly.graph_objs._figure import Figure import polars as pl import numpy as np +from polarstate import predict_aj_estimates, prepare_event_table # from rtichoke.helpers.send_post_request_to_r_rtichoke import send_requests_to_rtichoke_r @@ -106,27 +107,16 @@ def create_calibration_curve_times( competing events as censored, which calibration does not support. """ - 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 - ) 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: + if unsupported_competing_as_censored: raise ValueError( "Unsupported calibration heuristics: " "create_calibration_curve_times() does not support " - "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'." + "competing_heuristic='adjusted_as_censored'." ) calibration_curve_list_times = _create_calibration_curve_list_times( @@ -135,6 +125,7 @@ def create_calibration_curve_times( times, fixed_time_horizons=fixed_time_horizons, heuristics_sets=heuristics_sets, + calibration_type=calibration_type, size=size, color_values=color_values, ) @@ -943,6 +934,7 @@ def _build_initial_df_for_times( """Builds the initial DataFrame for time-dependent calibration curves.""" # Convert all inputs to dictionaries of arrays to unify processing + reals_was_dict = isinstance(reals, dict) if not isinstance(reals, dict): reals = {"single_population": np.asarray(reals)} if not isinstance(times, dict): @@ -979,9 +971,13 @@ def _build_initial_df_for_times( raise ValueError( f"Length of probabilities for model '{model_name}' does not match total number of observations." ) - return base_df.with_columns( - pl.Series("prob", prob_array), pl.lit(model_name).alias("model") - ) + expressions = [ + pl.Series("prob", prob_array), + pl.lit(model_name).alias("model"), + ] + if not reals_was_dict: + expressions.append(pl.lit(model_name).alias("reference_group")) + return base_df.with_columns(expressions) # Multiple models else: @@ -1062,12 +1058,96 @@ def _apply_heuristics_and_censoring( return df_adj +def _prepare_adjusted_event_data( + df: pl.DataFrame, horizon: float, competing_heuristic: str +) -> pl.DataFrame: + """Prepare event histories for Aalen-Johansen estimation.""" + event_data = df + if competing_heuristic == "excluded": + event_data = event_data.filter( + ~((pl.col("real") == 2) & (pl.col("time") <= horizon)) + ) + elif competing_heuristic == "adjusted_as_composite": + event_data = event_data.with_columns( + pl.when((pl.col("real") == 2) & (pl.col("time") <= horizon)) + .then(1) + .otherwise(pl.col("real")) + .alias("real") + ) + return event_data + + +def _aj_risk_at_horizon(df: pl.DataFrame, horizon: float) -> float: + """Estimate target-event cumulative incidence at one horizon.""" + event_table = prepare_event_table( + df.select(pl.col("time").alias("times"), pl.col("real").alias("reals")) + ) + estimate = predict_aj_estimates(event_table, pl.Series([horizon])) + return float(estimate["state_occupancy_probability_1"][0]) + + +def _make_adjusted_deciles_data( + df: pl.DataFrame, horizon: float, n_bins: int = 10 +) -> pl.DataFrame: + """Create calibration groups using within-group Aalen-Johansen risks.""" + grouped = df.with_columns( + ( + (pl.col("prob").rank("average").over("reference_group") - 1) + * n_bins + // pl.len().over("reference_group") + + 1 + ).alias("decile") + ) + rows = [] + for key, group_df in grouped.group_by(["reference_group", "decile"]): + reference_group, decile = key + estimate = _aj_risk_at_horizon(group_df, horizon) + n = group_df.height + rows.append( + { + "reference_group": reference_group, + "model": reference_group, + "decile": decile, + "n": n, + "x": cast(float, group_df["prob"].mean()), + "y": estimate, + "n_reals": estimate * n, + } + ) + return pl.DataFrame(rows).sort(["reference_group", "decile"]) + + +def _calculate_adjusted_pseudostates( + df: pl.DataFrame, horizon: float +) -> Dict[str, np.ndarray]: + """Calculate leave-one-out AJ pseudo-observations without a new dependency.""" + pseudo_by_group: Dict[str, np.ndarray] = {} + for key, group_df in df.group_by("reference_group", maintain_order=True): + reference_group = str(key[0]) + n = group_df.height + theta = _aj_risk_at_horizon(group_df, horizon) + if n == 1: + pseudo_by_group[reference_group] = np.array([theta]) + continue + leave_one_out = np.array( + [ + _aj_risk_at_horizon( + group_df.slice(0, i).vstack(group_df.slice(i + 1)), horizon + ) + for i in range(n) + ] + ) + pseudo_by_group[reference_group] = n * theta - (n - 1) * leave_one_out + return pseudo_by_group + + def _create_calibration_curve_list_times( probs: Dict[str, np.ndarray], reals: Union[np.ndarray, Dict[str, np.ndarray]], times: Union[np.ndarray, Dict[str, np.ndarray]], fixed_time_horizons: List[float], heuristics_sets: List[Dict[str, str]], + calibration_type: str = "discrete", size: int = 600, color_values: List[str] = [ "#1b9e77", @@ -1097,8 +1177,6 @@ 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 = [] all_smooth = [] @@ -1111,10 +1189,45 @@ def _create_calibration_curve_list_times( censoring_heuristic = heuristics["censoring_heuristic"] competing_heuristic = heuristics["competing_heuristic"] - if ( - (censoring_heuristic == "adjusted" and has_censoring) - or competing_heuristic == "adjusted_as_censored" - ): + if competing_heuristic == "adjusted_as_censored": + continue + + if censoring_heuristic == "adjusted": + df_adj = _prepare_adjusted_event_data( + initial_df, horizon, competing_heuristic + ) + if df_adj.height == 0: + continue + + deciles_data = _make_adjusted_deciles_data(df_adj, horizon) + probs_adj = { + key[0]: group_df["prob"].to_numpy() + for key, group_df in df_adj.group_by( + "reference_group", maintain_order=True + ) + } + if calibration_type == "smooth": + pseudo_by_group = _calculate_adjusted_pseudostates(df_adj, horizon) + smooth_data = _calculate_smooth_curve( + probs_adj, pseudo_by_group, performance_type + ) + else: + smooth_data = deciles_data.select("x", "y", "reference_group") + hist_data = _create_histogram_for_calibration(probs_adj) + + all_deciles.append( + deciles_data.with_columns( + pl.lit(horizon).alias("fixed_time_horizon") + ) + ) + all_smooth.append( + smooth_data.with_columns( + pl.lit(horizon).alias("fixed_time_horizon") + ) + ) + all_histograms.append( + hist_data.with_columns(pl.lit(horizon).alias("fixed_time_horizon")) + ) continue df_adj = _apply_heuristics_and_censoring( diff --git a/tests/test_calibration_times.py b/tests/test_calibration_times.py index 5be7d7e3..33ac4cc9 100644 --- a/tests/test_calibration_times.py +++ b/tests/test_calibration_times.py @@ -1,8 +1,14 @@ import numpy as np import pytest -from rtichoke import create_calibration_curve_times as create_calibration_curve_times_top_level +import polars as pl +from lifelines import AalenJohansenFitter, KaplanMeierFitter +from rtichoke import ( + create_calibration_curve_times as create_calibration_curve_times_top_level, +) from rtichoke.calibration import create_calibration_curve_times from rtichoke.calibration.calibration import ( + _calculate_adjusted_pseudostates, + _prepare_adjusted_event_data, create_calibration_curve_times as create_calibration_curve_times_direct, ) @@ -89,24 +95,141 @@ def test_create_calibration_curve_times_allows_adjusted_without_censoring(entry_ "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): +def test_create_calibration_curve_times_adjusts_independent_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", - } - ], + 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 + assert len(fig.data) == 3 + + +def test_adjusted_discrete_calibration_uses_kaplan_meier_risk(): + probs = {"model_1": np.repeat(0.5, 4)} + reals = np.array([1, 0, 1, 0]) + times = np.array([1.0, 2.0, 3.0, 4.0]) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[3.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + ) + + calibration_trace = fig.data[1] + reference = 1 - KaplanMeierFitter().fit(times, reals).predict(3.0) + assert calibration_trace.y[0] == pytest.approx(reference) + + +def test_adjusted_discrete_calibration_matches_aalen_johansen(): + probs = {"model_1": np.repeat(0.4, 5)} + reals = np.array([1, 2, 0, 2, 0]) + times = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[3.5], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + ) + + reference = ( + AalenJohansenFitter().fit(times, reals, event_of_interest=1).predict(3.5) + ) + assert fig.data[1].y[0] == pytest.approx(reference) + + +def test_excluded_competing_events_are_horizon_specific(): + data = pl.DataFrame( + { + "reference_group": ["model_1"] * 5, + "prob": [0.4] * 5, + "time": [1.0, 2.0, 3.0, 4.0, 5.0], + "real": [1, 2, 0, 2, 0], + } + ) + + adjusted = _prepare_adjusted_event_data(data, 3.5, "excluded") + + assert adjusted["time"].to_list() == [1.0, 3.0, 4.0, 5.0] + assert adjusted.filter(pl.col("time") == 4.0)["real"].item() == 2 + + +def test_adjusted_pseudostates_match_leave_one_out_aalen_johansen(): + times = np.array([1.0, 2.0, 3.0, 4.0]) + reals = np.array([1, 0, 2, 0]) + horizon = 3.5 + data = pl.DataFrame( + { + "reference_group": ["model_1"] * 4, + "prob": [0.1, 0.3, 0.6, 0.8], + "time": times, + "real": reals, + } + ) + + actual = _calculate_adjusted_pseudostates(data, horizon)["model_1"] + full = AalenJohansenFitter().fit(times, reals, event_of_interest=1).predict(horizon) + expected = [] + for index in range(len(times)): + keep = np.arange(len(times)) != index + leave_one_out = ( + AalenJohansenFitter() + .fit(times[keep], reals[keep], event_of_interest=1) + .predict(horizon) ) + expected.append(len(times) * full - (len(times) - 1) * leave_one_out) + + np.testing.assert_allclose(actual, expected) + + +def test_adjusted_smooth_calibration_uses_pseudo_observations(): + probs = {"model_1": np.array([0.1, 0.3, 0.6, 0.8])} + reals = np.array([1, 0, 1, 0]) + times = np.array([1.0, 2.0, 3.0, 4.0]) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[3.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + ) + + calibration_trace = fig.data[1] + assert len(calibration_trace.x) == 101 + assert np.isfinite(np.asarray(calibration_trace.y)).all() @pytest.mark.parametrize(