C09 · Does job training work? — Bayesian Modelling in Practice

The brief

A labour ministry ran a job-training programme for disadvantaged workers in the mid 1970s and wants to know: did it raise participants' earnings, and by enough to justify the cost? The ministry's analysts have only what ministries usually have - records of the participants, and a large national survey of people who did not take part. Their first comparison says participants earn $8,500 a year less than non-participants, and someone has already drafted the memo closing the programme.

You hold something they do not: the programme was in fact run as a randomised experiment, so you know the answer they should get. That makes this a rare chance to grade observational causal inference against the truth - the test Robert LaLonde set in 1986, which most methods of the day failed.

Your job: establish the experimental benchmark, then put yourself in the ministry's shoes, forget the experimental controls, and see how close a careful Bayesian analysis of the observational data can get - and how you would know, without the benchmark, whether to believe it.

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("task3", att_full=...) 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 warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
from scipy.special import expit

from pymc_challenges import Hints, data

RANDOM_SEED = 1986
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")

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

C09 - Job training: can observational data stand in for an experiment?

task0 - Look at the experiment (2 hints)
task1 - The experimental benchmark (3 hints, 2 check(s))
task2 - The ministry's dataset and the naive comparison (3 hints, 2 check(s))
task3 - DAG, generative model, pm.observe and pm.do (3 hints, 1 check(s))
task4 - Overlap, propensity scores and trimming (3 hints, 2 check(s))
task5 - Sensitivity to thresholds and adjustment sets (3 hints, 1 check(s))
task6 - What do you tell the ministry? (3 hints, 1 check(s))

The data

nsw is the experiment: men who applied to the National Supported Work programme and were randomly assigned to training (treat = 1) or not. cps is a national sample of men from the Current Population Survey, none of whom were in the programme (treat = 0).

column	meaning
`treat`	1 = received NSW training
`age`, `educ`	age and years of schooling
`black`, `hisp`, `marr`, `nodegree`	indicators: Black, Hispanic, married, no high-school degree
`re74`, `re75`	real earnings in 1974 and 1975 (USD) - before the programme
`re78`	real earnings in 1978 (USD) - the outcome

data.describe("nsw")
nsw = data.load("nsw")
nsw.head()

nsw: Randomised job-training experiment: 185 treated, 260 controls, earnings in 1978.
  source: LaLonde (1986) / Dehejia & Wahba (1999), National Supported Work experiment.
  url:    https://vincentarelbundock.github.io/Rdatasets/csv/causaldata/nsw_mixtape.csv

	data_id	treat	age	educ	black	hisp	marr	nodegree	re74	re75	re78
rownames
1	Dehejia-Wahba Sample	1	37	11	1	0	1	1	0.0	0.0	9930.045898
2	Dehejia-Wahba Sample	1	22	9	0	1	0	1	0.0	0.0	3595.894043
3	Dehejia-Wahba Sample	1	30	12	1	0	0	0	0.0	0.0	24909.449219
4	Dehejia-Wahba Sample	1	27	11	1	0	0	1	0.0	0.0	7506.145996
5	Dehejia-Wahba Sample	1	33	8	1	0	0	1	0.0	0.0	289.789886

data.describe("cps")
cps = data.load("cps")
cps.describe().round(1)

cps: 15,992 non-experimental 'controls' with the same covariates as `nsw`.
  source: Current Population Survey comparison sample used by LaLonde (1986).
  url:    https://vincentarelbundock.github.io/Rdatasets/csv/causaldata/cps_mixtape.csv

	treat	age	educ	black	hisp	marr	nodegree	re74	re75	re78
