Difference-in-Differences

Estimate treatment effects by comparing changes over time between treated and control groups using panel mode.

Author

Benjamin Vincent

Difference-in-differences (DiD) estimates a treatment effect by comparing changes over time between a treated group and a control group. pathmc’s panel mode fits this naturally: the structural model captures group and time effects, and the do() operator estimates the average treatment effect on the treated (ATT).

The causal structure

We model an outcome Y as a function of treatment group membership (treated), a continuous time trend (time), and a treatment activation indicator (treat_post).

treated: 1 for the treatment group, 0 for control
time: the time period (continuous), capturing the underlying trend shared by both groups
treat_post: the DiD interaction term (1 only for treated units after treatment)

Including time as a continuous predictor lets the model capture the underlying trend, so the treatment effect is identified as a level shift relative to that trend — the standard DiD design.

DiD DAG: treatment effect identified by the interaction term.

Simulate panel data

We generate data for 6 units (3 treated, 3 control) observed over 20 time periods. Treatment starts at period 11, with a true ATT of 2.0.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import arviz as az
import pathmc

rng = np.random.default_rng(42)

units_treated = ["T1", "T2", "T3"]
units_control = ["C1", "C2", "C3"]
all_units = units_treated + units_control
n_periods = 20
treatment_start = 11
true_att = 2.0

true_unit_effects = {
    "T1": 3.0,
    "T2": 4.0,
    "T3": 5.0,
    "C1": 2.0,
    "C2": 3.5,
    "C3": 4.5,
}

rows = []
for unit in all_units:
    is_treated = 1 if unit in units_treated else 0
    for t in range(1, n_periods + 1):
        post = 1 if t >= treatment_start else 0
        treat_post = is_treated * post
        y = (
            true_unit_effects[unit]
            + 0.3 * t
            + true_att * treat_post
            + rng.normal(scale=0.5)
        )
        rows.append({
            "unit": unit,
            "time": t,
            "treated": is_treated,
            "post": post,
            "treat_post": treat_post,
            "Y": y,
        })

df = pd.DataFrame(rows)
df.head(10)

	unit	time	treated	Y
0	T1	1	1	3.452359
1	T1	2	1	3.080008
2	T1	3	1	4.275226
3	T1	4	1	4.670282
4	T1	5	1	3.524482
5	T1	6	1	4.148910
6	T1	7	1	5.163920
7	T1	8	1	5.241879
8	T1	9	1	5.691599
9	T1	10	1	5.573478

Visualise the raw data

fig, ax = plt.subplots(figsize=(8, 4))
for unit in units_treated:
    d = df[df["unit"] == unit]
    ax.plot(d["time"], d["Y"], color="steelblue", alpha=0.6)
for unit in units_control:
    d = df[df["unit"] == unit]
    ax.plot(d["time"], d["Y"], color="coral", alpha=0.6)

ax.axvline(treatment_start - 0.5, ls="--", color="gray", label="Treatment onset")
ax.plot([], [], color="steelblue", label="Treated")
ax.plot([], [], color="coral", label="Control")
ax.set_xlabel("Time")
ax.set_ylabel("Y")
ax.legend()
ax.set_title("Raw outcome trajectories")
plt.tight_layout()
plt.show()

Outcome by group over time. Dashed line marks treatment onset.

Specify and fit the model

We use panel mode with partial pooling (random intercepts per unit) to account for baseline differences. Partial pooling already drops the fixed intercept, and the formula starts with 0 + so equations() does not print a phantom 1 + intercept that is not in the design matrix.

spec = "Y ~ 0 + treated + time + treat_post"

model = pathmc.model(
    spec,
    data=df,
    panel={"unit": "unit", "time": "time"},
    pooling="partial",
)
model.graph()

model.equations()

\begin{aligned} \beta_{Y} &\sim \text{Normal}(mu=0,\, sigma=10) \\ \sigma_{Y} &\sim \text{HalfNormal}(sigma=1) \\ \mu_{alpha,Y} &\sim \text{Normal}(mu=0,\, sigma=10) \\ \sigma_{alpha,Y} &\sim \text{HalfNormal}(sigma=1) \\ \alpha_{Y} &\sim \text{Normal}(mu\_alpha,\, sigma\_alpha) \\[6pt] \mu_{Y} &= \mathrm{treated} + \mathrm{time} + \mathrm{treat\_post} \\ \mathrm{Y} &\sim \text{Normal}(\mu_{Y},\, \sigma_{Y}) \end{aligned}

Sample

idata = model.fit(random_seed=42)

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [sigma_alpha_Y, mu_alpha_Y, alpha_Y, beta_Y, sigma_Y]

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

Results

The coefficient on treat_post is our estimate of the ATT.

model.summary()

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
mu_alpha_Y	3.288160	0.700282	2.146121	4.362130	1222.554445	1402.361969	1.003010	0.020253	0.016864
alpha_Y[C1]	2.011804	0.123242	1.813356	2.204900	1842.377558	2193.894835	1.002203	0.002871	0.002062
alpha_Y[C2]	3.370686	0.122059	3.179851	3.561686	1703.033196	2283.975489	1.004648	0.002958	0.002150
alpha_Y[C3]	4.483509	0.121855	4.290219	4.678462	1635.169359	2178.162928	1.002388	0.003014	0.002180
alpha_Y[T1]	2.228662	0.961314	0.654428	3.665343	789.216770	806.673397	1.004365	0.035162	0.030232
...	...	...	...	...	...	...	...	...	...
mu_Y[115]	9.213263	0.099242	9.058081	9.369133	4177.852632	2950.913503	1.000939	0.001541	0.001115
mu_Y[116]	9.508872	0.102984	9.347725	9.672497	3820.129187	2813.522532	1.001174	0.001673	0.001210
mu_Y[117]	9.804482	0.107184	9.635940	9.974467	3492.167631	2843.485723	1.001429	0.001818	0.001315
mu_Y[118]	10.100091	0.111789	9.924148	10.279614	3199.560027	2874.634430	1.001693	0.001981	0.001434
mu_Y[119]	10.395701	0.116752	10.212934	10.583437	2956.108900	2856.317772	1.001938	0.002155	0.001560

