diff --git a/pyproject.toml b/pyproject.toml index 77e16513..8cac906c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "great-tables>=0.18.0", ] name = "rtichoke" -version = "0.1.32" +version = "0.1.33" description = "interactive visualizations for performance of predictive models" readme = "README.md" diff --git a/src/rtichoke/calibration/calibration.py b/src/rtichoke/calibration/calibration.py index c2e05213..6f1a0ea3 100644 --- a/src/rtichoke/calibration/calibration.py +++ b/src/rtichoke/calibration/calibration.py @@ -141,7 +141,7 @@ def create_calibration_curve_times( 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), + ``"secondary_cox"`` (Austin, Harrell & McLernon secondary Cox regression with 3-knot restricted cubic splines on complementary log-log predictions), 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. @@ -1173,13 +1173,47 @@ def _make_adjusted_deciles_data( return pl.DataFrame(rows).sort(["reference_group", "decile"]) +def _calculate_rcs_basis_3knots( + x: np.ndarray, knots: Union[np.ndarray, None] = None +) -> tuple[np.ndarray, np.ndarray]: + """Calculate 3-knot restricted cubic spline basis matrix. + + Follows Harrell's RCS formulation (RMS Section 2.4.1) / rms::rcs in R. + For 3 knots (10th, 50th, 90th percentiles of x): + basis matrix has 2 columns: [x, u1(x)] + """ + x = np.asarray(x, dtype=float) + if knots is None: + knots = np.percentile(x, [10, 50, 90]) + knots = np.sort(np.asarray(knots, dtype=float)) + + t1, t2, t3 = knots[0], knots[1], knots[2] + + # Handle edge case where knots are duplicate / non-unique + if len(np.unique(knots)) < 3 or (t3 - t2) == 0 or (t2 - t1) == 0 or (t3 - t1) == 0: + return x[:, None], knots + + denom = (t3 - t1) ** 2 + + def pos_cube(val: np.ndarray) -> np.ndarray: + return np.maximum(val, 0) ** 3 + + u1 = ( + pos_cube(x - t1) + - ((t3 - t1) / (t3 - t2)) * pos_cube(x - t2) + + ((t2 - t1) / (t3 - t2)) * pos_cube(x - t3) + ) / denom + + basis = np.column_stack([x, u1]) + return basis, knots + + 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 + """Calculate smoothed calibration curve using secondary Cox regression (Austin et al. 2020 & McLernon et al. 2023 method).""" from lifelines import CoxPHFitter smooth_frames = [] @@ -1208,16 +1242,39 @@ def _calculate_secondary_cox_smooth( ) continue - fit_df = pd.DataFrame({"time": times, "event": events, "x": x}) + basis, knots = _calculate_rcs_basis_3knots(x) + + if basis.shape[1] == 2: + fit_df = pl.DataFrame( + { + "time": times, + "event": events, + "rcs_1": basis[:, 0], + "rcs_2": basis[:, 1], + } + ) + else: + fit_df = pl.DataFrame( + {"time": times, "event": events, "rcs_1": basis[:, 0]} + ) + try: cph = CoxPHFitter(penalizer=0.01) - cph.fit(fit_df, duration_col="time", event_col="event") + cph.fit(fit_df.to_pandas(), duration_col="time", event_col="event") xout = np.linspace(0.001, 0.999, 101) x_grid = np.log(-np.log(1 - xout)) + grid_basis, _ = _calculate_rcs_basis_3knots(x_grid, knots=knots) + + if grid_basis.shape[1] == 2 and "rcs_2" in fit_df.columns: + grid_df = pl.DataFrame( + {"rcs_1": grid_basis[:, 0], "rcs_2": grid_basis[:, 1]} + ) + else: + grid_df = pl.DataFrame({"rcs_1": grid_basis[:, 0]}) surv_at_t = cph.predict_survival_function( - pd.DataFrame({"x": x_grid}), times=[horizon] + grid_df.to_pandas(), times=[horizon] ).values.ravel() yout = np.clip(1.0 - surv_at_t, 0.0, 1.0) except Exception: diff --git a/tests/test_calibration_times.py b/tests/test_calibration_times.py index b6d2a047..bc2396b3 100644 --- a/tests/test_calibration_times.py +++ b/tests/test_calibration_times.py @@ -256,7 +256,9 @@ def test_create_calibration_curve_times_rejects_competing_as_censored(entry_poin ) -@pytest.mark.parametrize("smooth_method", ["local_aj", "secondary_cox", "pseudo_values"]) +@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)} @@ -307,7 +309,9 @@ def test_create_calibration_curve_times_invalid_smooth_method(): ) -@pytest.mark.parametrize("smooth_method", ["local_aj", "secondary_cox", "pseudo_values"]) +@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 @@ -334,3 +338,62 @@ def test_create_calibration_curve_times_competing_risks(smooth_method): 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) + + +def test_secondary_cox_rcs_3knots_parity(): + # Validation dataset with predictions, status, and survival times + probs = {"model_1": np.array([0.1, 0.2, 0.35, 0.5, 0.65, 0.8, 0.9])} + reals = np.array([0, 0, 1, 0, 1, 1, 1]) + times = np.array([12.0, 15.0, 8.0, 20.0, 6.0, 4.0, 3.0]) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[10.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + smooth_method="secondary_cox", + ) + + smooth_trace = fig.data[1] + y_vals = np.asarray(smooth_trace.y) + + assert len(smooth_trace.x) == 101 + assert np.all(y_vals >= 0.0) + assert np.all(y_vals <= 1.0) + # Higher predicted risk should correspond to higher actual risk + assert y_vals[-1] > y_vals[0] + + +def test_secondary_cox_duplicate_knots_fallback(): + # Data where probs contain duplicate values leading to duplicate knots + probs = {"model_1": np.array([0.2, 0.2, 0.2, 0.2, 0.8, 0.8])} + reals = np.array([0, 0, 1, 0, 1, 1]) + times = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + + fig = create_calibration_curve_times( + probs, + reals, + times, + fixed_time_horizons=[4.0], + heuristics_sets=[ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } + ], + calibration_type="smooth", + smooth_method="secondary_cox", + ) + + smooth_trace = fig.data[1] + y_vals = np.asarray(smooth_trace.y) + + assert len(smooth_trace.x) == 101 + assert np.all(np.isfinite(y_vals)) diff --git a/uv.lock b/uv.lock index e83f7972..799a2027 100644 --- a/uv.lock +++ b/uv.lock @@ -5119,7 +5119,7 @@ wheels = [ [[package]] name = "rtichoke" -version = "0.1.32" +version = "0.1.33" source = { editable = "." } dependencies = [ { name = "great-tables", version = "0.21.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },