Return an observation-level DataFrame for a fitted model.
augment(
model,
data=None,
**kwargs,
)
Produces one row per observation, appending model-derived columns (e.g., fitted values, residuals, or predicted survival probabilities) to the original data. This is the Python equivalent of R’s broom::augment().
Parameters
model: object-
A fitted Greenwood estimator (e.g., CoxPH).
data: Any = None-
The original data used to fit the model. Required by some adapters (e.g., Cox residuals need the covariate matrix); optional for others.
**kwargs: Any-
Forwarded to the registered adapter. Common options include
format=.
Returns
DataFrame-
An observation-level summary with predictions or residuals.
Examples
Register a custom augment adapter, then call augment(). The adapter receives the fitted model and the original data and returns a DataFrame with observation-level diagnostics:
import greenwood as gw
from greenwood.summaries import register_augment
lung = gw.load_dataset("lung", backend="polars")
cox = gw.CoxPH().fit("Surv(time, status) ~ age + sex", data=lung)
def _augment_cox(model, data=None, *, format=None, **kwargs):
resid = model.residuals(type="martingale")
return gw._backends.to_dataframe(
{"martingale": resid}, format=format,
)
register_augment("greenwood._cox.CoxPH", _augment_cox)
gw.augment(cox, data=lung, format="polars")
shape: (228, 1)| martingale |
|---|
| f64 |
| 0.00439 |
| -0.505762 |
| -3.129819 |
| 0.532754 |
| -2.350657 |
| … |
| -0.54877 |
| -0.292516 |
| -0.133129 |
| -0.366567 |
| -0.203602 |