From 5261ca08b8f62484f99e3cd96860777e8b3ffea2 Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Wed, 19 Aug 2026 10:51:43 +0300 Subject: [PATCH 1/4] Add smoothstate integration dependency --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 0ad82ea7..d3b6c03b 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 = "main" } + [tool.uv.dependency-groups] docs = {requires-python = ">=3.11"} From f08a0f13956e8b97348d7f990af5d0a933811558 Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Wed, 19 Aug 2026 10:52:21 +0300 Subject: [PATCH 2/4] Route secondary Cox smoothing through smoothstate --- src/rtichoke/calibration/_secondary_cox.py | 84 ++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 src/rtichoke/calibration/_secondary_cox.py 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) From 1f132cd3b4643ae605fdff2d2d52b3f8a5c2e611 Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Wed, 19 Aug 2026 10:52:33 +0300 Subject: [PATCH 3/4] Use smoothstate for secondary Cox calibration --- src/rtichoke/calibration/__init__.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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( From 25484c34cf7b223a4e95aa336fdcaca8f673139b Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Wed, 19 Aug 2026 11:04:59 +0300 Subject: [PATCH 4/4] Test against Python 3.9-compatible smoothstate branch --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d3b6c03b..a97648e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ docs = [ members = ["rtichoke"] [tool.uv.sources] -smoothstate = { git = "https://github.com/uriahf/smoothstate.git", branch = "main" } +smoothstate = { git = "https://github.com/uriahf/smoothstate.git", branch = "agent/python-39-compat" } [tool.uv.dependency-groups] docs = {requires-python = ">=3.11"}