132 rows × 9 columns

Posterior distribution of the ATT

The coefficient on treat_post is the ATT. We can visualise its full posterior distribution and compare it to the true value used in the simulation.

az.plot_dist(
    idata,
    var_names=["beta_Y"],
    coords={"Y_predictors": ["treat_post"]},
    ci_kind="hdi",
    ci_prob=0.94,
    visuals={"title": False, "point_estimate_text": False},
    figure_kwargs={"figsize": (8, 3)},
)
plt.gca().axvline(true_att, color="k", ls="--", lw=1.5)
plt.tight_layout()
plt.show()

Posterior distribution of the ATT. The dashed line marks the true value (2.0).

Causal effect via do()

We can also estimate the ATT using the do() operator by comparing outcomes with and without treatment.

r_treated = model.do(set={"treat_post": 1.0, "treated": 1.0, "time": 15.0})
r_control = model.do(set={"treat_post": 0.0, "treated": 1.0, "time": 15.0})

att = r_treated - r_control
att

DoResult — 4000 draws, 4 variables

variable	mean	94% HDI
treated	0.00	[0.00, 0.00]
time	0.00	[0.00, 0.00]
treat_post	1.00	[1.00, 1.00]
Y	2.08	[1.84, 2.32]

Methods: .draws() .mean() .hdi() .by_time() .plot() .dataset

The posterior mean should be close to the true ATT of 2.0.

Model predictions and counterfactual

With time as a continuous predictor, the model fits trend lines through each group. The counterfactual extrapolates the treated group’s pre-treatment trend forward — the gap between the counterfactual and the actual treated fit is the ATT.

stacked = idata.posterior.to_dataset().stack(sample=("chain", "draw"))
beta = stacked["beta_Y"]
alpha_draws = stacked["alpha_Y"]

# alpha_Y is the unit intercept; its mean is mu_alpha_Y. The formula uses
# 0 +, so there is no separate fixed Intercept to add.
b_treated = beta.sel(Y_predictors="treated").values
b_time = beta.sel(Y_predictors="time").values
b_tp = beta.sel(Y_predictors="treat_post").values


def group_mu(units, t, treat_post_val):
    """Posterior draws of group-averaged mu at a single time point."""
    preds = []
    for u in units:
        alpha_u = alpha_draws.sel(unit=u).values
        is_treated = 1.0 if u in units_treated else 0.0
        mu = alpha_u + b_treated * is_treated + b_time * t + b_tp * treat_post_val
        preds.append(mu)
    return np.mean(preds, axis=0)


times = np.arange(1, n_periods + 1, dtype=float)

treated_fit = np.array([
    group_mu(units_treated, t, 1.0 if t >= treatment_start else 0.0).mean()
    for t in times
])
treated_cf = np.array([group_mu(units_treated, t, 0.0).mean() for t in times])
control_fit = np.array([group_mu(units_control, t, 0.0).mean() for t in times])

fig, ax = plt.subplots(figsize=(9, 5))

treated_obs = df[df["treated"] == 1].groupby("time")["Y"].mean()
control_obs = df[df["treated"] == 0].groupby("time")["Y"].mean()
ax.scatter(
    treated_obs.index,
    treated_obs.values,
    color="steelblue",
    alpha=0.4,
    s=25,
    zorder=5,
    label="Treated (observed)",
)
ax.scatter(
    control_obs.index,
    control_obs.values,
    color="coral",
    alpha=0.4,
    s=25,
    zorder=5,
    label="Control (observed)",
)

ax.plot(times, treated_fit, color="steelblue", lw=2.5, label="Treated (fitted)")
ax.plot(times, control_fit, color="coral", lw=2.5, label="Control (fitted)")
ax.plot(
    times,
    treated_cf,
    color="steelblue",
    ls="--",
    lw=2,
    label="Treated (counterfactual)",
)

post_mask = times >= treatment_start
ax.fill_between(
    times[post_mask],
    treated_cf[post_mask],
    treated_fit[post_mask],
    color="steelblue",
    alpha=0.15,
    label="Causal effect (ATT)",
)

ax.axvline(treatment_start - 0.5, ls=":", color="gray", alpha=0.7)
ax.annotate(
    "Treatment onset",
    xy=(treatment_start - 0.5, ax.get_ylim()[1]),
    fontsize=9,
    color="gray",
    ha="right",
    xytext=(treatment_start - 1, ax.get_ylim()[1] - 0.3),
)

ax.set_xlabel("Time")
ax.set_ylabel("Y")
ax.set_title("Difference-in-Differences: observed, fitted, and counterfactual")
ax.legend(loc="upper left", fontsize=8, framealpha=0.9)
plt.tight_layout()
plt.show()

DiD results: observed group means, model trend lines, and counterfactual for the treated group.

The solid lines show the model’s fitted trends for each group. Pre-treatment, the treated fit and counterfactual overlap — confirming parallel trends. Post-treatment, the dashed line continues the treated group’s trend without treatment, and the shaded gap is the estimated ATT (\approx 2.0).