Survival — PyMC cross-check (Weibull PH with censoring)¶

Companion to survival_python.ipynb (from-scratch MLE + RW-Metropolis). PyMC samples $(\beta,\log k)$ with NUTS; we encode the censored Weibull log-likelihood directly with a pm.Potential — events contribute $\log h + \log S$, censored cases only $\log S$ — matching the from-scratch sampler exactly.

Model¶

$$\log L=\sum_i \delta_i\big[x_i'\beta+\log k+(k-1)\log t_i\big]-\sum_i e^{x_i'\beta}t_i^{k},\qquad \beta\sim N(0,10^2),\ \log k\sim N(0,1).$$ The HR for control vs 6-MP is $e^{\beta_1}$.

In [1]:
import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt, arviz as az
from survival_mcmc import weibull_rwm

d = pd.read_csv('leuk.csv')
trt = (d['treat'] == 'control').astype(float).to_numpy()
t = d['time'].to_numpy(float); delta = d['cens'].to_numpy(float)
X = np.column_stack([np.ones(len(d)), trt]); logt = np.log(t)

with pm.Model() as m:
    beta = pm.Normal('beta', 0, 10, shape=2)
    logk = pm.Normal('logk', 0, 1); k = pm.math.exp(logk)
    xb = pt.dot(X, beta)
    log_h = xb + logk + (k - 1) * logt          # log hazard (events only)
    log_S = -pt.exp(xb) * t ** k                 # log survival (everyone)
    pm.Potential('lik', (delta * log_h + log_S).sum())
    pm.Deterministic('HR', pm.math.exp(beta[1]))
    pm.Deterministic('k', k)
    idata = pm.sample(2000, tune=2000, chains=4, target_accept=0.95, random_seed=1, progressbar=False)

po = idata.posterior; hr = po['HR'].values.ravel()
mxr = float(az.summary(idata, var_names=['beta','logk'])['r_hat'].max())
print('beta_control = %.3f  95%% CI [%.2f, %.2f]' % (float(po['beta'][:,:,1].mean()),
      np.percentile(po['beta'].values[:,:,1],2.5), np.percentile(po['beta'].values[:,:,1],97.5)))
print('HR (control vs 6-MP) = %.2f  95%% CI [%.2f, %.2f]   P(HR>1) = %.3f' % (hr.mean(), np.percentile(hr,2.5), np.percentile(hr,97.5), (hr>1).mean()))
print('shape k = %.3f   max r_hat %.3f' % (float(po['k'].mean()), mxr))

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, logk]

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

beta_control = 1.754  95% CI [0.97, 2.62]
HR (control vs 6-MP) = 6.31  95% CI [2.64, 13.74]   P(HR>1) = 1.000
shape k = 1.356   max r_hat 1.000
In [2]:
# side-by-side with the from-scratch sampler
g = weibull_rwm(X, t, delta, seed=1)['draws']
print('%-22s %12s %12s' % ('parameter','from-scratch','PyMC'))
print('%-22s %12.3f %12.3f' % ('beta_control', g[:,1].mean(), float(po['beta'].values[:,:,1].mean())))
print('%-22s %12.2f %12.2f' % ('HR control/6-MP', np.exp(g[:,1]).mean(), hr.mean()))
print('%-22s %12.3f %12.3f' % ('shape k', np.exp(g[:,2]).mean(), float(po['k'].mean())))

parameter              from-scratch         PyMC
beta_control                  1.745        1.754
HR control/6-MP                6.31         6.31
shape k                       1.345        1.356

Results¶

  • PyMC (NUTS on the censored Weibull likelihood) reproduces the from-scratch fit: HR(control vs 6-MP) ≈ 5–6, P(HR>1) ≈ 1.00, shape k ≈ 1.35, r̂ ≈ 1.00.
  • Encoding survival in PyMC is just a pm.Potential with the two-piece log-likelihood ($\delta\log h+\log S$); no special distribution needed once you write $S(t)=e^{-\lambda t^k}$ explicitly.
  • Same treatment conclusion whether the latent event indicator is handled by the explicit censored likelihood (here) or by a partial likelihood / Cox model (R) — the drug roughly 5×-lowers the relapse hazard.