Extend smooth_method="secondary_cox" with 3-knot Restricted Cubic Splines

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"
 

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:

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