Survival inference with pymc

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import pymc as pm
import arviz as az
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler

from survivalpredict.estimators import CoxProportionalHazard
from survivalpredict.datasets import load_iranian_telecom_churn

az.style.use("arviz-darkgrid")

In cases of inference or prediction on a trained models, it might be valuable to gauge uncertainty. PyMC is a Bayesian/‘Markov chain Monte Carlo’ framework. Simply put, PyMC allows us to build models that can capture the uncertainty and possible variability of our model, given the observed data and model assumptions. There is much literature on PyMC/MCMC; we will assume some familiarity with the topic.

Curretly surivivalpredict supports generating equivalent PyMC models for two of its estimator classes. The CoxProportionalHazards get_pymc_model method will return an equivalent pymc model. Models with either ‘efron’ or ‘breslow’ ties, stratification, L2 regularization, and left-censorship are all supported. The ParametricDiscreteTimePH actually uses PyMC under the hood, calling the get_pymc_model will return the underlying pymc model.

#loading some sample data
iranian_telecom_churn_dict = load_iranian_telecom_churn()
X = iranian_telecom_churn_dict['X']
#in this case is important to scale our data
X = StandardScaler().fit_transform(X)
times = iranian_telecom_churn_dict['times'].astype(np.int64)
events = iranian_telecom_churn_dict['events'].astype(np.bool_)
column_names = iranian_telecom_churn_dict['column_names']

cox = CoxProportionalHazard(ties='efron').fit(X,times,events)
print(cox.coef_)

[ 0.50777218  0.48238747 -0.44392135  0.42167482 -2.31404596 -2.99677386
 -0.37073225 -0.417258    0.26621988 -0.03414636  0.30532322  1.86236737]

The CoxProportionalHazard class comes with a pymc model generator that builds pymc models that are equivalent to the survivalpredict model. It should be noted that adding time_start/’left-censorship’ will drastically reduce the samples per second of the pymc class, as it triggers more computation. Stratification will also somewhat slow down sampling.

When calling get_pymc_model, if the empirical_bayes argument is set to True, and our estimator class has already been trained, the initial values of our PyMC model coefficients will be the coefficients trained by our estimator class.

#Getting our pymc model, setting empirical_bayes to False as part of our demonstration.
pymc_model = cox.get_pymc_model(X, times, events,labes_names=column_names,empirical_bayes=False)

pm.model_to_graphviz(pymc_model)

A quick demonstration to show that this PyMC model is equivalent to our survivalpredict model would be to generate the ‘maximum likelihood estimation’/’maximum a posteriori’ estimates and compare them to our coefficients.

with pymc_model:
    mle = pm.find_MAP()

mle

{'coefs': array([ 0.50733508,  0.48238074, -0.44546389,  0.43406345, -2.31074651,
        -2.9503162 , -0.37137833, -0.4172489 ,  0.26588864, -0.03518892,
         0.30293834,  1.81416004]),
 'n_log_likelihood': array(3052.59817273)}

print('coefs: ',cox.coef_)
print('n_log_likelihood: ',cox.n_log_likelihood)
#looks pretty close to me

coefs:  [ 0.50777218  0.48238747 -0.44392135  0.42167482 -2.31404596 -2.99677386
 -0.37073225 -0.417258    0.26621988 -0.03414636  0.30532322  1.86236737]
n_log_likelihood:  3052.597200555113

#generating the trace/sampling
#this might take a bit, it can be faster with a gpu.
with pymc_model:
    idata = pm.sample()

NUTS[nutpie]: [coefs]

az.summary(idata)

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
coefs[call_failure]	0.51	0.078	0.38	0.64	2108	2427	1.00	0.0017	0.0012
coefs[complains]	0.482	0.0289	0.43	0.53	5043	3338	1.00	0.00041	0.00029
coefs[charge_amount]	-0.456	0.142	-0.69	-0.23	2603	2461	1.00	0.0028	0.002
coefs[seconds_of_use]	0.41	0.48	-0.35	1.2	888	1283	1.01	0.016	0.012
coefs[frequency_of_use]	-2.32	0.34	-2.9	-1.8	1650	2063	1.00	0.0084	0.0062
coefs[frequency_of_sms]	-3.04	1.07	-4.8	-1.4	828	1112	1.01	0.037	0.025
coefs[distinct_called_numbers]	-0.366	0.14	-0.59	-0.14	3186	2430	1.00	0.0025	0.0017
coefs[age_group]	-0.41	0.18	-0.69	-0.12	1411	1802	1.00	0.0048	0.0034
coefs[tariff_plan]	0.24	0.126	0.031	0.43	5427	3313	1.00	0.0017	0.0013
coefs[status]	-0.034	0.068	-0.14	0.074	2733	2441	1.00	0.0013	0.00091
coefs[age]	0.3	0.197	-0.011	0.61	1221	1672	1.00	0.0056	0.004
coefs[customer_value]	1.86	1.14	0.1	3.7	814	1001	1.01	0.04	0.027
n_log_likelihood	3058.6	2.5	3100	3100	1529	1939	1.00	0.065	0.062

#obtaining our 'highest density intervals'/credible intervals for the probable effect of different variables.
az.plot_dist(idata,var_names='coefs')
plt.tight_layout()

It is common for MCMC practitioners to generate ‘Posterior Probability Statements’ by asking binary questions across posterior traces and obtaining their mean as a gauge of certainty. Below is simply a comparison of the effects of parameters. There are a lot of good resources on more advanced bayesian inference techniques on MCMC traces.

coef1 = 'complains'
coef2 = 'seconds_of_use'

baysian_p_value = (idata['posterior']['coefs'].sel(labes=coef1) > idata['posterior']['coefs'].sel(labes=coef2)).mean().values

print(f'the baysian p value that {coef1} have more effect on {coef2} is :{baysian_p_value}')
#we should not be very certain with such middling number

the baysian p value that complains have more effect on seconds_of_use is :0.54975

Below is a taste of what is possible with PyMC traces. We can take a look at possible coefficients and build a range of possible predictions for our survival curves. It should be noted that our solution would look different for doing the same thing with a stratified model. Some what of is here may be packaged into tooling within survivalpredict in the future.

#poping out an internal function, ;)
from survivalpredict._base_hazard import _get_breslow_base_hazard

#turning coefs from the pymc traces into breslow base hazards
coefs_from_trace = idata.posterior.coefs.values.reshape(-1,X.shape[1])
risk_per_coef_per_row = np.exp(np.dot(coefs_from_trace,X.T))
max_time = np.max(times)
hazard_per_coef = np.apply_along_axis(lambda risk_array : _get_breslow_base_hazard(risk_array,times,events,max_time),1,risk_per_coef_per_row)
base_survival_per_coef = np.exp(- hazard_per_coef.cumsum(axis=1))

#the row we wants to build prediction intervals for. This row can also be something out of sample.
row_to_predict = X[17]

#building the survival_curves
predict_risk = np.exp(np.dot(coefs_from_trace,row_to_predict.T))
survival_curve_per_coef = base_survival_per_coef**predict_risk[:,None]

means = np.quantile(survival_curve_per_coef, .5,axis=0)
#we are looking at the 5% and 95% intervals
quantiles = np.quantile(survival_curve_per_coef, [0.01,0.99],axis=0)

plt.plot(np.arange(1,len(means)+1),means)
plt.fill_between(np.arange(1,len(means)+1),quantiles[0,:],quantiles[1,:],alpha=0.3)
ax = plt.gca()
ax.set_ylim(0.0,1.05)
plt.xlabel("Time Interval")
plt.ylabel("Probability of Survival")

Text(0, 0.5, 'Probability of Survival')