count	15992.0	15992.0	15992.0	15992.0	15992.0	15992.0	15992.0	15992.0	15992.0	15992.0
mean	0.0	33.2	12.0	0.1	0.1	0.7	0.3	14016.8	13650.8	14846.7
std	0.0	11.0	2.9	0.3	0.3	0.5	0.5	9569.8	9270.4	9647.4
min	0.0	16.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
25%	0.0	24.0	11.0	0.0	0.0	0.0	0.0	4403.5	4398.8	5669.3
50%	0.0	31.0	12.0	0.0	0.0	1.0	0.0	15123.6	14557.1	16422.0
75%	0.0	42.0	13.0	0.0	0.0	1.0	1.0	23584.2	22923.7	25564.7
max	0.0	55.0	18.0	1.0	1.0	1.0	1.0	25862.3	25243.6	25564.7

Task 0 · Look at the experiment

Deliver

A plot of the distribution of re78 in each experimental arm, and the numbers a likelihood would have to reproduce: share of exact zeros, median, mean, maximum.
A table showing whether randomisation balanced the pre-treatment covariates.
One sentence: which textbook likelihoods are ruled out for re78, and why?

Your turn Work this task in the notebook.

Task 1 · The experimental benchmark

Because of randomisation, the difference between the arms is the causal effect of training for men like these. Estimate it properly.

Deliver

A Bayesian model of re78 by arm with the most obvious likelihood, and a posterior predictive check that shows specifically what it gets wrong.
A better likelihood, with priors you can defend on the dollar scale, clean diagnostics, and a posterior predictive check that passes where the first failed.
The posterior of the average treatment effect in dollars (difference in mean earnings, zeros included), with a 94% interval and the probability that it exceeds $1,000 a year. This is the benchmark for everything that follows.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

A Normal likelihood will happily estimate two means. Ask its posterior predictive two questions: what share of replicated men earn exactly zero, and what share earn a negative amount? Then think of earnings as two outcomes - does he work at all, and how much if he does - and look for a distribution built that way. The quantity the ministry cares about is mean earnings including the zeros.

2 Approach

pm.HurdleGamma("re78", psi=..., alpha=..., beta=..., observed=...) is a point mass at zero with probability 1 - psi and a Gamma otherwise (pm.HurdleLogNormal also exists). Give psi and the Gamma mean an arm dimension and index them with the treatment indicator. With shape alpha and mean mu, the rate is beta = alpha / mu.

Mean earnings in an arm are psi * mu; wrap the difference between arms in a pm.Deterministic("ate", ...). Work in $1000s so that priors are easy to state: a full-time minimum-wage job paid about $5,500 in 1978.

For the PPC, az.extract(idata, group="posterior_predictive", var_names="re78", num_samples=200) gives replicated datasets; compare np.mean(rep == 0) and np.mean(rep < 0) with the data.

3 Skeleton

t = nsw.treat.values
y = nsw.re78.values / 1000
with pm.Model(coords={"arm": ["control", "treated"]}) as hurdle_model:
    psi = pm.Beta("psi", ___, ___, dims="arm")
    mu_pos = pm.Gamma("mu_pos", mu=___, sigma=___, dims="arm")
    shape = pm.Gamma("shape", mu=1.5, sigma=1)
    mean = pm.Deterministic("mean", ___, dims="arm")        # zeros included
    pm.Deterministic("ate", mean[1] - mean[0])
    pm.HurdleGamma("re78", psi=psi[t], alpha=shape, beta=___, observed=y)
    idata = pm.sample(random_seed=RANDOM_SEED)
    pm.sample_posterior_predictive(idata, extend_inferencedata=True, random_seed=RANDOM_SEED)

ate = idata.posterior["ate"]
float((ate > 1).mean())          # P(ATE > $1,000)

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

ate — posterior mean of the experimental ATE in DOLLARS, roughly 1,800 (did you convert from $1000s?)
p_ate_gt_1000 — posterior probability that the experimental ATE exceeds $1,000, roughly 0.85-0.9

Task 2 · The ministry's dataset

Now forget the experimental controls. The ministry has the 185 participants and the CPS.

Deliver

