diff --git a/pyproject.toml b/pyproject.toml index 0ad82ea7..a97648e5 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/rtichoke/calibration/__init__.py b/src/rtichoke/calibration/__init__.py index b04d51eb..87beddab 100644 --- a/src/rtichoke/calibration/__init__.py +++ b/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( diff --git a/src/rtichoke/calibration/_secondary_cox.py b/src/rtichoke/calibration/_secondary_cox.py new file mode 100644 index 00000000..b42fbaea --- /dev/null +++ b/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)