Mastectomy — PyMC cross-check (Cox via Poisson)

Companion to mastectomy_python.ipynb. Same person-time expansion (coxpois.py), fit with pm.Poisson (interval baseline + metastasis effect, log-exposure offset). This is essentially the model from PyMC's own "Bayesian Survival Analysis" example.

import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt, arviz as az
from coxpois import survsplit, cutpoints, build_design, pois_rwm
d = pd.read_csv('mastectomy.csv'); met=(d['metastasized']=='yes').astype(float).to_numpy()
t=d['time'].to_numpy(float); ev=d['event'].to_numpy(float)
cut=cutpoints(t,ev); G=len(cut)+1
subj,itv,expo,evt=survsplit(t,ev,cut); off=np.log(expo); x=met[subj]
with pm.Model() as m:
    lam=pm.Normal('lam',0,10,shape=G); bt=pm.Normal('bt',0,10)
    mu=pm.math.exp(lam[itv] + bt*x + off)
    pm.Poisson('y', mu=mu, observed=evt)
    idata=pm.sample(2000,tune=2000,chains=4,target_accept=0.95,random_seed=1,progressbar=False)
po=idata.posterior; b=po['bt'].values.ravel()
print('PyMC: metastasis coef %.3f [%.2f,%.2f]  HR %.2f  P(HR>1) %.3f  r_hat %.3f'
      %(b.mean(),np.percentile(b,2.5),np.percentile(b,97.5),np.exp(b.mean()),(b>0).mean(),float(az.summary(idata,var_names=['bt'])['r_hat'].iloc[0])))
g=pois_rwm(evt,build_design(itv,G,[x]),off,seed=1)['beta'][:,-1]
print('from-scratch Poisson: coef %.3f  HR %.2f'%(g.mean(),np.exp(g.mean())))

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [lam, bt]

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

PyMC: metastasis coef 0.643 [-0.23,1.63]  HR 1.90  P(HR>1) 0.926  r_hat 1.000

from-scratch Poisson: coef 0.658  HR 1.93

Results

PyMC reproduces the from-scratch Poisson fit: metastasis HR ≈ 1.9 (from-scratch 1.93), P(HR>1) ≈ 0.93, r̂ ≈ 1.0 — the count↔survival bridge on the mastectomy data, sampled by NUTS.