The observational dataset (NSW treated + all CPS men) and the naive estimate: the difference in mean 1978 earnings between participants and non-participants.
A covariate balance table - means by group and standardised mean differences - side by side with the same table for the experiment.
A short paragraph for the ministry explaining why the naive number says nothing about the programme. Look closely at the CPS earnings columns too: is there anything about how they were recorded that a model should know?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Stack the NSW treated rows on top of the CPS rows - the experimental controls are not available to the ministry. A raw difference in means compares two groups; before believing it, ask whether the groups were comparable before the programme. A difference in means is hard to judge across variables with different units - divide by a standard deviation.

2 Approach

pd.concat([nsw[nsw.treat == 1], cps], ignore_index=True). The standardised mean difference is (mean_treated - mean_control) / sqrt((var_treated + var_control) / 2); values beyond about 0.1 are usually flagged. Compute it for treated vs experimental controls and treated vs CPS and put the two columns side by side.

For the recording issue: look at cps.re78.max() and how many men sit exactly on it.

3 Skeleton

treated = nsw[nsw.treat == 1]
obs = pd.concat([treated, cps], ignore_index=True)
naive = obs.re78[obs.treat == 1].mean() - ___

def smd(a, b):
    return (a.mean() - b.mean()) / np.sqrt(___)

balance = pd.DataFrame({
    "SMD vs NSW controls": [smd(treated[c], nsw[nsw.treat == 0][c]) for c in covariates],
    "SMD vs CPS": [smd(treated[c], ___) for c in covariates],
}, index=covariates)

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

naive_diff — mean re78 of participants minus mean re78 of CPS men, in dollars - about -8,500
smd_black — standardised mean difference of black, participants vs CPS, about 2.4

Task 3 · A causal model, conditioned and intervened on

Deliver

The causal DAG you are assuming, drawn in a markdown cell, with the adjustment set it implies and the assumptions that are not testable from these data.
One generative PyMC model of the observational data - treatment assignment and 1978 earnings given the covariates - written without any observed data. Condition it on the data with pm.observe, fit it, and then use pm.do to set treat to 0 and to 1 for the participants. From those two interventions compute the posterior of the ATT (average effect on the 185 participants) by g-computation.
Compare with the benchmark. How far off is it, and in which direction?

Use all 16,177 rows here. Sampling takes around a minute; storing a 16,177-long deterministic for every draw is what you need to avoid.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Which variables were fixed before anyone was assigned to training, and could affect both selection into the programme and later earnings? Those are the back-door paths you need to block. Which assumption can you never check from these data alone?

For the model: write down how the data could have been generated - first who gets treated, then what they earn - without telling PyMC what was observed. Conditioning and intervening are then operations on that model, not new models.

2 Approach

Build the model with treat = pm.Bernoulli("treat", logit_p=...) and the earnings likelihood depending on treat, neither with observed=. Then

pm.observe(model, {"treat": t, "re78": y}) returns a new model with those variables observed - sample it as usual;
pm.do(model, {"treat": np.ones(n, dtype=int)}) returns a new model in which treat is a constant. pm.sample_posterior_predictive(idata, var_names=[...]) inside it pushes the fitted posterior through the intervened graph.

For the ATT you only need potential outcomes for the 185 participants. If the model is built by a function of a data frame, you can build it again on just those rows and apply pm.do there. Put E[re78 | treat, X] = psi * mu in a pm.Deterministic and average do(1) - do(0) over participants within each posterior draw.

With 16,177 rows, do not store that deterministic while sampling: pm.sample(var_names=[...parameter names...]) keeps only what you list. nuts_sampler="nutpie" helps.

3 Skeleton

def generative(df, names):
    X = features(df, names).values                 # your scaled design matrix
    with pm.Model(coords={"obs": df.index.values, "cov": names}) as gen:
        g0 = pm.Normal("g0", 0, 3); g = pm.Normal("g", 0, 2, dims="cov")
        treat = pm.Bernoulli("treat", logit_p=___, dims="obs")
        # hurdle part: P(re78 > 0)
        psi = pm.math.sigmoid(a_z + tau_z * treat + pm.math.dot(X, b_z))
        # level part: mean of positive earnings
        mu = pm.math.exp(___)
        pm.Deterministic("expected", psi * mu, dims="obs")
        pm.HurdleGamma("re78", psi=psi, alpha=shape, beta=shape / mu, dims="obs")
    return gen

with pm.observe(generative(obs, names), {"treat": ___, "re78": ___}):
    idata = pm.sample(var_names=PARAMETERS, nuts_sampler="nutpie", random_seed=RANDOM_SEED)

def expected_under(value):
    with pm.do(generative(treated, names), {"treat": np.full(len(treated), value)}):
        pp = pm.sample_posterior_predictive(idata, var_names=["expected"])
    return pp.posterior_predictive["expected"]

att = (expected_under(1) - expected_under(0)).mean("obs")     # dims: chain, draw

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

att_full — posterior mean ATT in dollars from outcome regression on ALL CPS men - far above the naive -8,500 but well short of the benchmark (somewhere between about -1,000 and +1,300 depending on your functional form)

Task 4 · Where do the counterfactuals come from?

The ATT needs, for each participant, a prediction of what he would have earned without training. Those predictions are only as good as the non-participants who resemble him.

Deliver

A diagnosis: estimate every man's probability of being a participant given his covariates (the propensity score - you may already have it) and compare its distribution in the two groups. What share of the CPS has a propensity below 1%?
Evidence, not using the experimental controls, that the Task 3 outcome model is wrong in the region of covariate space where the participants live.
A fix, a refit with clean diagnostics, and the new ATT next to the benchmark. State exactly which population your new estimate refers to.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

For a typical participant - young, Black, unmarried, no earnings in 1975 - how many CPS men are there like him? If the answer is "a handful out of sixteen thousand", who is determining the regression coefficients that produce his counterfactual?

And a check you can run without any experiment: does the fitted outcome model predict the 1978 earnings of the CPS men who most resemble participants?

2 Approach

The propensity score is P(treat = 1 | X). If your generative model has a treatment equation you have already estimated it: push the posterior of its coefficients through scipy.special.expit in NumPy (a few hundred draws are plenty) and average. Plot it on the logit scale with a log count axis - on the raw scale the CPS is one spike at 0.

A standard fix is to restrict the comparison group to the region of common support: drop CPS men whose propensity is below a threshold (0.05, say). Keeping all participants keeps the estimand equal to the ATT for the 185 men. Refit the same generative model on the trimmed data. Richer functional forms (age squared, indicators for zero earnings in 1974/75) help both the propensity and the outcome model.

3 Skeleton

post = az.extract(idata, var_names=["g0", "g"], num_samples=400)
logit = post["g0"].values[None, :] + X @ post["g"].transpose("cov", "sample").values
obs["ps"] = expit(logit).mean(axis=1)

share_below = np.mean(obs.ps[obs.treat == 0] < 0.01)

lookalikes = obs[(obs.ps >= 0.05) & (obs.treat == 0)]
# model-predicted mean under do(treat=0) for these rows vs lookalikes.re78.mean()

trim_df = obs[(obs.ps >= ___) | (obs.treat == 1)]
trim_idata = fit(trim_df, names)
att_trim = att(trim_idata, names)

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

cps_share_below_1pct — share of CPS men whose estimated propensity score is below 0.01 - about nine in ten
att_trimmed — posterior mean ATT in dollars after restricting to common support - positive, of the order of 1,000 to 1,700

Task 5 · How fragile is the answer?

You made choices in Task 4 that the ministry cannot audit. Show what they are worth.

Deliver a table and a plot of the ATT posterior under

at least three trimming thresholds, and
at least three adjustment sets on your preferred trimmed sample: all covariates, without re74, and without any pre-programme earnings.

Then answer: which choice matters most, and what does that say about the DAG in Task 3?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Two kinds of choice went into your estimate: who is in the comparison group and what you adjust for. Vary each separately and see which one moves the answer by more than its posterior standard deviation.

2 Approach

