import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt, arviz as az
from frailty_surv import frailty_weibull_gibbs
d = pd.read_csv('kidney.csv')
age = d['age'].to_numpy(float)-d['age'].mean(); female=(d['sex']==2).astype(float).to_numpy()
GN=(d['disease']=='GN').astype(float); AN=(d['disease']=='AN').astype(float); PKD=(d['disease']=='PKD').astype(float)
X = np.column_stack([np.ones(len(d)),age,female,GN,AN,PKD]); t=d['time'].to_numpy(float); delta=d['status'].to_numpy(float)
pat=(d['id']-1).to_numpy(int); J=int(pat.max()+1); logt=np.log(t)
with pm.Model() as m:
    beta=pm.Normal('beta',0,10,shape=6); logk=pm.Normal('logk',0,2); k=pm.math.exp(logk)
    sw=pm.HalfNormal('sw',1.0); z=pm.Normal('z',0,1,shape=J); w=pm.Deterministic('w',sw*z)
    eta=pt.dot(X,beta)+w[pat]
    pm.Potential('lik',(delta*(eta+logk+(k-1)*logt)-pt.exp(eta)*t**k).sum())
    idata=pm.sample(2000,tune=2000,chains=4,target_accept=0.95,random_seed=1,progressbar=False)
po=idata.posterior; b=po['beta'].values.reshape(-1,6); nm=['intercept','age','female','GN','AN','PKD']
print('covariate hazard ratios:')
for j in (1,2,3,4,5): print('  %-8s HR %.2f [%.2f, %.2f]'%(nm[j],np.exp(b[:,j]).mean(),np.exp(np.percentile(b[:,j],2.5)),np.exp(np.percentile(b[:,j],97.5))))
print('shape k = %.3f   frailty sd = %.3f [%.2f, %.2f]   max r_hat %.3f'%(np.exp(po['logk'].values).mean(),float(po['sw'].mean()),float(po['sw'].quantile(.025)),float(po['sw'].quantile(.975)),float(az.summary(idata,var_names=['beta','sw'])['r_hat'].max())))

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, logk, sw, z]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 8 seconds.

covariate hazard ratios:
  age      HR 1.00 [0.97, 1.04]
  female   HR 0.16 [0.04, 0.37]
  GN       HR 1.39 [0.38, 3.77]
  AN       HR 2.33 [0.61, 6.48]
  PKD      HR 0.48 [0.06, 1.77]
shape k = 1.243   frailty sd = 0.758 [0.09, 1.41]   max r_hat 1.000

g = frailty_weibull_gibbs(t, X, delta, pat, R=25000, burn=5000, seed=1)
print('%-14s %12s %12s'%('param','from-scratch','PyMC'))
for j,l in [(2,'HR female'),(5,'HR PKD'),(4,'HR AN')]:
    print('%-14s %12.2f %12.2f'%(l, np.exp(g['beta'][:,j]).mean(), np.exp(po['beta'].values[:,:,j]).mean()))
print('%-14s %12.3f %12.3f'%('shape k', g['k'].mean(), np.exp(po['logk'].values).mean()))
print('%-14s %12.3f %12.3f'%('frailty sd', g['sigma_w'].mean(), float(po['sw'].mean())))

param          from-scratch         PyMC
HR female              0.19         0.16
HR PKD                 0.57         0.48
HR AN                  1.97         2.33
shape k               1.108        1.243
frailty sd            0.655        0.758

Frailty survival — PyMC cross-check (Weibull + patient frailty)¶

Results¶