PySurvival Documentation

repository·master·Indexed 18 days ago

An open-source Python package for Survival Analysis modeling to predict the timing of events. Built on NumPy, SciPy, and PyTorch, it provides over 10 models including Cox Proportional Hazard (CoxPHModel), Neural Multi-Task Logistic Regression (LinearMultiTaskModel), and Random Survival Forest. It includes utilities for data handling via the Dataset class, performance metrics such as the concordance index (c-index) and integrated Brier score, and visualization tools for risk groups and feature importance.

Tokens: 4.4K
Snippets: 18
Records: 20
Agent score: 61%

What's inside PySurvival

Overview of PySurvival features
master
PySurvival is an open-source Python package for Survival Analysis modeling. It provides:
- Models: Over 10 models including Cox Proportional Hazard (CoxPHModel), Neural Multi-Task Logistic Regression (LinearMultiTaskModel), and Random Survival Forest.
- Metrics: Performance assessment tools like concordance_index (c-index) and brier_score.
- Data Handling: Utilities to load datasets and split them into training/testing sets via Dataset.load_train_test().
- Model Management: Capabilities to load and save models.
https://github.com/square/pysurvival/blob/master/README.md

Quickstart: Build and evaluate survival models

master

This example demonstrates the standard workflow in PySurvival: loading a dataset, splitting it into training and testing sets, fitting a Cox Proportional Hazard (CoxPH) model and a Linear Multi-Task Logistic Regression (MTLR) model, and evaluating them using the concordance index.

# Loading the modules
from pysurvival.models.semi_parametric import CoxPHModel
from pysurvival.models.multi_task import LinearMultiTaskModel
from pysurvival.datasets import Dataset
from pysurvival.utils.metrics import concordance_index

# Loading and splitting a simple example into train/test sets
X_train, T_train, E_train, X_test, T_test, E_test = \
	Dataset('simple_example').load_train_test()

# Building a CoxPH model
coxph_model = CoxPHModel()
coxph_model.fit(X=X_train, T=T_train, E=E_train, init_method='he_uniform', 
                l2_reg = 1e-4, lr = .4, tol = 1e-4)

# Building a MTLR model
mtlr = LinearMultiTaskModel()
mtlr.fit(X=X_train, T=T_train, E=E_train, init_method = 'glorot_uniform', 
           optimizer ='adam', lr = 8e-4)

# Checking the model performance
c_index1 = concordance_index(model=coxph_model, X=X_test, T=T_test, E=E_test )
print("CoxPH model c-index = {:.2f}".format(c_index1))

c_index2 = concordance_index(model=mtlr, X=X_test, T=T_test, E=E_test )
print("MTLR model c-index = {:.2f}".format(c_index2))

Install PySurvival via pip
master
If you have a working version of gcc installed, you can install PySurvival using pip. PySurvival is compatible with Python 2.7 through 3.7 and is built on NumPy, SciPy, and PyTorch.
```
pip install pysurvival
```
https://github.com/square/pysurvival/blob/master/README.md
Train a Random Survival Forest model
master
The RandomSurvivalForestModel can be used to fit survival data. The .fit() method requires three primary inputs:
- X: Feature matrix (e.g., a pandas DataFrame or numpy array).
- T: Time column (durations).
- E: Event column (binary indicators of whether the event occurred).
Common hyperparameters include num_trees, max_features, max_depth, and min_node_size.
```
from pysurvival.models.survival_forest import RandomSurvivalForestModel

# Fitting the model
rsf = RandomSurvivalForestModel(num_trees=100) 
rsf.fit(X_train, T_train, E_train, max_features='log2', 
        max_depth=2, min_node_size=5)
```
https://github.com/square/pysurvival/blob/master/notebooks/Employee%20retention%20-%20Knowing%20when%20your%20employees%20will%20quit.ipynb
Check feature importance in Survival Forests
master
For models like ExtraSurvivalTreesModel, you can access the feature importance via the variable_importance_table attribute. This provides a ranking of which features most significantly impact the model's predictions.
```
# Assuming xst is a fitted ExtraSurvivalTreesModel
xst.variable_importance_table.head(5)
```
https://github.com/square/pysurvival/blob/master/notebooks/Churn%20Prediction%20-%20Predicting%20when%20your%20customers%20will%20churn.ipynb
Save a trained model
master
You can persist a trained model to a file using the save_model utility from pysurvival.utils.
```
from pysurvival.utils import save_model
save_model(xst, 'path/to/model_filename.zip')
```
https://github.com/square/pysurvival/blob/master/notebooks/Churn%20Prediction%20-%20Predicting%20when%20your%20customers%20will%20churn.ipynb

