Difference in Differences Mock in Python

Generate mock panel data with a known treatment effect and compute a difference-in-differences estimate using group and period means.

Easy Python 3.9+ Aug 9, 2026 A/B testing & experimentation 16 views 0 copies

Requires third-party packages — install first
pip install numpy pandas

Python code

26 lines
Python 3.9+
import numpy as np
import pandas as pd

# Generate mock panel data: 2 groups (control=0, treatment=1) × 2 periods (pre=0, post=1)
rng = np.random.default_rng(42)
n_per_cell = 50

data = []
for group in [0, 1]:
    for period in [0, 1]:
        # True effect: treatment increases outcome by 5 in the post period
        base = 10 + period * 2 + group * 3
        effect = 5 if (group == 1 and period == 1) else 0
        outcome = base + effect + rng.normal(0, 1, n_per_cell)
        for y in outcome:
            data.append({"group": group, "period": period, "outcome": y})

df = pd.DataFrame(data)

# Compute difference-in-differences manually
means = df.groupby(["group", "period"])["outcome"].mean()
did = (means[1, 1] - means[1, 0]) - (means[0, 1] - means[0, 0])

print(f"DiD estimate: {did:.3f}")
print("Cell means:")
print(means.unstack().round(2))

Output

stdout
DiD estimate: 5.031
Cell means:
period     0     1
group            
0       10.11 12.22
1       13.02 20.14

How it works

The script simulates two groups across two time periods with random noise around known means. The true effect is encoded as an additive 5.0 boost for the treatment group in the post period. Grouping with pandas and averaging within each cell produces empirical means, then the did estimate subtracts the control group's change from the treatment group's change. The formula (mean_trt_post - mean_trt_pre) - (mean_ctrl_post - mean_ctrl_pre) cancels out shared time trends and group baselines. Running with a fixed random seed gives reproducible output close to the true effect.

Common mistakes

  • Forgetting to set a random seed, which makes results non-reproducible across runs
  • Confusing period and group indexes when extracting cell means from the grouped object
  • Applying the treatment effect to the pre-period or mixing up which group receives the boost
  • Assuming the DiD estimate is exact — it is an estimate with sampling noise from the random draws

Variations

  1. Use `scipy.stats.linregress` or statsmodels OLS with interaction terms for a regression-based DiD
  2. Replace the manual loop with `np.repeat` and `np.concatenate` to build the data array faster
  3. Bootstrap the DiD estimate by resampling outcomes within each cell to get confidence intervals

Real-world use cases

  • Validating new pricing rollouts by comparing sales changes across regions with and without the change.
  • Measuring the impact of a new onboarding flow in an A/B test with pre- and post-measurements.
  • Estimating policy effects in company experiments where you compare treated stores against control stores over time.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from A/B testing & experimentation

Related tutorials and quizzes for this topic.