Feature Request: Survival Analysis Support (similar to scikit-survival)

[HumbleHumbert]

Hi there,

First off, thank you for the incredible work on polars-ds. It has been a game-changer for integrating performant data science workflows directly into the Polars ecosystem, avoiding the costly round-trip to pandas/numpy for many standard tasks.

I often work with time-to-event data (survival analysis) and currently have to break out of the Polars lazy/expression API to use libraries like scikit-survival or lifelines. This involves materializing the data and converting it to pandas/numpy, which becomes a bottleneck with larger datasets that polars-ds would otherwise handle gracefully.

Ideally, this would mirror some of the core functionality found in scikit-survival

Specific features that would be high-value include:

1. Kaplan-Meier Estimator: A way to compute survival probabilities over time groups.

2. Nelson-Aalen Estimator: For cumulative hazard functions.

3. Cox Proportional Hazards (Metrics): Even if full training is complex, having the ability to calculate the Concordance Index (C-index) or Log-Rank Test statistics directly as expressions would be incredibly useful for evaluating models or comparing cohorts.

I understand this is a significant addition, but I believe it fits well with the library's goal of extending Polars for general DS use cases.

[abstractqqq]

I will take a look when I get time. But I cannot promise any meaningful progress until next year..

[abstractqqq]

This seems to be simpler than I thought. I added query_kaplan_meier_prob in this PR.

Here is an exampling calling the new implementation and comparing it against sk survival

import polars as pl
import polars_ds as pds
import pytest
from polars.testing import assert_frame_equal

def test_kaplan_meier():

    from sksurv.datasets import load_veterans_lung_cancer
    from sksurv.nonparametric import kaplan_meier_estimator    

    _, y = load_veterans_lung_cancer()
    time, prob_surv, conf_int = kaplan_meier_estimator(
        y["Status"], y["Survival_in_days"], conf_type="log-log"
    )

    df_result = pl.from_dict({
        "t": time
        , "prob": prob_surv
    })

    df = pl.from_dict({
        "status": y['Status']
        , "survival_time": y["Survival_in_days"]
    })

    df_pds_result = df.select(
        pds.query_kaplan_meier_prob(
            "status"
            , "survival_time"
        ).alias("estimate")
    ).unnest("estimate")

    assert_frame_equal(df_result, df_pds_result)

Because I don't do a lot of survival analysis and this is new, this API is experimental and is subject to breaking changes in the future. Please let me know what you think and what optional parameters should I include? E.g conf_level, reverse, etc. I don't want to over-complicate it but if there is a strong need for those parameters, then I will add them.

[HumbleHumbert]

@abstractqqq Thank you so much for the quick action on this!

I sincerely apologize for the lack of feedback on my end recently—I’ve been incredibly busy graduating and haven't had the chance to dig in as much as I wanted.

Regarding the implementation, it would be sufficient for my use case if it works the same as scikit-survival.

Thanks again!