Visualize cohort vs actual repayment using compare_to_actual

master

The compare_to_actual utility allows you to compare the predicted time series of events (e.g., loans repaid) or at-risk populations against the actual data.

from pysurvival.utils.display import compare_to_actual

# Compare predicted number of loans fully repaid
results = compare_to_actual(neural_mtlr, X_test, T_test, E_test,
                            is_at_risk=False, figure_size=(16, 6),
                            metrics=['rmse', 'mean', 'median'])

# Compare predicted number of loans still active
results = compare_to_actual(neural_mtlr, X_test, T_test, E_test,
                            is_at_risk=True, figure_size=(16, 6), 
                            metrics=['rmse', 'mean', 'median'])

Create risk groups for individual predictions
master
To segment individuals into risk categories (e.g., low, medium, high), use create_risk_groups from pysurvival.utils.display. This helps in visualizing survival functions for different risk profiles.
Parameters:
- model: The fitted model.
- X: Feature matrix.
- use_log: Whether to use log-transformed risk scores.
- num_bins: Number of bins for grouping.
- low, medium, high: Dictionaries defining the lower_bound, upper_bound, and color for each risk group.
```
from pysurvival.utils.display import create_risk_groups

risk_groups = create_risk_groups(model=xst, X=X_test,
    use_log = True, num_bins=30, figure_size=(20, 4),
    low={'lower_bound':0, 'upper_bound':1.65, 'color':'red'},
    medium={'lower_bound':1.65, 'upper_bound':2.2,'color':'green'},
    high={'lower_bound':2.2, 'upper_bound':3,  'color':'blue'}
    )
```
https://github.com/square/pysurvival/blob/master/notebooks/Churn%20Prediction%20-%20Predicting%20when%20your%20customers%20will%20churn.ipynb
Analyze feature importance and save models
master
After training a model like RandomSurvivalForestModel, you can inspect which features were most predictive using the variable_importance_table attribute. To persist your trained model, use the save_model utility.
```
# View feature importance
print(rsf.variable_importance_table.head())

# Save the model to a file
from pysurvival.utils import save_model
save_model(rsf, 'path/to/model.zip')
```
https://github.com/square/pysurvival/blob/master/notebooks/Employee%20retention%20-%20Knowing%20when%20your%20employees%20will%20quit.ipynb
Load datasets from pysurvival.datasets
master
You can load built-in datasets using the Dataset class from pysurvival.datasets. The .load() method returns the dataset, typically as a pandas DataFrame.
Example:
```
from pysurvival.datasets import Dataset
raw_dataset = Dataset('churn').load()
```
https://github.com/square/pysurvival/blob/master/notebooks/Churn%20Prediction%20-%20Predicting%20when%20your%20customers%20will%20churn.ipynb
Visualize cohort and individual survival predictions
master
PySurvival provides utilities for visualizing predictions:
- Cohort Predictions: Use compare_to_actual(model, X, T, E, ...) to compare the predicted number of events against actual observed counts over time.
- Individual Risk Groups: Use create_risk_groups(model, X, ...) to bin individuals into risk categories based on their risk scores. This is useful for comparing survival functions between high-risk and low-risk groups.
- Individual Survival Function: Use model.predict_survival(X_individual) to get the survival probability for a specific unit across all time points.
```
# Compare cohort predictions to actuals
from pysurvival.utils.display import compare_to_actual
results = compare_to_actual(rsf, X_test, T_test, E_test, is_at_risk=False)

# Create risk groups for visualization
from pysurvival.utils.display import create_risk_groups
risk_groups = create_risk_groups(model=rsf, X=X_test, num_bins=30)

# Predict survival for a single individual
survival_curve = rsf.predict_survival(X_test.iloc[0, :]).flatten()
```
https://github.com/square/pysurvival/blob/master/notebooks/Employee%20retention%20-%20Knowing%20when%20your%20employees%20will%20quit.ipynb
Load built-in datasets with Dataset.load()
master
You can access pre-configured datasets using the pysurvival.datasets.Dataset class. For example, to load the 'credit_risk' dataset, pass the name to the constructor and call .load().
```
from pysurvival.datasets import Dataset
raw_dataset = Dataset('credit_risk').load()
```
https://github.com/square/pysurvival/blob/master/notebooks/Credit%20Risk%20-%20Computing%20the%20speed%20of%20repayment%20of%20loans.ipynb