Cox via Poisson — PyMC cross-check

Companion to coxpois_python.ipynb. Same person-time expansion, fit with pm.Poisson (interval baseline + treatment, log-exposure offset) under NUTS.

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('leuk.csv'); control=(d['treat']=='control').astype(float).to_numpy()
t=d['time'].to_numpy(float); status=d['cens'].to_numpy(float)
cut=cutpoints(t,status); G=len(cut)+1
subj,itv,expo,ev=survsplit(t,status,cut); off=np.log(expo); trt=control[subj]
with pm.Model() as m:
    lam=pm.Normal('lam',0,10,shape=G)          # interval log baseline hazards
    bt=pm.Normal('bt',0,10)                     # log hazard ratio (control vs 6-MP)
    mu=pm.math.exp(lam[itv] + bt*trt + off)
    pm.Poisson('y', mu=mu, observed=ev)
    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: treat 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(ev,build_design(itv,G,[trt]),off,seed=1)['beta'][:,-1]
print('from-scratch Poisson: treat 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: treat coef 1.596 [0.79,2.48]  HR 4.94  P(HR>1) 1.000  r_hat 1.000
from-scratch Poisson: treat coef 1.615  HR 5.03

Results

PyMC's pm.Poisson on the person-time data gives the same treatment hazard ratio (HR ≈ 5, r̂ ≈ 1.0) as the from-scratch Poisson and as coxph — the count↔survival bridge, sampled by NUTS.