From d5a543dd2feabfcd0a71ee9afd318407482ea4d1 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 16 Aug 2026 14:14:43 +0000 Subject: [PATCH 1/2] feat: support local Aalen-Johansen and secondary Cox calibration smoothing - Added `smooth_method` parameter to `create_calibration_curve_times()` supporting `"local_aj"` (Gerds' preferred local Aalen-Johansen method, default), `"secondary_cox"` (Austin, Harrell & McLernon secondary Cox regression method), and `"pseudo_values"` (jackknife pseudo-value lowess). - Added `bandwidth` parameter for `local_aj` neighborhood window tuning. - Updated documentation and added comprehensive unit tests for time-dependent calibration smoothing methods across single, multiple model, and competing risk scenarios. Co-authored-by: uriahf <11351434+uriahf@users.noreply.github.com> --- src/rtichoke/calibration/calibration.py | 241 +++++++++++++++++++++- tests/test_calibration_times.py | 80 +++++++ user_guide/02-curve-api-compatibility.qmd | 11 +- user_guide/03-common-errors.qmd | 10 +- 4 files changed, 323 insertions(+), 19 deletions(-) diff --git a/src/rtichoke/calibration/calibration.py b/src/rtichoke/calibration/calibration.py index cb3af9b5..c2e05213 100644 --- a/src/rtichoke/calibration/calibration.py +++ b/src/rtichoke/calibration/calibration.py @@ -89,6 +89,8 @@ def create_calibration_curve_times( fixed_time_horizons: List[float], heuristics_sets: List[Dict[str, str]], calibration_type: str = "discrete", + smooth_method: str = "local_aj", + bandwidth: Union[float, None] = None, size: int = 600, color_values: List[str] = [ "#1b9e77", @@ -115,9 +117,48 @@ def create_calibration_curve_times( ) -> Figure: """Create a time-dependent calibration curve across fixed horizons. - Raises: - ValueError: If a heuristic set requests adjusted censoring or treats - competing events as censored, which calibration does not support. + This function generates time-dependent calibration curves evaluating predicted + probabilities against observed outcomes over specified prediction horizons. + + Parameters + ---------- + probs : Dict[str, np.ndarray] + A dictionary mapping model or dataset names to 1-D numpy arrays of + predicted probabilities. + reals : Union[np.ndarray, Dict[str, np.ndarray]] + True outcome indicators (0 for censored, 1 for event of interest, 2 for + competing risk). + times : Union[np.ndarray, Dict[str, np.ndarray]] + Follow-up times corresponding to `reals`. + fixed_time_horizons : List[float] + List of prediction horizons (times) at which to evaluate calibration. + heuristics_sets : List[Dict[str, str]] + List of heuristic dictionaries defining censoring and competing risk + adjustments. + calibration_type : str, optional + Type of calibration plot, either ``"discrete"`` (binned) or ``"smooth"``. + Defaults to ``"discrete"``. + smooth_method : str, optional + Smoothing method when `calibration_type="smooth"`. Supported options are + ``"local_aj"`` (Gerds' local Aalen-Johansen/KM neighborhood estimation), + ``"secondary_cox"`` (Austin, Harrell & McLernon secondary Cox regression method), + or ``"pseudo_values"`` (jackknife pseudo-values lowess). Defaults to ``"local_aj"``. + bandwidth : Union[float, None], optional + Bandwidth fraction for ``"local_aj"`` neighborhood smoothing. Defaults to None. + size : int, optional + Width and height of the Plotly figure in pixels. Defaults to 600. + color_values : List[str], optional + List of hex color strings for traces. + + Returns + ------- + Figure + A Plotly ``Figure`` object representing the time-dependent calibration curve. + + Raises + ------ + ValueError + If a heuristic set requests `competing_heuristic='adjusted_as_censored'`. """ unsupported_competing_as_censored = any( @@ -139,6 +180,8 @@ def create_calibration_curve_times( fixed_time_horizons=fixed_time_horizons, heuristics_sets=heuristics_sets, calibration_type=calibration_type, + smooth_method=smooth_method, + bandwidth=bandwidth, size=size, color_values=color_values, ) @@ -731,7 +774,7 @@ def process_single_array(p, r, group_name): # lowess returns a 2D array where the first column is x and the second is y smoothed = lowess(r, p, it=0) xout = np.linspace(0, 1, 101) - yout = np.interp(xout, smoothed[:, 0], smoothed[:, 1]) + yout = np.clip(np.interp(xout, smoothed[:, 0], smoothed[:, 1]), 0.0, 1.0) return pl.DataFrame( {"x": xout, "y": yout, "reference_group": [group_name] * len(xout)} ) @@ -1130,6 +1173,147 @@ def _make_adjusted_deciles_data( return pl.DataFrame(rows).sort(["reference_group", "decile"]) +def _calculate_secondary_cox_smooth( + df_adj: pl.DataFrame, + horizon: float, + performance_type: str, +) -> pl.DataFrame: + """Calculate smoothed calibration curve using secondary Cox regression (Austin et al. 2020 method).""" + import pandas as pd + from lifelines import CoxPHFitter + + smooth_frames = [] + + for key, group_df in df_adj.group_by("reference_group", maintain_order=True): + group_name = str(key[0]) + probs = group_df["prob"].to_numpy() + reals = group_df["real"].to_numpy() + times = group_df["time"].to_numpy() + + p_clipped = np.clip(probs, 1e-6, 1 - 1e-6) + x = np.log(-np.log(1 - p_clipped)) + events = (reals == 1).astype(int) + + if len(np.unique(x)) <= 1 or events.sum() == 0: + y_est = _aj_risk_at_horizon(group_df, horizon) + xout = np.linspace(0, 1, 101) + smooth_frames.append( + pl.DataFrame( + { + "x": xout, + "y": [y_est] * len(xout), + "reference_group": [group_name] * len(xout), + } + ) + ) + continue + + fit_df = pd.DataFrame({"time": times, "event": events, "x": x}) + try: + cph = CoxPHFitter(penalizer=0.01) + cph.fit(fit_df, duration_col="time", event_col="event") + + xout = np.linspace(0.001, 0.999, 101) + x_grid = np.log(-np.log(1 - xout)) + + surv_at_t = cph.predict_survival_function( + pd.DataFrame({"x": x_grid}), times=[horizon] + ).values.ravel() + yout = np.clip(1.0 - surv_at_t, 0.0, 1.0) + except Exception: + y_est = _aj_risk_at_horizon(group_df, horizon) + xout = np.linspace(0, 1, 101) + yout = np.array([y_est] * len(xout)) + + smooth_frames.append( + pl.DataFrame( + { + "x": xout, + "y": yout, + "reference_group": [group_name] * len(xout), + } + ) + ) + + if not smooth_frames: + return pl.DataFrame( + schema={ + "x": pl.Float64, + "y": pl.Float64, + "reference_group": pl.Utf8, + } + ) + + smooth_dat = pl.concat(smooth_frames) + return smooth_dat + + +def _calculate_local_aj_smooth( + df_adj: pl.DataFrame, + horizon: float, + performance_type: str, + bandwidth: Union[float, None] = None, +) -> pl.DataFrame: + """Calculate smoothed calibration curve using local Aalen-Johansen estimation (Gerds' method).""" + smooth_frames = [] + + for key, group_df in df_adj.group_by("reference_group", maintain_order=True): + group_name = str(key[0]) + n = group_df.height + probs = group_df["prob"].to_numpy() + + if len(np.unique(probs)) == 1: + y_est = _aj_risk_at_horizon(group_df, horizon) + xout = np.linspace(0, 1, 101) + smooth_frames.append( + pl.DataFrame( + { + "x": xout, + "y": [y_est] * len(xout), + "reference_group": [group_name] * len(xout), + } + ) + ) + continue + + xout = np.linspace(0, 1, 101) + yout = [] + + if bandwidth is not None: + k = max(5, int(bandwidth * n)) + else: + k = max(10, min(n, int(0.2 * n))) + + for p0 in xout: + distances = np.abs(probs - p0) + idx = np.argsort(distances, kind="stable")[:k] + sub_df = group_df[idx] + y_est = _aj_risk_at_horizon(sub_df, horizon) + yout.append(y_est) + + smooth_frames.append( + pl.DataFrame( + { + "x": xout, + "y": np.array(yout), + "reference_group": [group_name] * len(xout), + } + ) + ) + + if not smooth_frames: + return pl.DataFrame( + schema={ + "x": pl.Float64, + "y": pl.Float64, + "reference_group": pl.Utf8, + } + ) + + smooth_dat = pl.concat(smooth_frames) + return smooth_dat + + def _calculate_adjusted_pseudostates( df: pl.DataFrame, horizon: float ) -> Dict[str, np.ndarray]: @@ -1161,6 +1345,8 @@ def _create_calibration_curve_list_times( fixed_time_horizons: List[float], heuristics_sets: List[Dict[str, str]], calibration_type: str = "discrete", + smooth_method: str = "local_aj", + bandwidth: Union[float, None] = None, size: int = 600, color_values: List[str] = [ "#1b9e77", @@ -1220,10 +1406,26 @@ def _create_calibration_curve_list_times( ) } 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 - ) + if smooth_method == "local_aj": + smooth_data = _calculate_local_aj_smooth( + df_adj, horizon, performance_type, bandwidth=bandwidth + ) + elif smooth_method == "secondary_cox": + smooth_data = _calculate_secondary_cox_smooth( + df_adj, horizon, performance_type + ) + elif smooth_method == "pseudo_values": + pseudo_by_group = _calculate_adjusted_pseudostates( + df_adj, horizon + ) + smooth_data = _calculate_smooth_curve( + probs_adj, pseudo_by_group, performance_type + ) + else: + raise ValueError( + f"Unsupported smooth_method: '{smooth_method}'. " + "Supported options are 'local_aj', 'secondary_cox', and 'pseudo_values'." + ) else: smooth_data = deciles_data.select("x", "y", "reference_group") hist_data = _create_histogram_for_calibration(probs_adj) @@ -1270,9 +1472,26 @@ def _create_calibration_curve_list_times( ) # Smooth curve - smooth_data = _calculate_smooth_curve( - probs_adj, reals_adj, performance_type - ) + if calibration_type == "smooth": + if smooth_method == "local_aj": + smooth_data = _calculate_local_aj_smooth( + df_adj, horizon, performance_type, bandwidth=bandwidth + ) + elif smooth_method == "secondary_cox": + smooth_data = _calculate_secondary_cox_smooth( + df_adj, horizon, performance_type + ) + elif smooth_method == "pseudo_values": + smooth_data = _calculate_smooth_curve( + probs_adj, reals_adj, performance_type + ) + else: + raise ValueError( + f"Unsupported smooth_method: '{smooth_method}'. " + "Supported options are 'local_aj', 'secondary_cox', and 'pseudo_values'." + ) + else: + smooth_data = deciles_data.select("x", "y", "reference_group") all_smooth.append( smooth_data.with_columns(pl.lit(horizon).alias("fixed_time_horizon")) ) diff --git a/tests/test_calibration_times.py b/tests/test_calibration_times.py index 33ac4cc9..b6d2a047 100644 --- a/tests/test_calibration_times.py +++ b/tests/test_calibration_times.py @@ -254,3 +254,83 @@ def test_create_calibration_curve_times_rejects_competing_as_censored(entry_poin } ], ) + + +@pytest.mark.parametrize("smooth_method", ["local_aj", "secondary_cox", "pseudo_values"]) +def test_create_calibration_curve_times_smooth_methods(smooth_method): + np.random.seed(42) + probs = {"model_1": np.linspace(0.1, 0.9, 20)} + reals = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]) + times = np.linspace(1.0, 10.0, 20) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[5.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + smooth_method=smooth_method, + ) + + assert fig is not None + # Check that trace contains smooth points + smooth_trace = fig.data[1] + assert len(smooth_trace.x) == 101 + assert np.isfinite(np.asarray(smooth_trace.y)).all() + + +def test_create_calibration_curve_times_invalid_smooth_method(): + probs = {"model_1": np.array([0.1, 0.2, 0.3, 0.4])} + reals = np.array([1, 0, 1, 0]) + times = np.array([1.0, 2.0, 3.0, 4.0]) + + with pytest.raises(ValueError, match="Unsupported smooth_method"): + create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[2.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + smooth_method="unknown_method", + ) + + +@pytest.mark.parametrize("smooth_method", ["local_aj", "secondary_cox", "pseudo_values"]) +def test_create_calibration_curve_times_competing_risks(smooth_method): + probs = {"model_1": np.linspace(0.05, 0.95, 30)} + # 0 = censored, 1 = event of interest, 2 = competing event + reals = np.array([0, 1, 2] * 10) + times = np.linspace(1.0, 15.0, 30) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[8.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + smooth_method=smooth_method, + ) + + assert fig is not None + smooth_trace = fig.data[1] + assert len(smooth_trace.x) == 101 + assert np.all(np.asarray(smooth_trace.y) >= 0.0) + assert np.all(np.asarray(smooth_trace.y) <= 1.0) diff --git a/user_guide/02-curve-api-compatibility.qmd b/user_guide/02-curve-api-compatibility.qmd index faf8efd8..e603b3fd 100644 --- a/user_guide/02-curve-api-compatibility.qmd +++ b/user_guide/02-curve-api-compatibility.qmd @@ -51,19 +51,24 @@ These statements describe the effect of the heuristics on the estimates. Functio `create_calibration_curve_times()` differs from its ROC, precision-recall, Gains, Lift, and decision-curve siblings in two important ways: 1. `heuristics_sets` is currently required rather than defaulted. -2. Calibration explicitly rejects unsupported heuristic combinations, including `censoring_heuristic="adjusted"` and `competing_heuristic="adjusted_as_censored"`, with an `Unsupported calibration heuristics` error instead of silently skipping every requested horizon. +2. Calibration explicitly rejects unsupported heuristic combinations (specifically `competing_heuristic="adjusted_as_censored"`) with an `Unsupported calibration heuristics` error instead of silently skipping requested horizons. -Pass the calibration heuristic explicitly. For the currently working exclusion-based path: +Pass the calibration heuristic explicitly. For adjusted or exclusion-based paths: ```python heuristics_sets = [ { - "censoring_heuristic": "excluded", + "censoring_heuristic": "adjusted", "competing_heuristic": "adjusted_as_negative", } ] ``` +When `calibration_type="smooth"`, you can also specify the `smooth_method`: +- `"local_aj"` (default): Gerds' local Aalen-Johansen/KM neighborhood estimation. +- `"secondary_cox"`: Secondary Cox regression method (Austin, Harrell & McLernon 2020). +- `"pseudo_values"`: Leave-one-out Aalen-Johansen pseudo-observations lowess. + Then call: ```python diff --git a/user_guide/03-common-errors.qmd b/user_guide/03-common-errors.qmd index fca632d8..d5efdd7b 100644 --- a/user_guide/03-common-errors.qmd +++ b/user_guide/03-common-errors.qmd @@ -32,18 +32,18 @@ See [Curve API Compatibility](curve-api-compatibility.html) for the family-by-fa ### Why it happens -Calibration does not currently implement `censoring_heuristic="adjusted"` or -`competing_heuristic="adjusted_as_censored"`. These inputs are rejected -before curve construction rather than silently skipping every requested horizon. +Calibration does not support `competing_heuristic="adjusted_as_censored"`. +This input is rejected before curve construction rather than silently +skipping every requested horizon. ### Fix -Pass a supported calibration heuristic explicitly. For the exclusion-based path: +Pass a supported calibration heuristic explicitly. For example: ```python heuristics_sets = [ { - "censoring_heuristic": "excluded", + "censoring_heuristic": "adjusted", "competing_heuristic": "adjusted_as_negative", } ] From 459add8e3cf0f60e7d2b09454d0b92b57d2ea1d6 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 16 Aug 2026 14:49:17 +0000 Subject: [PATCH 2/2] feat: support local Aalen-Johansen and secondary Cox calibration smoothing - Added `smooth_method` parameter to `create_calibration_curve_times()` supporting `"local_aj"` (Gerds' preferred local Aalen-Johansen method, default), `"secondary_cox"` (Austin, Harrell & McLernon secondary Cox regression method), and `"pseudo_values"` (jackknife pseudo-value lowess). - Added `bandwidth` parameter for `local_aj` neighborhood window tuning. - Updated documentation and added comprehensive unit tests for time-dependent calibration smoothing methods across single, multiple model, and competing risk scenarios. Co-authored-by: uriahf <11351434+uriahf@users.noreply.github.com>