Integrate smoothstate secondary Cox smoothing

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "pandas>=2.2.3",
     "typing>=3.7.4.3",
     "polarstate==0.1.8",
+    "smoothstate>=0.1.0",
     "marimo>=0.17.0",
     "pyarrow>=21.0.0",
     "statsmodels>=0.14.0",
@@ -47,6 +48,9 @@ docs = [
 [tool.uv.workspace]
 members = ["rtichoke"]
 
+[tool.uv.sources]
+smoothstate = { git = "https://github.com/uriahf/smoothstate.git", branch = "agent/python-39-compat" }
+
 [tool.uv.dependency-groups]
 docs = {requires-python = ">=3.11"}
 

src/rtichoke/calibration/__init__.py

@@ -4,11 +4,26 @@
 
 from . import calibration as _calibration
 from ._interactive_aspect import enforce_square_calibration_panel
+from ._secondary_cox import calculate_secondary_cox_smooth
 
 _original_create_calibration_curve = _calibration.create_calibration_curve
 _original_create_calibration_curve_times = _calibration.create_calibration_curve_times
 
 
+# Route the existing private secondary-Cox hook through smoothstate while
+# preserving rtichoke's Aalen-Johansen fallback behavior.
+def _smoothstate_secondary_cox(df_adj, horizon, performance_type):
+    return calculate_secondary_cox_smooth(
+        df_adj,
+        horizon,
+        performance_type,
+        aj_risk_at_horizon=_calibration._aj_risk_at_horizon,
+    )
+
+
+_calibration._calculate_secondary_cox_smooth = _smoothstate_secondary_cox
+
+
 def create_calibration_curve(*args, **kwargs):
     """Create an interactive calibration plot with a square main panel."""
     return enforce_square_calibration_panel(

src/rtichoke/calibration/_secondary_cox.py

@@ -0,0 +1,84 @@
+"""Secondary Cox calibration smoothing backed by smoothstate."""
+
+from collections.abc import Callable
+
+import numpy as np
+import polars as pl
+from smoothstate import smooth_state_cox
+
+
+def calculate_secondary_cox_smooth(
+    df_adj: pl.DataFrame,
+    horizon: float,
+    performance_type: str,
+    *,
+    aj_risk_at_horizon: Callable[[pl.DataFrame, float], float],
+) -> pl.DataFrame:
+    """Calculate the secondary-Cox calibration curve with ``smoothstate``.
+
+    The degenerate/error fallback intentionally mirrors rtichoke's previous
+    implementation: return a constant Aalen-Johansen risk curve on [0, 1].
+    ``performance_type`` remains in the signature for API compatibility with
+    the existing private helper.
+    """
+    del performance_type
+    smooth_frames: list[pl.DataFrame] = []
+
+    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()
+        events = (reals == 1).astype(int)
+
+        p_clipped = np.clip(probs, 1e-6, 1 - 1e-6)
+        transformed = np.log(-np.log(1 - p_clipped))
+
+        if len(np.unique(transformed)) <= 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": np.full(len(xout), y_est),
+                        "reference_group": [group_name] * len(xout),
+                    }
+                )
+            )
+            continue
+
+        try:
+            smoothed = smooth_state_cox(
+                probs=probs,
+                times=times,
+                events=events,
+                horizon=horizon,
+                penalizer=0.01,
+            )
+            smooth_frames.append(
+                smoothed.with_columns(pl.lit(group_name).alias("reference_group"))
+            )
+        except Exception:
+            y_est = aj_risk_at_horizon(group_df, horizon)
+            xout = np.linspace(0, 1, 101)
+            smooth_frames.append(
+                pl.DataFrame(
+                    {
+                        "x": xout,
+                        "y": np.full(len(xout), y_est),
+                        "reference_group": [group_name] * len(xout),
+                    }
+                )
+            )
+
+    if not smooth_frames:
+        return pl.DataFrame(
+            schema={
+                "x": pl.Float64,
+                "y": pl.Float64,
+                "reference_group": pl.Utf8,
+            }
+        )
+
+    return pl.concat(smooth_frames)