coxrmstsimpson 0.2.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- coxrmstsimpson-0.2.2/.gitignore +7 -0
- coxrmstsimpson-0.2.2/LICENSE +17 -0
- coxrmstsimpson-0.2.2/PKG-INFO +99 -0
- coxrmstsimpson-0.2.2/README.md +82 -0
- coxrmstsimpson-0.2.2/pyproject.toml +28 -0
- coxrmstsimpson-0.2.2/src/coxrmstsimpson/__init__.py +6 -0
- coxrmstsimpson-0.2.2/src/coxrmstsimpson/core.py +197 -0
- coxrmstsimpson-0.2.2/src/coxrmstsimpson/plotting.py +80 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Ian Forde
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: coxrmstsimpson
|
|
3
|
+
Version: 0.2.2
|
|
4
|
+
Summary: Restricted mean survival time from Cox models using composite Simpson's rule
|
|
5
|
+
Project-URL: Homepage, https://github.com/fordeia/CoxRMSTSimpson
|
|
6
|
+
Project-URL: Repository, https://github.com/fordeia/CoxRMSTSimpson
|
|
7
|
+
Author-email: Ian Forde <fordeia@appstate.edu>
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Requires-Dist: lifelines>=0.28
|
|
12
|
+
Requires-Dist: matplotlib>=3.6
|
|
13
|
+
Requires-Dist: numpy>=1.23
|
|
14
|
+
Requires-Dist: pandas>=1.5
|
|
15
|
+
Requires-Dist: scipy>=1.10
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# CoxRMSTSimpson (Python)
|
|
19
|
+
|
|
20
|
+
A Python companion to the R package [CoxRMSTSimpson](https://github.com/fordeia/CoxRMSTSimpson). It predicts survival curves from a fitted `lifelines.CoxPHFitter`, estimates restricted mean survival time (RMST) using composite Simpson's rule, and plots survival curves with shaded areas.
|
|
21
|
+
|
|
22
|
+
**Development status:** This package is under development. Validate numerical results against the R implementation before assuming numerical equivalence.
|
|
23
|
+
|
|
24
|
+
## Install
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
python -m pip install coxrmstsimpson
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Example
|
|
31
|
+
|
|
32
|
+
The following example simulates survival data, fits a Cox proportional hazards model, estimates RMST up to time \(\tau=10\), and plots two predicted survival curves.
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
import numpy as np
|
|
36
|
+
import pandas as pd
|
|
37
|
+
from lifelines import CoxPHFitter
|
|
38
|
+
from coxrmstsimpson import cox_rmst, plot_cox_rmst
|
|
39
|
+
|
|
40
|
+
# Simulate reproducible survival data.
|
|
41
|
+
np.random.seed(42)
|
|
42
|
+
n = 200
|
|
43
|
+
|
|
44
|
+
group = np.random.binomial(1, 0.5, n)
|
|
45
|
+
event_time = np.random.exponential(
|
|
46
|
+
scale=1 / (0.10 * np.exp(0.4 * group))
|
|
47
|
+
)
|
|
48
|
+
censor_time = np.random.exponential(scale=15, size=n)
|
|
49
|
+
|
|
50
|
+
df = pd.DataFrame({
|
|
51
|
+
"time": np.minimum(event_time, censor_time),
|
|
52
|
+
"status": (event_time <= censor_time).astype(int),
|
|
53
|
+
"group": group,
|
|
54
|
+
})
|
|
55
|
+
|
|
56
|
+
# Fit the Cox proportional hazards model.
|
|
57
|
+
cph = CoxPHFitter()
|
|
58
|
+
cph.fit(df, duration_col="time", event_col="status")
|
|
59
|
+
|
|
60
|
+
# Define the two covariate profiles to compare.
|
|
61
|
+
profiles = pd.DataFrame({"group": [0, 1]})
|
|
62
|
+
|
|
63
|
+
# Estimate RMST up to tau = 10.
|
|
64
|
+
result = cox_rmst(cph, profiles, tau=10, n_grid=101)
|
|
65
|
+
print(result["results"])
|
|
66
|
+
|
|
67
|
+
# Plot survival curves with RMST estimates in the legend.
|
|
68
|
+
ax = plot_cox_rmst(
|
|
69
|
+
result,
|
|
70
|
+
labels=["Control", "Treatment"],
|
|
71
|
+
show_rmst=True,
|
|
72
|
+
)
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
In this example, `status` is 1 when the event occurs and 0 for right-censored observations. For your own data, use the covariates and encodings used to fit your Cox model. Each row of `profiles` specifies a covariate profile to evaluate, and the labels are supplied in the same order.
|
|
76
|
+
|
|
77
|
+
## Functions
|
|
78
|
+
|
|
79
|
+
- `simpson_rmst(time, survival)`: Integrates a supplied survival curve using composite Simpson's rule.
|
|
80
|
+
- `cox_rmst(model, newdata, tau, n_grid=101)`: Predicts survival curves from a fitted Cox model and estimates RMST up to `tau`.
|
|
81
|
+
- `plot_cox_rmst(result, ax=None, labels=None, show_rmst=False)`: Plots predicted survival curves and shades the areas beneath them. Use `labels` for custom profile labels and `show_rmst=True` to include RMST estimates in the legend.
|
|
82
|
+
|
|
83
|
+
## Definition
|
|
84
|
+
|
|
85
|
+
Restricted mean survival time up to time \(\tau\) is the area under the survival function from 0 to \(\tau\):
|
|
86
|
+
|
|
87
|
+
\[
|
|
88
|
+
\operatorname{RMST}(\tau)=\int_0^\tau S(t)\,dt.
|
|
89
|
+
\]
|
|
90
|
+
|
|
91
|
+
Here, \(S(t)\) is the predicted survival probability at time \(t\).
|
|
92
|
+
|
|
93
|
+
## Numerical note
|
|
94
|
+
|
|
95
|
+
`lifelines` interpolates predicted survival probabilities at requested times not already on its prediction timeline. This can differ from the step-function evaluation commonly used by R's `survival::survfit`. Cross-language validation is therefore needed before claiming exact numerical equivalence.
|
|
96
|
+
|
|
97
|
+
## Requirements
|
|
98
|
+
|
|
99
|
+
Python 3.10 or later. Dependencies, including NumPy, pandas, SciPy, Matplotlib, and lifelines, are installed automatically by pip.
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
# CoxRMSTSimpson (Python)
|
|
2
|
+
|
|
3
|
+
A Python companion to the R package [CoxRMSTSimpson](https://github.com/fordeia/CoxRMSTSimpson). It predicts survival curves from a fitted `lifelines.CoxPHFitter`, estimates restricted mean survival time (RMST) using composite Simpson's rule, and plots survival curves with shaded areas.
|
|
4
|
+
|
|
5
|
+
**Development status:** This package is under development. Validate numerical results against the R implementation before assuming numerical equivalence.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
python -m pip install coxrmstsimpson
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Example
|
|
14
|
+
|
|
15
|
+
The following example simulates survival data, fits a Cox proportional hazards model, estimates RMST up to time \(\tau=10\), and plots two predicted survival curves.
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
import numpy as np
|
|
19
|
+
import pandas as pd
|
|
20
|
+
from lifelines import CoxPHFitter
|
|
21
|
+
from coxrmstsimpson import cox_rmst, plot_cox_rmst
|
|
22
|
+
|
|
23
|
+
# Simulate reproducible survival data.
|
|
24
|
+
np.random.seed(42)
|
|
25
|
+
n = 200
|
|
26
|
+
|
|
27
|
+
group = np.random.binomial(1, 0.5, n)
|
|
28
|
+
event_time = np.random.exponential(
|
|
29
|
+
scale=1 / (0.10 * np.exp(0.4 * group))
|
|
30
|
+
)
|
|
31
|
+
censor_time = np.random.exponential(scale=15, size=n)
|
|
32
|
+
|
|
33
|
+
df = pd.DataFrame({
|
|
34
|
+
"time": np.minimum(event_time, censor_time),
|
|
35
|
+
"status": (event_time <= censor_time).astype(int),
|
|
36
|
+
"group": group,
|
|
37
|
+
})
|
|
38
|
+
|
|
39
|
+
# Fit the Cox proportional hazards model.
|
|
40
|
+
cph = CoxPHFitter()
|
|
41
|
+
cph.fit(df, duration_col="time", event_col="status")
|
|
42
|
+
|
|
43
|
+
# Define the two covariate profiles to compare.
|
|
44
|
+
profiles = pd.DataFrame({"group": [0, 1]})
|
|
45
|
+
|
|
46
|
+
# Estimate RMST up to tau = 10.
|
|
47
|
+
result = cox_rmst(cph, profiles, tau=10, n_grid=101)
|
|
48
|
+
print(result["results"])
|
|
49
|
+
|
|
50
|
+
# Plot survival curves with RMST estimates in the legend.
|
|
51
|
+
ax = plot_cox_rmst(
|
|
52
|
+
result,
|
|
53
|
+
labels=["Control", "Treatment"],
|
|
54
|
+
show_rmst=True,
|
|
55
|
+
)
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
In this example, `status` is 1 when the event occurs and 0 for right-censored observations. For your own data, use the covariates and encodings used to fit your Cox model. Each row of `profiles` specifies a covariate profile to evaluate, and the labels are supplied in the same order.
|
|
59
|
+
|
|
60
|
+
## Functions
|
|
61
|
+
|
|
62
|
+
- `simpson_rmst(time, survival)`: Integrates a supplied survival curve using composite Simpson's rule.
|
|
63
|
+
- `cox_rmst(model, newdata, tau, n_grid=101)`: Predicts survival curves from a fitted Cox model and estimates RMST up to `tau`.
|
|
64
|
+
- `plot_cox_rmst(result, ax=None, labels=None, show_rmst=False)`: Plots predicted survival curves and shades the areas beneath them. Use `labels` for custom profile labels and `show_rmst=True` to include RMST estimates in the legend.
|
|
65
|
+
|
|
66
|
+
## Definition
|
|
67
|
+
|
|
68
|
+
Restricted mean survival time up to time \(\tau\) is the area under the survival function from 0 to \(\tau\):
|
|
69
|
+
|
|
70
|
+
\[
|
|
71
|
+
\operatorname{RMST}(\tau)=\int_0^\tau S(t)\,dt.
|
|
72
|
+
\]
|
|
73
|
+
|
|
74
|
+
Here, \(S(t)\) is the predicted survival probability at time \(t\).
|
|
75
|
+
|
|
76
|
+
## Numerical note
|
|
77
|
+
|
|
78
|
+
`lifelines` interpolates predicted survival probabilities at requested times not already on its prediction timeline. This can differ from the step-function evaluation commonly used by R's `survival::survfit`. Cross-language validation is therefore needed before claiming exact numerical equivalence.
|
|
79
|
+
|
|
80
|
+
## Requirements
|
|
81
|
+
|
|
82
|
+
Python 3.10 or later. Dependencies, including NumPy, pandas, SciPy, Matplotlib, and lifelines, are installed automatically by pip.
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling>=1.26"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "coxrmstsimpson"
|
|
7
|
+
version = "0.2.2"
|
|
8
|
+
description = "Restricted mean survival time from Cox models using composite Simpson's rule"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
authors = [{ name = "Ian Forde", email = "fordeia@appstate.edu" }]
|
|
13
|
+
dependencies = ["numpy>=1.23", "pandas>=1.5", "scipy>=1.10", "matplotlib>=3.6", "lifelines>=0.28"]
|
|
14
|
+
|
|
15
|
+
[project.urls]
|
|
16
|
+
Homepage = "https://github.com/fordeia/CoxRMSTSimpson"
|
|
17
|
+
Repository = "https://github.com/fordeia/CoxRMSTSimpson"
|
|
18
|
+
|
|
19
|
+
[tool.hatch.build.targets.wheel]
|
|
20
|
+
packages = ["src/coxrmstsimpson"]
|
|
21
|
+
|
|
22
|
+
[tool.hatch.build.targets.sdist]
|
|
23
|
+
only-include = [
|
|
24
|
+
"src/coxrmstsimpson",
|
|
25
|
+
"README.md",
|
|
26
|
+
"pyproject.toml",
|
|
27
|
+
"LICENSE",
|
|
28
|
+
]
|
|
@@ -0,0 +1,197 @@
|
|
|
1
|
+
"""Core restricted mean survival time calculations."""
|
|
2
|
+
import numpy as np
|
|
3
|
+
import pandas as pd
|
|
4
|
+
from scipy.integrate import simpson
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def simpson_rmst(time, survival) -> float:
|
|
8
|
+
"""Integrate survival probabilities over time using composite Simpson's rule."""
|
|
9
|
+
t = np.asarray(time, dtype=float)
|
|
10
|
+
s = np.asarray(survival, dtype=float)
|
|
11
|
+
if t.ndim != 1 or s.ndim != 1 or len(t) != len(s):
|
|
12
|
+
raise ValueError("time and survival must be 1D arrays of equal length.")
|
|
13
|
+
if len(t) < 3:
|
|
14
|
+
raise ValueError("At least three time points are required.")
|
|
15
|
+
if not np.all(np.isfinite(t)) or not np.all(np.isfinite(s)):
|
|
16
|
+
raise ValueError("time and survival must contain only finite values.")
|
|
17
|
+
if np.any(np.diff(t) <= 0):
|
|
18
|
+
raise ValueError("time must be strictly increasing.")
|
|
19
|
+
if np.any((s < 0) | (s > 1)):
|
|
20
|
+
raise ValueError("survival probabilities must lie between 0 and 1.")
|
|
21
|
+
return float(simpson(s, x=t))
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _validate_model(model):
|
|
25
|
+
"""Validate that model is a fitted lifelines CoxPHFitter."""
|
|
26
|
+
from lifelines import CoxPHFitter
|
|
27
|
+
|
|
28
|
+
if not isinstance(model, CoxPHFitter):
|
|
29
|
+
raise TypeError("model must be a fitted lifelines.CoxPHFitter.")
|
|
30
|
+
if not hasattr(model, "params_"):
|
|
31
|
+
raise ValueError("model must be fitted before calling cox_rmst().")
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _mean_profile(model, covariates):
|
|
35
|
+
"""Return the model's fitted covariate means for prediction."""
|
|
36
|
+
means = getattr(model, "_norm_mean", None)
|
|
37
|
+
if means is not None:
|
|
38
|
+
profile = means.to_dict()
|
|
39
|
+
else:
|
|
40
|
+
profile = {}
|
|
41
|
+
for name in covariates:
|
|
42
|
+
if name not in profile or not np.isfinite(profile[name]):
|
|
43
|
+
profile[name] = 0.0
|
|
44
|
+
return profile
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def cox_rmst(
|
|
48
|
+
model,
|
|
49
|
+
newdata: pd.DataFrame = None,
|
|
50
|
+
tau: float = None,
|
|
51
|
+
n_grid: int = 101,
|
|
52
|
+
*,
|
|
53
|
+
treatment: str = None,
|
|
54
|
+
labels=None,
|
|
55
|
+
title: str = None,
|
|
56
|
+
show_plot: bool = True,
|
|
57
|
+
show_rmst: bool = True,
|
|
58
|
+
) -> dict:
|
|
59
|
+
"""Estimate RMST from a fitted lifelines CoxPHFitter model.
|
|
60
|
+
|
|
61
|
+
Two usage patterns are supported.
|
|
62
|
+
|
|
63
|
+
Automatic binary-treatment mode:
|
|
64
|
+
cox_rmst(model, treatment="group", tau=10,
|
|
65
|
+
labels=["Control", "Treatment"], title="RMST by group")
|
|
66
|
+
|
|
67
|
+
Existing explicit-profile mode:
|
|
68
|
+
cox_rmst(model, newdata, tau=10)
|
|
69
|
+
|
|
70
|
+
In automatic mode, ``treatment`` must be a numeric 0/1 covariate included
|
|
71
|
+
in the fitted model. Other model covariates are set to their fitted means
|
|
72
|
+
(lifelines' model normalization means). The package plots the curves by
|
|
73
|
+
default. Set ``show_plot=False`` to calculate without displaying a plot.
|
|
74
|
+
|
|
75
|
+
Parameters
|
|
76
|
+
----------
|
|
77
|
+
model : lifelines.CoxPHFitter
|
|
78
|
+
A fitted Cox proportional hazards model.
|
|
79
|
+
newdata : pandas.DataFrame, optional
|
|
80
|
+
Explicit covariate profiles for the existing API.
|
|
81
|
+
tau : float
|
|
82
|
+
Positive restriction time.
|
|
83
|
+
n_grid : int, default 101
|
|
84
|
+
Odd number of time-grid points, at least 3.
|
|
85
|
+
treatment : str, optional
|
|
86
|
+
Name of a numeric 0/1 covariate in the fitted model. When supplied,
|
|
87
|
+
profiles for treatment=0 and treatment=1 are created automatically.
|
|
88
|
+
labels : sequence of str, optional
|
|
89
|
+
Two legend/result labels in automatic mode, or one label per row of
|
|
90
|
+
``newdata`` in explicit-profile mode.
|
|
91
|
+
title : str, optional
|
|
92
|
+
Custom plot title.
|
|
93
|
+
show_plot : bool, default True
|
|
94
|
+
Display the plot before returning.
|
|
95
|
+
show_rmst : bool, default True
|
|
96
|
+
Include RMST values in the plot legend.
|
|
97
|
+
|
|
98
|
+
Returns
|
|
99
|
+
-------
|
|
100
|
+
dict
|
|
101
|
+
Includes ``curves`` (predicted survival curves), ``results`` (RMST
|
|
102
|
+
table), ``newdata``, ``tau``, and model metadata.
|
|
103
|
+
"""
|
|
104
|
+
_validate_model(model)
|
|
105
|
+
|
|
106
|
+
if tau is None or not np.isscalar(tau) or not np.isfinite(tau) or tau <= 0:
|
|
107
|
+
raise ValueError("tau must be a positive finite number.")
|
|
108
|
+
if not isinstance(n_grid, (int, np.integer)) or isinstance(n_grid, bool) or n_grid < 3 or n_grid % 2 == 0:
|
|
109
|
+
raise ValueError("n_grid must be an odd integer >= 3.")
|
|
110
|
+
|
|
111
|
+
covariates = list(model.params_.index)
|
|
112
|
+
|
|
113
|
+
if treatment is not None:
|
|
114
|
+
if newdata is not None:
|
|
115
|
+
raise ValueError("Provide either treatment= or newdata, not both.")
|
|
116
|
+
if not isinstance(treatment, str) or treatment not in covariates:
|
|
117
|
+
raise ValueError(
|
|
118
|
+
"treatment must name a covariate included in the fitted model "
|
|
119
|
+
f"({covariates}). For automatic mode, use a numeric 0/1 covariate."
|
|
120
|
+
)
|
|
121
|
+
profile = _mean_profile(model, covariates)
|
|
122
|
+
try:
|
|
123
|
+
profile0 = dict(profile)
|
|
124
|
+
profile1 = dict(profile)
|
|
125
|
+
profile0[treatment] = 0.0
|
|
126
|
+
profile1[treatment] = 1.0
|
|
127
|
+
newdata = pd.DataFrame([profile0, profile1], columns=covariates)
|
|
128
|
+
except Exception as exc:
|
|
129
|
+
raise ValueError("Could not construct the treatment profiles.") from exc
|
|
130
|
+
|
|
131
|
+
if labels is None:
|
|
132
|
+
labels = ["Control", "Treatment"]
|
|
133
|
+
else:
|
|
134
|
+
labels = list(labels)
|
|
135
|
+
if len(labels) != 2:
|
|
136
|
+
raise ValueError("labels must contain exactly two labels.")
|
|
137
|
+
newdata.index = [str(labels[0]), str(labels[1])]
|
|
138
|
+
else:
|
|
139
|
+
if not isinstance(newdata, pd.DataFrame) or newdata.empty:
|
|
140
|
+
raise ValueError(
|
|
141
|
+
"Provide a nonempty newdata DataFrame or use treatment='binary_covariate'."
|
|
142
|
+
)
|
|
143
|
+
missing = set(covariates) - set(newdata.columns)
|
|
144
|
+
if missing:
|
|
145
|
+
raise ValueError(f"newdata is missing covariates: {sorted(missing)}")
|
|
146
|
+
newdata = newdata.loc[:, covariates].copy()
|
|
147
|
+
if labels is not None:
|
|
148
|
+
labels = list(labels)
|
|
149
|
+
if len(labels) != len(newdata):
|
|
150
|
+
raise ValueError(
|
|
151
|
+
f"Provide one custom label for each profile ({len(newdata)} required)."
|
|
152
|
+
)
|
|
153
|
+
newdata.index = [str(label) for label in labels]
|
|
154
|
+
|
|
155
|
+
grid = np.linspace(0.0, float(tau), int(n_grid))
|
|
156
|
+
predicted = model.predict_survival_function(newdata, times=grid)
|
|
157
|
+
profile_labels = [str(x) if str(x).strip() else f"Profile {i+1}"
|
|
158
|
+
for i, x in enumerate(newdata.index)]
|
|
159
|
+
curves, values = [], []
|
|
160
|
+
for i, label in enumerate(profile_labels):
|
|
161
|
+
surv = predicted.iloc[:, i].to_numpy(dtype=float)
|
|
162
|
+
curves.append(pd.DataFrame({
|
|
163
|
+
"time": grid,
|
|
164
|
+
"survival": surv,
|
|
165
|
+
"profile": label,
|
|
166
|
+
}))
|
|
167
|
+
values.append(simpson_rmst(grid, surv))
|
|
168
|
+
|
|
169
|
+
curve_df = pd.concat(curves, ignore_index=True)
|
|
170
|
+
results = pd.DataFrame({
|
|
171
|
+
"profile": profile_labels,
|
|
172
|
+
"tau": float(tau),
|
|
173
|
+
"RMST": values,
|
|
174
|
+
})
|
|
175
|
+
result = {
|
|
176
|
+
"model": model,
|
|
177
|
+
"newdata": newdata.copy(),
|
|
178
|
+
"tau": float(tau),
|
|
179
|
+
"time_grid": grid,
|
|
180
|
+
"curves": curve_df,
|
|
181
|
+
"results": results,
|
|
182
|
+
"treatment": treatment,
|
|
183
|
+
}
|
|
184
|
+
|
|
185
|
+
if show_plot:
|
|
186
|
+
from .plotting import plot_cox_rmst
|
|
187
|
+
ax = plot_cox_rmst(
|
|
188
|
+
result,
|
|
189
|
+
labels=profile_labels,
|
|
190
|
+
title=title,
|
|
191
|
+
show_rmst=show_rmst,
|
|
192
|
+
show=True,
|
|
193
|
+
)
|
|
194
|
+
result["ax"] = ax
|
|
195
|
+
result["figure"] = ax.figure
|
|
196
|
+
|
|
197
|
+
return result
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
"""Plot predicted survival curves and shaded RMST areas."""
|
|
2
|
+
import matplotlib.pyplot as plt
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def plot_cox_rmst(
|
|
6
|
+
result,
|
|
7
|
+
ax=None,
|
|
8
|
+
labels=None,
|
|
9
|
+
show_rmst=False,
|
|
10
|
+
title=None,
|
|
11
|
+
show=True,
|
|
12
|
+
):
|
|
13
|
+
"""Plot predicted survival curves and shade their RMST areas.
|
|
14
|
+
|
|
15
|
+
Parameters
|
|
16
|
+
----------
|
|
17
|
+
result : dict
|
|
18
|
+
Result returned by cox_rmst().
|
|
19
|
+
ax : matplotlib.axes.Axes, optional
|
|
20
|
+
Existing axes on which to draw the plot.
|
|
21
|
+
labels : sequence of str, optional
|
|
22
|
+
Custom labels for profiles in the order plotted.
|
|
23
|
+
show_rmst : bool, default False
|
|
24
|
+
Include each profile's RMST estimate in the legend.
|
|
25
|
+
title : str, optional
|
|
26
|
+
Custom title. Defaults to a title containing tau.
|
|
27
|
+
show : bool, default True
|
|
28
|
+
Call matplotlib.pyplot.show() before returning.
|
|
29
|
+
"""
|
|
30
|
+
if not isinstance(result, dict) or "curves" not in result or "tau" not in result:
|
|
31
|
+
raise TypeError("result must be returned by cox_rmst().")
|
|
32
|
+
|
|
33
|
+
curves = result["curves"]
|
|
34
|
+
profiles = list(curves.groupby("profile", sort=False))
|
|
35
|
+
|
|
36
|
+
if labels is not None:
|
|
37
|
+
labels = list(labels)
|
|
38
|
+
if len(labels) != len(profiles):
|
|
39
|
+
raise ValueError(
|
|
40
|
+
f"Provide one custom label for each profile ({len(profiles)} required)."
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
results_table = result.get("results")
|
|
44
|
+
if show_rmst:
|
|
45
|
+
if results_table is None or not {"profile", "RMST"}.issubset(results_table.columns):
|
|
46
|
+
raise ValueError("RMST results are unavailable in result['results'].")
|
|
47
|
+
|
|
48
|
+
if ax is None:
|
|
49
|
+
_, ax = plt.subplots()
|
|
50
|
+
|
|
51
|
+
for i, (profile, data) in enumerate(profiles):
|
|
52
|
+
data = data.sort_values("time")
|
|
53
|
+
curve_label = str(labels[i]) if labels is not None else str(profile)
|
|
54
|
+
|
|
55
|
+
if show_rmst:
|
|
56
|
+
matching = results_table.loc[
|
|
57
|
+
results_table["profile"] == profile, "RMST"
|
|
58
|
+
]
|
|
59
|
+
if matching.empty:
|
|
60
|
+
raise ValueError(f"Could not find an RMST estimate for profile {profile!r}.")
|
|
61
|
+
curve_label = f"{curve_label} (RMST = {matching.iloc[0]:.3f})"
|
|
62
|
+
|
|
63
|
+
ax.plot(data["time"], data["survival"], label=curve_label)
|
|
64
|
+
ax.fill_between(data["time"], data["survival"], 0, alpha=0.18)
|
|
65
|
+
|
|
66
|
+
default_title = f"Restricted mean survival time (tau = {result['tau']:g})"
|
|
67
|
+
ax.set(
|
|
68
|
+
xlim=(0, result["tau"]),
|
|
69
|
+
ylim=(0, 1.02),
|
|
70
|
+
xlabel="Time",
|
|
71
|
+
ylabel="Predicted survival probability",
|
|
72
|
+
title=title if title is not None else default_title,
|
|
73
|
+
)
|
|
74
|
+
ax.legend()
|
|
75
|
+
ax.grid(alpha=0.25)
|
|
76
|
+
|
|
77
|
+
if show:
|
|
78
|
+
plt.show()
|
|
79
|
+
|
|
80
|
+
return ax
|