scikit-survival 0.28.0 released

I am pleased to announce the release of scikit-survival 0.28.0.

A highlight of this release is the support for Polars DataFrames alongside pandas DataFrames via the Narwhals dataframe abstraction layer. In addition, this release adds support for scikit-learn 1.9.

Support for Polars

Polars is a data frame library similar to pandas, but with its core written in Rust instead of Python, which often gives Polars an advantage in terms of performance.

All datasets shipped with scikit-survival can now be loaded as a Polars DataFrame by specifying the output_type argument.

from sksurv.datasets import load_gbsg2

# return X as a polars DataFrame
X, y = load_gbsg2(output_type="polars")

X.head()

shape: (5, 8)
┌──────┬────────┬───────┬──────────┬────────┬─────────┬────────┬───────┐
│ age  ┆ estrec ┆ horTh ┆ menostat ┆ pnodes ┆ progrec ┆ tgrade ┆ tsize │
│ ---  ┆ ---    ┆ ---   ┆ ---      ┆ ---    ┆ ---     ┆ ---    ┆ ---   │
│ f64  ┆ f64    ┆ enum  ┆ enum     ┆ f64    ┆ f64     ┆ enum   ┆ f64   │
╞══════╪════════╪═══════╪══════════╪════════╪═════════╪════════╪═══════╡
│ 70.0 ┆ 66.0   ┆ no    ┆ Post     ┆ 3.0    ┆ 48.0    ┆ II     ┆ 21.0  │
│ 56.0 ┆ 77.0   ┆ yes   ┆ Post     ┆ 7.0    ┆ 61.0    ┆ II     ┆ 12.0  │
│ 58.0 ┆ 271.0  ┆ yes   ┆ Post     ┆ 9.0    ┆ 52.0    ┆ II     ┆ 35.0  │
│ 59.0 ┆ 29.0   ┆ yes   ┆ Post     ┆ 4.0    ┆ 60.0    ┆ II     ┆ 17.0  │
│ 73.0 ┆ 65.0   ┆ no    ┆ Post     ┆ 1.0    ┆ 26.0    ┆ II     ┆ 35.0  │
└──────┴────────┴───────┴──────────┴────────┴─────────┴────────┴───────┘

Transforming DataFrames

scikit-learn enables transformers to return Polars data frames via the set_output API.

from sklearn import set_config
from sklearn.preprocessing import StandardScaler

set_config(transform_output="polars")
# standarize the columns of a polars DataFrame
X_standarized = StandardScaler().fit_transform(
  X.select("age", "tsize")
)

scikit-surival has two transformers that now accept polars and pandas data frames as input: ClinicalKernelTransform, and OneHotEncoder.

ClinicalKernelTransform

ClinicalKernelTransform computes a kernel matrix, so the output will be a numpy array, as before. With scikit-survival 0.28.0, it is aware of the polars colum types String, Categorical, Enum and Object. However, polars does not have a concept similar to ordered categories in pandas: pd.Categorical([…], ordered=True).

When computing the clinical kernel for a polars data frame, you can specify the order of categories with the ordinal_categories argument, otherwise all non-numeric columns will be treated as nominal columns, where values have no specific order (e.g. the column horTh with values “yes” and “no” in the example below).

from sksurv.kernels import ClinicalKernelTransform

K = ClinicalKernelTransform(
  ordinal_categories={"tgrade": ["I", "II", "III"]},
).fit_transform(
  X.select("age", "horTh", "tgrade")
)

OneHotEncoder

OneHotEncoder encodes the string-type columns to numeric columns. It automatically returns a data frame of the same type as the input:

from sksurv.preprocessing import OneHotEncoder

X_onehot = OneHotEncoder().fit_transform(X.select("horTh", "tgrade"))

X_onehot.head()

shape: (5, 3)
┌───────────┬───────────┬────────────┐
│ horTh=yes ┆ tgrade=II ┆ tgrade=III │
│ ---       ┆ ---       ┆ ---        │
│ f64       ┆ f64       ┆ f64        │
╞═══════════╪═══════════╪════════════╡
│ 0.0       ┆ 1.0       ┆ 0.0        │
│ 1.0       ┆ 1.0       ┆ 0.0        │
│ 1.0       ┆ 1.0       ┆ 0.0        │
│ 1.0       ┆ 1.0       ┆ 0.0        │
│ 0.0       ┆ 1.0       ┆ 0.0        │
└───────────┴───────────┴────────────┘

Fitting survival models

Finally, you can pass a polars data frame to the fit and predict functions of any estimator.

from sklearn.pipeline import make_pipeline
from sksurv.linear_model import CoxPHSurvivalAnalysis

pipe = make_pipeline(
  OneHotEncoder(), StandardScaler(), CoxPHSurvivalAnalysis()
)

pipe.fit(X[:500], y[:500])

risk_scores = pipe.predict(X[500:])

Note that this only works for eager data frames, lazy data frames will give an error.

X_lazy = X.lazy()

pipe.fit(X_lazy[:500], y[:500])

Traceback (most recent call last):
  File "example.py", line 21, in <module>
    pipe.fit(X_lazy[:500], y[:500])
  File "…/site-packages/sklearn/base.py", line 1403, in wrapper
    return fit_method(estimator, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  …
  File "…/site-packages/sksurv/_dataframe/_input.py", line 96, in ensure_eager_dataframe
    _reject_polars_lazyframe(obj)
  File "…/site-packages/sksurv/_dataframe/_input.py", line 84, in _reject_polars_lazyframe
    raise TypeError(_LAZYFRAME_NOT_SUPPORTED_MSG)
TypeError: polars.LazyFrame is not supported; call .collect() before passing to scikit-survival.

New Contributors

A big shoutout to our two new contributors:

• cakedev0 for adding support for scikit-learn 1.9
• 55Kamiryo for adding support for Polars data frames.

Updated Dependencies

With this release, the minimum supported version are:

Install

scikit-survival is available for Linux, macOS, and Windows and can be installed either

via pip:

pip install scikit-survival

or via conda

 conda install -c conda-forge scikit-survival