survival_analysis_tools

Kaplan–Meier plotting and Cox HR helpers for survival analysis workflows.

Install

pip install survival_analysis_tools

Usage

Example from tests/test_plot_km.ipynb:

import numpy as np
import pandas as pd

from survival_analysis_tools import plot_km

# Dummy survival data: two groups with different survival patterns
rng = np.random.default_rng(42)
n_per_group = 40

control = pd.DataFrame({
    "survival": rng.exponential(scale=8, size=n_per_group).clip(0.5, 36),
    "event": rng.binomial(1, 0.7, size=n_per_group),
    "treatment": "control",
})
treated = pd.DataFrame({
    "survival": rng.exponential(scale=14, size=n_per_group).clip(0.5, 36),
    "event": rng.binomial(1, 0.55, size=n_per_group),
    "treatment": "treated",
})

df = pd.concat([control, treated], ignore_index=True)
df.head()

    survival  event treatment
0  19.233669      1   control
1  18.689517      0   control
2  19.078088      0   control
3   2.238354      1   control
4   0.691499      0   control

summary, hr, pval = plot_km(
    df,
    time_col="survival",
    event_col="event",
    group_col="treatment",
    group_order=["control", "treated"],
    title="Dummy KM: control vs treated",
    font_scale = 1.1,

    # default parameters:
    xlabel='Time [months]',
    ylabel="Survival probability",
    plot_height=5,
    xlim=False,
    flip_legend_order=False,

    show_censors=True,
    show_ci=True,
    show_at_risk=True,
    show_stats=True,

    weights_col=None,
    cluster_col=None
)

Parameters

plot_km plots Kaplan–Meier curves grouped by group_col. If group_col is None, it plots KM for all data.

Returns a DataFrame with columns: group_name, n, median, median_ci_low, median_ci_high (plus HR and p-value).

Development

pip install -e ".[dev]"
python -m pytest
python -m build

Publishing

See PUBLISHING.md for TestPyPI and production PyPI upload steps with twine.