Loop over (threshold, covariate list) pairs, refit the same generative model and recompute the ATT - the trimmed fits take seconds. Collect mean, sd, a 94% interval and P(ATT > 1) in a DataFrame and draw it as a horizontal errorbar plot with the benchmark on top. Remember to remove the matching zero-earnings indicators when you drop re74 / re75.

3 Skeleton

specs = {
    "ps >= 0.01": (0.01, FLEX),
    "ps >= 0.10": (0.10, FLEX),
    "no re74": (0.05, [c for c in FLEX if c not in ("re74", "u74")]),
    "no earnings": (0.05, [c for c in FLEX if c not in (___)]),
}
out = {}
for label, (thr, names) in specs.items():
    df = obs[(obs.ps >= thr) | (obs.treat == 1)]
    out[label] = att(fit(df, names), names)

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

att_no_earnings — posterior mean ATT in dollars on the trimmed sample when neither re74 nor re75 (nor their zero indicators) is adjusted for - clearly lower than with them

Task 6 · What do you tell the ministry?

The ministry's economists say the programme pays for itself if it raises participants' earnings by more than $1,000 a year (an illustrative threshold).

Deliver

$P(\text{effect} > \$1{,}000)$ under the experiment, under the naive full-sample adjustment, and under your preferred observational analysis.
A short note to the ministry: what the observational data can and cannot support, what you would need to see before trusting an observational answer when there is no experiment to check it against, and whether the memo closing the programme should go out.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

A posterior probability answers "how sure are we that the effect clears the bar if the model's assumptions hold". Which assumptions differ between the experiment and the observational analyses, and which of them are reflected in the width of the posterior?

2 Approach

(att > 1).mean() on each posterior (in $1000s) gives the probabilities. For the note, separate statistical uncertainty (in the posterior) from specification uncertainty (the spread across Task 5's rows) and identification uncertainty (unmeasured confounding, different data sources - in no posterior at all). Without an experiment, the things you can show are: overlap, balance after trimming, predictive checks on look-alike controls, stability across reasonable specifications, and placebo outcomes (e.g. "effect" of training on 1975 earnings should be zero).

3 Skeleton

table = pd.DataFrame({
    "experiment": [float((ate_exp > 1).mean())],
    "full CPS": [float((att_full > 1).mean())],
    "trimmed": [float((___ > 1).mean())],
}, index=["P(effect > $1,000)"])

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

p_observational — P(ATT > $1,000) from the trimmed observational analysis - around a coin flip to two in three, lower than the experiment's ~0.88

Going further

Respect the top-coding. A quarter of CPS men have re78 recorded at the cap. Model that explicitly with a censored likelihood for positive earnings (you will have to split the hurdle into its two parts to do it). Does the full-sample estimate move? Does the trimmed one?
Weight instead of trim. Use the propensity scores to weight CPS men by $e(x) / (1 - e(x))$ in the outcome likelihood (pm.Potential), or combine weighting and outcome modelling into a doubly-robust estimate. Propagate the uncertainty in $e(x)$.
A more flexible response surface. Replace the linear predictors by an HSGP or BART-like additive model in re74, re75 and age. Does flexibility substitute for trimming, or only hide the extrapolation?
Unmeasured confounding. Add a sensitivity parameter: a latent binary $U$ that raises both the odds of participation and earnings. How strong must it be to push $P(\text{effect} > \$1{,}000)$ below one half?

h.progress()

C09: 0/20 hints used, 9/9 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 experiment

#Task 1 · The experimental benchmark

#Task 2 · The ministry's dataset

#Task 3 · A causal model, conditioned and intervened on

#Task 4 · Where do the counterfactuals come from?

#Task 5 · How fragile is the answer?

#Task 6 · What do you tell the ministry?

#Going further

Reference solution

The brief

How this notebook works

The data

Task 0 · Look at the experiment

Task 1 · The experimental benchmark

Task 2 · The ministry's dataset

Task 3 · A causal model, conditioned and intervened on

Task 4 · Where do the counterfactuals come from?

Task 5 · How fragile is the answer?

Task 6 · What do you tell the ministry?

Going further