Panel DID simulation with true value (genriesz)
We implement DID as ATT on the differenced outcome
\[\Delta Y = Y_1 - Y_0,\]
where:
\(Y0\) is the pre-period outcome,
\(Y1\) is the post-period outcome,
the same units are observed in both periods (panel),
\(D\) is a binary treatment indicator (treatment happens in the post period).
With a standard panel DID setup:
\[Y_{0} = \mu(Z) + u + \varepsilon_0, \qquad
Y_{1} = \mu(Z) + \text{trend}(Z) + u + \tau D + \varepsilon_1,\]
the DID effect equals the constant treatment effect \(\tau\), provided the parallel trends condition holds after conditioning on \(Z\).
This notebook:
simulates a large population to compute an approximate “true” DID effect,
samples a dataset and calls
genriesz.grr_did(X, Y0=..., Y1=...).
[1]:
import numpy as np
from genriesz import (
grr_did,
SquaredGenerator,
UKLGenerator,
BPGenerator,
PolynomialBasis,
TreatmentInteractionBasis,
RBFRandomFourierBasis,
KNNCatchmentBasis,
)
rng = np.random.default_rng(0)
DGP
[7]:
def draw_panel(n: int, d_z: int, tau: float, seed: int = 0):
rng = np.random.default_rng(seed)
Z = rng.normal(size=(n, d_z))
logits = 0.6 * Z[:, 0] - 0.25 * Z[:, 1]
e = 1.0 / (1.0 + np.exp(-logits))
D = rng.binomial(1, e, size=n).astype(int)
mu = 0.5 * Z[:, 0] - 0.2 * Z[:, 1] ** 2
trend = 0.5 + 0.1 * Z[:, 0] # common trend that depends on Z
u = rng.normal(scale=1.0, size=n) # unit fixed effect
Y0 = mu + u + rng.normal(scale=1.0, size=n)
Y1 = mu + trend + u + tau * D + rng.normal(scale=1.0, size=n)
X = np.column_stack([D.astype(float), Z])
return X, Y0, Y1, D
tau_true = 1.0
# Large population for an approximate truth
X_pop, Y0_pop, Y1_pop, D_pop = draw_panel(n=200_000, d_z=5, tau=tau_true, seed=1)
true_did = np.mean((Y1_pop - Y0_pop)[D_pop == 1]) - np.mean((Y1_pop - Y0_pop)[D_pop == 0]) # naive DID
# Our target here is "ATT on ΔY", whose true value equals tau_true by construction.
print("True tau (by construction):", tau_true)
print("Naive DID (difference in mean ΔY):", true_did)
True tau (by construction): 1.0
Naive DID (difference in mean ΔY): 1.0477956893531037
Example 1: Polynomial basis + treatment interactions
[9]:
# Sample a dataset from the same DGP
X, Y0, Y1, D = draw_panel(n=6000, d_z=5, tau=tau_true, seed=0)
psi = PolynomialBasis(degree=2, include_bias=True)
phi = TreatmentInteractionBasis(base_basis=psi)
gen = SquaredGenerator(C=0.0).as_generator()
res = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=phi,
generator=gen,
cross_fit=True,
folds=5,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
riesz_penalty="l2",
riesz_lam=1e-3,
max_iter=300,
tol=1e-8,
)
print(res.summary_text())
DID estimates (n=6000)
alpha=0.05 | null=0.0
diagnostics: max_abs_smd_unweighted=0.5468841195092563, max_abs_smd_weighted=0.002194706023289376, ess_treated=3017.871638373246, ess_control=2036.5721815918366
Estimator Estimate SE CI p-value
---------------------------------------------------------------------------------
RA 0.983907 0.0129607 [ 0.958505, 1.00931] 0
RW 0.976091 0.0495386 [ 0.878997, 1.07318] 0
ARW 0.98479 0.0421389 [ 0.9022, 1.06738] 0
TMLE 0.984776 0.0421433 [ 0.902176, 1.06738] 0
Example 2: RKHS basis (RBF random Fourier features)
[ ]:
psi_rff = RBFRandomFourierBasis(
n_features=500,
sigma=1.0,
standardize=True,
random_state=0,
)
phi_rff = TreatmentInteractionBasis(base_basis=psi_rff)
res_phi_rff = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=phi_rff,
generator=gen,
cross_fit=True,
folds=5,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
riesz_penalty="l2",
riesz_lam=1e-3,
max_iter=300,
tol=1e-8,
)
print(res_phi_rff.summary_text())
Example 3: KNN catchment basis (nearest-neighbor matching)
Nearest-neighbor matching as a special case of squared-loss Riesz regression.
[ ]:
basis_knn = KNNCatchmentBasis(n_neighbors=5, include_bias=False)
phi_knn = TreatmentInteractionBasis(base_basis=basis_knn)
res_phi_knn = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=phi_knn,
generator=gen,
cross_fit=True,
folds=5,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
riesz_penalty="l2",
riesz_lam=1e-3,
max_iter=300,
tol=1e-8,
)
print(res_phi_knn.summary_text())
Example 4: Random forest leaf basis (optional)
[ ]:
from sklearn.ensemble import RandomForestRegressor
from genriesz.sklearn_basis import RandomForestLeafBasis
rf = RandomForestRegressor(n_estimators=200, max_depth=6, random_state=0)
leaf_basis = RandomForestLeafBasis(rf).fit(X, Y1 - Y0)
res_leaf_basis = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=leaf_basis,
generator=gen,
cross_fit=True,
folds=5,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
riesz_penalty="l2",
riesz_lam=1e-3,
max_iter=300,
tol=1e-8,
)
print(res_leaf_basis.summary_text())
Example 5: Neural network embedding basis (optional)
[ ]:
import torch
from genriesz.torch_basis import MLPEmbeddingNet, TorchEmbeddingBasis
torch.manual_seed(0)
net = MLPEmbeddingNet(input_dim=X.shape[1], hidden_dims=(64,), output_dim=32)
nn_basis = TorchEmbeddingBasis(net, include_bias=True, device="cpu")
res_nn_basis = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=nn_basis,
generator=gen,
cross_fit=True,
folds=5,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
riesz_penalty="l2",
riesz_lam=1e-3,
max_iter=300,
tol=1e-8,
)
print(res_nn_basis.summary_text())
Generator / regularization sweep (SQ / UKL / BP)
We repeat the DID estimation (implemented as ATT on the differenced outcome) under SQ-Riesz / UKL-Riesz / BP-Riesz, multiple regularization norms, and multiple regularization strengths.
For UKL/BP we set a branch function to match the treatment/control sign pattern. As with ATT, the KL-type generators use a zero shift (\(C=0\)), not the ATE default \(C=1\): the DID (ATT-on-\(\Delta Y\)) control-branch Riesz representer can have magnitude below one, so \(C=1\) excludes part of its range and can drive the fit toward the generator boundary. If the internal clip binds, the fitted representer targets a modified, clipped estimand and weights can be extreme.
[ ]:
branch = lambda x: int(x[0] == 1.0)
generator_grid = [
("SQ", SquaredGenerator(C=0.0).as_generator()),
("UKL (C=0)", UKLGenerator(C=0.0, branch_fn=branch).as_generator()),
("BP (omega=0.1, C=0)", BPGenerator(C=0.0, omega=0.1, branch_fn=branch).as_generator()),
("BP (omega=0.2, C=0)", BPGenerator(C=0.0, omega=0.2, branch_fn=branch).as_generator()),
("BP (omega=0.5, C=0)", BPGenerator(C=0.0, omega=0.5, branch_fn=branch).as_generator()),
]
penalty_grid = [
{"penalty": "l2", "lam": 1e-4, "p_norm": None},
{"penalty": "l2", "lam": 1e-3, "p_norm": None},
{"penalty": "l1", "lam": 1e-4, "p_norm": None},
{"penalty": "lp", "lam": 1e-3, "p_norm": 1.5},
]
rows = []
for gname, gen_i in generator_grid:
for cfg in penalty_grid:
res_i = grr_did(
X=X,
Y0=Y0,
Y1=Y1,
basis=phi,
generator=gen_i,
cross_fit=True,
folds=3,
random_state=0,
estimators=("ra", "rw", "arw", "tmle"),
outcome_models="shared",
outcome_link="identity",
riesz_penalty=cfg["penalty"],
riesz_lam=cfg["lam"],
riesz_p_norm=cfg.get("p_norm"),
max_iter=250,
tol=1e-8,
)
row = {
"generator": gname,
"penalty": cfg["penalty"],
"lam": cfg["lam"],
}
for k in ("ra", "rw", "arw", "tmle"):
e = res_i.estimates[k]
row[f"{k}"] = e.estimate
row[f"{k}_se"] = e.se
row[f"{k}_err"] = e.estimate - tau_true
rows.append(row)
import pandas as pd
df = pd.DataFrame(rows)
df = df.sort_values(by="arw_err", key=lambda s: np.abs(s))
display(df)
[ ]: