C11 · Who needs the second drug? — Bayesian Modelling in Practice

The brief

It is 1996. ACTG 175 has just shown that adding didanosine (ddI) to zidovudine (ZDV) slows HIV disease better than ZDV alone. But ddI is not free: it causes pancreatitis and painful peripheral neuropathy in a minority of patients, and a clinic has to decide whom to put on the combination. The HIV clinic you work for asks:

"Does everybody benefit, or is there a subgroup of patients for whom ddI does little and who could be spared the toxicity? Our rule of thumb: ddI is worth it for a patient if it raises their CD4 count after 20 weeks by at least 15%, and we want to be at least 80% sure of that."

A statistician on the team has already fitted a "clustering regression" and says it found responders and non-responders. Your job: build the model properly - a regression whose coefficients may differ between hidden subgroups, with a Dirichlet-process prior so that the data decide how many subgroups there are - find out what it can and cannot say about who benefits, and turn it into a rule the clinic can use.

This is the regression sequel to example E26 (Dirichlet-process mixtures of densities). Do E26 first: this challenge assumes you know stick-breaking, truncation checks and why component labels are meaningless.

How this notebook works

Each task states what to deliver, not how. Write your code in the YOUR CODE HERE cells.
Stuck? h.hint("task3") reveals hints one level at a time: nudge → approach → code skeleton. Try to get by on nudges.
h.check("task2", ate=...) compares your numbers with the reference solution.
A full worked solution lives in notebooks/solutions/. Open it only when you are done (or truly stuck).

import logging
import time

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import xarray as xr
from scipy import stats

from pymc_challenges import Hints, data

RANDOM_SEED = 1996
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # several fits: no sampler banner per fit

h = Hints("C11")
h.tasks()

C11 - Who needs the second drug? Hidden subgroups with their own regressions

task0 - Look at the trial (2 hints)
task1 - One regression for everybody (3 hints, 2 check(s))
task2 - A Dirichlet-process mixture of regressions (3 hints, 2 check(s))
task3 - The colleague's responder list (3 hints, 2 check(s))
task4 - Let the baseline profile choose the subgroup (3 hints, 2 check(s))
task5 - The clinic's rule (3 hints, 3 check(s))

The data

One row per patient. We use the two arms the clinic is choosing between: trt = 0 (zidovudine alone) and trt = 1 (zidovudine + didanosine). CD4 cells are the immune cells HIV destroys; their count per mm³ of blood is the standard marker of disease progression (below 200 means AIDS). The trial enrolled patients with 200-500 cells/mm³.

column	meaning
`cd40`, `cd420`	CD4 count at baseline and at 20 ± 5 weeks - the outcome is `cd420`
`cd80`, `cd820`	CD8 count at baseline and 20 weeks (do not use `cd820`: it is post-treatment)
`age`, `wtkg`	age (years), weight (kg)
`str2`	prior antiretroviral therapy (1 = experienced, 0 = naive)
`symptom`	1 = symptomatic HIV infection at baseline
`karnof`	Karnofsky performance score (100 = normal activity)
`gender`, `race`	1 = male; 1 = non-white
`homo`, `drugs`, `hemo`	homosexual activity, injection drug use, haemophilia (1 = yes)
`offtrt`, `time`, `cid`	post-randomisation: stopped treatment, follow-up days, clinical event - do not use

data.describe("actg175")
actg = data.load("actg175")
trial = actg[actg.trt.isin([0, 1])].copy()
print(trial.trt.value_counts().sort_index().rename({0: "ZDV", 1: "ZDV + ddI"}))
trial.describe().T.round(1)

actg175: 2139 HIV-infected adults with 200-500 CD4 cells/mm3, randomised to zidovudine (ZDV, trt 0), ZDV + didanosine (trt 1), ZDV + zalcitabine (trt 2) or didanosine alone (trt 3). One row per patient: baseline covariates (age, wtkg, hemo, homo, drugs, karnof, race, gender, str2 = prior antiretroviral therapy, symptom, ...), CD4/CD8 counts at baseline (cd40, cd80) and at 20 +- 5 weeks (cd420, cd820), and time to a clinical event (time, cid).
  source: Hammer et al. (1996, New England Journal of Medicine 335:1081), AIDS Clinical Trials Group Study 175; UCI Machine Learning Repository dataset 890 (doi:10.24432/C5ZG8F), the same data as the R package `speff2trial` (Juraska et al.).
  url:    https://archive.ics.uci.edu/static/public/890/data.csv
trt
ZDV          532
ZDV + ddI    522
Name: count, dtype: int64

	count	mean	std	min	25%	50%	75%	max
cid	1054.0	0.3	0.4	0.0	0.0	0.0	1.0	1.0
time	1054.0	858.2	302.9	33.0	647.2	989.0	1090.0	1231.0
trt	1054.0	0.5	0.5	0.0	0.0	0.0	1.0	1.0
age	1054.0	35.2	8.8	12.0	29.0	34.0	40.0	70.0
wtkg	1054.0	75.5	13.4	31.0	67.1	74.5	83.0	159.9
hemo	1054.0	0.1	0.3	0.0	0.0	0.0	0.0	1.0
homo	1054.0	0.7	0.5	0.0	0.0	1.0	1.0	1.0
drugs	1054.0	0.1	0.3	0.0	0.0	0.0	0.0	1.0
karnof	1054.0	95.5	5.9	70.0	90.0	100.0	100.0	100.0
oprior	1054.0	0.0	0.2	0.0	0.0	0.0	0.0	1.0
z30	1054.0	0.5	0.5	0.0	0.0	1.0	1.0	1.0
zprior	1054.0	1.0	0.0	1.0	1.0	1.0	1.0	1.0
preanti	1054.0	376.8	468.5	0.0	0.0	123.5	728.0	2500.0
race	1054.0	0.3	0.4	0.0	0.0	0.0	1.0	1.0
gender	1054.0	0.8	0.4	0.0	1.0	1.0	1.0	1.0
str2	1054.0	0.6	0.5	0.0	0.0	1.0	1.0	1.0
strat	1054.0	2.0	0.9	1.0	1.0	2.0	3.0	3.0
symptom	1054.0	0.2	0.4	0.0	0.0	0.0	0.0	1.0
treat	1054.0	0.5	0.5	0.0	0.0	0.0	1.0	1.0
offtrt	1054.0	0.4	0.5	0.0	0.0	0.0	1.0	1.0
cd40	1054.0	351.0	122.3	0.0	264.2	340.0	423.8	1199.0
cd420	1054.0	369.3	147.9	49.0	264.2	352.5	459.0	1119.0
cd80	1054.0	995.7	481.4	40.0	654.0	894.5	1214.5	4255.0
cd820	1054.0	947.9	439.3	124.0	637.0	866.0	1179.8	3407.0

Task 0 · Look at the trial

Deliver

A plot of CD4 at 20 weeks against baseline CD4, by arm, on the scale you will model. Why that scale?
The records that cannot be right, and what you do with them.
The distribution of the change from baseline in each arm: centre, spread, and shape. What would a regression's error distribution have to look like?

Your turn Work this task in the notebook.

Task 1 · One regression for everybody

Deliver

A linear regression of log CD4 at 20 weeks on log baseline CD4, treatment and the other baseline covariates, with priors you can defend on the log scale, and clean diagnostics.
The treatment effect as a percentage change in CD4, with a 94% interval.
A posterior predictive check, per arm, that shows what this model gets wrong - pick the statistic that matters for a clinic worried about patients who deteriorate.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

On the log scale a coefficient is a relative change: what is a plausible effect of a drug on CD4 within 20 weeks - 5%? 50%? 500%? For the check, think about what a clinic worried about deterioration would count, and about a statistic a Normal error distribution is not free to fit.

2 Approach

Regress log(cd420) on centred log(cd40), the treatment indicator and the other baseline covariates standardised (never cd820, offtrt, time, cid: they are measured after randomisation). exp(c) - 1 is the percentage effect.

For the PPC, pm.sample_posterior_predictive gives replicated log counts; subtract log(cd40) and compute, per replicated dataset and per arm, the share falling by more than 25% and the skewness. Compare the distribution of each with the observed value.

3 Skeleton

with pm.Model(coords={"covariate": list(Z.columns)}) as single:
    a = pm.Normal("a", ___, 1.0)                     # log CD4 of an average patient
    b = pm.Normal("b", 0.7, 0.5)                     # elasticity on baseline CD4
    c = pm.Normal("c", 0.0, ___)                     # treatment effect, log scale
    g = pm.Normal("g", 0.0, 0.2, dims="covariate")
    sigma = pm.HalfNormal("sigma", 1.0)
    pm.Normal("y", a + b * lx + c * t + Z.to_numpy() @ g, sigma, observed=y)
    idata = pm.sample(random_seed=RANDOM_SEED)
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)

rep = az.extract(idata, group="posterior_predictive", var_names="y", num_samples=500)
rep = rep.transpose("sample", ...).to_numpy() - log_cd4        # (draws, patients)
fall = (rep[:, t == 0] < np.log(0.75)).mean(axis=___)          # one number per replicate

Self-checks in the notebook (h.check("task1", …)):

effect — posterior mean of the treatment coefficient on the LOG scale, about 0.20 (i.e. about +22%)
sigma — posterior mean of the residual sd on the log scale, about 0.33

Task 2 · A Dirichlet-process mixture of regressions

Suppose the patients fall into hidden subgroups, each with its own regression: its own intercept, slope on baseline CD4, treatment effect and residual spread (the other covariates may keep shared coefficients). Nobody knows how many subgroups there are.

Deliver

That model, with a truncated Dirichlet-process (stick-breaking) prior on the subgroup weights and the subgroup memberships marginalised. Show what your prior on the concentration implies for the number of subgroups among these patients, and check that the truncation does not bite.
Evidence that the sampler worked. Decide which quantities it is meaningful to diagnose.
Did it fix what Task 1 got wrong? The same PPC, and a LOO comparison with Task 1 and with a single regression that merely has heavier tails. How much of the improvement needs subgroups?
The average treatment effect, and how different the subgroups' treatment effects are.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

This is E26's model with a regression inside each component: the component means now differ per patient. Before fitting, simulate how many occupied subgroups your prior on the concentration implies for about 1000 patients. After fitting, ask which quantities have a meaning that does not depend on what the components are called. And to judge "subgroups", compare against a single regression that only has heavier tails.

2 Approach

pm.StickBreakingWeights("w", alpha=alpha, K=T - 1) gives T weights (K counts the breaks). Build a (patients, T) matrix of means a + b * lx[:, None] + c * t[:, None] + (Z @ g)[:, None] with a, b, c, sigma of shape T, and pass it to pm.NormalMixture("y", w=w, mu=mu, sigma=sigma, observed=y): that marginalises the memberships. Give the sds a prior with no mass near zero (pm.InverseGamma).

Label-free quantities: the average effect (w * c).sum(), the spread of subgroup effects sqrt(sum w (c - ate)^2), alpha, and sample_stats["logp"]: put them in an xr.Dataset and call az.rhat / az.ess. Check that the last weight is negligible. For the heavy-tail comparison use pm.StudentT with a prior on nu, and compare the three models with az.compare.

3 Skeleton

T = 12
with pm.Model() as dp:
    alpha = pm.Gamma("alpha", ___, ___)
    w = pm.StickBreakingWeights("w", alpha=alpha, K=___)
    a = pm.Normal("a", 5.8, 1.0, shape=T)
    b = pm.Normal("b", 0.7, 0.5, shape=T)
    c = pm.Normal("c", 0.0, 0.5, shape=T)
    sigma = pm.InverseGamma("sigma", 3.0, 0.5, shape=T)
    g = pm.Normal("g", 0.0, 0.2, shape=Z.shape[1])
    mu = ___                                                    # (patients, T)
    ate = pm.Deterministic("ate", ___)
    pm.NormalMixture("y", w=w, mu=mu, sigma=sigma, observed=y)
    idata_dp = pm.sample(random_seed=RANDOM_SEED)

Self-checks in the notebook (h.check("task2", …)):

ate — posterior mean of the average treatment effect (sum of w_k c_k) on the LOG scale, about 0.19
elpd_gain — LOO elpd of the mixture minus that of the Task 1 normal regression, about 60

Task 3 · The colleague's responder list

The statistician on the team took your Task 2 posterior and computed, for every patient $i$, "their own treatment effect":

$$\widehat{\text{effect}}_i = \sum_k P(\text{patient } i \in \text{subgroup } k \mid \text{data})\; c_k,$$

averaged over posterior draws. The spread across patients is real, she says: rank them, give ddI to the top half (the "responders"), and spare the bottom half.

Deliver

Her numbers, for every patient in both arms.
What is wrong with using them to decide who gets ddI. There are two separate problems; show each with a number or a plot, not only an argument.
What your Task 2 model predicts as the treatment effect for a new patient walking into the clinic, as a function of their baseline covariates. Explain the result.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Write down what her probability conditions on. Which of those quantities does a clinician know on the day they decide the treatment? Then compute her score for the patients in the ZDV arm too - they never received ddI.

2 Approach

For each posterior draw, the membership probability is proportional to w_k * Normal(y_i | mu_ik, sigma_k); normalise over k (work on the log scale), multiply by c_k, sum over k and average over draws. Plot the score against the observed log(cd420 / cd40), colour by arm, and compute the correlation per arm.

For a new patient there is no y: the membership probabilities fall back to the prior ones. In this model, what are those, and do they depend on the covariates?

3 Skeleton

post = az.extract(idata_dp, var_names=["w", "a", "b", "c", "sigma", "g"], num_samples=500)
W, A, B, C, S = (post[v].transpose("sample", ...).to_numpy() for v in ["w", "a", "b", "c", "sigma"])
G = post["g"].transpose("sample", ...).to_numpy()
MU = A[:, None, :] + B[:, None, :] * lx[None, :, None] + ___ + (G @ Z.to_numpy().T)[:, :, None]
log_r = np.log(W)[:, None, :] + stats.norm.logpdf(y[None, :, None], MU, S[:, None, :])
resp = np.exp(log_r - log_r.max(-1, keepdims=True))
resp /= resp.sum(___, keepdims=True)
score = (resp * C[:, None, :]).sum(-1).mean(0)
np.corrcoef(score[t == 0], ___)[0, 1]
cate_new = (W * C).sum(-1)       # effect for a new patient: does it depend on anything?

Self-checks in the notebook (h.check("task3", …)):

corr_control — correlation, in the ZDV arm, between her score and the observed log(cd420 / cd40), about -0.7 (negative!)
corr_treated — the same correlation in the ZDV + ddI arm, about +0.5

Task 4 · Let the baseline profile choose the subgroup

Deliver

A mixture of regressions whose subgroup probabilities depend on baseline covariates (and on nothing measured after randomisation). Say which covariates, and why.
Evidence that the quantities you will base decisions on have converged - for every patient, not just on average - and that your conclusions do not depend on which chain you happen to look at.
A plot of each patient's conditional average treatment effect (CATE) against baseline CD4, with uncertainty, next to what Task 2 said.
A LOO comparison with Task 2. What does it tell you - and what can LOO not see here?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Task 3 showed that a treatment rule needs the prior membership to depend on what is known at baseline. How can the stick-breaking fractions become functions of covariates? And when a per-patient quantity is label-free, is it automatically well identified? Finally: how much of a patient's predictive density does the treatment effect move?

2 Approach

Logistic stick-breaking: v_ik = sigmoid(g0_k + x_i @ gamma_k) for k < T, weights w_ik = v_ik * prod_{l<k} (1 - v_il), the last component takes the rest. Build the (patients, T) weight matrix with pt.cumprod and pt.concatenate and pass it to pm.NormalMixture as w. The CATE of patient i is (w[i] * c).sum(): store it as a Deterministic and run az.rhat on it - all 1051 values, not only the mean. If some mix poorly, look at which patients, and compare the chains' average CATE within bands of baseline CD4. Keep the gate small: every covariate you add is another way to route thin corners of covariate space. For the last component, count per draw how many components carry more than 5% for any patient.

3 Skeleton

Xg = (lx / lx.std())[:, None]                     # gate covariates (baseline only!)
with pm.Model() as gated:
    ...                                           # a, b, c, sigma, g as in Task 2, T = 6
    g0 = pm.Normal("g0", 0.0, 1.5, shape=T - 1)
    gam = pm.Normal("gamma", 0.0, 1.0, shape=(Xg.shape[1], T - 1))
    v = pm.math.sigmoid(___)                      # (patients, T - 1)
    w = pt.concatenate([v, pt.ones((N, 1))], axis=1) * pt.concatenate(
        [pt.ones((N, 1)), pt.cumprod(___, axis=1)], axis=1)
    pm.Deterministic("cate", (w * c).sum(-1))
    pm.NormalMixture("y", w=w, mu=mu, sigma=sigma, observed=y)
    idata_gated = pm.sample(tune=1000, draws=2000, random_seed=RANDOM_SEED)

Self-checks in the notebook (h.check("task4", …)):

cate_low — mean CATE (log scale) over patients in the lowest quartile of baseline CD4 (<= 265 cells), about 0.26
cate_high — mean CATE (log scale) over patients in the highest quartile of baseline CD4 (>= 424 cells), about 0.16

Task 5 · The clinic's rule

Recall the rule: ddI is worth it for a patient if it raises their CD4 count at 20 weeks by at least 15%, with probability at least 0.8.

Deliver

For every patient, the probability that their expected gain exceeds 15% under your Task 4 model, and the share of patients the rule would put on ddI. Who is left out? Does the answer depend on the chain?
The same from two plain regressions: one with a treatment × baseline CD4 interaction, one with treatment × every baseline covariate. Where do the three disagree, and why?
Your recommendation to the clinic in a few sentences, including what the analysis cannot tell them.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

The rule is a statement about each patient's CATE: P(CATE_i > log 1.15) >= 0.8. When three models give three shares, ask separately where they disagree (which CD4 values) and why (functional form? number of modifiers?). Randomisation gives you a model-free referee within bands of baseline CD4.

2 Approach

From cate draws of shape (draws, patients): p = (cate > np.log(1.15)).mean(0), the share treated is (p >= 0.8).mean(); per chain, do the same on idata.posterior["cate"].mean("draw")-style slices. For the interaction regressions add c_cd4 * lx_std (and Z @ c_z) to the treatment coefficient, store the resulting CATE as a Deterministic, and apply the same rule. The referee: within quartiles of cd40 (pd.qcut), the difference in mean log(cd420 / cd40) between the arms, with its standard error.

3 Skeleton

p_mix = (cate_draws > np.log(1.15)).mean(0)
share_mix = (p_mix >= ___).mean()

with pm.Model(coords={"covariate": list(Z.columns)}) as interaction:
    ...                                            # Task 1 model, plus:
    c_cd4 = pm.Normal("c_cd4", 0.0, 0.2)
    cate = c + c_cd4 * lx / lx.std()               # + Z @ c_z for the "all" version
    cate = pm.Deterministic("cate", cate)
    mu = a + b * lx + ___ * t + Z.to_numpy() @ g

band = pd.qcut(df.cd40, 4)
g = pd.DataFrame({"change": change, "arm": t, "band": band}).groupby(["band", "arm"])["change"]
g.mean().unstack()[1] - g.mean().unstack()[0]      # model-free effect per band

Self-checks in the notebook (h.check("task5", …)):

share_mixture — share of the 1051 patients the rule puts on ddI under the CD4-gated mixture, about 0.88
share_interaction_cd4 — share under a regression with a treatment x log baseline CD4 interaction, about 0.69
share_interaction_all — share under a regression with treatment x every baseline covariate, about 0.5

h.progress()

C11: 0/17 hints used, 11/11 checks passed.

Reference solution

A complete, executed solution with commentary on every task — including the ways the obvious model fails. Try the hint ladder first.

Open the worked solution →

#The brief

#How this notebook works

#The data

#Task 0 · Look at the trial

#Task 1 · One regression for everybody

#Task 2 · A Dirichlet-process mixture of regressions

#Task 3 · The colleague's responder list

#Task 4 · Let the baseline profile choose the subgroup

#Task 5 · The clinic's rule

Reference solution

The brief

How this notebook works

The data

Task 0 · Look at the trial

Task 1 · One regression for everybody

Task 2 · A Dirichlet-process mixture of regressions

Task 3 · The colleague's responder list

Task 4 · Let the baseline profile choose the subgroup

Task 5 · The clinic's rule