diff --git a/tests/conftest.py b/tests/conftest.py
index 32bbfe75..06c54118 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -19,7 +19,6 @@
import pytest
-
# =============================================================================
# R Availability Fixtures (Lazy Loading)
# =============================================================================
diff --git a/tests/test_aliases.py b/tests/test_aliases.py
index a6d51112..a1eb6621 100644
--- a/tests/test_aliases.py
+++ b/tests/test_aliases.py
@@ -22,8 +22,18 @@ def test_alias_identity():
def test_aliases_in_all():
"""All aliases are listed in __all__."""
aliases = [
- "DiD", "TWFE", "EventStudy", "SDiD", "CS", "CDiD",
- "SA", "BJS", "Gardner", "DDD", "Stacked", "Bacon",
+ "DiD",
+ "TWFE",
+ "EventStudy",
+ "SDiD",
+ "CS",
+ "CDiD",
+ "SA",
+ "BJS",
+ "Gardner",
+ "DDD",
+ "Stacked",
+ "Bacon",
]
for alias in aliases:
assert alias in diff_diff.__all__, f"{alias} missing from __all__"
diff --git a/tests/test_bandwidth_selector.py b/tests/test_bandwidth_selector.py
index 6326e93c..197e59a0 100644
--- a/tests/test_bandwidth_selector.py
+++ b/tests/test_bandwidth_selector.py
@@ -37,8 +37,7 @@ def golden():
not copied alongside tests/)."""
if not GOLDEN_PATH.exists():
pytest.skip(
- "Golden values file not found; "
- "run: Rscript benchmarks/R/generate_nprobust_golden.R"
+ "Golden values file not found; " "run: Rscript benchmarks/R/generate_nprobust_golden.R"
)
with GOLDEN_PATH.open() as f:
return json.load(f)
@@ -139,9 +138,9 @@ def test_R_parity(self, dgp_case, stage):
if expected == 0:
assert actual == pytest.approx(0, abs=1e-10), f"{name} {stage}"
else:
- assert actual == pytest.approx(expected, rel=_PARITY_TOL), (
- f"{name} {stage}: py={actual!r} R={expected!r}"
- )
+ assert actual == pytest.approx(
+ expected, rel=_PARITY_TOL
+ ), f"{name} {stage}: py={actual!r} R={expected!r}"
# =============================================================================
@@ -257,9 +256,7 @@ def test_public_wrapper_fixes_vce_nn_nnmatch_3(self):
G = 2000
d = rng.uniform(0, 1, size=G)
y = d + d**2 + rng.normal(0, 0.5, size=G)
- via_wrapper = mse_optimal_bandwidth(
- d, y, kernel="epanechnikov", return_diagnostics=True
- )
+ via_wrapper = mse_optimal_bandwidth(d, y, kernel="epanechnikov", return_diagnostics=True)
via_port_nn = lpbwselect_mse_dpi(
y,
d,
@@ -477,9 +474,7 @@ def test_wrapper_rank_deficient_raises_valueerror(self):
# 0.3, 0.4). The B1 design matrix has o_B+1 = 5 columns but
# only 4 independent rows -> rank-deficient X'X -> Cholesky
# fails in qrXXinv.
- d = np.concatenate(
- [[0.1], np.full(50, 0.15), np.full(50, 0.3), np.full(50, 0.4)]
- )
+ d = np.concatenate([[0.1], np.full(50, 0.15), np.full(50, 0.3), np.full(50, 0.4)])
y = d + rng.normal(0, 0.01, size=d.size)
with pytest.raises(ValueError, match="qrXXinv"):
mse_optimal_bandwidth(d, y, boundary=float(d.min()))
diff --git a/tests/test_bias_corrected_lprobust.py b/tests/test_bias_corrected_lprobust.py
index 404fb299..5e46933d 100644
--- a/tests/test_bias_corrected_lprobust.py
+++ b/tests/test_bias_corrected_lprobust.py
@@ -22,10 +22,7 @@
)
GOLDEN_PATH = (
- Path(__file__).resolve().parents[1]
- / "benchmarks"
- / "data"
- / "nprobust_lprobust_golden.json"
+ Path(__file__).resolve().parents[1] / "benchmarks" / "data" / "nprobust_lprobust_golden.json"
)
@@ -51,7 +48,7 @@ def _smoke_data(self, seed=7):
rng = np.random.default_rng(seed)
G = 800
d = rng.uniform(0.0, 1.0, G)
- y = d + d ** 2 + rng.normal(0, 0.4, G)
+ y = d + d**2 + rng.normal(0, 0.4, G)
return d, y
def test_returns_bias_corrected_fit(self):
@@ -124,52 +121,34 @@ def test_dgp1_all_six_outputs(self, golden):
d, y, g = self._dgp(golden, "dgp1")
fit = bias_corrected_local_linear(d, y, boundary=0.0, kernel="epanechnikov")
# DGP 1 hits bit-parity through the Cholesky path for tau_cl / se_cl.
- np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_classical, g["se_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_robust, g["se_rb"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_low, g["ci_low"],
- atol=1e-13, rtol=1e-13)
- np.testing.assert_allclose(fit.ci_high, g["ci_high"],
- atol=1e-13, rtol=1e-13)
+ np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_classical, g["se_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_robust, g["se_rb"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_low, g["ci_low"], atol=1e-13, rtol=1e-13)
+ np.testing.assert_allclose(fit.ci_high, g["ci_high"], atol=1e-13, rtol=1e-13)
# Bandwidth also should match (Phase 1b bit-parity path).
np.testing.assert_allclose(fit.h, g["h"], atol=1e-12, rtol=1e-12)
def test_dgp2_all_six_outputs(self, golden):
d, y, g = self._dgp(golden, "dgp2")
fit = bias_corrected_local_linear(d, y, boundary=0.0, kernel="epanechnikov")
- np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_classical, g["se_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_robust, g["se_rb"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_low, g["ci_low"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_high, g["ci_high"],
- atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_classical, g["se_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_robust, g["se_rb"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_low, g["ci_low"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_high, g["ci_high"], atol=1e-12, rtol=1e-12)
def test_dgp3_all_six_outputs(self, golden):
d, y, g = self._dgp(golden, "dgp3")
fit = bias_corrected_local_linear(d, y, boundary=0.0, kernel="epanechnikov")
- np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_classical, g["se_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_robust, g["se_rb"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_low, g["ci_low"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_high, g["ci_high"],
- atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_classical, g["se_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_robust, g["se_rb"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_low, g["ci_low"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_high, g["ci_high"], atol=1e-12, rtol=1e-12)
def test_clustered_parity_dgp_4(self, golden):
"""DGP 4 uses manual h=b=0.3 to sidestep an nprobust-internal
@@ -179,26 +158,25 @@ def test_clustered_parity_dgp_4(self, golden):
y = np.asarray(g["y"], dtype=np.float64)
cluster = np.asarray(g["cluster"])
fit = bias_corrected_local_linear(
- d, y, boundary=0.0, kernel="epanechnikov",
- h=g["h"], b=g["b"], cluster=cluster,
+ d,
+ y,
+ boundary=0.0,
+ kernel="epanechnikov",
+ h=g["h"],
+ b=g["b"],
+ cluster=cluster,
)
# Cluster IDs are sequential (1..50), so np.unique order matches R's
# unique(); bit-parity is achievable.
- np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"],
- atol=1e-14, rtol=1e-14)
- np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"],
- atol=1e-14, rtol=1e-14)
- np.testing.assert_allclose(fit.se_classical, g["se_cl"],
- atol=1e-14, rtol=1e-14)
- np.testing.assert_allclose(fit.se_robust, g["se_rb"],
- atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"], atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"], atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.se_classical, g["se_cl"], atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.se_robust, g["se_rb"], atol=1e-14, rtol=1e-14)
# CI bounds at bit-parity too (Python's scipy ppf and R's qnorm
# agree to ULP; the remaining drift is pure tau.bc +/- z * se.rb
# floating-point arithmetic).
- np.testing.assert_allclose(fit.ci_low, g["ci_low"],
- atol=1e-14, rtol=1e-14)
- np.testing.assert_allclose(fit.ci_high, g["ci_high"],
- atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.ci_low, g["ci_low"], atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(fit.ci_high, g["ci_high"], atol=1e-14, rtol=1e-14)
def test_shifted_boundary_parity_dgp_5(self, golden):
"""Design 1 continuous-near-d_lower: boundary = d.min() > 0."""
@@ -207,22 +185,19 @@ def test_shifted_boundary_parity_dgp_5(self, golden):
y = np.asarray(g["y"], dtype=np.float64)
eval_point = float(g["eval_point_override"])
fit = bias_corrected_local_linear(
- d, y, boundary=eval_point, kernel="epanechnikov",
+ d,
+ y,
+ boundary=eval_point,
+ kernel="epanechnikov",
)
- np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_classical, g["se_cl"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.se_robust, g["se_rb"],
- atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_classical, g["tau_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.estimate_bias_corrected, g["tau_bc"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_classical, g["se_cl"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.se_robust, g["se_rb"], atol=1e-12, rtol=1e-12)
# CI bounds at the same tolerance (Python scipy ppf vs R qnorm
# agree to ULP; tau.bc +/- z * se.rb inherits se_rb's drift).
- np.testing.assert_allclose(fit.ci_low, g["ci_low"],
- atol=1e-12, rtol=1e-12)
- np.testing.assert_allclose(fit.ci_high, g["ci_high"],
- atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_low, g["ci_low"], atol=1e-12, rtol=1e-12)
+ np.testing.assert_allclose(fit.ci_high, g["ci_high"], atol=1e-12, rtol=1e-12)
# =============================================================================
@@ -242,7 +217,10 @@ def test_ci_symmetric_about_tau_bc(self):
fit = bias_corrected_local_linear(d, y, boundary=0.0, h=0.2, b=0.2)
midpoint = 0.5 * (fit.ci_low + fit.ci_high)
np.testing.assert_allclose(
- midpoint, fit.estimate_bias_corrected, atol=1e-14, rtol=1e-14,
+ midpoint,
+ fit.estimate_bias_corrected,
+ atol=1e-14,
+ rtol=1e-14,
)
def test_alpha_01_wider_than_alpha_05(self):
@@ -310,9 +288,7 @@ def test_off_support_boundary_raises(self):
def test_mass_point_raises(self):
# d.min() > 0 with > 2% at the mode.
rng = np.random.default_rng(3)
- d = np.concatenate(
- [np.full(100, 0.1), rng.uniform(0.1, 1.0, 400)]
- )
+ d = np.concatenate([np.full(100, 0.1), rng.uniform(0.1, 1.0, 400)])
y = d + rng.normal(0, 0.3, 500)
with pytest.raises(NotImplementedError, match="mass-point"):
bias_corrected_local_linear(d, y, boundary=0.1)
@@ -350,9 +326,7 @@ def test_cluster_object_none_raises(self):
catches it (CI review PR #340 follow-up P1)."""
d = np.linspace(0.0, 1.0, 100)
y = d.copy()
- cluster = np.array(
- [i // 10 if i != 5 else None for i in range(100)], dtype=object
- )
+ cluster = np.array([i // 10 if i != 5 else None for i in range(100)], dtype=object)
with pytest.raises(ValueError, match="cluster contains missing"):
bias_corrected_local_linear(d, y, h=0.3, cluster=cluster)
@@ -360,9 +334,7 @@ def test_cluster_object_nan_raises(self):
"""Object-dtype cluster with np.nan is rejected."""
d = np.linspace(0.0, 1.0, 100)
y = d.copy()
- cluster = np.array(
- [i // 10 if i != 5 else np.nan for i in range(100)], dtype=object
- )
+ cluster = np.array([i // 10 if i != 5 else np.nan for i in range(100)], dtype=object)
with pytest.raises(ValueError, match="cluster contains missing"):
bias_corrected_local_linear(d, y, h=0.3, cluster=cluster)
@@ -387,7 +359,11 @@ def test_cluster_with_manual_bandwidth(self):
y = d + rng.normal(0, 0.3, G)
cluster = np.repeat(np.arange(20), 20)
fit = bias_corrected_local_linear(
- d, y, h=0.3, b=0.3, cluster=cluster,
+ d,
+ y,
+ h=0.3,
+ b=0.3,
+ cluster=cluster,
)
assert np.isfinite(fit.estimate_bias_corrected)
assert fit.se_robust > 0
@@ -498,7 +474,7 @@ def _smoke_data(self, seed=33):
rng = np.random.default_rng(seed)
G = 600
d = rng.uniform(0.0, 1.0, G)
- y = d + d ** 2 + rng.normal(0, 0.3, G)
+ y = d + d**2 + rng.normal(0, 0.3, G)
return d, y
def test_auto_cluster_returns_finite(self):
@@ -591,11 +567,7 @@ class TestWeightedBiasCorrectedLocalLinear:
def _panel(self, seed=42, G=300, boundary=0.0):
rng = np.random.default_rng(seed)
d = rng.uniform(boundary, boundary + 1.0, G)
- y = (
- 2.0 * (d - boundary)
- + 0.3 * (d - boundary) ** 2
- + rng.normal(0, 0.2, G)
- )
+ y = 2.0 * (d - boundary) + 0.3 * (d - boundary) ** 2 + rng.normal(0, 0.2, G)
return d, y
def test_uniform_weights_bit_parity_full_struct(self):
@@ -603,22 +575,18 @@ def test_uniform_weights_bit_parity_full_struct(self):
d, y = self._panel()
base = bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3)
- w1 = bias_corrected_local_linear(
- d, y, boundary=0.0, h=0.3, b=0.3, weights=np.ones(d.size)
- )
+ w1 = bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3, weights=np.ones(d.size))
np.testing.assert_allclose(
w1.estimate_classical, base.estimate_classical, atol=1e-14, rtol=1e-14
)
np.testing.assert_allclose(
- w1.estimate_bias_corrected, base.estimate_bias_corrected,
- atol=1e-14, rtol=1e-14,
- )
- np.testing.assert_allclose(
- w1.se_classical, base.se_classical, atol=1e-14, rtol=1e-14
- )
- np.testing.assert_allclose(
- w1.se_robust, base.se_robust, atol=1e-14, rtol=1e-14
+ w1.estimate_bias_corrected,
+ base.estimate_bias_corrected,
+ atol=1e-14,
+ rtol=1e-14,
)
+ np.testing.assert_allclose(w1.se_classical, base.se_classical, atol=1e-14, rtol=1e-14)
+ np.testing.assert_allclose(w1.se_robust, base.se_robust, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(w1.ci_low, base.ci_low, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(w1.ci_high, base.ci_high, atol=1e-14, rtol=1e-14)
assert w1.n_used == base.n_used
@@ -632,14 +600,14 @@ def test_uniform_weights_bit_parity_auto_bandwidth(self):
d, y = self._panel()
base = bias_corrected_local_linear(d, y, boundary=0.0)
- w1 = bias_corrected_local_linear(
- d, y, boundary=0.0, weights=np.ones(d.size)
- )
+ w1 = bias_corrected_local_linear(d, y, boundary=0.0, weights=np.ones(d.size))
np.testing.assert_allclose(w1.h, base.h, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(w1.b, base.b, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(
- w1.estimate_bias_corrected, base.estimate_bias_corrected,
- atol=1e-14, rtol=1e-14,
+ w1.estimate_bias_corrected,
+ base.estimate_bias_corrected,
+ atol=1e-14,
+ rtol=1e-14,
)
def test_nontrivial_weights_change_estimate(self):
@@ -648,12 +616,8 @@ def test_nontrivial_weights_change_estimate(self):
d, y = self._panel()
base = bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3)
w = np.exp(-d * 3.0) # informative: concentrate near boundary
- r_w = bias_corrected_local_linear(
- d, y, boundary=0.0, h=0.3, b=0.3, weights=w
- )
- assert not np.isclose(
- r_w.estimate_bias_corrected, base.estimate_bias_corrected, atol=1e-6
- )
+ r_w = bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3, weights=w)
+ assert not np.isclose(r_w.estimate_bias_corrected, base.estimate_bias_corrected, atol=1e-6)
def test_negative_weights_reject(self):
from diff_diff.local_linear import bias_corrected_local_linear
@@ -662,9 +626,7 @@ def test_negative_weights_reject(self):
w = np.ones(d.size)
w[0] = -0.5
with pytest.raises(ValueError, match="non-negative"):
- bias_corrected_local_linear(
- d, y, boundary=0.0, h=0.3, b=0.3, weights=w
- )
+ bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3, weights=w)
def test_nan_weights_reject(self):
from diff_diff.local_linear import bias_corrected_local_linear
@@ -673,18 +635,14 @@ def test_nan_weights_reject(self):
w = np.ones(d.size)
w[0] = np.nan
with pytest.raises(ValueError, match="non-finite"):
- bias_corrected_local_linear(
- d, y, boundary=0.0, h=0.3, b=0.3, weights=w
- )
+ bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3, weights=w)
def test_zero_sum_weights_reject(self):
from diff_diff.local_linear import bias_corrected_local_linear
d, y = self._panel()
with pytest.raises(ValueError, match="sum to zero"):
- bias_corrected_local_linear(
- d, y, boundary=0.0, h=0.3, b=0.3, weights=np.zeros(d.size)
- )
+ bias_corrected_local_linear(d, y, boundary=0.0, h=0.3, b=0.3, weights=np.zeros(d.size))
def test_weights_length_mismatch_reject(self):
from diff_diff.local_linear import bias_corrected_local_linear
diff --git a/tests/test_chaisemartin_dhaultfoeuille_parity.py b/tests/test_chaisemartin_dhaultfoeuille_parity.py
index 01a98feb..247e013c 100644
--- a/tests/test_chaisemartin_dhaultfoeuille_parity.py
+++ b/tests/test_chaisemartin_dhaultfoeuille_parity.py
@@ -387,9 +387,7 @@ def _golden_to_df_with_covariates(data_dict: dict) -> pd.DataFrame:
return pd.DataFrame(cols)
-def _golden_to_df_with_extra(
- data_dict: dict, extra_cols: List[str]
-) -> pd.DataFrame:
+def _golden_to_df_with_extra(data_dict: dict, extra_cols: List[str]) -> pd.DataFrame:
"""Reconstruct a panel DataFrame including arbitrary extra columns.
Used for scenarios that ship non-covariate side columns (e.g.,
@@ -423,12 +421,18 @@ class TestDCDHDynRParityPhase3:
# Trends-only: exact (0.0000%) at both horizons after cumulation fix.
# Combined: observed gap 0.30%-0.59% (from OLS residualization only).
# SE: 3-5% from cell-count weighting; 12-18% for cumulated SEs.
- POINT_RTOL = 0.01 # 1% for controls (observed: 0.26%)
- SE_RTOL = 0.20 # 20% for SE (cell-count weighting + cumulation)
+ POINT_RTOL = 0.01 # 1% for controls (observed: 0.26%)
+ SE_RTOL = 0.20 # 20% for SE (cell-count weighting + cumulation)
def _check_phase3_scenario(
- self, golden_values, scenario_name, L_max, controls=None,
- trends_linear=None, point_rtol=None, se_rtol=None,
+ self,
+ golden_values,
+ scenario_name,
+ L_max,
+ controls=None,
+ trends_linear=None,
+ point_rtol=None,
+ se_rtol=None,
):
scenario = golden_values.get(scenario_name)
if scenario is None:
@@ -437,9 +441,14 @@ def _check_phase3_scenario(
df = _golden_to_df_with_covariates(scenario["data"])
est = ChaisemartinDHaultfoeuille()
results = est.fit(
- df, outcome="outcome", group="group", time="period",
- treatment="treatment", L_max=L_max,
- controls=controls, trends_linear=trends_linear,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=L_max,
+ controls=controls,
+ trends_linear=trends_linear,
)
r_results = scenario["results"]
rtol = point_rtol or self.POINT_RTOL
@@ -458,9 +467,7 @@ def _check_phase3_scenario(
# Check per-horizon effects
for h_str, r_eff in r_results.get("effects", {}).items():
h = int(h_str)
- assert h in py_effects, (
- f"Horizon {h} missing from Python results"
- )
+ assert h in py_effects, f"Horizon {h} missing from Python results"
py_eff = py_effects[h]["effect"]
assert py_eff == pytest.approx(
r_eff["overall_att"], rel=rtol
@@ -482,7 +489,9 @@ def test_parity_joiners_only_controls(self, golden_values):
vs obs-count weighting deviation in REGISTRY.md.
"""
self._check_phase3_scenario(
- golden_values, "joiners_only_controls", L_max=2,
+ golden_values,
+ "joiners_only_controls",
+ L_max=2,
controls=["X1"],
point_rtol=self.POINT_RTOL,
)
@@ -493,7 +502,9 @@ def test_parity_joiners_only_trends_lin(self, golden_values):
Exact match (0.0000%) at both horizons after per-group cumulation fix.
"""
self._check_phase3_scenario(
- golden_values, "joiners_only_trends_lin", L_max=2,
+ golden_values,
+ "joiners_only_trends_lin",
+ L_max=2,
trends_linear=True,
point_rtol=1e-4, # exact match
)
@@ -505,8 +516,11 @@ def test_parity_joiners_only_controls_trends_lin(self, golden_values):
the trends cumulation is now exact after per-group cumulation fix).
"""
self._check_phase3_scenario(
- golden_values, "joiners_only_controls_trends_lin", L_max=2,
- controls=["X1"], trends_linear=True,
+ golden_values,
+ "joiners_only_controls_trends_lin",
+ L_max=2,
+ controls=["X1"],
+ trends_linear=True,
point_rtol=self.POINT_RTOL,
)
@@ -623,8 +637,7 @@ def _compare_by_path(self, scenario, by_path, L_max, point_rtol, se_rtol):
)
assert py_h["effect"] == pytest.approx(r_h["effect"], rel=point_rtol), (
- f"path={path_key} h={h}: "
- f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
+ f"path={path_key} h={h}: " f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
)
# Assert matching finite/missing state BEFORE the numeric
@@ -646,8 +659,7 @@ def _compare_by_path(self, scenario, by_path, L_max, point_rtol, se_rtol):
)
if py_finite_positive and r_finite_positive:
assert py_se == pytest.approx(r_se, rel=se_rtol), (
- f"path={path_key} h={h} SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ f"path={path_key} h={h} SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
def test_parity_mixed_single_switch_by_path(self, golden_values):
@@ -720,14 +732,10 @@ def test_parity_multi_path_reversible_by_path_placebo(self, golden_values):
scenario = golden_values.get("multi_path_reversible_by_path_placebo")
if scenario is None:
- pytest.skip(
- "scenario 'multi_path_reversible_by_path_placebo' not in golden values"
- )
+ pytest.skip("scenario 'multi_path_reversible_by_path_placebo' not in golden values")
df = _golden_to_df(scenario["data"])
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, by_path=3, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -776,11 +784,8 @@ def test_parity_multi_path_reversible_by_path_placebo(self, golden_values):
f"before comparing SE."
)
- assert py_h["effect"] == pytest.approx(
- r_h["effect"], rel=self.POINT_RTOL
- ), (
- f"path={path_key} lag={h}: "
- f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
+ assert py_h["effect"] == pytest.approx(r_h["effect"], rel=self.POINT_RTOL), (
+ f"path={path_key} lag={h}: " f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
)
py_se = py_h["se"]
@@ -793,8 +798,7 @@ def test_parity_multi_path_reversible_by_path_placebo(self, golden_values):
)
if py_finite_positive and r_finite_positive:
assert py_se == pytest.approx(r_se, rel=self.SE_RTOL), (
- f"path={path_key} lag={h} placebo SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ f"path={path_key} lag={h} placebo SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
@@ -849,9 +853,7 @@ def test_parity_multi_path_reversible_by_path_controls(self, golden_values):
scenario = golden_values.get("multi_path_reversible_by_path_controls")
if scenario is None:
- pytest.skip(
- "scenario 'multi_path_reversible_by_path_controls' not in golden values"
- )
+ pytest.skip("scenario 'multi_path_reversible_by_path_controls' not in golden values")
df = _golden_to_df_with_covariates(scenario["data"])
est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3)
@@ -899,11 +901,8 @@ def test_parity_multi_path_reversible_by_path_controls(self, golden_values):
f"py={py_h['n_obs']} vs r={int(r_h['n_switchers'])}"
)
- assert py_h["effect"] == pytest.approx(
- r_h["effect"], rel=self.POINT_RTOL
- ), (
- f"path={path_key} h={h}: "
- f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
+ assert py_h["effect"] == pytest.approx(r_h["effect"], rel=self.POINT_RTOL), (
+ f"path={path_key} h={h}: " f"py={py_h['effect']:.4f} vs r={r_h['effect']:.4f}"
)
py_se = py_h["se"]
@@ -911,13 +910,11 @@ def test_parity_multi_path_reversible_by_path_controls(self, golden_values):
py_finite_positive = math.isfinite(py_se) and py_se > 0.0
r_finite_positive = math.isfinite(r_se) and r_se > 0.0
assert py_finite_positive == r_finite_positive, (
- f"path={path_key} h={h} SE state mismatch "
- f"(py_se={py_se}, r_se={r_se})"
+ f"path={path_key} h={h} SE state mismatch " f"(py_se={py_se}, r_se={r_se})"
)
if py_finite_positive and r_finite_positive:
assert py_se == pytest.approx(r_se, rel=self.SE_RTOL), (
- f"path={path_key} h={h} SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ f"path={path_key} h={h} SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
@@ -992,19 +989,14 @@ def test_parity_single_baseline_multi_path_trends_lin(self, golden_values):
import math
import warnings
- scenario = golden_values.get(
- "single_baseline_multi_path_by_path_trends_lin"
- )
+ scenario = golden_values.get("single_baseline_multi_path_by_path_trends_lin")
if scenario is None:
pytest.skip(
- "scenario 'single_baseline_multi_path_by_path_trends_lin' "
- "not in golden values"
+ "scenario 'single_baseline_multi_path_by_path_trends_lin' " "not in golden values"
)
df = _golden_to_df(scenario["data"])
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, by_path=3, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -1019,8 +1011,7 @@ def test_parity_single_baseline_multi_path_trends_lin(self, golden_values):
r_by_path = scenario["results"]["by_path"]
assert results.path_cumulated_event_study is not None, (
- "path_cumulated_event_study should be populated under "
- "by_path + trends_linear=True"
+ "path_cumulated_event_study should be populated under " "by_path + trends_linear=True"
)
py_cum = results.path_cumulated_event_study
@@ -1043,8 +1034,7 @@ def test_parity_single_baseline_multi_path_trends_lin(self, golden_values):
if h <= 0:
continue
assert h in py_path_cum, (
- f"path={path_key}: horizon {h} missing from "
- f"path_cumulated_event_study"
+ f"path={path_key}: horizon {h} missing from " f"path_cumulated_event_study"
)
py_h = py_path_cum[h]
@@ -1065,15 +1055,11 @@ def test_parity_single_baseline_multi_path_trends_lin(self, golden_values):
py_finite_positive = math.isfinite(py_se) and py_se > 0.0
r_finite_positive = math.isfinite(r_se) and r_se > 0.0
assert py_finite_positive == r_finite_positive, (
- f"path={path_key} h={h} SE state mismatch "
- f"(py_se={py_se}, r_se={r_se})"
+ f"path={path_key} h={h} SE state mismatch " f"(py_se={py_se}, r_se={r_se})"
)
if py_finite_positive and r_finite_positive:
- assert py_se == pytest.approx(
- r_se, rel=self.CUM_SE_RTOL
- ), (
- f"path={path_key} h={h} cumulated SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ assert py_se == pytest.approx(r_se, rel=self.CUM_SE_RTOL), (
+ f"path={path_key} h={h} cumulated SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
@@ -1113,26 +1099,19 @@ class TestDCDHDynRParityByPathTrendsNonparam:
def _path_key_from_r_label(self, r_label: str):
return tuple(int(x) for x in r_label.split(","))
- def test_parity_multi_path_reversible_by_path_trends_nonparam(
- self, golden_values
- ):
+ def test_parity_multi_path_reversible_by_path_trends_nonparam(self, golden_values):
"""3-path case with state-set trends: by_path=3, trends_nonparam='state'."""
import math
import warnings
- scenario = golden_values.get(
- "multi_path_reversible_by_path_trends_nonparam"
- )
+ scenario = golden_values.get("multi_path_reversible_by_path_trends_nonparam")
if scenario is None:
pytest.skip(
- "scenario 'multi_path_reversible_by_path_trends_nonparam' "
- "not in golden values"
+ "scenario 'multi_path_reversible_by_path_trends_nonparam' " "not in golden values"
)
df = _golden_to_df_with_extra(scenario["data"], extra_cols=["state"])
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, by_path=3, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -1174,8 +1153,7 @@ def test_parity_multi_path_reversible_by_path_trends_nonparam(
h = int(h_str)
if h > 0:
assert h in py_path["horizons"], (
- f"path={path_key}: horizon {h} missing from "
- f"Python path_effects"
+ f"path={path_key}: horizon {h} missing from " f"Python path_effects"
)
py_h = py_path["horizons"][h]
else:
@@ -1190,11 +1168,8 @@ def test_parity_multi_path_reversible_by_path_trends_nonparam(
f"py={py_h['n_obs']} vs r={int(r_h['n_switchers'])}"
)
- assert py_h["effect"] == pytest.approx(
- r_h["effect"], rel=self.POINT_RTOL
- ), (
- f"path={path_key} h={h}: "
- f"py={py_h['effect']:.6f} vs r={r_h['effect']:.6f}"
+ assert py_h["effect"] == pytest.approx(r_h["effect"], rel=self.POINT_RTOL), (
+ f"path={path_key} h={h}: " f"py={py_h['effect']:.6f} vs r={r_h['effect']:.6f}"
)
py_se = py_h["se"]
@@ -1202,13 +1177,11 @@ def test_parity_multi_path_reversible_by_path_trends_nonparam(
py_finite_positive = math.isfinite(py_se) and py_se > 0.0
r_finite_positive = math.isfinite(r_se) and r_se > 0.0
assert py_finite_positive == r_finite_positive, (
- f"path={path_key} h={h} SE state mismatch "
- f"(py_se={py_se}, r_se={r_se})"
+ f"path={path_key} h={h} SE state mismatch " f"(py_se={py_se}, r_se={r_se})"
)
if py_finite_positive and r_finite_positive:
assert py_se == pytest.approx(r_se, rel=self.SE_RTOL), (
- f"path={path_key} h={h} SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ f"path={path_key} h={h} SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
@@ -1248,26 +1221,19 @@ class TestDCDHDynRParityByPathNonBinary:
def _path_key_from_r_label(self, r_label: str):
return tuple(int(x) for x in r_label.split(","))
- def test_parity_multi_path_reversible_by_path_non_binary(
- self, golden_values
- ):
+ def test_parity_multi_path_reversible_by_path_non_binary(self, golden_values):
"""3-path case with D in {0, 1, 2}: by_path=3, placebo=1."""
import math
import warnings
- scenario = golden_values.get(
- "multi_path_reversible_by_path_non_binary"
- )
+ scenario = golden_values.get("multi_path_reversible_by_path_non_binary")
if scenario is None:
pytest.skip(
- "scenario 'multi_path_reversible_by_path_non_binary' "
- "not in golden values"
+ "scenario 'multi_path_reversible_by_path_non_binary' " "not in golden values"
)
df = _golden_to_df(scenario["data"])
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, by_path=3, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -1308,8 +1274,7 @@ def test_parity_multi_path_reversible_by_path_non_binary(
h = int(h_str)
if h > 0:
assert h in py_path["horizons"], (
- f"path={path_key}: horizon {h} missing from "
- f"Python path_effects"
+ f"path={path_key}: horizon {h} missing from " f"Python path_effects"
)
py_h = py_path["horizons"][h]
else:
@@ -1326,23 +1291,18 @@ def test_parity_multi_path_reversible_by_path_non_binary(
assert py_h["effect"] == pytest.approx(
r_h["effect"], rel=self.POINT_RTOL, abs=self.POINT_ATOL
- ), (
- f"path={path_key} h={h}: "
- f"py={py_h['effect']:.6f} vs r={r_h['effect']:.6f}"
- )
+ ), (f"path={path_key} h={h}: " f"py={py_h['effect']:.6f} vs r={r_h['effect']:.6f}")
py_se = py_h["se"]
r_se = r_h["se"]
py_finite_positive = math.isfinite(py_se) and py_se > 0.0
r_finite_positive = math.isfinite(r_se) and r_se > 0.0
assert py_finite_positive == r_finite_positive, (
- f"path={path_key} h={h} SE state mismatch "
- f"(py_se={py_se}, r_se={r_se})"
+ f"path={path_key} h={h} SE state mismatch " f"(py_se={py_se}, r_se={r_se})"
)
if py_finite_positive and r_finite_positive:
assert py_se == pytest.approx(r_se, rel=self.SE_RTOL), (
- f"path={path_key} h={h} SE: "
- f"py={py_se:.4f} vs r={r_se:.4f}"
+ f"path={path_key} h={h} SE: " f"py={py_se:.4f} vs r={r_se:.4f}"
)
@@ -1373,10 +1333,7 @@ def test_parity_multi_path_reversible_predict_het(self, golden_values):
scenario = golden_values.get("multi_path_reversible_predict_het")
if scenario is None:
- pytest.skip(
- "scenario 'multi_path_reversible_predict_het' not in "
- "golden values"
- )
+ pytest.skip("scenario 'multi_path_reversible_predict_het' not in " "golden values")
df = _golden_to_df_with_extra(scenario["data"], extra_cols=["het_x"])
est = ChaisemartinDHaultfoeuille(drop_larger_lower=False)
@@ -1397,25 +1354,25 @@ def test_parity_multi_path_reversible_predict_het(self, golden_values):
for h_str, r_h in r_predict_het.items():
h = int(h_str)
- assert h in results.heterogeneity_effects, (
- f"horizon {h} missing from Python heterogeneity_effects"
- )
+ assert (
+ h in results.heterogeneity_effects
+ ), f"horizon {h} missing from Python heterogeneity_effects"
py_h = results.heterogeneity_effects[h]
assert py_h["beta"] == pytest.approx(
r_h["beta"], rel=self.BETA_RTOL, abs=self.BETA_ATOL
), f"h={h} beta: py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}"
- assert py_h["se"] == pytest.approx(r_h["se"], rel=self.SE_RTOL), (
- f"h={h} se: py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
- )
+ assert py_h["se"] == pytest.approx(
+ r_h["se"], rel=self.SE_RTOL
+ ), f"h={h} se: py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
# `t_stat = beta / se` is invariant to the Wald-test
# critical-value distribution; pin it at SE_RTOL so a
# regression in beta or se surfaces here too.
- assert py_h["t_stat"] == pytest.approx(r_h["t"], rel=self.SE_RTOL), (
- f"h={h} t_stat: py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
- )
- assert int(py_h["n_obs"]) == int(r_h["n_obs"]), (
- f"h={h} n_obs: py={py_h['n_obs']} vs r={r_h['n_obs']}"
- )
+ assert py_h["t_stat"] == pytest.approx(
+ r_h["t"], rel=self.SE_RTOL
+ ), f"h={h} t_stat: py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
+ assert int(py_h["n_obs"]) == int(
+ r_h["n_obs"]
+ ), f"h={h} n_obs: py={py_h['n_obs']} vs r={r_h['n_obs']}"
# `p_value` and `conf_int` parity (post-2026-05-15 df threading).
# `_compute_heterogeneity_test` now passes
# `df = n_obs - rank(design)` to `safe_inference`, matching
@@ -1427,17 +1384,11 @@ def test_parity_multi_path_reversible_predict_het(self, golden_values):
assert py_h["p_value"] == pytest.approx(
r_h["p_value"], rel=self.INFERENCE_RTOL
), f"h={h} p_value: py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
- assert py_h["conf_int"][0] == pytest.approx(
- r_h["ci_lo"], rel=self.INFERENCE_RTOL
- ), (
- f"h={h} ci_lo: py={py_h['conf_int'][0]:.6f} "
- f"vs r={r_h['ci_lo']:.6f}"
+ assert py_h["conf_int"][0] == pytest.approx(r_h["ci_lo"], rel=self.INFERENCE_RTOL), (
+ f"h={h} ci_lo: py={py_h['conf_int'][0]:.6f} " f"vs r={r_h['ci_lo']:.6f}"
)
- assert py_h["conf_int"][1] == pytest.approx(
- r_h["ci_hi"], rel=self.INFERENCE_RTOL
- ), (
- f"h={h} ci_hi: py={py_h['conf_int'][1]:.6f} "
- f"vs r={r_h['ci_hi']:.6f}"
+ assert py_h["conf_int"][1] == pytest.approx(r_h["ci_hi"], rel=self.INFERENCE_RTOL), (
+ f"h={h} ci_hi: py={py_h['conf_int'][1]:.6f} " f"vs r={r_h['ci_hi']:.6f}"
)
@@ -1466,19 +1417,14 @@ class TestDCDHDynRParityByPathHeterogeneity:
def _path_key_from_r_label(self, r_label: str):
return tuple(int(x) for x in r_label.split(","))
- def test_parity_multi_path_reversible_by_path_predict_het(
- self, golden_values
- ):
+ def test_parity_multi_path_reversible_by_path_predict_het(self, golden_values):
"""Per-path heterogeneity coefficient parity per (path, horizon)."""
import warnings
- scenario = golden_values.get(
- "multi_path_reversible_by_path_predict_het"
- )
+ scenario = golden_values.get("multi_path_reversible_by_path_predict_het")
if scenario is None:
pytest.skip(
- "scenario 'multi_path_reversible_by_path_predict_het' "
- "not in golden values"
+ "scenario 'multi_path_reversible_by_path_predict_het' " "not in golden values"
)
df = _golden_to_df_with_extra(scenario["data"], extra_cols=["het_x"])
@@ -1519,33 +1465,21 @@ def test_parity_multi_path_reversible_by_path_predict_het(
py_h = py_horizons[h]
assert py_h["beta"] == pytest.approx(
r_h["beta"], rel=self.BETA_RTOL, abs=self.BETA_ATOL
- ), (
- f"path={path_key} h={h} beta: "
- f"py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}"
- )
- assert py_h["se"] == pytest.approx(
- r_h["se"], rel=self.SE_RTOL
- ), (
- f"path={path_key} h={h} se: "
- f"py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
+ ), (f"path={path_key} h={h} beta: " f"py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}")
+ assert py_h["se"] == pytest.approx(r_h["se"], rel=self.SE_RTOL), (
+ f"path={path_key} h={h} se: " f"py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
)
# `t_stat = beta / se` is invariant to the Wald-test
# critical-value distribution; pin at SE_RTOL.
- assert py_h["t_stat"] == pytest.approx(
- r_h["t"], rel=self.SE_RTOL
- ), (
- f"path={path_key} h={h} t_stat: "
- f"py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
+ assert py_h["t_stat"] == pytest.approx(r_h["t"], rel=self.SE_RTOL), (
+ f"path={path_key} h={h} t_stat: " f"py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
)
assert int(py_h["n_obs"]) == int(r_h["n_obs"]), (
- f"path={path_key} h={h} n_obs: "
- f"py={py_h['n_obs']} vs r={r_h['n_obs']}"
+ f"path={path_key} h={h} n_obs: " f"py={py_h['n_obs']} vs r={r_h['n_obs']}"
)
# `p_value` and `conf_int` parity — same df-threading
# rationale as the global heterogeneity class.
- assert py_h["p_value"] == pytest.approx(
- r_h["p_value"], rel=self.INFERENCE_RTOL
- ), (
+ assert py_h["p_value"] == pytest.approx(r_h["p_value"], rel=self.INFERENCE_RTOL), (
f"path={path_key} h={h} p_value: "
f"py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
)
@@ -1594,19 +1528,14 @@ class TestDCDHDynRParityByPathHeterogeneityWithPlacebo:
def _path_key_from_r_label(self, r_label: str):
return tuple(int(x) for x in r_label.split(","))
- def test_parity_multi_path_reversible_predict_het_with_placebo(
- self, golden_values
- ):
+ def test_parity_multi_path_reversible_predict_het_with_placebo(self, golden_values):
"""Per-path heterogeneity on forward + backward horizons."""
import warnings
- scenario = golden_values.get(
- "multi_path_reversible_predict_het_with_placebo"
- )
+ scenario = golden_values.get("multi_path_reversible_predict_het_with_placebo")
if scenario is None:
pytest.skip(
- "scenario 'multi_path_reversible_predict_het_with_placebo' "
- "not in golden values"
+ "scenario 'multi_path_reversible_predict_het_with_placebo' " "not in golden values"
)
df = _golden_to_df_with_extra(scenario["data"], extra_cols=["het_x"])
@@ -1615,9 +1544,7 @@ def test_parity_multi_path_reversible_predict_het_with_placebo(
# subset (placebo=2 in scenario 22 -> backward horizons -1, -2).
# Python's extra backward horizons (e.g., -3 if eligible) are
# allowed and not parity-checked.
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, by_path=3, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, by_path=3, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -1649,18 +1576,14 @@ def test_parity_multi_path_reversible_predict_het_with_placebo(
for h_str, r_h in r_path_entry["horizons"].items():
h = int(h_str)
assert h > 0
- assert h in py_horizons, (
- f"path={path_key}: forward horizon {h} missing"
- )
+ assert h in py_horizons, f"path={path_key}: forward horizon {h} missing"
py_h = py_horizons[h]
self._assert_horizon_parity(path_key, h, py_h, r_h)
# Placebo horizons (negative int keys). R-verified: scenario 22
# has 2 placebo rows per path (-1, -2); Python mirrors with
# negative-int keys in path_heterogeneity_effects.
- for h_str, r_h in _as_dict(
- r_path_entry.get("placebo_horizons")
- ).items():
+ for h_str, r_h in _as_dict(r_path_entry.get("placebo_horizons")).items():
h = int(h_str)
assert h < 0
assert h in py_horizons, (
@@ -1672,43 +1595,26 @@ def test_parity_multi_path_reversible_predict_het_with_placebo(
def _assert_horizon_parity(self, path_key, h, py_h, r_h):
"""Pin all 6 inference fields against R."""
- assert py_h["beta"] == pytest.approx(
- r_h["beta"], rel=self.BETA_RTOL, abs=self.BETA_ATOL
- ), (
- f"path={path_key} h={h} beta: "
- f"py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}"
+ assert py_h["beta"] == pytest.approx(r_h["beta"], rel=self.BETA_RTOL, abs=self.BETA_ATOL), (
+ f"path={path_key} h={h} beta: " f"py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}"
)
assert py_h["se"] == pytest.approx(r_h["se"], rel=self.SE_RTOL), (
- f"path={path_key} h={h} se: "
- f"py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
+ f"path={path_key} h={h} se: " f"py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
)
- assert py_h["t_stat"] == pytest.approx(
- r_h["t"], rel=self.SE_RTOL
- ), (
- f"path={path_key} h={h} t_stat: "
- f"py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
+ assert py_h["t_stat"] == pytest.approx(r_h["t"], rel=self.SE_RTOL), (
+ f"path={path_key} h={h} t_stat: " f"py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
)
assert int(py_h["n_obs"]) == int(r_h["n_obs"]), (
- f"path={path_key} h={h} n_obs: "
- f"py={py_h['n_obs']} vs r={r_h['n_obs']}"
+ f"path={path_key} h={h} n_obs: " f"py={py_h['n_obs']} vs r={r_h['n_obs']}"
)
- assert py_h["p_value"] == pytest.approx(
- r_h["p_value"], rel=self.INFERENCE_RTOL
- ), (
- f"path={path_key} h={h} p_value: "
- f"py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
+ assert py_h["p_value"] == pytest.approx(r_h["p_value"], rel=self.INFERENCE_RTOL), (
+ f"path={path_key} h={h} p_value: " f"py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
)
- assert py_h["conf_int"][0] == pytest.approx(
- r_h["ci_lo"], rel=self.INFERENCE_RTOL
- ), (
- f"path={path_key} h={h} ci_lo: "
- f"py={py_h['conf_int'][0]:.6f} vs r={r_h['ci_lo']:.6f}"
+ assert py_h["conf_int"][0] == pytest.approx(r_h["ci_lo"], rel=self.INFERENCE_RTOL), (
+ f"path={path_key} h={h} ci_lo: " f"py={py_h['conf_int'][0]:.6f} vs r={r_h['ci_lo']:.6f}"
)
- assert py_h["conf_int"][1] == pytest.approx(
- r_h["ci_hi"], rel=self.INFERENCE_RTOL
- ), (
- f"path={path_key} h={h} ci_hi: "
- f"py={py_h['conf_int'][1]:.6f} vs r={r_h['ci_hi']:.6f}"
+ assert py_h["conf_int"][1] == pytest.approx(r_h["ci_hi"], rel=self.INFERENCE_RTOL), (
+ f"path={path_key} h={h} ci_hi: " f"py={py_h['conf_int'][1]:.6f} vs r={r_h['ci_hi']:.6f}"
)
@@ -1736,15 +1642,11 @@ class TestDCDHDynRParityHeterogeneityWithPlacebo:
SE_RTOL = 1e-5
INFERENCE_RTOL = 1e-4
- def test_parity_multi_path_reversible_predict_het_with_placebo_global(
- self, golden_values
- ):
+ def test_parity_multi_path_reversible_predict_het_with_placebo_global(self, golden_values):
"""Global heterogeneity on forward + backward horizons."""
import warnings
- scenario = golden_values.get(
- "multi_path_reversible_predict_het_with_placebo_global"
- )
+ scenario = golden_values.get("multi_path_reversible_predict_het_with_placebo_global")
if scenario is None:
pytest.skip(
"scenario 'multi_path_reversible_predict_het_with_placebo_global' "
@@ -1757,9 +1659,7 @@ def test_parity_multi_path_reversible_predict_het_with_placebo_global(
# R's emitted subset (placebo=2 in scenario 23 -> backward
# horizons -1, -2). Python's extra backward horizons (e.g., -3
# if eligible) are allowed and not parity-checked.
- est = ChaisemartinDHaultfoeuille(
- drop_larger_lower=False, placebo=True
- )
+ est = ChaisemartinDHaultfoeuille(drop_larger_lower=False, placebo=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
results = est.fit(
@@ -1788,9 +1688,7 @@ def test_parity_multi_path_reversible_predict_het_with_placebo_global(
for h_str, r_h in r_placebo_het.items():
h = int(h_str)
assert h < 0
- assert h in py_het, (
- f"placebo horizon {h} missing from Python heterogeneity_effects"
- )
+ assert h in py_het, f"placebo horizon {h} missing from Python heterogeneity_effects"
self._assert_horizon_parity(h, py_het[h], r_h)
def _assert_horizon_parity(self, h, py_h, r_h):
@@ -1798,20 +1696,18 @@ def _assert_horizon_parity(self, h, py_h, r_h):
assert py_h["beta"] == pytest.approx(
r_h["beta"], rel=self.BETA_RTOL, abs=self.BETA_ATOL
), f"h={h} beta: py={py_h['beta']:.6f} vs r={r_h['beta']:.6f}"
- assert py_h["se"] == pytest.approx(r_h["se"], rel=self.SE_RTOL), (
- f"h={h} se: py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
- )
- assert py_h["t_stat"] == pytest.approx(r_h["t"], rel=self.SE_RTOL), (
- f"h={h} t_stat: py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
- )
- assert int(py_h["n_obs"]) == int(r_h["n_obs"]), (
- f"h={h} n_obs: py={py_h['n_obs']} vs r={r_h['n_obs']}"
- )
+ assert py_h["se"] == pytest.approx(
+ r_h["se"], rel=self.SE_RTOL
+ ), f"h={h} se: py={py_h['se']:.6f} vs r={r_h['se']:.6f}"
+ assert py_h["t_stat"] == pytest.approx(
+ r_h["t"], rel=self.SE_RTOL
+ ), f"h={h} t_stat: py={py_h['t_stat']:.6f} vs r={r_h['t']:.6f}"
+ assert int(py_h["n_obs"]) == int(
+ r_h["n_obs"]
+ ), f"h={h} n_obs: py={py_h['n_obs']} vs r={r_h['n_obs']}"
assert py_h["p_value"] == pytest.approx(
r_h["p_value"], rel=self.INFERENCE_RTOL
- ), (
- f"h={h} p_value: py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
- )
+ ), f"h={h} p_value: py={py_h['p_value']:.6e} vs r={r_h['p_value']:.6e}"
assert py_h["conf_int"][0] == pytest.approx(
r_h["ci_lo"], rel=self.INFERENCE_RTOL
), f"h={h} ci_lo: py={py_h['conf_int'][0]:.6f} vs r={r_h['ci_lo']:.6f}"
diff --git a/tests/test_conley_vcov.py b/tests/test_conley_vcov.py
index f66f8bc3..f983c76a 100644
--- a/tests/test_conley_vcov.py
+++ b/tests/test_conley_vcov.py
@@ -3705,9 +3705,7 @@ def test_singular_gram_rank_reduces_not_raises(self):
dropped_idx = int(np.flatnonzero(nan_diag)[0])
assert np.all(np.isnan(V[dropped_idx, :])) and np.all(np.isnan(V[:, dropped_idx]))
msgs = [str(w.message) for w in caught]
- assert any(
- "Conley spatial HAC variance" in m and "rank-deficient" in m for m in msgs
- ), msgs
+ assert any("Conley spatial HAC variance" in m and "rank-deficient" in m for m in msgs), msgs
def test_rank_zero_gram_returns_nan(self):
X, resid, coords = self._cross_section(1, dup="indep")
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index a0ec05bb..28f2cf2a 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -5,9 +5,6 @@
including both the download/cache mechanism and the fallback data generation.
"""
-import os
-import tempfile
-from pathlib import Path
from unittest.mock import patch
import numpy as np
@@ -15,7 +12,6 @@
import pytest
from diff_diff.datasets import (
- _CACHE_DIR,
_construct_card_krueger_data,
_construct_castle_doctrine_data,
_construct_divorce_laws_data,
@@ -23,10 +19,7 @@
clear_cache,
list_datasets,
load_card_krueger,
- load_castle_doctrine,
load_dataset,
- load_divorce_laws,
- load_mpdta,
)
@@ -263,9 +256,7 @@ def test_card_krueger_with_did(self):
# Should be able to fit DiD
did = DifferenceInDifferences()
- results = did.fit(
- df_long, outcome="employment", treatment="treated", time="post"
- )
+ results = did.fit(df_long, outcome="employment", treatment="treated", time="post")
assert hasattr(results, "att")
assert hasattr(results, "se")
diff --git a/tests/test_dcdh_bootstrap_cell_period_coverage.py b/tests/test_dcdh_bootstrap_cell_period_coverage.py
index b31283c3..9934c2e0 100644
--- a/tests/test_dcdh_bootstrap_cell_period_coverage.py
+++ b/tests/test_dcdh_bootstrap_cell_period_coverage.py
@@ -70,21 +70,17 @@ def _simulate_panel(
# two global PSUs reused across groups.
psu_id = int(g) * 2 + parity
d = 1 if (treated[g] and t >= first_treated_period) else 0
- y = (
- group_fe[g]
- + 0.1 * t
- + tau * d
- + psu_fe[g, parity]
- + rng.normal(0.0, obs_sigma)
+ y = group_fe[g] + 0.1 * t + tau * d + psu_fe[g, parity] + rng.normal(0.0, obs_sigma)
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": float(y),
+ "psu": psu_id,
+ "pw": 1.0,
+ }
)
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": float(y),
- "psu": psu_id,
- "pw": 1.0,
- })
return pd.DataFrame(rows)
@@ -140,12 +136,16 @@ def test_bootstrap_cell_period_coverage_varying_psu():
# broadcast is exercised) is finite-sample-sensitive
# and B=500 would risk a spurious edge-of-band miss.
res = ChaisemartinDHaultfoeuille(
- n_bootstrap=1000, seed=r + 1,
+ n_bootstrap=1000,
+ seed=r + 1,
).fit(
df,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=2,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=2,
)
except Exception:
failed += 1
@@ -196,8 +196,7 @@ def test_bootstrap_cell_period_coverage_varying_psu():
completed = n_reps - failed
assert completed >= int(0.95 * n_reps), (
- f"MC simulation had {failed}/{n_reps} fit failures, above "
- f"the 5% tolerance."
+ f"MC simulation had {failed}/{n_reps} fit failures, above " f"the 5% tolerance."
)
coverage_overall = covered_overall / completed
coverage_h1 = covered_h1 / completed
diff --git a/tests/test_dcdh_cell_period_coverage.py b/tests/test_dcdh_cell_period_coverage.py
index b3ba0203..719b1fa4 100644
--- a/tests/test_dcdh_cell_period_coverage.py
+++ b/tests/test_dcdh_cell_period_coverage.py
@@ -60,21 +60,17 @@ def _simulate_panel(
# trivially satisfied.
psu_id = int(g) * 2 + parity
d = 1 if (treated[g] and t >= first_treated_period) else 0
- y = (
- group_fe[g]
- + 0.1 * t
- + tau * d
- + psu_fe[g, parity]
- + rng.normal(0.0, obs_sigma)
+ y = group_fe[g] + 0.1 * t + tau * d + psu_fe[g, parity] + rng.normal(0.0, obs_sigma)
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": float(y),
+ "psu": psu_id,
+ "pw": 1.0,
+ }
)
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": float(y),
- "psu": psu_id,
- "pw": 1.0,
- })
return pd.DataFrame(rows)
@@ -111,9 +107,12 @@ def test_cell_period_allocator_coverage_within_group_varying_psu():
warnings.simplefilter("ignore")
res = ChaisemartinDHaultfoeuille(seed=r + 1).fit(
df,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
except Exception:
failed += 1
@@ -129,8 +128,7 @@ def test_cell_period_allocator_coverage_within_group_varying_psu():
completed = n_reps - failed
assert completed >= int(0.95 * n_reps), (
- f"MC simulation had {failed}/{n_reps} fit failures, above "
- f"the 5% tolerance."
+ f"MC simulation had {failed}/{n_reps} fit failures, above " f"the 5% tolerance."
)
coverage = covered / completed
assert 0.925 <= coverage <= 0.975, (
diff --git a/tests/test_diagnostics.py b/tests/test_diagnostics.py
index d019c1e6..ae53085e 100644
--- a/tests/test_diagnostics.py
+++ b/tests/test_diagnostics.py
@@ -48,13 +48,15 @@ def panel_data_parallel_trends():
y += np.random.normal(0, 0.5)
- data.append({
- "unit": unit,
- "period": period,
- "treated": int(is_treated),
- "post": int(period >= 3),
- "outcome": y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": int(period >= 3),
+ "outcome": y,
+ }
+ )
return pd.DataFrame(data)
@@ -88,13 +90,15 @@ def panel_data_violated_trends():
y += np.random.normal(0, 0.5)
- data.append({
- "unit": unit,
- "period": period,
- "treated": int(is_treated),
- "post": int(period >= 3),
- "outcome": y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": int(period >= 3),
+ "outcome": y,
+ }
+ )
return pd.DataFrame(data)
@@ -117,13 +121,15 @@ def simple_panel_data():
if is_treated and period >= 2:
y += 5.0 # Treatment effect
- data.append({
- "unit": unit,
- "period": period,
- "treated": int(is_treated),
- "post": int(period >= 2),
- "outcome": y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": int(period >= 2),
+ "outcome": y,
+ }
+ )
return pd.DataFrame(data)
@@ -215,15 +221,27 @@ def test_significance_stars(self):
"""Test significance stars property."""
# Very significant
r1 = PlaceboTestResults(
- test_type="test", placebo_effect=0, se=1, t_stat=0,
- p_value=0.0001, conf_int=(0, 0), n_obs=100, is_significant=True
+ test_type="test",
+ placebo_effect=0,
+ se=1,
+ t_stat=0,
+ p_value=0.0001,
+ conf_int=(0, 0),
+ n_obs=100,
+ is_significant=True,
)
assert r1.significance_stars == "***"
# Not significant
r2 = PlaceboTestResults(
- test_type="test", placebo_effect=0, se=1, t_stat=0,
- p_value=0.5, conf_int=(0, 0), n_obs=100, is_significant=False
+ test_type="test",
+ placebo_effect=0,
+ se=1,
+ t_stat=0,
+ p_value=0.5,
+ conf_int=(0, 0),
+ n_obs=100,
+ is_significant=False,
)
assert r2.significance_stars == ""
@@ -319,9 +337,7 @@ class TestPlaceboGroupTest:
def test_basic_fake_group(self, simple_panel_data):
"""Test basic fake group test."""
# Use some control units as fake treated
- control_units = simple_panel_data[
- simple_panel_data["treated"] == 0
- ]["unit"].unique()
+ control_units = simple_panel_data[simple_panel_data["treated"] == 0]["unit"].unique()
fake_treated = list(control_units[:3])
results = placebo_group_test(
@@ -428,8 +444,7 @@ def test_permutation_reproducibility(self, simple_panel_data):
assert results1.p_value == results2.p_value
np.testing.assert_array_equal(
- results1.permutation_distribution,
- results2.permutation_distribution
+ results1.permutation_distribution, results2.permutation_distribution
)
def test_permutation_detects_true_effect(self, simple_panel_data):
@@ -482,9 +497,7 @@ def test_loo_returns_all_treated_units(self, simple_panel_data):
)
# Get number of treated units
- n_treated = simple_panel_data[
- simple_panel_data["treated"] == 1
- ]["unit"].nunique()
+ n_treated = simple_panel_data[simple_panel_data["treated"] == 1]["unit"].nunique()
assert len(results.leave_one_out_effects) == n_treated
@@ -509,13 +522,15 @@ def test_loo_detects_influential_unit(self):
else:
y += 2.0 # Normal treatment effect
- data.append({
- "unit": unit,
- "period": period,
- "treated": int(is_treated),
- "post": period,
- "outcome": y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": period,
+ "outcome": y,
+ }
+ )
df = pd.DataFrame(data)
@@ -581,12 +596,14 @@ def test_loo_single_valid_effect_nan_inference(self):
# Only 1 valid LOO effect → SE should be NaN (not 0.0)
assert np.isnan(result.se), f"SE should be NaN with 1 valid effect, got {result.se}"
- assert_nan_inference({
- "se": result.se,
- "t_stat": result.t_stat,
- "p_value": result.p_value,
- "conf_int": result.conf_int,
- })
+ assert_nan_inference(
+ {
+ "se": result.se,
+ "t_stat": result.t_stat,
+ "p_value": result.p_value,
+ "conf_int": result.conf_int,
+ }
+ )
# =============================================================================
@@ -730,16 +747,19 @@ def test_permutation_test_tstat_nan_when_se_zero(self):
y = 5.0
if is_treated and post == 1:
y += 2.0
- data.append({
- "unit": unit,
- "post": post,
- "outcome": y,
- "treated": int(is_treated),
- })
+ data.append(
+ {
+ "unit": unit,
+ "post": post,
+ "outcome": y,
+ "treated": int(is_treated),
+ }
+ )
df = pd.DataFrame(data)
import warnings
+
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
result = permutation_test(
@@ -756,14 +776,11 @@ def test_permutation_test_tstat_nan_when_se_zero(self):
t_stat = result.t_stat
if not np.isfinite(se) or se == 0:
- assert np.isnan(t_stat), (
- f"permutation t_stat should be NaN when SE={se}, got {t_stat}"
- )
+ assert np.isnan(t_stat), f"permutation t_stat should be NaN when SE={se}, got {t_stat}"
else:
expected = result.original_effect / se
assert np.isclose(t_stat, expected), (
- f"permutation t_stat should be effect/SE, "
- f"expected {expected}, got {t_stat}"
+ f"permutation t_stat should be effect/SE, " f"expected {expected}, got {t_stat}"
)
def test_leave_one_out_tstat_nan_when_se_zero(self):
@@ -780,16 +797,19 @@ def test_leave_one_out_tstat_nan_when_se_zero(self):
y = 5.0
if is_treated and post == 1:
y += 2.0
- data.append({
- "unit": unit,
- "post": post,
- "outcome": y,
- "treated": int(is_treated),
- })
+ data.append(
+ {
+ "unit": unit,
+ "post": post,
+ "outcome": y,
+ "treated": int(is_treated),
+ }
+ )
df = pd.DataFrame(data)
import warnings
+
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
result = leave_one_out_test(
@@ -804,13 +824,11 @@ def test_leave_one_out_tstat_nan_when_se_zero(self):
t_stat = result.t_stat
if not np.isfinite(se) or se == 0:
- assert np.isnan(t_stat), (
- f"LOO t_stat should be NaN when SE={se}, got {t_stat}"
- )
+ assert np.isnan(t_stat), f"LOO t_stat should be NaN when SE={se}, got {t_stat}"
ci = result.conf_int
- assert np.isnan(ci[0]) and np.isnan(ci[1]), (
- f"LOO conf_int should be (NaN, NaN) when SE={se}, got {ci}"
- )
+ assert np.isnan(ci[0]) and np.isnan(
+ ci[1]
+ ), f"LOO conf_int should be (NaN, NaN) when SE={se}, got {ci}"
def test_permutation_tstat_consistency(self, simple_panel_data):
"""permutation_test t_stat = effect/SE when SE is valid."""
@@ -828,11 +846,9 @@ def test_permutation_tstat_consistency(self, simple_panel_data):
t_stat = result.t_stat
if not np.isfinite(se) or se == 0:
- assert np.isnan(t_stat), (
- f"t_stat should be NaN when SE={se}, got {t_stat}"
- )
+ assert np.isnan(t_stat), f"t_stat should be NaN when SE={se}, got {t_stat}"
else:
expected = result.original_effect / se
- assert np.isclose(t_stat, expected), (
- f"t_stat should be effect/SE, expected {expected}, got {t_stat}"
- )
+ assert np.isclose(
+ t_stat, expected
+ ), f"t_stat should be effect/SE, expected {expected}, got {t_stat}"
diff --git a/tests/test_doc_deps_integrity.py b/tests/test_doc_deps_integrity.py
index a1490b31..6a49214a 100644
--- a/tests/test_doc_deps_integrity.py
+++ b/tests/test_doc_deps_integrity.py
@@ -172,7 +172,8 @@ def test_group_primaries_have_sources_entry():
"""Each ``groups:`` block must be non-empty and its primary (first) member must have a
``sources:`` entry. Group members resolve to the primary module for doc lookup, so a group
whose primary lacks a ``sources:`` entry dead-ends — /docs-impact would yield no docs for the
- whole group even though every member is nominally 'mapped' by ``test_all_public_sources_mapped``."""
+ whole group even though every member is nominally 'mapped' by ``test_all_public_sources_mapped``.
+ """
broken = []
for name, members in _GROUPS.items():
if not members:
diff --git a/tests/test_efficient_did_validation.py b/tests/test_efficient_did_validation.py
index 9a29ec22..a20321eb 100644
--- a/tests/test_efficient_did_validation.py
+++ b/tests/test_efficient_did_validation.py
@@ -13,9 +13,9 @@
import numpy as np
import pandas as pd
import pytest
+from edid_dgp import make_compustat_dgp, true_es_avg, true_overall_att
from diff_diff import CallawaySantAnna, EfficientDiD
-from edid_dgp import make_compustat_dgp, true_es_avg, true_overall_att
# =============================================================================
# Data Loaders & Helpers
@@ -83,8 +83,7 @@ def _assert_close(actual, expected, label, se=None, se_frac=0.05):
tol = max(0.05 * abs(expected), 50)
diff = abs(actual - expected)
assert diff < tol, (
- f"{label}: expected {expected}, got {actual:.1f} "
- f"(diff={diff:.1f}, tol={tol:.1f})"
+ f"{label}: expected {expected}, got {actual:.1f} " f"(diff={diff:.1f}, tol={tol:.1f})"
)
@@ -92,11 +91,7 @@ def _compute_es_avg(result):
"""Compute ES_avg (Eq 2.3): uniform average over post-treatment horizons."""
if result.event_study_effects is None:
raise ValueError("No event study effects; use aggregate='all'")
- es = {
- int(e): d["effect"]
- for e, d in result.event_study_effects.items()
- if int(e) >= 0
- }
+ es = {int(e): d["effect"] for e, d in result.event_study_effects.items() if int(e) >= 0}
return np.mean(list(es.values()))
@@ -116,8 +111,12 @@ def _run_mc_simulation(n_sims, rho, seed=1000, also_cs=False):
edid = EfficientDiD(pt_assumption="all")
res = edid.fit(
- data, outcome="y", unit="unit", time="time",
- first_treat="first_treat", aggregate="all",
+ data,
+ outcome="y",
+ unit="unit",
+ time="time",
+ first_treat="first_treat",
+ aggregate="all",
)
edid_estimates.append(_compute_es_avg(res))
edid_overall_att.append(res.overall_att)
@@ -127,8 +126,12 @@ def _run_mc_simulation(n_sims, rho, seed=1000, also_cs=False):
if also_cs:
cs = CallawaySantAnna(control_group="never_treated")
cs_res = cs.fit(
- data, outcome="y", unit="unit", time="time",
- first_treat="first_treat", aggregate="event_study",
+ data,
+ outcome="y",
+ unit="unit",
+ time="time",
+ first_treat="first_treat",
+ aggregate="event_study",
)
cs_estimates_list.append(_compute_es_avg(cs_res))
@@ -159,8 +162,12 @@ def edid_hrs_result(hrs_data):
"""Fit EDiD on HRS data (shared across tests)."""
edid = EfficientDiD(pt_assumption="all")
return edid.fit(
- hrs_data, outcome="outcome", unit="unit", time="time",
- first_treat="first_treat", aggregate="all",
+ hrs_data,
+ outcome="outcome",
+ unit="unit",
+ time="time",
+ first_treat="first_treat",
+ aggregate="all",
)
@@ -185,9 +192,12 @@ def test_sample_selection_yields_expected_counts(self, hrs_data):
n_inf = groups.apply(np.isinf).sum()
assert n_inf == 65, f"G=inf: expected 65, got {n_inf}"
- assert sorted(hrs_data["time"].unique()) == [7, 8, 9, 10], (
- f"Expected waves [7,8,9,10], got {sorted(hrs_data['time'].unique())}"
- )
+ assert sorted(hrs_data["time"].unique()) == [
+ 7,
+ 8,
+ 9,
+ 10,
+ ], f"Expected waves [7,8,9,10], got {sorted(hrs_data['time'].unique())}"
def test_group_time_effects_match_table6(self, edid_hrs_result):
for (g, t), (expected_effect, se) in TABLE6_EDID.items():
@@ -213,9 +223,9 @@ def test_se_diagnostic_comparison(self, edid_hrs_result):
for (g, t), (_, cluster_se) in TABLE6_EDID.items():
info = _get_effect(edid_hrs_result.group_time_effects, g, t)
analytical_se = info["se"]
- assert np.isfinite(analytical_se) and analytical_se > 0, (
- f"ATT({g},{t}): analytical SE should be finite positive, got {analytical_se}"
- )
+ assert (
+ np.isfinite(analytical_se) and analytical_se > 0
+ ), f"ATT({g},{t}): analytical SE should be finite positive, got {analytical_se}"
ratio = analytical_se / cluster_se
assert 0.3 < ratio < 3.0, (
f"ATT({g},{t}): SE ratio (analytical/cluster) = {ratio:.2f} "
@@ -227,17 +237,28 @@ def test_cs_cross_validation(self, hrs_data):
"""Cross-validate data loading using CallawaySantAnna."""
cs = CallawaySantAnna(control_group="never_treated")
cs_result = cs.fit(
- hrs_data, outcome="outcome", unit="unit", time="time",
+ hrs_data,
+ outcome="outcome",
+ unit="unit",
+ time="time",
first_treat="first_treat",
)
# CS-SA paper SEs from Table 6
- cs_ses = {(8,8): 1035, (8,9): 909, (8,10): 1008,
- (9,9): 702, (9,10): 651, (10,10): 995}
+ cs_ses = {
+ (8, 8): 1035,
+ (8, 9): 909,
+ (8, 10): 1008,
+ (9, 9): 702,
+ (9, 10): 651,
+ (10, 10): 995,
+ }
for (g, t), expected_effect in TABLE6_CS_SA.items():
info = _get_effect(cs_result.group_time_effects, g, t)
_assert_close(
- info["effect"], expected_effect,
- f"CS ATT({g},{t})", se=cs_ses[(g, t)],
+ info["effect"],
+ expected_effect,
+ f"CS ATT({g},{t})",
+ se=cs_ses[(g, t)],
)
def test_pretreatment_effects_near_zero(self, edid_hrs_result):
@@ -289,26 +310,24 @@ def test_unbiasedness(self, ci_params, rho):
def test_edid_has_lower_rmse_than_cs(self, ci_params, rho):
n_sims = ci_params.bootstrap(150, min_n=49)
mc = _run_mc_simulation(
- n_sims, rho=rho, seed=3000 + int(rho * 100), also_cs=True,
+ n_sims,
+ rho=rho,
+ seed=3000 + int(rho * 100),
+ also_cs=True,
)
- rmse_edid = np.sqrt(
- np.mean((mc["edid_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2)
- )
- rmse_cs = np.sqrt(
- np.mean((mc["cs_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2)
- )
+ rmse_edid = np.sqrt(np.mean((mc["edid_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2))
+ rmse_cs = np.sqrt(np.mean((mc["cs_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2))
# EDiD should not be meaningfully worse than CS
- assert rmse_edid <= rmse_cs * 1.15, (
- f"rho={rho}: RMSE_edid={rmse_edid:.4f} > RMSE_cs={rmse_cs:.4f} * 1.15"
- )
+ assert (
+ rmse_edid <= rmse_cs * 1.15
+ ), f"rho={rho}: RMSE_edid={rmse_edid:.4f} > RMSE_cs={rmse_cs:.4f} * 1.15"
# For negative rho, efficiency gain should be clear
if rho == -0.5:
assert rmse_edid < rmse_cs, (
- f"rho={rho}: Expected RMSE_edid < RMSE_cs, "
- f"got {rmse_edid:.4f} >= {rmse_cs:.4f}"
+ f"rho={rho}: Expected RMSE_edid < RMSE_cs, " f"got {rmse_edid:.4f} >= {rmse_cs:.4f}"
)
def test_efficiency_gain_increases_with_serial_correlation(self, ci_params):
@@ -317,12 +336,8 @@ def test_efficiency_gain_increases_with_serial_correlation(self, ci_params):
mc_neg = _run_mc_simulation(n_sims, rho=-0.5, seed=4500, also_cs=True)
def rel_rmse(mc):
- rmse_e = np.sqrt(
- np.mean((mc["edid_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2)
- )
- rmse_c = np.sqrt(
- np.mean((mc["cs_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2)
- )
+ rmse_e = np.sqrt(np.mean((mc["edid_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2))
+ rmse_c = np.sqrt(np.mean((mc["cs_estimates"] - _TRUE_ES_AVG_COMPUSTAT) ** 2))
return rmse_c / rmse_e if rmse_e > 0 else 1.0
rel_zero = rel_rmse(mc_zero)
@@ -338,20 +353,17 @@ def test_coverage_approximately_correct(self, ci_params):
mc = _run_mc_simulation(n_sims, rho=0, seed=5000)
true_overall = true_overall_att()
- covered = sum(
- ci[0] <= true_overall <= ci[1]
- for ci in mc["edid_overall_ci"]
- )
+ covered = sum(ci[0] <= true_overall <= ci[1] for ci in mc["edid_overall_ci"])
coverage = covered / n_sims
if n_sims >= 200:
- assert 0.88 <= coverage <= 0.99, (
- f"Coverage={coverage:.2f}, expected 0.88-0.99 (n_sims={n_sims})"
- )
+ assert (
+ 0.88 <= coverage <= 0.99
+ ), f"Coverage={coverage:.2f}, expected 0.88-0.99 (n_sims={n_sims})"
else:
- assert 0.80 <= coverage <= 1.00, (
- f"Coverage={coverage:.2f}, expected 0.80-1.00 (n_sims={n_sims})"
- )
+ assert (
+ 0.80 <= coverage <= 1.00
+ ), f"Coverage={coverage:.2f}, expected 0.80-1.00 (n_sims={n_sims})"
def test_analytical_se_calibration(self, ci_params):
n_sims = ci_params.bootstrap(200, min_n=49)
diff --git a/tests/test_honest_did.py b/tests/test_honest_did.py
index f5229a2e..02a16e2a 100644
--- a/tests/test_honest_did.py
+++ b/tests/test_honest_did.py
@@ -299,7 +299,7 @@ class TestParameterExtraction:
def test_extract_from_multiperiod(self, mock_multiperiod_results):
"""Test extraction from MultiPeriodDiDResults."""
- (beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df) = (
+ beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df = (
_extract_event_study_params(mock_multiperiod_results)
)
@@ -687,7 +687,7 @@ def test_multiperiod_sub_vcov_extraction(self, simple_panel_data):
reference_period=3,
)
- (beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df) = (
+ beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df = (
_extract_event_study_params(results)
)
@@ -1363,13 +1363,15 @@ def _fit_dcdh(n_groups=40, n_periods=6, seed=42, L_max=2):
from diff_diff import ChaisemartinDHaultfoeuille
from diff_diff.prep import generate_reversible_did_data
- df = generate_reversible_did_data(
- n_groups=n_groups, n_periods=n_periods, seed=seed
- )
+ df = generate_reversible_did_data(n_groups=n_groups, n_periods=n_periods, seed=seed)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
return ChaisemartinDHaultfoeuille(seed=1).fit(
- df, "outcome", "group", "period", "treatment",
+ df,
+ "outcome",
+ "group",
+ "period",
+ "treatment",
L_max=L_max,
)
@@ -1385,9 +1387,7 @@ def test_dcdh_integration(self):
def test_dcdh_extraction(self):
"""_extract_event_study_params returns correct shapes for dCDH."""
results = self._fit_dcdh()
- beta_hat, sigma, n_pre, n_post, pre_t, post_t, df_s = (
- _extract_event_study_params(results)
- )
+ beta_hat, sigma, n_pre, n_post, pre_t, post_t, df_s = _extract_event_study_params(results)
assert n_pre >= 1
assert n_post >= 1
assert beta_hat.shape == (n_pre + n_post,)
@@ -1407,7 +1407,11 @@ def test_dcdh_no_placebos_raises(self):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
r = ChaisemartinDHaultfoeuille(seed=1, placebo=False).fit(
- df, "outcome", "group", "period", "treatment",
+ df,
+ "outcome",
+ "group",
+ "period",
+ "treatment",
)
with pytest.raises(ValueError, match="placebo_event_study"):
compute_honest_did(r)
@@ -1421,13 +1425,13 @@ def test_dcdh_emits_placebo_warning(self):
warnings.simplefilter("always")
compute_honest_did(results)
placebo_warnings = [
- x for x in w
- if "placebo" in str(x.message).lower()
- and "pre-period" in str(x.message).lower()
+ x
+ for x in w
+ if "placebo" in str(x.message).lower() and "pre-period" in str(x.message).lower()
]
- assert len(placebo_warnings) >= 1, (
- "Expected a UserWarning about placebo-based pre-period inputs"
- )
+ assert (
+ len(placebo_warnings) >= 1
+ ), "Expected a UserWarning about placebo-based pre-period inputs"
def test_dcdh_empty_consecutive_block_raises(self):
"""ValueError when all placebos have NaN SE (no valid pre-periods)."""
@@ -1472,13 +1476,8 @@ def test_dcdh_interior_gap_triggers_trimming_warning(self):
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
bounds = compute_honest_did(results)
- trim_warns = [
- x for x in w
- if "dropping non-consecutive" in str(x.message).lower()
- ]
- assert len(trim_warns) >= 1, (
- "Expected a warning about dropping non-consecutive horizons"
- )
+ trim_warns = [x for x in w if "dropping non-consecutive" in str(x.message).lower()]
+ assert len(trim_warns) >= 1, "Expected a warning about dropping non-consecutive horizons"
# Retained pre should be [-1] only (h=-3 dropped due to gap at -2)
assert bounds.pre_periods_used == [-1]
diff --git a/tests/test_local_linear.py b/tests/test_local_linear.py
index 37e7d7e2..241bf677 100644
--- a/tests/test_local_linear.py
+++ b/tests/test_local_linear.py
@@ -16,7 +16,6 @@
uniform_kernel,
)
-
# =============================================================================
# Kernel support and shape
# =============================================================================
@@ -67,7 +66,7 @@ def _numeric_kappa(kernel_name: str, k: int) -> float:
kfun = KERNELS[kernel_name]
def integrand(t: float) -> float:
- return (t ** k) * kfun(np.array([t]))[0]
+ return (t**k) * kfun(np.array([t]))[0]
val, _ = integrate.quad(integrand, 0.0, 1.0, limit=200)
return val
@@ -114,9 +113,9 @@ def test_closed_form_kappa_matches_numerical_integration(self, kernel, k):
def test_C_matches_formula(self, kernel):
"""C = (kappa_2^2 - kappa_1 kappa_3) / (kappa_0 kappa_2 - kappa_1^2)."""
moms = kernel_moments(kernel)
- expected = (
- moms["kappa_2"] ** 2 - moms["kappa_1"] * moms["kappa_3"]
- ) / (moms["kappa_0"] * moms["kappa_2"] - moms["kappa_1"] ** 2)
+ expected = (moms["kappa_2"] ** 2 - moms["kappa_1"] * moms["kappa_3"]) / (
+ moms["kappa_0"] * moms["kappa_2"] - moms["kappa_1"] ** 2
+ )
assert moms["C"] == pytest.approx(expected, abs=1e-15)
def test_kstar_L2_norm_matches_direct_integration(self):
@@ -128,7 +127,7 @@ def test_kstar_L2_norm_matches_direct_integration(self):
def integrand(t: float) -> float:
kt = epanechnikov_kernel(np.array([t]))[0]
w = (k2 - k1 * t) / denom
- return (w ** 2) * (kt ** 2)
+ return (w**2) * (kt**2)
expected, _ = integrate.quad(integrand, 0.0, 1.0, limit=200)
assert moms["kstar_L2_norm"] == pytest.approx(expected, abs=1e-12)
@@ -153,9 +152,7 @@ def test_recovers_intercept_from_linear_dgp(self):
d = rng.uniform(0.0, 1.0, size=n)
y = a_true + b_true * d + rng.normal(0.0, 0.01, size=n)
- fit = local_linear_fit(
- d, y, bandwidth=0.3, boundary=0.0, kernel="epanechnikov"
- )
+ fit = local_linear_fit(d, y, bandwidth=0.3, boundary=0.0, kernel="epanechnikov")
# Tolerance is several noise-sigmas given the effective sample size.
assert fit.intercept == pytest.approx(a_true, abs=0.01)
assert fit.slope == pytest.approx(b_true, abs=0.05)
@@ -168,9 +165,7 @@ def test_intercept_unbiased_at_exact_linear_data(self):
"""With noiseless linear data, local-linear recovers intercept exactly."""
d = np.linspace(0.01, 0.5, 50)
y = 1.5 + 2.0 * d
- fit = local_linear_fit(
- d, y, bandwidth=0.4, boundary=0.0, kernel="epanechnikov"
- )
+ fit = local_linear_fit(d, y, bandwidth=0.4, boundary=0.0, kernel="epanechnikov")
assert fit.intercept == pytest.approx(1.5, abs=1e-10)
assert fit.slope == pytest.approx(2.0, abs=1e-10)
@@ -186,9 +181,7 @@ def test_matches_weighted_ols_directly(self):
fit = local_linear_fit(d, y, bandwidth=h, boundary=0.0, kernel="uniform")
retain = (d >= 0.0) & (d <= h)
- X_manual = np.column_stack(
- [np.ones(retain.sum()), d[retain] - 0.0]
- )
+ X_manual = np.column_stack([np.ones(retain.sum()), d[retain] - 0.0])
w_manual = np.ones(retain.sum())
coef_manual, _, _ = solve_ols( # type: ignore[call-overload]
X_manual,
@@ -209,7 +202,11 @@ def test_weights_composed_with_kernel(self):
y = 1.0 + 0.5 * d + rng.normal(0.0, 0.05, size=n)
user_w = rng.uniform(0.5, 2.0, size=n)
fit = local_linear_fit(
- d, y, bandwidth=0.4, boundary=0.0, kernel="epanechnikov",
+ d,
+ y,
+ bandwidth=0.4,
+ boundary=0.0,
+ kernel="epanechnikov",
weights=user_w,
)
# Just a smoke test that weights don't error and produce a close-to-1
@@ -221,16 +218,12 @@ def test_weights_composed_with_kernel(self):
def test_returns_dataclass_fields(self):
d = np.linspace(0.01, 0.5, 30)
y = np.random.default_rng(0).normal(size=30)
- fit = local_linear_fit(
- d, y, bandwidth=0.4, boundary=0.0, kernel="triangular"
- )
+ fit = local_linear_fit(d, y, bandwidth=0.4, boundary=0.0, kernel="triangular")
# Dataclass invariants
assert fit.n_effective == len(fit.residuals) == len(fit.kernel_weights)
assert fit.design_matrix.shape == (fit.n_effective, 2)
# The first column of the design is the intercept column.
- np.testing.assert_array_equal(
- fit.design_matrix[:, 0], np.ones(fit.n_effective)
- )
+ np.testing.assert_array_equal(fit.design_matrix[:, 0], np.ones(fit.n_effective))
def test_n_effective_counts_positive_kernel_weights(self):
"""Observations outside [d0, d0 + h] are excluded."""
@@ -238,9 +231,7 @@ def test_n_effective_counts_positive_kernel_weights(self):
d = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 1.5, 2.0, 2.5, 3.0, 3.5])
y = np.zeros_like(d)
y[:5] = 1.0
- fit = local_linear_fit(
- d, y, bandwidth=0.6, boundary=0.0, kernel="uniform"
- )
+ fit = local_linear_fit(d, y, bandwidth=0.6, boundary=0.0, kernel="uniform")
assert fit.n_effective == 5
def test_bandwidth_too_narrow_raises(self):
@@ -261,25 +252,19 @@ def test_negative_bandwidth_raises(self):
d = np.array([0.1, 0.2])
y = np.array([1.0, 2.0])
with pytest.raises(ValueError, match="bandwidth must be positive"):
- local_linear_fit(
- d, y, bandwidth=-0.1, boundary=0.0, kernel="uniform"
- )
+ local_linear_fit(d, y, bandwidth=-0.1, boundary=0.0, kernel="uniform")
def test_zero_bandwidth_raises(self):
d = np.array([0.1, 0.2])
y = np.array([1.0, 2.0])
with pytest.raises(ValueError, match="bandwidth must be positive"):
- local_linear_fit(
- d, y, bandwidth=0.0, boundary=0.0, kernel="uniform"
- )
+ local_linear_fit(d, y, bandwidth=0.0, boundary=0.0, kernel="uniform")
def test_unknown_kernel_raises(self):
d = np.array([0.1, 0.2])
y = np.array([1.0, 2.0])
with pytest.raises(ValueError, match="Unknown kernel"):
- local_linear_fit(
- d, y, bandwidth=0.5, boundary=0.0, kernel="my_kernel"
- )
+ local_linear_fit(d, y, bandwidth=0.5, boundary=0.0, kernel="my_kernel")
def test_mismatched_shapes_raise(self):
d = np.array([0.1, 0.2, 0.3])
@@ -292,18 +277,14 @@ def test_mismatched_weights_shape_raises(self):
y = np.array([1.0, 2.0, 3.0])
w = np.array([1.0, 1.0]) # wrong length
with pytest.raises(ValueError, match="weights must have"):
- local_linear_fit(
- d, y, bandwidth=0.5, boundary=0.0, weights=w
- )
+ local_linear_fit(d, y, bandwidth=0.5, boundary=0.0, weights=w)
def test_negative_weights_raise(self):
d = np.array([0.1, 0.2, 0.3])
y = np.array([1.0, 2.0, 3.0])
w = np.array([1.0, -0.5, 1.0])
with pytest.raises(ValueError, match="nonnegative"):
- local_linear_fit(
- d, y, bandwidth=0.5, boundary=0.0, weights=w
- )
+ local_linear_fit(d, y, bandwidth=0.5, boundary=0.0, weights=w)
def test_non_finite_d_raises(self):
d = np.array([0.1, np.nan, 0.3, 0.4])
@@ -331,9 +312,7 @@ def test_nonzero_boundary(self):
d = rng.uniform(1.0, 2.0, size=n) # Support starts at d_lower = 1.0
y = 3.0 + 0.4 * (d - 1.0) + rng.normal(0.0, 0.02, size=n)
- fit = local_linear_fit(
- d, y, bandwidth=0.3, boundary=1.0, kernel="epanechnikov"
- )
+ fit = local_linear_fit(d, y, bandwidth=0.3, boundary=1.0, kernel="epanechnikov")
# Boundary estimate should recover the intercept at d0=1.0.
assert fit.intercept == pytest.approx(3.0, abs=0.02)
assert fit.boundary == 1.0
diff --git a/tests/test_methodology_callaway.py b/tests/test_methodology_callaway.py
index 1db5dffa..0aa3c91a 100644
--- a/tests/test_methodology_callaway.py
+++ b/tests/test_methodology_callaway.py
@@ -22,12 +22,11 @@
import pytest
from diff_diff import CallawaySantAnna, SurveyDesign
-from diff_diff.linalg import solve_logit
-from diff_diff.prep import generate_staggered_data
from diff_diff.bootstrap_utils import (
generate_bootstrap_weights_batch as _generate_bootstrap_weights_batch,
)
-
+from diff_diff.linalg import solve_logit
+from diff_diff.prep import generate_staggered_data
# =============================================================================
# Test Fixtures and Helpers
@@ -50,23 +49,41 @@ def generate_hand_calculable_data() -> Tuple[pd.DataFrame, float]:
# - Baseline effect varies by unit
# - Time trend: +1 per period for all units
# - Treatment effect: +3 for treated units at t=2
- data = pd.DataFrame({
- 'unit': [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8],
- 'period': [0, 1, 2] * 8,
- 'first_treat': [2] * 6 + [2] * 6 + [0] * 6 + [0] * 6, # 4 treated at g=2, 4 never
- 'outcome': [
- # Treated units (g=2): base + time trend + treatment at t=2
- 10, 11, 15, # unit 1: Y[0]=10, Y[1]=11, Y[2]=15 (effect=15-11-(12-11)=3)
- 12, 13, 17, # unit 2: Y[0]=12, Y[1]=13, Y[2]=17
- 11, 12, 16, # unit 3
- 13, 14, 18, # unit 4
- # Control units: base + time trend only
- 10, 11, 12, # unit 5
- 12, 13, 14, # unit 6
- 11, 12, 13, # unit 7
- 13, 14, 15, # unit 8
- ]
- })
+ data = pd.DataFrame(
+ {
+ "unit": [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8],
+ "period": [0, 1, 2] * 8,
+ "first_treat": [2] * 6 + [2] * 6 + [0] * 6 + [0] * 6, # 4 treated at g=2, 4 never
+ "outcome": [
+ # Treated units (g=2): base + time trend + treatment at t=2
+ 10,
+ 11,
+ 15, # unit 1: Y[0]=10, Y[1]=11, Y[2]=15 (effect=15-11-(12-11)=3)
+ 12,
+ 13,
+ 17, # unit 2: Y[0]=12, Y[1]=13, Y[2]=17
+ 11,
+ 12,
+ 16, # unit 3
+ 13,
+ 14,
+ 18, # unit 4
+ # Control units: base + time trend only
+ 10,
+ 11,
+ 12, # unit 5
+ 12,
+ 13,
+ 14, # unit 6
+ 11,
+ 12,
+ 13, # unit 7
+ 13,
+ 14,
+ 15, # unit 8
+ ],
+ }
+ )
# Hand calculation for ATT(g=2, t=2):
# Base period = g-1 = 1 (for post-treatment effect)
@@ -99,71 +116,57 @@ def test_att_gt_basic_formula_hand_calculation(self):
"""
data, expected_att = generate_hand_calculable_data()
- cs = CallawaySantAnna(estimation_method='reg', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="reg", n_bootstrap=0)
results = cs.fit(
- data,
- outcome='outcome',
- unit='unit',
- time='period',
- first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# ATT(g=2, t=2) should match hand calculation exactly
- actual = results.group_time_effects[(2, 2)]['effect']
- assert np.isclose(actual, expected_att, rtol=1e-10), \
- f"ATT(2,2) expected {expected_att}, got {actual}"
+ actual = results.group_time_effects[(2, 2)]["effect"]
+ assert np.isclose(
+ actual, expected_att, rtol=1e-10
+ ), f"ATT(2,2) expected {expected_att}, got {actual}"
def test_att_gt_with_outcome_regression(self):
"""Test outcome regression produces consistent ATT(g,t)."""
data, expected_att = generate_hand_calculable_data()
- cs = CallawaySantAnna(estimation_method='reg', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="reg", n_bootstrap=0)
results = cs.fit(
- data,
- outcome='outcome',
- unit='unit',
- time='period',
- first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Outcome regression without covariates should match simple DID
- actual = results.group_time_effects[(2, 2)]['effect']
+ actual = results.group_time_effects[(2, 2)]["effect"]
assert np.isclose(actual, expected_att, rtol=1e-10)
def test_att_gt_with_ipw(self):
"""Test IPW produces consistent ATT(g,t) without covariates."""
data, expected_att = generate_hand_calculable_data()
- cs = CallawaySantAnna(estimation_method='ipw', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="ipw", n_bootstrap=0)
results = cs.fit(
- data,
- outcome='outcome',
- unit='unit',
- time='period',
- first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# IPW without covariates should approximate simple DID
# (may differ slightly due to unconditional propensity weighting)
- actual = results.group_time_effects[(2, 2)]['effect']
- assert np.isclose(actual, expected_att, rtol=0.01), \
- f"ATT(2,2) expected ~{expected_att}, got {actual}"
+ actual = results.group_time_effects[(2, 2)]["effect"]
+ assert np.isclose(
+ actual, expected_att, rtol=0.01
+ ), f"ATT(2,2) expected ~{expected_att}, got {actual}"
def test_att_gt_with_doubly_robust(self):
"""Test doubly robust produces consistent ATT(g,t)."""
data, expected_att = generate_hand_calculable_data()
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
results = cs.fit(
- data,
- outcome='outcome',
- unit='unit',
- time='period',
- first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# DR without covariates should match simple DID
- actual = results.group_time_effects[(2, 2)]['effect']
+ actual = results.group_time_effects[(2, 2)]["effect"]
assert np.isclose(actual, expected_att, rtol=1e-10)
@@ -183,28 +186,27 @@ def test_base_period_varying_vs_universal_post_treatment(self):
cohort_periods=[4],
never_treated_frac=0.3,
treatment_effect=2.0,
- seed=42
+ seed=42,
)
cs_varying = CallawaySantAnna(base_period="varying", n_bootstrap=0)
cs_universal = CallawaySantAnna(base_period="universal", n_bootstrap=0)
results_v = cs_varying.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
results_u = cs_universal.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Post-treatment effects (t >= 4) should match exactly
for t in [4, 5, 6, 7]:
if (4, t) in results_v.group_time_effects and (4, t) in results_u.group_time_effects:
- eff_v = results_v.group_time_effects[(4, t)]['effect']
- eff_u = results_u.group_time_effects[(4, t)]['effect']
- assert np.isclose(eff_v, eff_u, rtol=1e-10), \
- f"Post-treatment ATT(4,{t}) should match: varying={eff_v}, universal={eff_u}"
+ eff_v = results_v.group_time_effects[(4, t)]["effect"]
+ eff_u = results_u.group_time_effects[(4, t)]["effect"]
+ assert np.isclose(
+ eff_v, eff_u, rtol=1e-10
+ ), f"Post-treatment ATT(4,{t}) should match: varying={eff_v}, universal={eff_u}"
def test_base_period_varying_pre_treatment_uses_consecutive(self):
"""
@@ -218,20 +220,18 @@ def test_base_period_varying_pre_treatment_uses_consecutive(self):
cohort_periods=[5],
never_treated_frac=0.3,
treatment_effect=2.0,
- seed=42
+ seed=42,
)
cs = CallawaySantAnna(base_period="varying", n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Pre-treatment periods should exist (varying computes them)
# With g=5, pre-treatment would be t in {1,2,3,4} (if anticipation=0)
pre_treatment_exists = any(
- (g, t) in results.group_time_effects
- for g in [5] for t in [1, 2, 3, 4]
+ (g, t) in results.group_time_effects for g in [5] for t in [1, 2, 3, 4]
)
assert pre_treatment_exists, "Varying base period should produce pre-treatment effects"
@@ -248,26 +248,29 @@ def test_base_period_universal_includes_reference_period(self):
cohort_periods=[4],
never_treated_frac=0.3,
treatment_effect=2.0,
- seed=42
+ seed=42,
)
cs = CallawaySantAnna(base_period="universal", anticipation=0, n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
# Reference period e=-1 should exist with effect=0
- assert results.event_study_effects is not None, \
- "Event study effects should be computed"
- assert -1 in results.event_study_effects, \
- "Universal base period should include e=-1 in event study"
+ assert results.event_study_effects is not None, "Event study effects should be computed"
+ assert (
+ -1 in results.event_study_effects
+ ), "Universal base period should include e=-1 in event study"
ref = results.event_study_effects[-1]
- assert ref['effect'] == 0.0, "Reference period effect should be 0"
- assert np.isnan(ref['se']), "Reference period SE should be NaN"
- assert ref['n_groups'] == 0, "Reference period n_groups should be 0"
+ assert ref["effect"] == 0.0, "Reference period effect should be 0"
+ assert np.isnan(ref["se"]), "Reference period SE should be NaN"
+ assert ref["n_groups"] == 0, "Reference period n_groups should be 0"
class TestDoublyRobustEstimator:
@@ -287,19 +290,19 @@ def test_doubly_robust_recovers_true_effect(self):
cohort_periods=[3],
treatment_effect=2.5,
never_treated_frac=0.3,
- seed=42
+ seed=42,
)
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Should recover approximately 2.5 treatment effect
# Allow wider tolerance due to dynamic effects and noise
- assert abs(results.overall_att - 2.5) < 1.0, \
- f"DR should recover ~2.5 effect, got {results.overall_att}"
+ assert (
+ abs(results.overall_att - 2.5) < 1.0
+ ), f"DR should recover ~2.5 effect, got {results.overall_att}"
def test_estimation_methods_produce_similar_results(self):
"""
@@ -309,26 +312,22 @@ def test_estimation_methods_produce_similar_results(self):
reg, ipw, and dr should all produce very similar ATT estimates.
"""
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[4],
- treatment_effect=3.0,
- seed=123
+ n_units=200, n_periods=8, cohort_periods=[4], treatment_effect=3.0, seed=123
)
results = {}
- for method in ['reg', 'ipw', 'dr']:
+ for method in ["reg", "ipw", "dr"]:
cs = CallawaySantAnna(estimation_method=method, n_bootstrap=0)
results[method] = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# All methods should produce similar overall ATT
- atts = [results[m].overall_att for m in ['reg', 'ipw', 'dr']]
+ atts = [results[m].overall_att for m in ["reg", "ipw", "dr"]]
max_diff = max(atts) - min(atts)
- assert max_diff < 0.5, \
- f"Estimation methods differ by {max_diff}: reg={atts[0]}, ipw={atts[1]}, dr={atts[2]}"
+ assert (
+ max_diff < 0.5
+ ), f"Estimation methods differ by {max_diff}: reg={atts[0]}, ipw={atts[1]}, dr={atts[2]}"
class TestDRNoCovariateSEUniformity:
@@ -469,7 +468,7 @@ def _run_r_estimation(
estimation_method: str = "dr",
control_group: str = "nevertreated",
anticipation: int = 0,
- base_period: str = "varying"
+ base_period: str = "varying",
) -> Dict[str, Any]:
"""
Run R's did::att_gt() and return results as dictionary.
@@ -494,7 +493,7 @@ def _run_r_estimation(
# Escape path for cross-platform compatibility (Windows backslashes, spaces)
escaped_path = data_path.replace("\\", "/")
- r_script = f'''
+ r_script = f"""
suppressMessages(library(did))
suppressMessages(library(jsonlite))
@@ -532,27 +531,25 @@ def _run_r_estimation(
)
cat(toJSON(output, pretty = TRUE))
- '''
+ """
result = subprocess.run(
- ["Rscript", "-e", r_script],
- capture_output=True,
- text=True,
- timeout=60
+ ["Rscript", "-e", r_script], capture_output=True, text=True, timeout=60
)
if result.returncode != 0:
raise RuntimeError(f"R script failed: {result.stderr}")
import json
+
parsed = json.loads(result.stdout)
# Handle R's JSON serialization quirks
# Extract scalar values from single-element lists if needed
- if isinstance(parsed.get('overall_att'), list):
- parsed['overall_att'] = parsed['overall_att'][0]
- if isinstance(parsed.get('overall_se'), list):
- parsed['overall_se'] = parsed['overall_se'][0]
+ if isinstance(parsed.get("overall_att"), list):
+ parsed["overall_att"] = parsed["overall_att"][0]
+ if isinstance(parsed.get("overall_se"), list):
+ parsed["overall_se"] = parsed["overall_se"][0]
return parsed
@@ -565,7 +562,7 @@ def benchmark_data(self, tmp_path):
cohort_periods=[4, 6],
treatment_effect=2.0,
never_treated_frac=0.3,
- seed=12345
+ seed=12345,
)
csv_path = tmp_path / "benchmark_data.csv"
data.to_csv(csv_path, index=False)
@@ -580,10 +577,9 @@ def test_overall_att_matches_r_dr(self, require_r, benchmark_data):
data, csv_path = benchmark_data
# Python estimation
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# R estimation
@@ -591,8 +587,9 @@ def test_overall_att_matches_r_dr(self, require_r, benchmark_data):
# Compare overall ATT - use 20% tolerance for aggregation differences
# The discrepancy is primarily in aggregation weights, not ATT(g,t) values
- assert np.isclose(py_results.overall_att, r_results['overall_att'], rtol=0.20), \
- f"ATT mismatch: Python={py_results.overall_att}, R={r_results['overall_att']}"
+ assert np.isclose(
+ py_results.overall_att, r_results["overall_att"], rtol=0.20
+ ), f"ATT mismatch: Python={py_results.overall_att}, R={r_results['overall_att']}"
def test_overall_att_matches_r_reg(self, require_r, benchmark_data):
"""Test overall ATT matches R with outcome regression.
@@ -602,16 +599,16 @@ def test_overall_att_matches_r_reg(self, require_r, benchmark_data):
"""
data, csv_path = benchmark_data
- cs = CallawaySantAnna(estimation_method='reg', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="reg", n_bootstrap=0)
py_results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
r_results = self._run_r_estimation(csv_path, estimation_method="reg")
- assert np.isclose(py_results.overall_att, r_results['overall_att'], rtol=0.20), \
- f"ATT mismatch: Python={py_results.overall_att}, R={r_results['overall_att']}"
+ assert np.isclose(
+ py_results.overall_att, r_results["overall_att"], rtol=0.20
+ ), f"ATT mismatch: Python={py_results.overall_att}, R={r_results['overall_att']}"
def test_group_time_effects_match_r(self, require_r, benchmark_data):
"""Test individual ATT(g,t) values match R for post-treatment periods.
@@ -626,29 +623,28 @@ def test_group_time_effects_match_r(self, require_r, benchmark_data):
"""
data, csv_path = benchmark_data
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
r_results = self._run_r_estimation(csv_path, estimation_method="dr")
# Compare each ATT(g,t) for post-treatment only
- r_gt = r_results['group_time']
+ r_gt = r_results["group_time"]
n_comparisons = 0
mismatches = []
- for i in range(len(r_gt['group'])):
- g = int(r_gt['group'][i])
- t = int(r_gt['time'][i])
- r_att = r_gt['att'][i]
+ for i in range(len(r_gt["group"])):
+ g = int(r_gt["group"][i])
+ t = int(r_gt["time"][i])
+ r_att = r_gt["att"][i]
# Only compare post-treatment effects (t >= g)
if t < g:
continue
if (g, t) in py_results.group_time_effects:
- py_att = py_results.group_time_effects[(g, t)]['effect']
+ py_att = py_results.group_time_effects[(g, t)]["effect"]
# Post-treatment effects should match within 20% or 0.5 abs
# Wider tolerance accounts for differences in dynamic effect handling
if not np.isclose(py_att, r_att, rtol=0.20, atol=0.5):
@@ -656,10 +652,12 @@ def test_group_time_effects_match_r(self, require_r, benchmark_data):
n_comparisons += 1
# Should have made at least some comparisons
- assert n_comparisons > 0, "No post-treatment group-time effects matched between Python and R"
+ assert (
+ n_comparisons > 0
+ ), "No post-treatment group-time effects matched between Python and R"
# Report mismatches if any
- assert len(mismatches) == 0, f"Post-treatment ATT mismatches:\n" + "\n".join(mismatches)
+ assert len(mismatches) == 0, "Post-treatment ATT mismatches:\n" + "\n".join(mismatches)
# =============================================================================
@@ -678,17 +676,12 @@ def test_single_obs_group_produces_valid_result(self):
"""
# Create data with one group having very few units
data = generate_staggered_data(
- n_units=50,
- n_periods=6,
- cohort_periods=[3, 5],
- never_treated_frac=0.4,
- seed=42
+ n_units=50, n_periods=6, cohort_periods=[3, 5], never_treated_frac=0.4, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Should produce valid results
@@ -711,33 +704,32 @@ def test_no_post_treatment_returns_nan_with_warning(self):
periods = np.tile(np.arange(n_periods), n_units)
# 15 never-treated (first_treat=0), 35 treated at period 10 (after data ends)
- first_treat_by_unit = np.concatenate([
- np.zeros(15), # Never treated
- np.full(35, 10) # Treated at period 10 (after data ends)
- ]).astype(int)
+ first_treat_by_unit = np.concatenate(
+ [
+ np.zeros(15), # Never treated
+ np.full(35, 10), # Treated at period 10 (after data ends)
+ ]
+ ).astype(int)
first_treat = np.repeat(first_treat_by_unit, n_periods)
outcomes = np.random.randn(len(units)) + units * 0.1 + periods * 0.5
- data = pd.DataFrame({
- 'unit': units,
- 'period': periods,
- 'first_treat': first_treat,
- 'outcome': outcomes
- })
+ data = pd.DataFrame(
+ {"unit": units, "period": periods, "first_treat": first_treat, "outcome": outcomes}
+ )
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Check warning was emitted
warning_messages = [str(warning.message) for warning in w]
- assert any("post-treatment" in msg.lower() for msg in warning_messages), \
- f"Expected post-treatment warning, got: {warning_messages}"
+ assert any(
+ "post-treatment" in msg.lower() for msg in warning_messages
+ ), f"Expected post-treatment warning, got: {warning_messages}"
# All overall inference fields should be NaN
assert np.isnan(results.overall_att), "overall_att should be NaN"
@@ -753,31 +745,26 @@ def test_nonfinite_se_produces_nan_tstat(self):
not 0.0 which would be misleading.
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=8,
- cohort_periods=[4],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=8, cohort_periods=[4], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Check that t_stat handling is consistent
for (_g, _t), effect in results.group_time_effects.items():
- se = effect['se']
- t_stat = effect['t_stat']
+ se = effect["se"]
+ t_stat = effect["t_stat"]
if not np.isfinite(se) or se <= 0:
- assert np.isnan(t_stat), \
- f"t_stat should be NaN when SE={se}, got {t_stat}"
+ assert np.isnan(t_stat), f"t_stat should be NaN when SE={se}, got {t_stat}"
elif np.isfinite(se) and se > 0:
# t_stat should be effect/se
- expected_t = effect['effect'] / se
- assert np.isclose(t_stat, expected_t, rtol=1e-10), \
- f"t_stat should be effect/se when SE is valid"
+ expected_t = effect["effect"] / se
+ assert np.isclose(
+ t_stat, expected_t, rtol=1e-10
+ ), "t_stat should be effect/se when SE is valid"
def test_anticipation_shifts_reference_period(self):
"""
@@ -787,18 +774,13 @@ def test_anticipation_shifts_reference_period(self):
g-1-k for post-treatment effects.
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=10,
- cohort_periods=[5],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=10, cohort_periods=[5], treatment_effect=2.0, seed=42
)
# With anticipation=1, post-treatment starts at t >= g-1 = 4
cs = CallawaySantAnna(anticipation=1, n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# With anticipation=1, period 4 (= g-1 = 5-1) should be post-treatment
@@ -815,23 +797,17 @@ def test_not_yet_treated_excludes_cohort_g(self):
"""
# Two cohorts: g=4 and g=7
data = generate_staggered_data(
- n_units=100,
- n_periods=10,
- cohort_periods=[4, 7],
- never_treated_frac=0.3,
- seed=42
+ n_units=100, n_periods=10, cohort_periods=[4, 7], never_treated_frac=0.3, seed=42
)
- cs_nyt = CallawaySantAnna(control_group='not_yet_treated', n_bootstrap=0)
- cs_nt = CallawaySantAnna(control_group='never_treated', n_bootstrap=0)
+ cs_nyt = CallawaySantAnna(control_group="not_yet_treated", n_bootstrap=0)
+ cs_nt = CallawaySantAnna(control_group="never_treated", n_bootstrap=0)
results_nyt = cs_nyt.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
results_nt = cs_nt.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Both should produce valid results
@@ -839,8 +815,8 @@ def test_not_yet_treated_excludes_cohort_g(self):
assert results_nt.overall_att is not None
# Control group setting should be recorded
- assert results_nyt.control_group == 'not_yet_treated'
- assert results_nt.control_group == 'never_treated'
+ assert results_nyt.control_group == "not_yet_treated"
+ assert results_nt.control_group == "never_treated"
# Results should differ (not_yet_treated uses more controls)
# n_control for not_yet_treated should be >= never_treated for early periods
@@ -873,8 +849,7 @@ def test_missing_columns_raises_error(self):
cs = CallawaySantAnna()
with pytest.raises(ValueError, match="Missing columns"):
cs.fit(
- data, outcome='nonexistent', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="nonexistent", unit="unit", time="period", first_treat="first_treat"
)
def test_no_never_treated_raises_error(self):
@@ -884,15 +859,12 @@ def test_no_never_treated_raises_error(self):
n_periods=5,
cohort_periods=[3],
never_treated_frac=0.0, # No never-treated
- seed=42
+ seed=42,
)
cs = CallawaySantAnna()
with pytest.raises(ValueError, match="No never-treated units"):
- cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
- )
+ cs.fit(data, outcome="outcome", unit="unit", time="period", first_treat="first_treat")
class TestRankDeficiencyHandling:
@@ -944,33 +916,30 @@ def test_analytical_se_close_to_bootstrap_se(self, ci_params):
"""
n_boot = ci_params.bootstrap(499, min_n=199)
data = generate_staggered_data(
- n_units=300,
- n_periods=8,
- cohort_periods=[4],
- treatment_effect=2.0,
- seed=42
+ n_units=300, n_periods=8, cohort_periods=[4], treatment_effect=2.0, seed=42
)
cs_anal = CallawaySantAnna(n_bootstrap=0)
cs_boot = CallawaySantAnna(n_bootstrap=n_boot, seed=42)
results_anal = cs_anal.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
results_boot = cs_boot.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Check overall ATT SE (wider tolerance when min_n cap reduces
# bootstrap iterations in pure Python mode)
if results_boot.overall_se > 0:
- rel_diff = abs(results_anal.overall_se - results_boot.overall_se) / results_boot.overall_se
+ rel_diff = (
+ abs(results_anal.overall_se - results_boot.overall_se) / results_boot.overall_se
+ )
threshold = 0.40 if n_boot < 100 else 0.25
- assert rel_diff < threshold, \
- f"Analytical SE ({results_anal.overall_se}) differs from bootstrap SE " \
+ assert rel_diff < threshold, (
+ f"Analytical SE ({results_anal.overall_se}) differs from bootstrap SE "
f"({results_boot.overall_se}) by {rel_diff*100:.1f}%"
+ )
def test_bootstrap_weight_moments_rademacher(self):
"""
@@ -979,7 +948,7 @@ def test_bootstrap_weight_moments_rademacher(self):
These are the standard multiplier bootstrap weights.
"""
rng = np.random.default_rng(42)
- weights = _generate_bootstrap_weights_batch(10000, 100, 'rademacher', rng)
+ weights = _generate_bootstrap_weights_batch(10000, 100, "rademacher", rng)
# E[w] should be ~0
mean_w = np.mean(weights)
@@ -996,7 +965,7 @@ def test_bootstrap_weight_moments_mammen(self):
Mammen's two-point distribution matches skewness of Bernoulli.
"""
rng = np.random.default_rng(42)
- weights = _generate_bootstrap_weights_batch(10000, 100, 'mammen', rng)
+ weights = _generate_bootstrap_weights_batch(10000, 100, "mammen", rng)
mean_w = np.mean(weights)
assert abs(mean_w) < 0.02, f"Mammen E[w] should be ~0, got {mean_w}"
@@ -1017,7 +986,7 @@ def test_bootstrap_weight_moments_webb(self):
This matches R's `did` package behavior.
"""
rng = np.random.default_rng(42)
- weights = _generate_bootstrap_weights_batch(10000, 100, 'webb', rng)
+ weights = _generate_bootstrap_weights_batch(10000, 100, "webb", rng)
mean_w = np.mean(weights)
assert abs(mean_w) < 0.02, f"Webb E[w] should be ~0, got {mean_w}"
@@ -1034,17 +1003,12 @@ def test_bootstrap_produces_valid_inference(self, ci_params):
"""
n_boot = ci_params.bootstrap(99)
data = generate_staggered_data(
- n_units=100,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=3.0,
- seed=42
+ n_units=100, n_periods=6, cohort_periods=[3], treatment_effect=3.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=n_boot, seed=42)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Bootstrap results should exist
@@ -1057,7 +1021,7 @@ def test_bootstrap_produces_valid_inference(self, ci_params):
# Group-time p-values should be in [0, 1]
for effect in results.group_time_effects.values():
- assert 0 <= effect['p_value'] <= 1
+ assert 0 <= effect["p_value"] <= 1
class TestAggregationMethods:
@@ -1070,17 +1034,12 @@ def test_simple_aggregation_weights_by_group_size(self):
Overall ATT = Σ w_g * ATT(g,t) where w_g ∝ n_g
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=8,
- cohort_periods=[4, 6],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=8, cohort_periods=[4, 6], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# Overall ATT should be weighted average of post-treatment effects
@@ -1094,18 +1053,17 @@ def test_event_study_aggregation_by_relative_time(self):
ATT(e) = Σ_g w_g * ATT(g, g+e) for each event time e.
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=10,
- cohort_periods=[4, 6],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=10, cohort_periods=[4, 6], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.event_study_effects is not None
@@ -1123,18 +1081,17 @@ def test_group_aggregation_by_cohort(self):
ATT(g) = (1/T_g) Σ_t ATT(g,t) for t >= g - anticipation.
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=10,
- cohort_periods=[4, 6],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=10, cohort_periods=[4, 6], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='group'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="group",
)
assert results.group_effects is not None
@@ -1146,18 +1103,17 @@ def test_group_aggregation_by_cohort(self):
def test_all_aggregation_computes_everything(self):
"""Test aggregate='all' computes event_study and group effects."""
data = generate_staggered_data(
- n_units=100,
- n_periods=8,
- cohort_periods=[4],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=8, cohort_periods=[4], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='all'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="all",
)
assert results.event_study_effects is not None
@@ -1170,28 +1126,28 @@ class TestGetSetParams:
def test_get_params_returns_all_parameters(self):
"""Test that get_params returns all constructor parameters."""
cs = CallawaySantAnna(
- control_group='not_yet_treated',
+ control_group="not_yet_treated",
anticipation=1,
- estimation_method='ipw',
+ estimation_method="ipw",
alpha=0.10,
n_bootstrap=100,
- bootstrap_weights='mammen',
+ bootstrap_weights="mammen",
seed=42,
- rank_deficient_action='silent',
- base_period='universal'
+ rank_deficient_action="silent",
+ base_period="universal",
)
params = cs.get_params()
- assert params['control_group'] == 'not_yet_treated'
- assert params['anticipation'] == 1
- assert params['estimation_method'] == 'ipw'
- assert params['alpha'] == 0.10
- assert params['n_bootstrap'] == 100
- assert params['bootstrap_weights'] == 'mammen'
- assert params['seed'] == 42
- assert params['rank_deficient_action'] == 'silent'
- assert params['base_period'] == 'universal'
+ assert params["control_group"] == "not_yet_treated"
+ assert params["anticipation"] == 1
+ assert params["estimation_method"] == "ipw"
+ assert params["alpha"] == 0.10
+ assert params["n_bootstrap"] == 100
+ assert params["bootstrap_weights"] == "mammen"
+ assert params["seed"] == 42
+ assert params["rank_deficient_action"] == "silent"
+ assert params["base_period"] == "universal"
def test_set_params_modifies_attributes(self):
"""Test that set_params modifies estimator attributes."""
@@ -1221,17 +1177,12 @@ class TestResultsObject:
def test_results_summary_contains_key_info(self):
"""Test that summary() output contains key information."""
data = generate_staggered_data(
- n_units=50,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=50, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
summary = results.summary()
@@ -1243,48 +1194,42 @@ def test_results_summary_contains_key_info(self):
def test_results_to_dataframe_group_time(self):
"""Test to_dataframe with level='group_time'."""
data = generate_staggered_data(
- n_units=50,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=50, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
- df = results.to_dataframe(level='group_time')
+ df = results.to_dataframe(level="group_time")
- assert 'group' in df.columns
- assert 'time' in df.columns
- assert 'effect' in df.columns
- assert 'se' in df.columns
+ assert "group" in df.columns
+ assert "time" in df.columns
+ assert "effect" in df.columns
+ assert "se" in df.columns
assert len(df) == len(results.group_time_effects)
def test_results_to_dataframe_event_study(self):
"""Test to_dataframe with level='event_study'."""
data = generate_staggered_data(
- n_units=50,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=50, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
- df = results.to_dataframe(level='event_study')
+ df = results.to_dataframe(level="event_study")
- assert 'relative_period' in df.columns
- assert 'effect' in df.columns
+ assert "relative_period" in df.columns
+ assert "effect" in df.columns
def test_results_significance_properties(self):
"""Test is_significant and significance_stars properties."""
@@ -1293,13 +1238,12 @@ def test_results_significance_properties(self):
n_periods=8,
cohort_periods=[4],
treatment_effect=5.0, # Large effect for significance
- seed=42
+ seed=42,
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
# With large effect, should be significant
@@ -1328,24 +1272,23 @@ def test_event_study_analytical_se_includes_wif(self):
so it can only add (or maintain) variance, never reduce it.
"""
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.event_study_effects is not None
for e, eff_data in results.event_study_effects.items():
- se_with_wif = eff_data['se']
+ se_with_wif = eff_data["se"]
if np.isfinite(se_with_wif) and se_with_wif > 0:
# WIF-adjusted SE should be positive
assert se_with_wif > 0, f"SE for e={e} should be positive"
@@ -1356,27 +1299,29 @@ def test_event_study_bootstrap_consistent_with_analytical(self, ci_params):
"""
n_boot = ci_params.bootstrap(999)
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
# Analytical SEs
cs_analytical = CallawaySantAnna(n_bootstrap=0)
results_analytical = cs_analytical.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
# Bootstrap SEs
cs_boot = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=False)
results_boot = cs_boot.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results_analytical.event_study_effects is not None
@@ -1385,15 +1330,16 @@ def test_event_study_bootstrap_consistent_with_analytical(self, ci_params):
threshold = 0.40 if n_boot < 100 else 0.20
n_compared = 0
for e in results_analytical.event_study_effects:
- se_a = results_analytical.event_study_effects[e]['se']
+ se_a = results_analytical.event_study_effects[e]["se"]
if e not in results_boot.event_study_effects:
continue
- se_b = results_boot.event_study_effects[e]['se']
+ se_b = results_boot.event_study_effects[e]["se"]
if np.isfinite(se_a) and se_a > 0 and np.isfinite(se_b) and se_b > 0:
rel_diff = abs(se_a - se_b) / se_a
- assert rel_diff < threshold, \
- f"e={e}: analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} " \
+ assert rel_diff < threshold, (
+ f"e={e}: analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} "
f"(diff={rel_diff*100:.1f}% > {threshold*100}%)"
+ )
n_compared += 1
assert n_compared > 0, "No event times had finite SEs for comparison"
@@ -1405,23 +1351,17 @@ def test_overall_att_bootstrap_uses_combined_if(self, ci_params):
"""
n_boot = ci_params.bootstrap(999)
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
cs_analytical = CallawaySantAnna(n_bootstrap=0)
results_a = cs_analytical.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
cs_boot = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=False)
results_b = cs_boot.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat'
+ data, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
)
se_a = results_a.overall_se
@@ -1430,8 +1370,9 @@ def test_overall_att_bootstrap_uses_combined_if(self, ci_params):
threshold = 0.40 if n_boot < 100 else 0.20
if np.isfinite(se_a) and se_a > 0 and np.isfinite(se_b) and se_b > 0:
rel_diff = abs(se_a - se_b) / se_a
- assert rel_diff < threshold, \
- f"Analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} (diff={rel_diff*100:.1f}%)"
+ assert (
+ rel_diff < threshold
+ ), f"Analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} (diff={rel_diff*100:.1f}%)"
def test_single_group_event_time_wif_effect(self):
"""
@@ -1439,27 +1380,25 @@ def test_single_group_event_time_wif_effect(self):
adjustment should add minimal variance (weight is near-deterministic).
"""
data = generate_staggered_data(
- n_units=100,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.event_study_effects is not None
# With a single cohort, all event times have only one group
for e, eff_data in results.event_study_effects.items():
- if eff_data['n_groups'] == 1:
+ if eff_data["n_groups"] == 1:
# SE should still be finite and positive
- assert np.isfinite(eff_data['se']), \
- f"Single-group SE for e={e} should be finite"
+ assert np.isfinite(eff_data["se"]), f"Single-group SE for e={e} should be finite"
def test_unbalanced_panel_bootstrap_uses_global_n(self, ci_params):
"""
@@ -1475,51 +1414,53 @@ def test_unbalanced_panel_bootstrap_uses_global_n(self, ci_params):
n_boot = ci_params.bootstrap(999)
# Generate balanced staggered data
data = generate_staggered_data(
- n_units=100,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
# Add extra never-treated units with NaN outcomes in all periods.
# These units will be in the panel (and precomputed['all_units'])
# but excluded from all IFs because NaN fails the validity check.
n_nan_units = 10
- max_unit = data['unit'].max()
- periods = sorted(data['period'].unique())
+ max_unit = data["unit"].max()
+ periods = sorted(data["period"].unique())
nan_rows = []
for i in range(1, n_nan_units + 1):
uid = max_unit + i
for p in periods:
- nan_rows.append({
- 'unit': uid,
- 'period': p,
- 'first_treat': 0, # never-treated
- 'outcome': np.nan,
- })
- data_unbalanced = pd.concat(
- [data, pd.DataFrame(nan_rows)], ignore_index=True
- )
-
- n_global = data_unbalanced['unit'].nunique()
- n_valid = data['unit'].nunique() # units that actually appear in IFs
+ nan_rows.append(
+ {
+ "unit": uid,
+ "period": p,
+ "first_treat": 0, # never-treated
+ "outcome": np.nan,
+ }
+ )
+ data_unbalanced = pd.concat([data, pd.DataFrame(nan_rows)], ignore_index=True)
+
+ n_global = data_unbalanced["unit"].nunique()
+ n_valid = data["unit"].nunique() # units that actually appear in IFs
assert n_global > n_valid, "Test setup: global N should exceed IF unit count"
# Analytical SEs (already fixed to use global N)
cs_analytical = CallawaySantAnna(n_bootstrap=0)
results_a = cs_analytical.fit(
- data_unbalanced, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data_unbalanced,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
# Bootstrap SEs (the fix under test)
cs_boot = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=False)
results_b = cs_boot.fit(
- data_unbalanced, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data_unbalanced,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results_a.event_study_effects is not None
@@ -1528,15 +1469,16 @@ def test_unbalanced_panel_bootstrap_uses_global_n(self, ci_params):
threshold = 0.40 if n_boot < 100 else 0.25
n_compared = 0
for e in results_a.event_study_effects:
- se_a = results_a.event_study_effects[e]['se']
+ se_a = results_a.event_study_effects[e]["se"]
if e not in results_b.event_study_effects:
continue
- se_b = results_b.event_study_effects[e]['se']
+ se_b = results_b.event_study_effects[e]["se"]
if np.isfinite(se_a) and se_a > 0 and np.isfinite(se_b) and se_b > 0:
rel_diff = abs(se_a - se_b) / se_a
- assert rel_diff < threshold, \
- f"e={e}: analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} " \
+ assert rel_diff < threshold, (
+ f"e={e}: analytical SE={se_a:.4f} vs bootstrap SE={se_b:.4f} "
f"(diff={rel_diff*100:.1f}% > {threshold*100}%)"
+ )
n_compared += 1
assert n_compared > 0, "No event times had finite SEs for comparison"
@@ -1552,72 +1494,72 @@ def test_cband_crit_value_exceeds_pointwise(self, ci_params):
"""
n_boot = ci_params.bootstrap(999, min_n=199)
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=True)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
- assert results.cband_crit_value is not None, \
- "cband_crit_value should be set when cband=True and n_bootstrap > 0"
- assert results.cband_crit_value > 1.96, \
- f"Simultaneous critical value ({results.cband_crit_value:.3f}) " \
+ assert (
+ results.cband_crit_value is not None
+ ), "cband_crit_value should be set when cband=True and n_bootstrap > 0"
+ assert results.cband_crit_value > 1.96, (
+ f"Simultaneous critical value ({results.cband_crit_value:.3f}) "
"should exceed pointwise z_0.025=1.96"
+ )
def test_cband_confidence_intervals_wider(self, ci_params):
"""Simultaneous CIs should be wider than pointwise CIs."""
n_boot = ci_params.bootstrap(999, min_n=199)
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=True)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.event_study_effects is not None
for e, eff_data in results.event_study_effects.items():
- if 'cband_conf_int' in eff_data:
- pw_ci = eff_data['conf_int']
- cb_ci = eff_data['cband_conf_int']
+ if "cband_conf_int" in eff_data:
+ pw_ci = eff_data["conf_int"]
+ cb_ci = eff_data["cband_conf_int"]
pw_width = pw_ci[1] - pw_ci[0]
cb_width = cb_ci[1] - cb_ci[0]
if np.isfinite(pw_width) and np.isfinite(cb_width) and pw_width > 0:
- assert cb_width >= pw_width, \
- f"e={e}: cband CI width ({cb_width:.4f}) should >= " \
+ assert cb_width >= pw_width, (
+ f"e={e}: cband CI width ({cb_width:.4f}) should >= "
f"pointwise CI width ({pw_width:.4f})"
+ )
def test_cband_false_disables_simultaneous_ci(self, ci_params):
"""When cband=False, cband_crit_value should be None."""
n_boot = ci_params.bootstrap(999, min_n=199)
data = generate_staggered_data(
- n_units=100,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=False)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.cband_crit_value is None
@@ -1625,18 +1567,17 @@ def test_cband_false_disables_simultaneous_ci(self, ci_params):
def test_cband_requires_bootstrap(self):
"""When n_bootstrap=0, cband has no effect (remains None)."""
data = generate_staggered_data(
- n_units=100,
- n_periods=6,
- cohort_periods=[3],
- treatment_effect=2.0,
- seed=42
+ n_units=100, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=0, cband=True)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.cband_crit_value is None
@@ -1658,19 +1599,22 @@ def test_cband_validity_threshold(self):
def mock_max(a, axis=None):
result = original_max(a, axis=axis)
- if axis == 0 and hasattr(result, '__len__') and len(result) > 10:
+ if axis == 0 and hasattr(result, "__len__") and len(result) > 10:
# Make >50% of sup-t draws NaN
result = result.copy()
n_nan = int(len(result) * 0.6)
result[:n_nan] = np.nan
return result
- with unittest.mock.patch('diff_diff.staggered_bootstrap.np.max', side_effect=mock_max):
+ with unittest.mock.patch("diff_diff.staggered_bootstrap.np.max", side_effect=mock_max):
with pytest.warns(RuntimeWarning, match="Too few valid sup-t"):
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study',
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
assert results.cband_crit_value is None
@@ -1679,30 +1623,29 @@ def test_cband_in_get_params(self):
"""Test that cband parameter appears in get_params."""
cs = CallawaySantAnna(cband=False)
params = cs.get_params()
- assert 'cband' in params
- assert params['cband'] is False
+ assert "cband" in params
+ assert params["cband"] is False
def test_cband_to_dataframe_columns(self, ci_params):
"""Test that to_dataframe includes cband columns."""
n_boot = ci_params.bootstrap(999, min_n=199)
data = generate_staggered_data(
- n_units=200,
- n_periods=8,
- cohort_periods=[3, 5],
- treatment_effect=2.0,
- seed=42
+ n_units=200, n_periods=8, cohort_periods=[3, 5], treatment_effect=2.0, seed=42
)
cs = CallawaySantAnna(n_bootstrap=n_boot, seed=42, cband=True)
results = cs.fit(
- data, outcome='outcome', unit='unit',
- time='period', first_treat='first_treat',
- aggregate='event_study'
+ data,
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ aggregate="event_study",
)
- df = results.to_dataframe(level='event_study')
- assert 'cband_lower' in df.columns
- assert 'cband_upper' in df.columns
+ df = results.to_dataframe(level="event_study")
+ assert "cband_lower" in df.columns
+ assert "cband_upper" in df.columns
class TestMPDTARComparison:
@@ -1736,7 +1679,7 @@ def _get_r_mpdta_and_results(self, tmp_path) -> Tuple[pd.DataFrame, dict]:
csv_path = tmp_path / "r_mpdta.csv"
escaped_path = str(csv_path).replace("\\", "/")
- r_script = f'''
+ r_script = f"""
suppressMessages(library(did))
suppressMessages(library(jsonlite))
@@ -1810,13 +1753,10 @@ def _get_r_mpdta_and_results(self, tmp_path) -> Tuple[pd.DataFrame, dict]:
)
cat(toJSON(output, pretty = TRUE))
- '''
+ """
result = subprocess.run(
- ["Rscript", "-e", r_script],
- capture_output=True,
- text=True,
- timeout=120
+ ["Rscript", "-e", r_script], capture_output=True, text=True, timeout=120
)
if result.returncode != 0:
@@ -1825,10 +1765,10 @@ def _get_r_mpdta_and_results(self, tmp_path) -> Tuple[pd.DataFrame, dict]:
parsed = json.loads(result.stdout)
# Handle R's JSON serialization quirks
- if isinstance(parsed.get('overall_att'), list):
- parsed['overall_att'] = parsed['overall_att'][0]
- if isinstance(parsed.get('overall_se'), list):
- parsed['overall_se'] = parsed['overall_se'][0]
+ if isinstance(parsed.get("overall_att"), list):
+ parsed["overall_att"] = parsed["overall_att"][0]
+ if isinstance(parsed.get("overall_se"), list):
+ parsed["overall_se"] = parsed["overall_se"][0]
# Read the exported CSV
mpdta = pd.read_csv(csv_path)
@@ -1845,20 +1785,19 @@ def test_mpdta_overall_att_matches_r_strict(self, require_r, tmp_path):
mpdta, r_results = self._get_r_mpdta_and_results(tmp_path)
# Python estimation using R's data
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
- mpdta,
- outcome='lemp',
- unit='countyreal',
- time='year',
- first_treat='first_treat'
+ mpdta, outcome="lemp", unit="countyreal", time="year", first_treat="first_treat"
)
# Compare overall ATT - strict 1% tolerance for MPDTA
- rel_diff = abs(py_results.overall_att - r_results['overall_att']) / abs(r_results['overall_att'])
- assert rel_diff < 0.01, \
- f"MPDTA ATT mismatch: Python={py_results.overall_att:.6f}, " \
+ rel_diff = abs(py_results.overall_att - r_results["overall_att"]) / abs(
+ r_results["overall_att"]
+ )
+ assert rel_diff < 0.01, (
+ f"MPDTA ATT mismatch: Python={py_results.overall_att:.6f}, "
f"R={r_results['overall_att']:.6f}, diff={rel_diff*100:.2f}%"
+ )
def test_mpdta_overall_se_matches_r_strict(self, require_r, tmp_path):
"""Test overall SE matches R within 1% using MPDTA dataset.
@@ -1867,20 +1806,17 @@ def test_mpdta_overall_se_matches_r_strict(self, require_r, tmp_path):
"""
mpdta, r_results = self._get_r_mpdta_and_results(tmp_path)
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
- mpdta,
- outcome='lemp',
- unit='countyreal',
- time='year',
- first_treat='first_treat'
+ mpdta, outcome="lemp", unit="countyreal", time="year", first_treat="first_treat"
)
# Compare overall SE - strict 1% tolerance for MPDTA
- rel_diff = abs(py_results.overall_se - r_results['overall_se']) / r_results['overall_se']
- assert rel_diff < 0.01, \
- f"MPDTA SE mismatch: Python={py_results.overall_se:.6f}, " \
+ rel_diff = abs(py_results.overall_se - r_results["overall_se"]) / r_results["overall_se"]
+ assert rel_diff < 0.01, (
+ f"MPDTA SE mismatch: Python={py_results.overall_se:.6f}, "
f"R={r_results['overall_se']:.6f}, diff={rel_diff*100:.2f}%"
+ )
def test_mpdta_group_time_effects_match_r_strict(self, require_r, tmp_path):
"""Test individual ATT(g,t) values match R within 1% for MPDTA.
@@ -1890,52 +1826,48 @@ def test_mpdta_group_time_effects_match_r_strict(self, require_r, tmp_path):
"""
mpdta, r_results = self._get_r_mpdta_and_results(tmp_path)
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
- mpdta,
- outcome='lemp',
- unit='countyreal',
- time='year',
- first_treat='first_treat'
+ mpdta, outcome="lemp", unit="countyreal", time="year", first_treat="first_treat"
)
# Compare each post-treatment ATT(g,t)
- r_gt = r_results['group_time']
+ r_gt = r_results["group_time"]
n_comparisons = 0
mismatches = []
- for i in range(len(r_gt['group'])):
- g = int(r_gt['group'][i])
- t = int(r_gt['time'][i])
- r_att = r_gt['att'][i]
+ for i in range(len(r_gt["group"])):
+ g = int(r_gt["group"][i])
+ t = int(r_gt["time"][i])
+ r_att = r_gt["att"][i]
# Skip pre-treatment effects
if t < g:
continue
if (g, t) in py_results.group_time_effects:
- py_att = py_results.group_time_effects[(g, t)]['effect']
+ py_att = py_results.group_time_effects[(g, t)]["effect"]
# Handle near-zero effects (use absolute tolerance)
if abs(r_att) < 0.001:
if abs(py_att - r_att) > 0.01:
- mismatches.append(
- f"ATT({g},{t}): Python={py_att:.6f}, R={r_att:.6f}"
- )
+ mismatches.append(f"ATT({g},{t}): Python={py_att:.6f}, R={r_att:.6f}")
else:
rel_diff = abs(py_att - r_att) / abs(r_att)
if rel_diff > 0.01: # 1% tolerance
mismatches.append(
- f"ATT({g},{t}): Python={py_att:.6f}, R={r_att:.6f}, " \
+ f"ATT({g},{t}): Python={py_att:.6f}, R={r_att:.6f}, "
f"diff={rel_diff*100:.2f}%"
)
n_comparisons += 1
- assert n_comparisons > 0, \
- "No post-treatment group-time effects matched between Python and R"
+ assert (
+ n_comparisons > 0
+ ), "No post-treatment group-time effects matched between Python and R"
- assert len(mismatches) == 0, \
- f"MPDTA post-treatment ATT mismatches (>1% diff):\n" + "\n".join(mismatches)
+ assert (
+ len(mismatches) == 0
+ ), "MPDTA post-treatment ATT mismatches (>1% diff):\n" + "\n".join(mismatches)
def test_mpdta_event_study_ses_match_r(self, require_r, tmp_path):
"""Test event study ATTs and SEs match R's aggte(type='dynamic').
@@ -1946,23 +1878,22 @@ def test_mpdta_event_study_ses_match_r(self, require_r, tmp_path):
"""
mpdta, r_results = self._get_r_mpdta_and_results(tmp_path)
- cs = CallawaySantAnna(estimation_method='dr', n_bootstrap=0)
+ cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=0)
py_results = cs.fit(
mpdta,
- outcome='lemp',
- unit='countyreal',
- time='year',
- first_treat='first_treat',
- aggregate='event_study'
+ outcome="lemp",
+ unit="countyreal",
+ time="year",
+ first_treat="first_treat",
+ aggregate="event_study",
)
- assert py_results.event_study_effects is not None, \
- "Python event_study_effects is None"
+ assert py_results.event_study_effects is not None, "Python event_study_effects is None"
- r_es = r_results['event_study']
- r_event_times = r_es['event_time']
- r_atts = r_es['att']
- r_ses = r_es['se']
+ r_es = r_results["event_study"]
+ r_event_times = r_es["event_time"]
+ r_atts = r_es["att"]
+ r_ses = r_es["se"]
n_compared = 0
mismatches = []
@@ -1981,35 +1912,28 @@ def test_mpdta_event_study_ses_match_r(self, require_r, tmp_path):
continue
py_eff = py_results.event_study_effects[e]
- py_att = py_eff['effect']
- py_se = py_eff['se']
+ py_att = py_eff["effect"]
+ py_se = py_eff["se"]
# Compare ATT: rtol=5%, atol=0.01
if not np.isclose(py_att, r_att, rtol=0.05, atol=0.01):
- mismatches.append(
- f"e={e} ATT: Python={py_att:.6f}, R={r_att:.6f}"
- )
+ mismatches.append(f"e={e} ATT: Python={py_att:.6f}, R={r_att:.6f}")
# Compare SE: rtol=10%, atol=0.005 (wider for WIF sensitivity)
if not np.isclose(py_se, r_se, rtol=0.10, atol=0.005):
- mismatches.append(
- f"e={e} SE: Python={py_se:.6f}, R={r_se:.6f}"
- )
+ mismatches.append(f"e={e} SE: Python={py_se:.6f}, R={r_se:.6f}")
n_compared += 1
if unmatched_r:
# Log but don't fail for unmatched event times
- warnings.warn(
- f"R event times not in Python: {unmatched_r}"
- )
+ warnings.warn(f"R event times not in Python: {unmatched_r}")
- assert n_compared > 0, \
- "No event times were compared between Python and R"
+ assert n_compared > 0, "No event times were compared between Python and R"
- assert len(mismatches) == 0, \
- f"Event study mismatches ({n_compared} compared):\n" + \
- "\n".join(mismatches)
+ assert (
+ len(mismatches) == 0
+ ), f"Event study mismatches ({n_compared} compared):\n" + "\n".join(mismatches)
def test_mpdta_cband_crit_value_vs_r(self, require_r, tmp_path):
"""Test simultaneous confidence band critical value matches R.
@@ -2021,21 +1945,21 @@ def test_mpdta_cband_crit_value_vs_r(self, require_r, tmp_path):
"""
mpdta, r_results = self._get_r_mpdta_and_results(tmp_path)
- r_crit = r_results['event_study_cband']['crit_val']
+ r_crit = r_results["event_study_cband"]["crit_val"]
cs = CallawaySantAnna(
- estimation_method='dr',
+ estimation_method="dr",
n_bootstrap=999,
seed=42,
cband=True,
)
py_results = cs.fit(
mpdta,
- outcome='lemp',
- unit='countyreal',
- time='year',
- first_treat='first_treat',
- aggregate='event_study'
+ outcome="lemp",
+ unit="countyreal",
+ time="year",
+ first_treat="first_treat",
+ aggregate="event_study",
)
py_crit = py_results.cband_crit_value
@@ -2047,9 +1971,10 @@ def test_mpdta_cband_crit_value_vs_r(self, require_r, tmp_path):
assert py_crit > 1.96, f"Python crit value {py_crit} <= 1.96"
# 30% tolerance for bootstrap variability across implementations
- assert np.isclose(py_crit, r_crit, rtol=0.30), \
- f"Cband crit value mismatch: Python={py_crit:.4f}, " \
+ assert np.isclose(py_crit, r_crit, rtol=0.30), (
+ f"Cband crit value mismatch: Python={py_crit:.4f}, "
f"R={r_crit:.4f}, rtol={abs(py_crit - r_crit) / r_crit:.2%}"
+ )
class TestCSCovariateScaleEquilibration:
@@ -2117,9 +2042,7 @@ def test_or_fit_offset_invariant(self):
# DRDID panel influence-function parity (numpy cross-checks)
# =============================================================================
-_GOLDEN_VALUES_PATH = (
- Path(__file__).parents[1] / "benchmarks" / "data" / "csdid_golden_values.json"
-)
+_GOLDEN_VALUES_PATH = Path(__file__).parents[1] / "benchmarks" / "data" / "csdid_golden_values.json"
def _reg_did_panel_ref(
@@ -2141,9 +2064,7 @@ def _reg_did_panel_ref(
Xi = np.column_stack([np.ones(n), X])
c = D == 0
Wc = w[c]
- beta, *_ = np.linalg.lstsq(
- np.sqrt(Wc)[:, None] * Xi[c], np.sqrt(Wc) * dY[c], rcond=None
- )
+ beta, *_ = np.linalg.lstsq(np.sqrt(Wc)[:, None] * Xi[c], np.sqrt(Wc) * dY[c], rcond=None)
pred = Xi @ beta
w_treat = w * D
mw_t = w_treat.mean()
@@ -2277,9 +2198,7 @@ class TestDRDIDPanelIFParity:
R by 4-13% / 3-20% (reg) and ~7x per-cell (ipw).
"""
- def _fit(
- self, df: pd.DataFrame, method: str, survey_design: Any = None, **cs_kwargs
- ) -> Any:
+ def _fit(self, df: pd.DataFrame, method: str, survey_design: Any = None, **cs_kwargs) -> Any:
fit_kwargs = {"survey_design": survey_design} if survey_design is not None else {}
with warnings.catch_warnings():
warnings.simplefilter("ignore")
diff --git a/tests/test_methodology_did.py b/tests/test_methodology_did.py
index d1baf5a1..0f999455 100644
--- a/tests/test_methodology_did.py
+++ b/tests/test_methodology_did.py
@@ -43,21 +43,35 @@ def generate_hand_calculable_data() -> Tuple[pd.DataFrame, float]:
expected_att : float
The hand-calculated ATT value (3.0)
"""
- data = pd.DataFrame({
- 'unit': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
- 'outcome': [
- # Treated, Pre-period (mean = 11.5)
- 10, 11, 12, 13,
- # Treated, Post-period (mean = 15.5) -> diff = 4.0
- 14, 15, 16, 17,
- # Control, Pre-period (mean = 9.5)
- 8, 9, 10, 11,
- # Control, Post-period (mean = 10.5) -> diff = 1.0
- 9, 10, 11, 12,
- ],
- 'treated': [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
- 'post': [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1],
- })
+ data = pd.DataFrame(
+ {
+ "unit": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+ "outcome": [
+ # Treated, Pre-period (mean = 11.5)
+ 10,
+ 11,
+ 12,
+ 13,
+ # Treated, Post-period (mean = 15.5) -> diff = 4.0
+ 14,
+ 15,
+ 16,
+ 17,
+ # Control, Pre-period (mean = 9.5)
+ 8,
+ 9,
+ 10,
+ 11,
+ # Control, Post-period (mean = 10.5) -> diff = 1.0
+ 9,
+ 10,
+ 11,
+ 12,
+ ],
+ "treated": [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+ "post": [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1],
+ }
+ )
# Hand calculation:
# Ȳ_treated_pre = (10+11+12+13)/4 = 11.5
@@ -71,10 +85,7 @@ def generate_hand_calculable_data() -> Tuple[pd.DataFrame, float]:
def generate_did_data_with_covariates(
- n_units: int = 200,
- treatment_effect: float = 2.5,
- covariate_effect: float = 1.5,
- seed: int = 42
+ n_units: int = 200, treatment_effect: float = 2.5, covariate_effect: float = 1.5, seed: int = 42
) -> pd.DataFrame:
"""
Generate 2x2 DiD data with a covariate.
@@ -110,13 +121,15 @@ def generate_did_data_with_covariates(
y += treatment_effect
y += np.random.normal(0, 1)
- data.append({
- 'unit': unit,
- 'outcome': y,
- 'treated': int(is_treated),
- 'post': period,
- 'x1': covariate,
- })
+ data.append(
+ {
+ "unit": unit,
+ "outcome": y,
+ "treated": int(is_treated),
+ "post": period,
+ "x1": covariate,
+ }
+ )
return pd.DataFrame(data)
@@ -126,7 +139,7 @@ def generate_clustered_did_data(
cluster_size: int = 10,
treatment_effect: float = 3.0,
icc: float = 0.3,
- seed: int = 42
+ seed: int = 42,
) -> pd.DataFrame:
"""
Generate DiD data with cluster structure.
@@ -165,21 +178,21 @@ def generate_clustered_did_data(
y += treatment_effect
y += np.random.normal(0, np.sqrt(1 - icc))
- data.append({
- 'cluster_id': cluster,
- 'unit': cluster * cluster_size + unit,
- 'outcome': y,
- 'treated': int(is_treated),
- 'post': period,
- })
+ data.append(
+ {
+ "cluster_id": cluster,
+ "unit": cluster * cluster_size + unit,
+ "outcome": y,
+ "treated": int(is_treated),
+ "post": period,
+ }
+ )
return pd.DataFrame(data)
def generate_heteroskedastic_did_data(
- n_units: int = 200,
- treatment_effect: float = 3.0,
- seed: int = 42
+ n_units: int = 200, treatment_effect: float = 3.0, seed: int = 42
) -> pd.DataFrame:
"""
Generate DiD data with heteroskedastic errors.
@@ -207,12 +220,14 @@ def generate_heteroskedastic_did_data(
y += treatment_effect
y += np.random.normal(0, error_scale)
- data.append({
- 'unit': unit,
- 'outcome': y,
- 'treated': int(is_treated),
- 'post': period,
- })
+ data.append(
+ {
+ "unit": unit,
+ "outcome": y,
+ "treated": int(is_treated),
+ "post": period,
+ }
+ )
return pd.DataFrame(data)
@@ -278,16 +293,12 @@ def test_att_equals_double_difference_of_means(self):
data, expected_att = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
# ATT should match hand calculation exactly (to numerical precision)
- assert np.isclose(results.att, expected_att, rtol=1e-10), \
- f"ATT expected {expected_att}, got {results.att}"
+ assert np.isclose(
+ results.att, expected_att, rtol=1e-10
+ ), f"ATT expected {expected_att}, got {results.att}"
def test_att_equals_regression_interaction_coefficient(self):
"""
@@ -298,71 +309,87 @@ def test_att_equals_regression_interaction_coefficient(self):
data, expected_att = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
# Check that the interaction coefficient matches ATT
assert results.coefficients is not None, "coefficients should not be None"
- interaction_key = 'treated:post'
- assert interaction_key in results.coefficients, \
- f"Expected '{interaction_key}' in coefficients"
- assert np.isclose(results.coefficients[interaction_key], results.att, rtol=1e-10), \
- "ATT should equal interaction coefficient"
- assert np.isclose(results.att, expected_att, rtol=1e-10), \
- f"ATT expected {expected_att}, got {results.att}"
+ interaction_key = "treated:post"
+ assert (
+ interaction_key in results.coefficients
+ ), f"Expected '{interaction_key}' in coefficients"
+ assert np.isclose(
+ results.coefficients[interaction_key], results.att, rtol=1e-10
+ ), "ATT should equal interaction coefficient"
+ assert np.isclose(
+ results.att, expected_att, rtol=1e-10
+ ), f"ATT expected {expected_att}, got {results.att}"
def test_att_with_covariates_close_to_true_effect(self):
"""
Verify ATT with covariates recovers approximately the true effect.
"""
data = generate_did_data_with_covariates(
- n_units=500,
- treatment_effect=2.5,
- covariate_effect=1.5,
- seed=42
+ n_units=500, treatment_effect=2.5, covariate_effect=1.5, seed=42
)
did = DifferenceInDifferences()
results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- covariates=['x1']
+ data, outcome="outcome", treatment="treated", time="post", covariates=["x1"]
)
# Should recover approximately 2.5 treatment effect
- assert abs(results.att - 2.5) < 0.5, \
- f"ATT expected ~2.5, got {results.att}"
+ assert abs(results.att - 2.5) < 0.5, f"ATT expected ~2.5, got {results.att}"
def test_att_sign_matches_data_direction(self):
"""
Verify ATT sign is correct: positive when treated group improves more.
"""
# Positive treatment effect
- data_pos = pd.DataFrame({
- 'outcome': [10, 10, 20, 20, 10, 10, 10, 10], # Treated improves by 10, control stays same
- 'treated': [1, 1, 1, 1, 0, 0, 0, 0],
- 'post': [0, 0, 1, 1, 0, 0, 1, 1],
- })
+ data_pos = pd.DataFrame(
+ {
+ "outcome": [
+ 10,
+ 10,
+ 20,
+ 20,
+ 10,
+ 10,
+ 10,
+ 10,
+ ], # Treated improves by 10, control stays same
+ "treated": [1, 1, 1, 1, 0, 0, 0, 0],
+ "post": [0, 0, 1, 1, 0, 0, 1, 1],
+ }
+ )
did = DifferenceInDifferences()
- results_pos = did.fit(data_pos, outcome='outcome', treatment='treated', time='post')
- assert np.isclose(results_pos.att, 10.0, rtol=1e-10), f"Expected ATT=10.0, got {results_pos.att}"
+ results_pos = did.fit(data_pos, outcome="outcome", treatment="treated", time="post")
+ assert np.isclose(
+ results_pos.att, 10.0, rtol=1e-10
+ ), f"Expected ATT=10.0, got {results_pos.att}"
# Negative treatment effect
- data_neg = pd.DataFrame({
- 'outcome': [20, 20, 10, 10, 10, 10, 10, 10], # Treated decreases by 10, control stays same
- 'treated': [1, 1, 1, 1, 0, 0, 0, 0],
- 'post': [0, 0, 1, 1, 0, 0, 1, 1],
- })
+ data_neg = pd.DataFrame(
+ {
+ "outcome": [
+ 20,
+ 20,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ ], # Treated decreases by 10, control stays same
+ "treated": [1, 1, 1, 1, 0, 0, 0, 0],
+ "post": [0, 0, 1, 1, 0, 0, 1, 1],
+ }
+ )
- results_neg = did.fit(data_neg, outcome='outcome', treatment='treated', time='post')
- assert np.isclose(results_neg.att, -10.0, rtol=1e-10), f"Expected ATT=-10.0, got {results_neg.att}"
+ results_neg = did.fit(data_neg, outcome="outcome", treatment="treated", time="post")
+ assert np.isclose(
+ results_neg.att, -10.0, rtol=1e-10
+ ), f"Expected ATT=-10.0, got {results_neg.att}"
# =============================================================================
@@ -374,11 +401,7 @@ class TestRBenchmarkDiD:
"""Tests comparing Python implementation to R's fixest::feols()."""
def _run_r_estimation(
- self,
- data_path: str,
- covariates: list = None,
- cluster: str = None,
- robust: bool = True
+ self, data_path: str, covariates: list = None, cluster: str = None, robust: bool = True
) -> Dict[str, Any]:
"""
Run R's fixest::feols() and return results as dictionary.
@@ -407,13 +430,13 @@ def _run_r_estimation(
# Build vcov specification
if cluster:
- vcov_spec = f'~{cluster}'
+ vcov_spec = f"~{cluster}"
elif robust:
vcov_spec = '"HC1"'
else:
vcov_spec = '"iid"'
- r_script = f'''
+ r_script = f"""
suppressMessages(library(fixest))
suppressMessages(library(jsonlite))
@@ -453,13 +476,10 @@ def _run_r_estimation(
)
cat(toJSON(output, pretty = TRUE))
- '''
+ """
result = subprocess.run(
- ["Rscript", "-e", r_script],
- capture_output=True,
- text=True,
- timeout=60
+ ["Rscript", "-e", r_script], capture_output=True, text=True, timeout=60
)
if result.returncode != 0:
@@ -468,7 +488,7 @@ def _run_r_estimation(
parsed = json.loads(result.stdout)
# Handle R's JSON serialization quirks
- for key in ['att', 'se', 'pvalue', 'ci_lower', 'ci_upper', 'r_squared']:
+ for key in ["att", "se", "pvalue", "ci_lower", "ci_upper", "r_squared"]:
if isinstance(parsed.get(key), list):
parsed[key] = parsed[key][0]
@@ -489,12 +509,14 @@ def benchmark_data(self, tmp_path):
y += 3.0 # True ATT
y += np.random.normal(0, 1)
- data.append({
- 'unit': unit,
- 'outcome': y,
- 'treated': int(is_treated),
- 'post': period,
- })
+ data.append(
+ {
+ "unit": unit,
+ "outcome": y,
+ "treated": int(is_treated),
+ "post": period,
+ }
+ )
df = pd.DataFrame(data)
csv_path = tmp_path / "benchmark_data.csv"
@@ -505,10 +527,7 @@ def benchmark_data(self, tmp_path):
def benchmark_data_with_covariate(self, tmp_path):
"""Generate benchmark data with covariate and save to CSV."""
df = generate_did_data_with_covariates(
- n_units=200,
- treatment_effect=2.5,
- covariate_effect=1.5,
- seed=12345
+ n_units=200, treatment_effect=2.5, covariate_effect=1.5, seed=12345
)
csv_path = tmp_path / "benchmark_data_cov.csv"
df.to_csv(csv_path, index=False)
@@ -518,11 +537,7 @@ def benchmark_data_with_covariate(self, tmp_path):
def benchmark_data_clustered(self, tmp_path):
"""Generate clustered benchmark data and save to CSV."""
df = generate_clustered_did_data(
- n_clusters=20,
- cluster_size=10,
- treatment_effect=3.0,
- icc=0.3,
- seed=12345
+ n_clusters=20, cluster_size=10, treatment_effect=3.0, icc=0.3, seed=12345
)
csv_path = tmp_path / "benchmark_data_clustered.csv"
df.to_csv(csv_path, index=False)
@@ -534,79 +549,63 @@ def test_att_matches_r_basic(self, require_fixest, benchmark_data):
# Python estimation
did = DifferenceInDifferences(robust=True)
- py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ py_results = did.fit(data, outcome="outcome", treatment="treated", time="post")
# R estimation
r_results = self._run_r_estimation(csv_path, robust=True)
# Compare ATT - R's JSON output is truncated to 4 decimal places
# so use 1e-3 tolerance
- assert np.isclose(py_results.att, r_results['att'], rtol=1e-3), \
- f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
+ assert np.isclose(
+ py_results.att, r_results["att"], rtol=1e-3
+ ), f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
def test_se_matches_r_hc1(self, require_fixest, benchmark_data):
"""Test HC1 standard errors match R within 5%."""
data, csv_path = benchmark_data
did = DifferenceInDifferences(robust=True)
- py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ py_results = did.fit(data, outcome="outcome", treatment="treated", time="post")
r_results = self._run_r_estimation(csv_path, robust=True)
# Compare SE - within 5% tolerance
- rel_diff = abs(py_results.se - r_results['se']) / r_results['se']
- assert rel_diff < 0.05, \
- f"SE mismatch: Python={py_results.se}, R={r_results['se']}, diff={rel_diff*100:.1f}%"
+ rel_diff = abs(py_results.se - r_results["se"]) / r_results["se"]
+ assert (
+ rel_diff < 0.05
+ ), f"SE mismatch: Python={py_results.se}, R={r_results['se']}, diff={rel_diff*100:.1f}%"
def test_pvalue_matches_r(self, require_fixest, benchmark_data):
"""Test p-value matches R within 0.01."""
data, csv_path = benchmark_data
did = DifferenceInDifferences(robust=True)
- py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ py_results = did.fit(data, outcome="outcome", treatment="treated", time="post")
r_results = self._run_r_estimation(csv_path, robust=True)
# Compare p-value - within 0.01 absolute
- assert abs(py_results.p_value - r_results['pvalue']) < 0.01, \
- f"P-value mismatch: Python={py_results.p_value}, R={r_results['pvalue']}"
+ assert (
+ abs(py_results.p_value - r_results["pvalue"]) < 0.01
+ ), f"P-value mismatch: Python={py_results.p_value}, R={r_results['pvalue']}"
def test_ci_overlaps_r(self, require_fixest, benchmark_data):
"""Test confidence intervals overlap with R."""
data, csv_path = benchmark_data
did = DifferenceInDifferences(robust=True, alpha=0.05)
- py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ py_results = did.fit(data, outcome="outcome", treatment="treated", time="post")
r_results = self._run_r_estimation(csv_path, robust=True)
py_lower, py_upper = py_results.conf_int
- r_lower = float(r_results['ci_lower'])
- r_upper = float(r_results['ci_upper'])
+ r_lower = float(r_results["ci_lower"])
+ r_upper = float(r_results["ci_upper"])
# CIs should overlap
- assert py_lower < r_upper and r_lower < py_upper, \
- f"CIs don't overlap: Python=[{py_lower}, {py_upper}], R=[{r_lower}, {r_upper}]"
+ assert (
+ py_lower < r_upper and r_lower < py_upper
+ ), f"CIs don't overlap: Python=[{py_lower}, {py_upper}], R=[{r_lower}, {r_upper}]"
def test_att_matches_r_with_covariates(self, require_fixest, benchmark_data_with_covariate):
"""Test ATT matches R when covariates are included."""
@@ -614,42 +613,36 @@ def test_att_matches_r_with_covariates(self, require_fixest, benchmark_data_with
did = DifferenceInDifferences(robust=True)
py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- covariates=['x1']
+ data, outcome="outcome", treatment="treated", time="post", covariates=["x1"]
)
- r_results = self._run_r_estimation(csv_path, covariates=['x1'], robust=True)
+ r_results = self._run_r_estimation(csv_path, covariates=["x1"], robust=True)
# Compare ATT - R's JSON output is truncated to 4 decimal places
# so use 1e-3 tolerance
- assert np.isclose(py_results.att, r_results['att'], rtol=1e-3), \
- f"ATT mismatch with covariates: Python={py_results.att}, R={r_results['att']}"
+ assert np.isclose(
+ py_results.att, r_results["att"], rtol=1e-3
+ ), f"ATT mismatch with covariates: Python={py_results.att}, R={r_results['att']}"
def test_se_matches_r_with_clustering(self, require_fixest, benchmark_data_clustered):
"""Test cluster-robust SEs match R."""
data, csv_path = benchmark_data_clustered
- did = DifferenceInDifferences(cluster='cluster_id')
- py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post'
- )
+ did = DifferenceInDifferences(cluster="cluster_id")
+ py_results = did.fit(data, outcome="outcome", treatment="treated", time="post")
- r_results = self._run_r_estimation(csv_path, cluster='cluster_id')
+ r_results = self._run_r_estimation(csv_path, cluster="cluster_id")
# Compare ATT - R's JSON output is truncated to 4 decimal places
- assert np.isclose(py_results.att, r_results['att'], rtol=1e-3), \
- f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
+ assert np.isclose(
+ py_results.att, r_results["att"], rtol=1e-3
+ ), f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
# Compare SE - within 10% for clustered (more variance expected)
- rel_diff = abs(py_results.se - r_results['se']) / r_results['se']
- assert rel_diff < 0.10, \
- f"Clustered SE mismatch: Python={py_results.se}, R={r_results['se']}, diff={rel_diff*100:.1f}%"
+ rel_diff = abs(py_results.se - r_results["se"]) / r_results["se"]
+ assert (
+ rel_diff < 0.10
+ ), f"Clustered SE mismatch: Python={py_results.se}, R={r_results['se']}, diff={rel_diff*100:.1f}%"
# =============================================================================
@@ -668,28 +661,23 @@ def test_hc1_vs_classical_se_differ(self):
different standard errors. The relationship depends on the pattern
of heteroskedasticity.
"""
- data = generate_heteroskedastic_did_data(
- n_units=200,
- treatment_effect=3.0,
- seed=42
- )
+ data = generate_heteroskedastic_did_data(n_units=200, treatment_effect=3.0, seed=42)
did_hc1 = DifferenceInDifferences(robust=True)
did_classical = DifferenceInDifferences(robust=False)
- results_hc1 = did_hc1.fit(
- data, outcome='outcome', treatment='treated', time='post'
- )
+ results_hc1 = did_hc1.fit(data, outcome="outcome", treatment="treated", time="post")
results_classical = did_classical.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ data, outcome="outcome", treatment="treated", time="post"
)
# SE methods should produce both valid and positive SEs
assert results_hc1.se > 0, "HC1 SE should be positive"
assert results_classical.se > 0, "Classical SE should be positive"
# ATT should be the same regardless of SE method
- assert np.isclose(results_hc1.att, results_classical.att, rtol=1e-10), \
- "ATT should be identical for both SE methods"
+ assert np.isclose(
+ results_hc1.att, results_classical.att, rtol=1e-10
+ ), "ATT should be identical for both SE methods"
def test_cluster_se_differs_from_robust(self):
"""
@@ -699,30 +687,23 @@ def test_cluster_se_differs_from_robust(self):
is within-cluster correlation.
"""
data = generate_clustered_did_data(
- n_clusters=20,
- cluster_size=10,
- treatment_effect=3.0,
- icc=0.3,
- seed=42
+ n_clusters=20, cluster_size=10, treatment_effect=3.0, icc=0.3, seed=42
)
did_robust = DifferenceInDifferences(robust=True)
- did_cluster = DifferenceInDifferences(cluster='cluster_id')
+ did_cluster = DifferenceInDifferences(cluster="cluster_id")
- results_robust = did_robust.fit(
- data, outcome='outcome', treatment='treated', time='post'
- )
- results_cluster = did_cluster.fit(
- data, outcome='outcome', treatment='treated', time='post'
- )
+ results_robust = did_robust.fit(data, outcome="outcome", treatment="treated", time="post")
+ results_cluster = did_cluster.fit(data, outcome="outcome", treatment="treated", time="post")
# Both SE methods should produce valid positive SEs
assert results_robust.se > 0, "Robust SE should be positive"
assert results_cluster.se > 0, "Cluster SE should be positive"
# ATT should be the same regardless of SE method
- assert np.isclose(results_robust.att, results_cluster.att, rtol=1e-10), \
- "ATT should be identical for both SE methods"
+ assert np.isclose(
+ results_robust.att, results_cluster.att, rtol=1e-10
+ ), "ATT should be identical for both SE methods"
def test_se_decreases_with_sample_size(self):
"""Verify SE decreases approximately with sqrt(n)."""
@@ -738,18 +719,17 @@ def test_se_decreases_with_sample_size(self):
if is_treated and period == 1:
y += 3.0
y += np.random.normal(0, 1)
- data.append({
- 'outcome': y,
- 'treated': int(is_treated),
- 'post': period,
- })
+ data.append(
+ {
+ "outcome": y,
+ "treated": int(is_treated),
+ "post": period,
+ }
+ )
did = DifferenceInDifferences(robust=True)
results = did.fit(
- pd.DataFrame(data),
- outcome='outcome',
- treatment='treated',
- time='post'
+ pd.DataFrame(data), outcome="outcome", treatment="treated", time="post"
)
ses.append(results.se)
@@ -766,9 +746,7 @@ def test_vcov_positive_semidefinite(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences(robust=True)
- results = did.fit(
- data, outcome='outcome', treatment='treated', time='post'
- )
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert results.vcov is not None, "vcov should not be None"
@@ -780,8 +758,9 @@ def test_vcov_positive_semidefinite(self):
# All eigenvalues should be >= 0 (with numerical tolerance)
eigenvalues = np.linalg.eigvalsh(vcov)
- assert np.all(eigenvalues >= -1e-10), \
- f"VCoV not positive semi-definite: min eigenvalue = {eigenvalues.min()}"
+ assert np.all(
+ eigenvalues >= -1e-10
+ ), f"VCoV not positive semi-definite: min eigenvalue = {eigenvalues.min()}"
class TestWildBootstrapInference:
@@ -791,97 +770,71 @@ def test_wild_bootstrap_produces_valid_se(self, ci_params):
"""Test wild bootstrap produces finite, positive SE."""
n_boot = ci_params.bootstrap(199)
data = generate_clustered_did_data(
- n_clusters=15,
- cluster_size=10,
- treatment_effect=3.0,
- seed=42
+ n_clusters=15, cluster_size=10, treatment_effect=3.0, seed=42
)
did = DifferenceInDifferences(
- inference='wild_bootstrap',
- cluster='cluster_id',
+ inference="wild_bootstrap",
+ cluster="cluster_id",
n_bootstrap=n_boot,
- bootstrap_weights='rademacher',
- seed=42
- )
- results = did.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ bootstrap_weights="rademacher",
+ seed=42,
)
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert np.isfinite(results.se), "Bootstrap SE should be finite"
assert results.se > 0, "Bootstrap SE should be positive"
- assert results.inference_method == 'wild_bootstrap'
+ assert results.inference_method == "wild_bootstrap"
def test_wild_bootstrap_pvalue_in_valid_range(self, ci_params):
"""Test wild bootstrap p-value is in [0, 1]."""
n_boot = ci_params.bootstrap(199)
data = generate_clustered_did_data(
- n_clusters=15,
- cluster_size=10,
- treatment_effect=3.0,
- seed=42
+ n_clusters=15, cluster_size=10, treatment_effect=3.0, seed=42
)
did = DifferenceInDifferences(
- inference='wild_bootstrap',
- cluster='cluster_id',
- n_bootstrap=n_boot,
- seed=42
- )
- results = did.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ inference="wild_bootstrap", cluster="cluster_id", n_bootstrap=n_boot, seed=42
)
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
- assert 0 <= results.p_value <= 1, \
- f"P-value {results.p_value} not in [0, 1]"
+ assert 0 <= results.p_value <= 1, f"P-value {results.p_value} not in [0, 1]"
def test_wild_bootstrap_ci_contains_point_estimate(self, ci_params):
"""Test wild bootstrap CI contains point estimate."""
n_boot = ci_params.bootstrap(199)
data = generate_clustered_did_data(
- n_clusters=15,
- cluster_size=10,
- treatment_effect=3.0,
- seed=42
+ n_clusters=15, cluster_size=10, treatment_effect=3.0, seed=42
)
did = DifferenceInDifferences(
- inference='wild_bootstrap',
- cluster='cluster_id',
- n_bootstrap=n_boot,
- seed=42
- )
- results = did.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ inference="wild_bootstrap", cluster="cluster_id", n_bootstrap=n_boot, seed=42
)
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
lower, upper = results.conf_int
# CI should contain point estimate or be very close
# (bootstrap CI may differ slightly from point estimate)
- assert lower < results.att + 0.5 and results.att - 0.5 < upper, \
- f"CI [{lower}, {upper}] should approximately contain ATT {results.att}"
+ assert (
+ lower < results.att + 0.5 and results.att - 0.5 < upper
+ ), f"CI [{lower}, {upper}] should approximately contain ATT {results.att}"
@pytest.mark.parametrize("weight_type", ["rademacher", "mammen", "webb"])
def test_wild_bootstrap_weight_types(self, weight_type, ci_params):
"""Test all wild bootstrap weight types work."""
n_boot = ci_params.bootstrap(99)
data = generate_clustered_did_data(
- n_clusters=15,
- cluster_size=10,
- treatment_effect=3.0,
- seed=42
+ n_clusters=15, cluster_size=10, treatment_effect=3.0, seed=42
)
did = DifferenceInDifferences(
- inference='wild_bootstrap',
- cluster='cluster_id',
+ inference="wild_bootstrap",
+ cluster="cluster_id",
n_bootstrap=n_boot,
bootstrap_weights=weight_type,
- seed=42
- )
- results = did.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ seed=42,
)
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert np.isfinite(results.se), f"SE should be finite for {weight_type}"
assert results.se > 0, f"SE should be positive for {weight_type}"
@@ -904,16 +857,18 @@ def test_empty_cell_produces_nan_or_warning(self):
error → ValueError, silent → NaN without warning).
"""
# No treated units in pre-period
- data = pd.DataFrame({
- 'outcome': [10, 11, 12, 13],
- 'treated': [0, 0, 1, 1],
- 'post': [0, 1, 1, 1], # No treated-pre observations
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 12, 13],
+ "treated": [0, 0, 1, 1],
+ "post": [0, 1, 1, 1], # No treated-pre observations
+ }
+ )
- did = DifferenceInDifferences(rank_deficient_action='warn')
+ did = DifferenceInDifferences(rank_deficient_action="warn")
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
# Should have rank deficiency warning due to empty cell
# causing collinearity between treated and post columns
@@ -922,8 +877,9 @@ def test_empty_cell_produces_nan_or_warning(self):
has_nan = results.coefficients is not None and any(
np.isnan(v) for v in results.coefficients.values()
)
- assert len(rank_warnings) > 0 or has_nan, \
- "Empty cell should cause rank deficiency warning or NaN coefficient"
+ assert (
+ len(rank_warnings) > 0 or has_nan
+ ), "Empty cell should cause rank deficiency warning or NaN coefficient"
def test_singleton_cluster_handling(self):
"""
@@ -932,16 +888,18 @@ def test_singleton_cluster_handling(self):
With clustering, singleton clusters may produce warnings.
"""
# Create data with one singleton cluster
- data = pd.DataFrame({
- 'outcome': list(range(20)) + [100],
- 'treated': [0]*10 + [1]*10 + [1],
- 'post': [0, 1]*10 + [1],
- 'cluster': list(range(10))*2 + [99], # Cluster 99 is singleton
- })
-
- did = DifferenceInDifferences(cluster='cluster')
+ data = pd.DataFrame(
+ {
+ "outcome": list(range(20)) + [100],
+ "treated": [0] * 10 + [1] * 10 + [1],
+ "post": [0, 1] * 10 + [1],
+ "cluster": list(range(10)) * 2 + [99], # Cluster 99 is singleton
+ }
+ )
+
+ did = DifferenceInDifferences(cluster="cluster")
# Should run without error (warning may be emitted)
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert results is not None
def test_rank_deficient_warn_mode(self):
@@ -950,23 +908,21 @@ def test_rank_deficient_warn_mode(self):
REGISTRY.md: "Rank-deficient design matrix: warns and sets NA for dropped coefficients"
"""
- data = pd.DataFrame({
- 'outcome': [10, 11, 15, 18, 9, 10, 12, 13],
- 'treated': [1, 1, 1, 1, 0, 0, 0, 0],
- 'post': [0, 0, 1, 1, 0, 0, 1, 1],
- 'collinear': [1, 1, 1, 1, 0, 0, 0, 0], # Perfectly collinear with treated
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 15, 18, 9, 10, 12, 13],
+ "treated": [1, 1, 1, 1, 0, 0, 0, 0],
+ "post": [0, 0, 1, 1, 0, 0, 1, 1],
+ "collinear": [1, 1, 1, 1, 0, 0, 0, 0], # Perfectly collinear with treated
+ }
+ )
- did = DifferenceInDifferences(rank_deficient_action='warn')
+ did = DifferenceInDifferences(rank_deficient_action="warn")
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- covariates=['collinear']
+ data, outcome="outcome", treatment="treated", time="post", covariates=["collinear"]
)
# Should have warning about rank deficiency
@@ -982,45 +938,41 @@ def test_rank_deficient_error_mode(self):
"""
Test rank_deficient_action='error' raises ValueError.
"""
- data = pd.DataFrame({
- 'outcome': [10, 11, 15, 18, 9, 10, 12, 13],
- 'treated': [1, 1, 1, 1, 0, 0, 0, 0],
- 'post': [0, 0, 1, 1, 0, 0, 1, 1],
- 'collinear': [1, 1, 1, 1, 0, 0, 0, 0],
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 15, 18, 9, 10, 12, 13],
+ "treated": [1, 1, 1, 1, 0, 0, 0, 0],
+ "post": [0, 0, 1, 1, 0, 0, 1, 1],
+ "collinear": [1, 1, 1, 1, 0, 0, 0, 0],
+ }
+ )
- did = DifferenceInDifferences(rank_deficient_action='error')
+ did = DifferenceInDifferences(rank_deficient_action="error")
with pytest.raises(ValueError, match="rank-deficient"):
did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- covariates=['collinear']
+ data, outcome="outcome", treatment="treated", time="post", covariates=["collinear"]
)
def test_rank_deficient_silent_mode(self):
"""
Test rank_deficient_action='silent' produces no warning.
"""
- data = pd.DataFrame({
- 'outcome': [10, 11, 15, 18, 9, 10, 12, 13],
- 'treated': [1, 1, 1, 1, 0, 0, 0, 0],
- 'post': [0, 0, 1, 1, 0, 0, 1, 1],
- 'collinear': [1, 1, 1, 1, 0, 0, 0, 0],
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 15, 18, 9, 10, 12, 13],
+ "treated": [1, 1, 1, 1, 0, 0, 0, 0],
+ "post": [0, 0, 1, 1, 0, 0, 1, 1],
+ "collinear": [1, 1, 1, 1, 0, 0, 0, 0],
+ }
+ )
- did = DifferenceInDifferences(rank_deficient_action='silent')
+ did = DifferenceInDifferences(rank_deficient_action="silent")
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- covariates=['collinear']
+ data, outcome="outcome", treatment="treated", time="post", covariates=["collinear"]
)
rank_warnings = [x for x in w if "rank" in str(x.message).lower()]
@@ -1035,65 +987,75 @@ def test_non_binary_treatment_raises_error(self):
REGISTRY.md: "Treatment and post indicators must be binary (0/1) with variation in both"
"""
- data = pd.DataFrame({
- 'outcome': [10, 11, 12, 13],
- 'treated': [0, 1, 2, 3], # Non-binary
- 'post': [0, 0, 1, 1],
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 12, 13],
+ "treated": [0, 1, 2, 3], # Non-binary
+ "post": [0, 0, 1, 1],
+ }
+ )
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="binary"):
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
def test_non_binary_time_raises_error(self):
"""
Test non-binary time indicator raises ValueError.
"""
- data = pd.DataFrame({
- 'outcome': [10, 11, 12, 13],
- 'treated': [0, 0, 1, 1],
- 'post': [0, 1, 2, 3], # Non-binary
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 12, 13],
+ "treated": [0, 0, 1, 1],
+ "post": [0, 1, 2, 3], # Non-binary
+ }
+ )
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="binary"):
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
def test_no_treatment_variation_raises_error(self):
"""Test error when treatment has no variation."""
- data = pd.DataFrame({
- 'outcome': [10, 11, 12, 13],
- 'treated': [1, 1, 1, 1], # All treated
- 'post': [0, 0, 1, 1],
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 12, 13],
+ "treated": [1, 1, 1, 1], # All treated
+ "post": [0, 0, 1, 1],
+ }
+ )
did = DifferenceInDifferences()
with pytest.raises(ValueError):
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
def test_no_time_variation_raises_error(self):
"""Test error when time has no variation."""
- data = pd.DataFrame({
- 'outcome': [10, 11, 12, 13],
- 'treated': [0, 0, 1, 1],
- 'post': [1, 1, 1, 1], # All post
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, 11, 12, 13],
+ "treated": [0, 0, 1, 1],
+ "post": [1, 1, 1, 1], # All post
+ }
+ )
did = DifferenceInDifferences()
with pytest.raises(ValueError):
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
def test_missing_values_raise_error(self):
"""Test error when data contains missing values."""
- data = pd.DataFrame({
- 'outcome': [10, np.nan, 12, 13],
- 'treated': [0, 0, 1, 1],
- 'post': [0, 0, 1, 1],
- })
+ data = pd.DataFrame(
+ {
+ "outcome": [10, np.nan, 12, 13],
+ "treated": [0, 0, 1, 1],
+ "post": [0, 0, 1, 1],
+ }
+ )
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="missing"):
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
# =============================================================================
@@ -1109,7 +1071,7 @@ def test_formula_star_syntax(self):
data, expected_att = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, formula='outcome ~ treated * post')
+ results = did.fit(data, formula="outcome ~ treated * post")
assert np.isclose(results.att, expected_att, rtol=1e-10)
@@ -1118,7 +1080,7 @@ def test_formula_explicit_interaction_syntax(self):
data, expected_att = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, formula='outcome ~ treated + post + treated:post')
+ results = did.fit(data, formula="outcome ~ treated + post + treated:post")
assert np.isclose(results.att, expected_att, rtol=1e-10)
@@ -1128,10 +1090,10 @@ def test_formula_with_covariates(self):
did = DifferenceInDifferences()
# Note: formula parsing for covariates after interaction
- results = did.fit(data, formula='outcome ~ treated * post + x1')
+ results = did.fit(data, formula="outcome ~ treated * post + x1")
assert results.coefficients is not None, "coefficients should not be None"
- assert 'x1' in results.coefficients
+ assert "x1" in results.coefficients
assert np.isfinite(results.att)
def test_formula_missing_tilde_raises_error(self):
@@ -1140,7 +1102,7 @@ def test_formula_missing_tilde_raises_error(self):
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="~"):
- did.fit(data, formula='outcome treated * post')
+ did.fit(data, formula="outcome treated * post")
def test_formula_missing_interaction_raises_error(self):
"""Test error when formula has no interaction term."""
@@ -1148,7 +1110,7 @@ def test_formula_missing_interaction_raises_error(self):
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="interaction"):
- did.fit(data, formula='outcome ~ treated + post')
+ did.fit(data, formula="outcome ~ treated + post")
def test_formula_missing_column_raises_error(self):
"""Test error when formula references non-existent column."""
@@ -1156,7 +1118,7 @@ def test_formula_missing_column_raises_error(self):
did = DifferenceInDifferences()
with pytest.raises(ValueError, match="not found"):
- did.fit(data, formula='outcome ~ nonexistent * post')
+ did.fit(data, formula="outcome ~ nonexistent * post")
def test_formula_and_params_conflict(self):
"""Test that formula overrides outcome/treatment/time params."""
@@ -1166,10 +1128,10 @@ def test_formula_and_params_conflict(self):
# Formula should override the (invalid) parameter values
results = did.fit(
data,
- formula='outcome ~ treated * post',
- outcome='wrong', # Would fail if used
- treatment='wrong',
- time='wrong'
+ formula="outcome ~ treated * post",
+ outcome="wrong", # Would fail if used
+ treatment="wrong",
+ time="wrong",
)
assert np.isclose(results.att, expected_att, rtol=1e-10)
@@ -1204,13 +1166,15 @@ def panel_data(self):
y += 3.0 # True ATT
y += np.random.normal(0, 0.5)
- data.append({
- 'unit': unit,
- 'period': period,
- 'treated': int(is_treated),
- 'post': post,
- 'outcome': y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": post,
+ "outcome": y,
+ }
+ )
return pd.DataFrame(data)
@@ -1218,31 +1182,21 @@ def test_absorb_produces_valid_att(self, panel_data):
"""Test that absorb option produces valid ATT estimate."""
did = DifferenceInDifferences()
results = did.fit(
- panel_data,
- outcome='outcome',
- treatment='treated',
- time='post',
- absorb=['unit']
+ panel_data, outcome="outcome", treatment="treated", time="post", absorb=["unit"]
)
# ATT should be close to 3.0
- assert abs(results.att - 3.0) < 1.0, \
- f"ATT with absorb expected ~3.0, got {results.att}"
+ assert abs(results.att - 3.0) < 1.0, f"ATT with absorb expected ~3.0, got {results.att}"
def test_fixed_effects_produces_valid_att(self, panel_data):
"""Test that fixed_effects option produces valid ATT estimate."""
did = DifferenceInDifferences()
results = did.fit(
- panel_data,
- outcome='outcome',
- treatment='treated',
- time='post',
- fixed_effects=['unit']
+ panel_data, outcome="outcome", treatment="treated", time="post", fixed_effects=["unit"]
)
# ATT should be close to 3.0
- assert abs(results.att - 3.0) < 1.0, \
- f"ATT with FE expected ~3.0, got {results.att}"
+ assert abs(results.att - 3.0) < 1.0, f"ATT with FE expected ~3.0, got {results.att}"
def test_absorb_produces_valid_results(self, panel_data):
"""Test that absorb option produces valid ATT estimate.
@@ -1252,11 +1206,7 @@ def test_absorb_produces_valid_results(self, panel_data):
"""
did = DifferenceInDifferences()
results = did.fit(
- panel_data,
- outcome='outcome',
- treatment='treated',
- time='post',
- absorb=['unit']
+ panel_data, outcome="outcome", treatment="treated", time="post", absorb=["unit"]
)
# Verify ATT is estimated (not NaN)
@@ -1277,7 +1227,7 @@ def _run_r_feols_fe(self, data_path: str) -> Dict[str, Any]:
"""
escaped_path = data_path.replace("\\", "/")
- r_script = f'''
+ r_script = f"""
suppressMessages(library(fixest))
suppressMessages(library(jsonlite))
@@ -1304,20 +1254,17 @@ def _run_r_feols_fe(self, data_path: str) -> Dict[str, Any]:
)
cat(toJSON(output, pretty = TRUE))
- '''
+ """
result = subprocess.run(
- ["Rscript", "-e", r_script],
- capture_output=True,
- text=True,
- timeout=60
+ ["Rscript", "-e", r_script], capture_output=True, text=True, timeout=60
)
if result.returncode != 0:
raise RuntimeError(f"R script failed: {result.stderr}")
parsed = json.loads(result.stdout)
- for key in ['att', 'se', 'r_squared']:
+ for key in ["att", "se", "r_squared"]:
if isinstance(parsed.get(key), list):
parsed[key] = parsed[key][0]
@@ -1344,13 +1291,15 @@ def panel_data_csv(self, tmp_path):
y += 3.0
y += np.random.normal(0, 0.5)
- data.append({
- 'unit': unit,
- 'period': period,
- 'treated': int(is_treated),
- 'post': post,
- 'outcome': y,
- })
+ data.append(
+ {
+ "unit": unit,
+ "period": period,
+ "treated": int(is_treated),
+ "post": post,
+ "outcome": y,
+ }
+ )
df = pd.DataFrame(data)
csv_path = tmp_path / "panel_data.csv"
@@ -1364,19 +1313,16 @@ def test_absorb_matches_r_feols_fe(self, require_fixest, panel_data_csv):
# Python with absorb
did = DifferenceInDifferences(robust=True)
py_results = did.fit(
- data,
- outcome='outcome',
- treatment='treated',
- time='post',
- absorb=['unit']
+ data, outcome="outcome", treatment="treated", time="post", absorb=["unit"]
)
# R with | unit syntax
r_results = self._run_r_feols_fe(csv_path)
# Compare ATT
- assert np.isclose(py_results.att, r_results['att'], rtol=0.01), \
- f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
+ assert np.isclose(
+ py_results.att, r_results["att"], rtol=0.01
+ ), f"ATT mismatch: Python={py_results.att}, R={r_results['att']}"
# =============================================================================
@@ -1391,25 +1337,25 @@ def test_get_params_returns_all_parameters(self):
"""Test that get_params returns all constructor parameters."""
did = DifferenceInDifferences(
robust=False,
- cluster='cluster_id',
+ cluster="cluster_id",
alpha=0.10,
- inference='wild_bootstrap',
+ inference="wild_bootstrap",
n_bootstrap=500,
- bootstrap_weights='webb',
+ bootstrap_weights="webb",
seed=123,
- rank_deficient_action='silent'
+ rank_deficient_action="silent",
)
params = did.get_params()
- assert params['robust'] is False
- assert params['cluster'] == 'cluster_id'
- assert params['alpha'] == 0.10
- assert params['inference'] == 'wild_bootstrap'
- assert params['n_bootstrap'] == 500
- assert params['bootstrap_weights'] == 'webb'
- assert params['seed'] == 123
- assert params['rank_deficient_action'] == 'silent'
+ assert params["robust"] is False
+ assert params["cluster"] == "cluster_id"
+ assert params["alpha"] == 0.10
+ assert params["inference"] == "wild_bootstrap"
+ assert params["n_bootstrap"] == 500
+ assert params["bootstrap_weights"] == "webb"
+ assert params["seed"] == 123
+ assert params["rank_deficient_action"] == "silent"
def test_set_params_modifies_attributes(self):
"""Test that set_params modifies estimator attributes."""
@@ -1440,11 +1386,9 @@ def test_params_affect_estimation(self):
did_classical = DifferenceInDifferences()
did_classical.set_params(robust=False)
- results_default = did_default.fit(
- data, outcome='outcome', treatment='treated', time='post'
- )
+ results_default = did_default.fit(data, outcome="outcome", treatment="treated", time="post")
results_classical = did_classical.fit(
- data, outcome='outcome', treatment='treated', time='post'
+ data, outcome="outcome", treatment="treated", time="post"
)
# VCoV should differ (robust vs classical)
@@ -1466,7 +1410,7 @@ def test_summary_contains_key_info(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
summary = results.summary()
@@ -1478,24 +1422,24 @@ def test_to_dict_contains_all_fields(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
result_dict = results.to_dict()
- assert 'att' in result_dict
- assert 'se' in result_dict
- assert 't_stat' in result_dict
- assert 'p_value' in result_dict
+ assert "att" in result_dict
+ assert "se" in result_dict
+ assert "t_stat" in result_dict
+ assert "p_value" in result_dict
# Note: conf_int is split into conf_int_lower and conf_int_upper
- assert 'conf_int_lower' in result_dict or 'conf_int' in result_dict
- assert 'n_obs' in result_dict
+ assert "conf_int_lower" in result_dict or "conf_int" in result_dict
+ assert "n_obs" in result_dict
def test_to_dataframe_has_correct_shape(self):
"""Test to_dataframe produces DataFrame with one row."""
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
df = results.to_dataframe()
@@ -1507,7 +1451,7 @@ def test_is_significant_property(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences(alpha=0.05)
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert isinstance(results.is_significant, bool)
@@ -1516,7 +1460,7 @@ def test_significance_stars_property(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert results.significance_stars in ["", ".", "*", "**", "***"]
@@ -1525,37 +1469,38 @@ def test_coefficients_dict(self):
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert results.coefficients is not None, "coefficients should not be None"
- assert 'const' in results.coefficients
- assert 'treated' in results.coefficients
- assert 'post' in results.coefficients
- assert 'treated:post' in results.coefficients
+ assert "const" in results.coefficients
+ assert "treated" in results.coefficients
+ assert "post" in results.coefficients
+ assert "treated:post" in results.coefficients
def test_residuals_and_fitted_values(self):
"""Test residuals and fitted values are correctly computed."""
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- results = did.fit(data, outcome='outcome', treatment='treated', time='post')
+ results = did.fit(data, outcome="outcome", treatment="treated", time="post")
assert results.residuals is not None, "residuals should not be None"
assert results.fitted_values is not None, "fitted_values should not be None"
# Residuals + fitted should equal original outcome
reconstructed = results.residuals + results.fitted_values
- original = data['outcome'].values.astype(float)
+ original = data["outcome"].values.astype(float)
- assert np.allclose(reconstructed, original), \
- "Residuals + fitted should equal original outcome"
+ assert np.allclose(
+ reconstructed, original
+ ), "Residuals + fitted should equal original outcome"
def test_predict_contract_points_to_fitted_values(self):
"""predict() is intentionally unsupported until post-estimation is designed."""
data, _ = generate_hand_calculable_data()
did = DifferenceInDifferences()
- did.fit(data, outcome='outcome', treatment='treated', time='post')
+ did.fit(data, outcome="outcome", treatment="treated", time="post")
with pytest.raises(
NotImplementedError,
@@ -1563,6 +1508,7 @@ def test_predict_contract_points_to_fitted_values(self):
):
did.predict(data)
+
# =============================================================================
# Multi-absorb (N>1 FE) iterative alternating-projection demeaning
# =============================================================================
@@ -1607,12 +1553,20 @@ def test_did_multi_absorb_matches_fixed_effects(self, weighted):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res_abs = DifferenceInDifferences().fit(
- df, outcome="outcome", treatment="treated", time="post",
- absorb=["a", "b"], **kw,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="post",
+ absorb=["a", "b"],
+ **kw,
)
res_fe = DifferenceInDifferences().fit(
- df, outcome="outcome", treatment="treated", time="post",
- fixed_effects=["a", "b"], **kw,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="post",
+ fixed_effects=["a", "b"],
+ **kw,
)
# Iterative MAP matches the full-dummy ATT to solver precision; the old
# single-pass sweep would differ by ~1e-2 (well outside this tolerance).
@@ -1627,7 +1581,10 @@ def test_did_collinear_regressors_dropped_to_nan(self):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res = DifferenceInDifferences().fit(
- df, outcome="outcome", treatment="treated", time="post",
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="post",
absorb=["a", "b"],
)
# 'treated' is constant within 'a' (treated = a>=3) -> absorbed -> NaN coef.
@@ -1652,8 +1609,11 @@ def test_mpd_multi_absorb_matches_fixed_effects(self, weighted):
df["treated"] = ((df["a"] >= 3) & (rng.integers(0, 2, size=n) == 1)).astype(int)
df["time"] = rng.integers(0, 3, size=n)
df["outcome"] = (
- 1.5 * df["treated"] + 0.4 * df["a"] - 0.3 * df["b"]
- + 0.5 * df["time"] + rng.normal(scale=0.5, size=n)
+ 1.5 * df["treated"]
+ + 0.4 * df["a"]
+ - 0.3 * df["b"]
+ + 0.5 * df["time"]
+ + rng.normal(scale=0.5, size=n)
)
df["w"] = rng.uniform(0.5, 2.0, size=n)
kw = {}
@@ -1662,12 +1622,22 @@ def test_mpd_multi_absorb_matches_fixed_effects(self, weighted):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res_abs = MultiPeriodDiD().fit(
- df, outcome="outcome", treatment="treated", time="time",
- post_periods=[2], absorb=["a", "b"], **kw,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="time",
+ post_periods=[2],
+ absorb=["a", "b"],
+ **kw,
)
res_fe = MultiPeriodDiD().fit(
- df, outcome="outcome", treatment="treated", time="time",
- post_periods=[2], fixed_effects=["a", "b"], **kw,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="time",
+ post_periods=[2],
+ fixed_effects=["a", "b"],
+ **kw,
)
a_eff = {p: pe.effect for p, pe in res_abs.period_effects.items()}
f_eff = {p: pe.effect for p, pe in res_fe.period_effects.items()}
@@ -1696,8 +1666,7 @@ def test_did_multi_absorb_zero_total_weight_group_is_inert(self):
df["treated"] = (df["a"] >= 3).astype(int)
df["post"] = rng.integers(0, 2, size=n)
df["outcome"] = (
- 2.0 * df["treated"] * df["post"] + 0.3 * df["a"] - 0.2 * df["b"]
- + rng.normal(size=n)
+ 2.0 * df["treated"] * df["post"] + 0.3 * df["a"] - 0.2 * df["b"] + rng.normal(size=n)
)
df["w"] = rng.uniform(0.5, 2.0, size=n)
df.loc[df["a"] == 0, "w"] = 0.0 # entire absorbed category a==0 is zero-weight
@@ -1705,12 +1674,20 @@ def test_did_multi_absorb_zero_total_weight_group_is_inert(self):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res_abs = DifferenceInDifferences().fit(
- df, outcome="outcome", treatment="treated", time="post",
- absorb=["a", "b"], survey_design=sd,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="post",
+ absorb=["a", "b"],
+ survey_design=sd,
)
res_fe = DifferenceInDifferences().fit(
- df, outcome="outcome", treatment="treated", time="post",
- fixed_effects=["a", "b"], survey_design=sd,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="post",
+ fixed_effects=["a", "b"],
+ survey_design=sd,
)
assert np.isfinite(res_abs.att) # no NaN/Inf from the zero-weight group
assert abs(res_abs.att - res_fe.att) < 1e-8
@@ -1728,8 +1705,11 @@ def test_mpd_multi_absorb_zero_total_weight_group_is_inert(self):
df["treated"] = ((df["a"] >= 3) & (rng.integers(0, 2, size=n) == 1)).astype(int)
df["time"] = rng.integers(0, 3, size=n)
df["outcome"] = (
- 1.5 * df["treated"] + 0.3 * df["a"] - 0.2 * df["b"]
- + 0.5 * df["time"] + rng.normal(size=n)
+ 1.5 * df["treated"]
+ + 0.3 * df["a"]
+ - 0.2 * df["b"]
+ + 0.5 * df["time"]
+ + rng.normal(size=n)
)
df["w"] = rng.uniform(0.5, 2.0, size=n)
df.loc[df["a"] == 0, "w"] = 0.0
@@ -1737,20 +1717,26 @@ def test_mpd_multi_absorb_zero_total_weight_group_is_inert(self):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res_abs = MultiPeriodDiD().fit(
- df, outcome="outcome", treatment="treated", time="time",
- post_periods=[2], absorb=["a", "b"], survey_design=sd,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="time",
+ post_periods=[2],
+ absorb=["a", "b"],
+ survey_design=sd,
)
res_fe = MultiPeriodDiD().fit(
- df, outcome="outcome", treatment="treated", time="time",
- post_periods=[2], fixed_effects=["a", "b"], survey_design=sd,
+ df,
+ outcome="outcome",
+ treatment="treated",
+ time="time",
+ post_periods=[2],
+ fixed_effects=["a", "b"],
+ survey_design=sd,
)
a_eff = {p: pe.effect for p, pe in res_abs.period_effects.items()}
f_eff = {p: pe.effect for p, pe in res_fe.period_effects.items()}
- compared = [
- (a_eff[p], f_eff[p])
- for p in a_eff
- if p in f_eff and np.isfinite(f_eff[p])
- ]
+ compared = [(a_eff[p], f_eff[p]) for p in a_eff if p in f_eff and np.isfinite(f_eff[p])]
assert compared, "no period effects to compare"
for ae, fe in compared:
assert np.isfinite(ae) # no NaN/Inf leaked from the zero-weight group
diff --git a/tests/test_methodology_honest_did.py b/tests/test_methodology_honest_did.py
index 1fdcae87..7a0efe59 100644
--- a/tests/test_methodology_honest_did.py
+++ b/tests/test_methodology_honest_did.py
@@ -5,7 +5,6 @@
equations, known analytical cases, and expected mathematical properties.
"""
-
import json
import os
import warnings
@@ -15,7 +14,6 @@
from diff_diff.honest_did import (
HonestDiD,
- _compute_flci,
_compute_optimal_flci,
_compute_pre_first_differences,
_construct_A_sd,
@@ -39,19 +37,22 @@ def test_row_count(self):
for T, Tbar in [(2, 2), (3, 3), (4, 2), (1, 1), (3, 1), (1, 3)]:
A = _construct_A_sd(T, Tbar)
expected_rows = T + Tbar - 1
- assert A.shape == (expected_rows, T + Tbar), (
- f"T={T}, Tbar={Tbar}: expected {expected_rows} rows, got {A.shape[0]}"
- )
+ assert A.shape == (
+ expected_rows,
+ T + Tbar,
+ ), f"T={T}, Tbar={Tbar}: expected {expected_rows} rows, got {A.shape[0]}"
def test_2pre_2post_hand_computed(self):
"""Hand-computed matrix for 2 pre + 2 post periods."""
# delta = [d_{-2}, d_{-1}, d_1, d_2]
A = _construct_A_sd(2, 2)
- expected = np.array([
- [1, -2, 0, 0], # t=-1: d_{-2} - 2*d_{-1} + 0
- [0, 1, 1, 0], # t= 0: d_{-1} + d_1 (bridge)
- [0, 0, -2, 1], # t= 1: 0 - 2*d_1 + d_2
- ])
+ expected = np.array(
+ [
+ [1, -2, 0, 0], # t=-1: d_{-2} - 2*d_{-1} + 0
+ [0, 1, 1, 0], # t= 0: d_{-1} + d_1 (bridge)
+ [0, 0, -2, 1], # t= 1: 0 - 2*d_1 + d_2
+ ]
+ )
np.testing.assert_array_equal(A, expected)
def test_bridge_constraint_present(self):
@@ -63,8 +64,8 @@ def test_bridge_constraint_present(self):
for row in A:
if row[T - 1] != 0 and row[T] != 0:
# This should be [0, ..., 1, 1, ..., 0]
- assert row[T - 1] == 1, f"Bridge row should have 1 at delta_{{-1}}"
- assert row[T] == 1, f"Bridge row should have 1 at delta_1"
+ assert row[T - 1] == 1, "Bridge row should have 1 at delta_{-1}"
+ assert row[T] == 1, "Bridge row should have 1 at delta_1"
bridge_found = True
assert bridge_found, f"Bridge constraint not found for T={T}, Tbar={Tbar}"
@@ -112,9 +113,9 @@ def test_bounds_widen_with_m(self):
A, b = _construct_constraints_sd(3, 1, M=M)
lb, ub = _solve_bounds_lp(beta_pre, beta_post, l_vec, A, b, 3)
width = ub - lb
- assert width >= prev_width - 1e-10, (
- f"Width should increase: M={M}, width={width}, prev={prev_width}"
- )
+ assert (
+ width >= prev_width - 1e-10
+ ), f"Width should increase: M={M}, width={width}, prev={prev_width}"
prev_width = width
def test_three_period_analytical(self):
@@ -134,10 +135,8 @@ def test_three_period_analytical(self):
lb, ub = _solve_bounds_lp(beta_pre, beta_post, np.array([1.0]), A, b, 1)
expected_lb = 3.0 + 0.5 - M
expected_ub = 3.0 + 0.5 + M
- np.testing.assert_allclose(lb, expected_lb, atol=1e-6,
- err_msg=f"M={M}: lb mismatch")
- np.testing.assert_allclose(ub, expected_ub, atol=1e-6,
- err_msg=f"M={M}: ub mismatch")
+ np.testing.assert_allclose(lb, expected_lb, atol=1e-6, err_msg=f"M={M}: lb mismatch")
+ np.testing.assert_allclose(ub, expected_ub, atol=1e-6, err_msg=f"M={M}: ub mismatch")
# =============================================================================
@@ -176,12 +175,12 @@ def test_rm_constraints_are_first_differences(self):
assert A.shape[1] == 5 # 2 pre + 3 post
# First pair: d_1 <= 0.1 and -d_1 <= 0.1
- assert A[0, 2] == 1 # d_1
+ assert A[0, 2] == 1 # d_1
assert A[1, 2] == -1 # -d_1
# Second pair: d_2 - d_1 <= 0.1
- assert A[2, 3] == 1 and A[2, 2] == -1 # d_2 - d_1
- assert A[3, 3] == -1 and A[3, 2] == 1 # -(d_2 - d_1)
+ assert A[2, 3] == 1 and A[2, 2] == -1 # d_2 - d_1
+ assert A[3, 3] == -1 and A[3, 2] == 1 # -(d_2 - d_1)
def test_mbar0_gives_point_estimate(self):
"""Mbar=0: all post first diffs = 0, theta = l'beta_post."""
@@ -220,6 +219,7 @@ class TestOptimalFLCI:
def test_cv_alpha_at_zero(self):
"""cv_alpha(0, alpha) = z_{alpha/2} (standard normal quantile)."""
from scipy.stats import norm
+
np.testing.assert_allclose(_cv_alpha(0, 0.05), norm.ppf(0.975), atol=1e-4)
np.testing.assert_allclose(_cv_alpha(0, 0.01), norm.ppf(0.995), atol=1e-4)
@@ -264,9 +264,7 @@ def test_m0_short_circuit(self):
l_vec = np.array([1.0])
with patch("diff_diff.honest_did.optimize.linprog") as mock_linprog:
- ci_lb, ci_ub = _compute_optimal_flci(
- beta_pre, beta_post, sigma, l_vec, 3, 1, M=0.0
- )
+ ci_lb, ci_ub = _compute_optimal_flci(beta_pre, beta_post, sigma, l_vec, 3, 1, M=0.0)
assert mock_linprog.call_count == 0, (
f"M=0 must skip the LP solver (fast path at "
@@ -274,9 +272,9 @@ def test_m0_short_circuit(self):
f"{mock_linprog.call_count} linprog call(s)."
)
# End-to-end correctness: M=0 CI is still well-defined.
- assert np.isfinite(ci_lb) and np.isfinite(ci_ub), (
- f"M=0 CI must be finite; got [{ci_lb}, {ci_ub}]"
- )
+ assert np.isfinite(ci_lb) and np.isfinite(
+ ci_ub
+ ), f"M=0 CI must be finite; got [{ci_lb}, {ci_ub}]"
assert ci_lb <= ci_ub, f"M=0 CI must be ordered; got [{ci_lb}, {ci_ub}]"
def test_smoothness_flci_with_survey_df(self):
@@ -295,34 +293,34 @@ def test_smoothness_flci_with_survey_df(self):
)
width_norm = ci_ub_norm - ci_lb_norm
width_t = ci_ub_t - ci_lb_t
- assert width_t > width_norm, (
- f"Survey df=2 should widen CI: norm={width_norm:.4f}, t={width_t:.4f}"
- )
+ assert (
+ width_t > width_norm
+ ), f"Survey df=2 should widen CI: norm={width_norm:.4f}, t={width_t:.4f}"
def test_m0_se_includes_pre_period_variance(self):
"""M=0 SE should account for pre-period variance, not just post."""
# Use off-diagonal covariance to make pre-period SE matter
- sigma = np.array([
- [0.04, 0.02, 0.01], # pre-1 has high variance
- [0.02, 0.01, 0.005],
- [0.01, 0.005, 0.01],
- ])
+ sigma = np.array(
+ [
+ [0.04, 0.02, 0.01], # pre-1 has high variance
+ [0.02, 0.01, 0.005],
+ [0.01, 0.005, 0.01],
+ ]
+ )
beta_pre = np.array([0.2, 0.1]) # linear pre-trend
beta_post = np.array([2.0])
l_vec = np.array([1.0])
- ci_lb, ci_ub = _compute_optimal_flci(
- beta_pre, beta_post, sigma, l_vec, 2, 1, M=0.0
- )
+ ci_lb, ci_ub = _compute_optimal_flci(beta_pre, beta_post, sigma, l_vec, 2, 1, M=0.0)
# CI should be finite and the width should reflect pre-period variance
assert np.isfinite(ci_lb) and np.isfinite(ci_ub), "M=0 CI should be finite"
width = ci_ub - ci_lb
# Compare to post-only SE: sqrt(l'Sigma_post l) = sqrt(0.01) = 0.1
post_only_width = 2 * 1.96 * np.sqrt(sigma[2, 2])
- assert width > post_only_width, (
- f"M=0 width ({width:.4f}) should exceed post-only ({post_only_width:.4f})"
- )
+ assert (
+ width > post_only_width
+ ), f"M=0 width ({width:.4f}) should exceed post-only ({post_only_width:.4f})"
def test_optimal_flci_width_increases_with_m_positive(self):
"""Regression for P0: smoothness CI width must increase with M for M > 0."""
@@ -368,8 +366,12 @@ def test_three_period_m0_flci_center(self):
)
center = (ci_lb + ci_ub) / 2
expected_center = 3.0 + 0.5 # beta_1 + beta_{-1}
- np.testing.assert_allclose(center, expected_center, atol=1e-4,
- err_msg="M=0 FLCI should be centered on beta_1 + beta_{-1}")
+ np.testing.assert_allclose(
+ center,
+ expected_center,
+ atol=1e-4,
+ err_msg="M=0 FLCI should be centered on beta_1 + beta_{-1}",
+ )
def test_multi_post_m0_finite(self):
"""Default l_vec with Tbar>1: M=0 gives finite CI."""
@@ -378,12 +380,10 @@ def test_multi_post_m0_finite(self):
sigma = np.eye(5) * 0.01
l_vec = np.array([0.5, 0.5]) # average of 2 post periods
- ci_lb, ci_ub = _compute_optimal_flci(
- beta_pre, beta_post, sigma, l_vec, 3, 2, M=0.0
- )
- assert np.isfinite(ci_lb) and np.isfinite(ci_ub), (
- f"Multi-post M=0 should give finite CI, got [{ci_lb}, {ci_ub}]"
- )
+ ci_lb, ci_ub = _compute_optimal_flci(beta_pre, beta_post, sigma, l_vec, 3, 2, M=0.0)
+ assert np.isfinite(ci_lb) and np.isfinite(
+ ci_ub
+ ), f"Multi-post M=0 should give finite CI, got [{ci_lb}, {ci_ub}]"
def test_multi_post_m_positive_finite(self):
"""Default l_vec with Tbar>1: M>0 gives finite CI."""
@@ -392,12 +392,10 @@ def test_multi_post_m_positive_finite(self):
sigma = np.eye(5) * 0.01
l_vec = np.array([0.5, 0.5])
- ci_lb, ci_ub = _compute_optimal_flci(
- beta_pre, beta_post, sigma, l_vec, 3, 2, M=0.5
- )
- assert np.isfinite(ci_lb) and np.isfinite(ci_ub), (
- f"Multi-post M=0.5 should give finite CI, got [{ci_lb}, {ci_ub}]"
- )
+ ci_lb, ci_ub = _compute_optimal_flci(beta_pre, beta_post, sigma, l_vec, 3, 2, M=0.5)
+ assert np.isfinite(ci_lb) and np.isfinite(
+ ci_ub
+ ), f"Multi-post M=0.5 should give finite CI, got [{ci_lb}, {ci_ub}]"
def test_infeasible_lp_returns_nan(self):
"""Regression for P1: infeasible LP should return NaN, not [-inf, inf]."""
@@ -408,9 +406,7 @@ def test_infeasible_lp_returns_nan(self):
lb, ub = _solve_bounds_lp(beta_pre, beta_post, np.array([1.0]), A, b, 3)
# M=0 with non-linear pre-trends: should be infeasible
- assert np.isnan(lb) and np.isnan(ub), (
- f"Infeasible LP should return NaN, got [{lb}, {ub}]"
- )
+ assert np.isnan(lb) and np.isnan(ub), f"Infeasible LP should return NaN, got [{lb}, {ub}]"
def test_infeasible_smoothness_fit_returns_flci_with_empty_idset(self):
"""Fit-level: an empty ESTIMATED identified set still yields a finite FLCI.
@@ -427,19 +423,31 @@ def test_infeasible_smoothness_fit_returns_flci_with_empty_idset(self):
# Non-linear pre-trends: inconsistent with Delta^SD(M=0.0)
period_effects = {
- 1: PeriodEffect(period=1, effect=1.0, se=0.1, t_stat=10.0,
- p_value=0.0, conf_int=(0.8, 1.2)),
- 2: PeriodEffect(period=2, effect=0.0, se=0.1, t_stat=0.0,
- p_value=1.0, conf_int=(-0.2, 0.2)),
- 3: PeriodEffect(period=3, effect=1.0, se=0.1, t_stat=10.0,
- p_value=0.0, conf_int=(0.8, 1.2)),
- 5: PeriodEffect(period=5, effect=2.0, se=0.1, t_stat=20.0,
- p_value=0.0, conf_int=(1.8, 2.2)),
+ 1: PeriodEffect(
+ period=1, effect=1.0, se=0.1, t_stat=10.0, p_value=0.0, conf_int=(0.8, 1.2)
+ ),
+ 2: PeriodEffect(
+ period=2, effect=0.0, se=0.1, t_stat=0.0, p_value=1.0, conf_int=(-0.2, 0.2)
+ ),
+ 3: PeriodEffect(
+ period=3, effect=1.0, se=0.1, t_stat=10.0, p_value=0.0, conf_int=(0.8, 1.2)
+ ),
+ 5: PeriodEffect(
+ period=5, effect=2.0, se=0.1, t_stat=20.0, p_value=0.0, conf_int=(1.8, 2.2)
+ ),
}
results = MultiPeriodDiDResults(
- avg_att=2.0, avg_se=0.1, avg_t_stat=20.0, avg_p_value=0.0,
- avg_conf_int=(1.8, 2.2), n_obs=500, n_treated=250, n_control=250,
- period_effects=period_effects, pre_periods=[1, 2, 3], post_periods=[5],
+ avg_att=2.0,
+ avg_se=0.1,
+ avg_t_stat=20.0,
+ avg_p_value=0.0,
+ avg_conf_int=(1.8, 2.2),
+ n_obs=500,
+ n_treated=250,
+ n_control=250,
+ period_effects=period_effects,
+ pre_periods=[1, 2, 3],
+ post_periods=[5],
vcov=np.eye(4) * 0.01,
interaction_indices={1: 0, 2: 1, 3: 2, 5: 3},
)
@@ -447,12 +455,14 @@ def test_infeasible_smoothness_fit_returns_flci_with_empty_idset(self):
honest = HonestDiD(method="smoothness", M=0.0)
r = honest.fit(results)
# Estimated identified set is empty -> NaN bounds ...
- assert np.isnan(r.lb) and np.isnan(r.ub), f"Expected NaN id-set bounds, got [{r.lb}, {r.ub}]"
+ assert np.isnan(r.lb) and np.isnan(
+ r.ub
+ ), f"Expected NaN id-set bounds, got [{r.lb}, {r.ub}]"
# ... but the FLCI is finite and matches R createSensitivityResults(M=0):
# [2.082644, 2.488866] (verified against HonestDiD 0.2.6).
- assert np.isfinite(r.ci_lb) and np.isfinite(r.ci_ub), (
- f"Expected finite FLCI, got [{r.ci_lb}, {r.ci_ub}]"
- )
+ assert np.isfinite(r.ci_lb) and np.isfinite(
+ r.ci_ub
+ ), f"Expected finite FLCI, got [{r.ci_lb}, {r.ci_ub}]"
assert abs(r.ci_lb - 2.082644) < 1e-3, f"ci_lb={r.ci_lb:.6f} vs R 2.082644"
assert abs(r.ci_ub - 2.488866) < 1e-3, f"ci_ub={r.ci_ub:.6f} vs R 2.488866"
# Finite CI excluding 0 -> significant, with a star.
@@ -499,13 +509,11 @@ def test_smoothness_flci_finite_and_matches_r_across_M_grid(self):
assert np.isfinite(ci_lb) and np.isfinite(ci_ub), f"M={M}: non-finite FLCI"
assert ci_ub > ci_lb, f"M={M}: degenerate CI"
center = (ci_lb + ci_ub) / 2
- assert abs(center - r_center[M]) < 1e-3, (
- f"M={M} center={center:.6f} vs R {r_center[M]:.6f}"
- )
+ assert (
+ abs(center - r_center[M]) < 1e-3
+ ), f"M={M} center={center:.6f} vs R {r_center[M]:.6f}"
# M=0: tight R parity on the endpoints too (R 0.2.6: [2.082880, 2.488631]).
- ci_lb, ci_ub = _compute_optimal_flci(
- beta_pre, beta_post, sigma, lvec, 3, 1, M=0.0, df=None
- )
+ ci_lb, ci_ub = _compute_optimal_flci(beta_pre, beta_post, sigma, lvec, 3, 1, M=0.0, df=None)
assert abs(ci_lb - 2.082880) < 1e-3, f"M=0 ci_lb={ci_lb:.6f} vs R 2.082880"
assert abs(ci_ub - 2.488631) < 1e-3, f"M=0 ci_ub={ci_ub:.6f} vs R 2.488631"
@@ -534,8 +542,13 @@ def test_signed_contrast_lvec_matches_r(self):
from diff_diff.honest_did import _compute_optimal_flci
lb, ub = _compute_optimal_flci(
- np.array([0.1, 0.0, 0.1]), np.array([1.2, 0.9]),
- np.eye(5) * 0.05, np.array([1.0, -1.0]), 3, 2, M=1.0,
+ np.array([0.1, 0.0, 0.1]),
+ np.array([1.2, 0.9]),
+ np.eye(5) * 0.05,
+ np.array([1.0, -1.0]),
+ 3,
+ 2,
+ M=1.0,
)
assert abs((lb + ub) / 2 - 0.200000) < 1e-3, f"center={(lb + ub) / 2:.6f} vs R 0.200000"
assert abs((ub - lb) / 2 - 4.637049) < 1e-3, f"half={(ub - lb) / 2:.6f} vs R 4.637049"
@@ -547,8 +560,13 @@ def test_negative_M_raises(self):
with pytest.raises(ValueError, match="non-negative"):
_compute_optimal_flci(
- np.array([0.1, 0.0, 0.1]), np.array([1.0]),
- np.eye(4) * 0.01, np.array([1.0]), 3, 1, M=-0.1,
+ np.array([0.1, 0.0, 0.1]),
+ np.array([1.0]),
+ np.eye(4) * 0.01,
+ np.array([1.0]),
+ 3,
+ 1,
+ M=-0.1,
)
@@ -603,8 +621,12 @@ def test_center_and_halflength_match_override_r(self):
lb, ub = _compute_optimal_flci(bp, bq, sigma, l_vec, npre, npost, c["M"], c["alpha"])
center, half = (lb + ub) / 2, (ub - lb) / 2
tag = f"npre={npre} npost={npost} M={c['M']}"
- assert abs(center - c["center"]) < 1e-3, f"{tag}: center {center:.6f} vs R {c['center']:.6f}"
- assert abs(half - c["half_length"]) < 1e-3, f"{tag}: half {half:.6f} vs R {c['half_length']:.6f}"
+ assert (
+ abs(center - c["center"]) < 1e-3
+ ), f"{tag}: center {center:.6f} vs R {c['center']:.6f}"
+ assert (
+ abs(half - c["half_length"]) < 1e-3
+ ), f"{tag}: half {half:.6f} vs R {c['half_length']:.6f}"
max_dc = max(max_dc, abs(center - c["center"]))
max_dw = max(max_dw, abs(half - c["half_length"]))
assert max_dc < 1e-3 and max_dw < 1e-3
@@ -618,9 +640,9 @@ def test_optimal_vec_matches_override_r(self):
bp, bq, sigma, l_vec, npre, npost = self._unpack(c)
_, _, v = _flci_solve(bp, bq, sigma, l_vec, npre, npost, c["M"], c["alpha"])
r_v = np.atleast_1d(np.asarray(c["optimal_vec"], float))
- assert v is not None and np.max(np.abs(v - r_v)) < 3e-3, (
- f"npre={npre} M={c['M']}: optvec gap {np.max(np.abs(v - r_v)):.2e}"
- )
+ assert (
+ v is not None and np.max(np.abs(v - r_v)) < 3e-3
+ ), f"npre={npre} M={c['M']}: optvec gap {np.max(np.abs(v - r_v)):.2e}"
def test_center_matches_stock_r_within_mc_noise(self):
for c in self._cases():
@@ -630,9 +652,9 @@ def test_center_matches_stock_r_within_mc_noise(self):
lb, ub = _compute_optimal_flci(bp, bq, sigma, l_vec, npre, npost, c["M"], c["alpha"])
# Loose tier: R's MC .qfoldednormal noise reaches ~1.4e-2 on flat-surface
# cases (max observed 1.39e-2); the tight parity is vs override-R (1e-3).
- assert abs((lb + ub) / 2 - c["stock_center"]) < 1.5e-2, (
- f"npre={npre} M={c['M']}: vs stock-R {abs((lb + ub) / 2 - c['stock_center']):.2e}"
- )
+ assert (
+ abs((lb + ub) / 2 - c["stock_center"]) < 1.5e-2
+ ), f"npre={npre} M={c['M']}: vs stock-R {abs((lb + ub) / 2 - c['stock_center']):.2e}"
def test_inner_solve_failure_nan_consistent(self, monkeypatch):
"""A failed / infeasible inner solve NaN-propagates the full FLCI (no silent
@@ -640,11 +662,18 @@ def test_inner_solve_failure_nan_consistent(self, monkeypatch):
import diff_diff.honest_did as hd
monkeypatch.setattr(
- hd, "_flci_min_bias_given_h",
+ hd,
+ "_flci_min_bias_given_h",
lambda P, h, x0_w=None: (np.zeros(P["num_pre"]), np.nan, False),
)
lb, ub = _compute_optimal_flci(
- np.array([0.1, -0.05, 0.1]), np.array([1.2]), np.eye(4) * 0.01, np.array([1.0]), 3, 1, M=0.1
+ np.array([0.1, -0.05, 0.1]),
+ np.array([1.2]),
+ np.eye(4) * 0.01,
+ np.array([1.0]),
+ 3,
+ 1,
+ M=0.1,
)
assert np.isnan(lb) and np.isnan(ub)
@@ -663,17 +692,28 @@ def test_breakdown_monotonicity(self):
# Use a weak effect so breakdown is reachable at moderate M
period_effects = {
- 1: PeriodEffect(period=1, effect=0.1, se=0.05, t_stat=2.0,
- p_value=0.05, conf_int=(0.0, 0.2)),
- 2: PeriodEffect(period=2, effect=0.05, se=0.05, t_stat=1.0,
- p_value=0.32, conf_int=(-0.05, 0.15)),
- 4: PeriodEffect(period=4, effect=0.15, se=0.05, t_stat=3.0,
- p_value=0.003, conf_int=(0.05, 0.25)),
+ 1: PeriodEffect(
+ period=1, effect=0.1, se=0.05, t_stat=2.0, p_value=0.05, conf_int=(0.0, 0.2)
+ ),
+ 2: PeriodEffect(
+ period=2, effect=0.05, se=0.05, t_stat=1.0, p_value=0.32, conf_int=(-0.05, 0.15)
+ ),
+ 4: PeriodEffect(
+ period=4, effect=0.15, se=0.05, t_stat=3.0, p_value=0.003, conf_int=(0.05, 0.25)
+ ),
}
results = MultiPeriodDiDResults(
- avg_att=0.15, avg_se=0.05, avg_t_stat=3.0, avg_p_value=0.003,
- avg_conf_int=(0.05, 0.25), n_obs=500, n_treated=250, n_control=250,
- period_effects=period_effects, pre_periods=[1, 2], post_periods=[4],
+ avg_att=0.15,
+ avg_se=0.05,
+ avg_t_stat=3.0,
+ avg_p_value=0.003,
+ avg_conf_int=(0.05, 0.25),
+ n_obs=500,
+ n_treated=250,
+ n_control=250,
+ period_effects=period_effects,
+ pre_periods=[1, 2],
+ post_periods=[4],
vcov=np.eye(3) * 0.0025,
interaction_indices={1: 0, 2: 1, 4: 2},
)
@@ -739,7 +779,8 @@ def test_enumerate_vertices_quiet_on_healthy_enumeration(self):
warnings.simplefilter("always", RuntimeWarning)
vertices = _enumerate_vertices(X_tilde, sigma_tilde_diag, n_moments=4)
diag_warnings = [
- w for w in caught
+ w
+ for w in caught
if "exhausted" in str(w.message) or "heavily constrained" in str(w.message)
]
assert not diag_warnings, (
diff --git a/tests/test_methodology_staggered_triple_diff.py b/tests/test_methodology_staggered_triple_diff.py
index 9c2a3454..4a345589 100644
--- a/tests/test_methodology_staggered_triple_diff.py
+++ b/tests/test_methodology_staggered_triple_diff.py
@@ -665,21 +665,29 @@ def test_constant_covariate_finite_se_matches_drop_one(self, method):
drop_one = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1"],
)
with_const = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data_const, "outcome", "unit", "period", "first_treat", "eligibility",
+ data_const,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "xc"],
)
assert np.isfinite(with_const.overall_se)
assert with_const.overall_se < 1.0 # was 30-100x inflated
- np.testing.assert_allclose(
- with_const.overall_se, drop_one.overall_se, rtol=1e-9
- )
+ np.testing.assert_allclose(with_const.overall_se, drop_one.overall_se, rtol=1e-9)
def test_constant_covariate_emits_rank_guard_warning(self):
# reg uses the OR bread only -> exactly one OR-path rank-guard warning.
@@ -687,10 +695,13 @@ def test_constant_covariate_emits_rank_guard_warning(self):
data["xc"] = 5.0
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
- StaggeredTripleDifference(
- estimation_method="reg", control_group="nevertreated"
- ).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ StaggeredTripleDifference(estimation_method="reg", control_group="nevertreated").fit(
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "xc"],
)
rank_guard = [w for w in caught if "rank-guarded inverse" in str(w.message)]
@@ -703,7 +714,12 @@ def test_well_conditioned_no_rank_guard_warning(self):
res = StaggeredTripleDifference(
estimation_method="dr", control_group="nevertreated"
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "x2"],
)
assert not any("rank-guarded inverse" in str(w.message) for w in caught)
@@ -728,20 +744,30 @@ def test_survey_weighted_constant_covariate_finite_se(self, method):
drop_one = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
- covariates=["x1"], survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
+ covariates=["x1"],
+ survey_design=sd,
)
with_const = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data_const, "outcome", "unit", "period", "first_treat", "eligibility",
- covariates=["x1", "xc"], survey_design=sd,
+ data_const,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
+ covariates=["x1", "xc"],
+ survey_design=sd,
)
assert np.isfinite(with_const.overall_se)
assert with_const.overall_se < 1.0
- np.testing.assert_allclose(
- with_const.overall_se, drop_one.overall_se, rtol=1e-9
- )
+ np.testing.assert_allclose(with_const.overall_se, drop_one.overall_se, rtol=1e-9)
@pytest.mark.parametrize("method", ["reg", "ipw", "dr"])
def test_notyettreated_silent_constant_covariate(self, method):
@@ -758,7 +784,12 @@ def test_notyettreated_silent_constant_covariate(self, method):
control_group="notyettreated",
rank_deficient_action="silent",
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "xc"],
)
assert np.isfinite(res.overall_se)
@@ -782,7 +813,12 @@ def test_error_mode_raises_before_rank_guard(self, method):
control_group="nevertreated",
rank_deficient_action="error",
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "x1c"],
)
@@ -800,27 +836,33 @@ def test_control_cell_aliasing_close_to_drop_one(self, method):
rng = np.random.default_rng(0)
nt = (data["first_treat"] == 0).to_numpy()
x2d = np.where(nt, 2.0 * data["x1"].to_numpy(), rng.normal(size=len(data)))
- data["x2d"] = pd.Series(x2d, index=data.index).groupby(
- data["unit"]
- ).transform("first")
+ data["x2d"] = pd.Series(x2d, index=data.index).groupby(data["unit"]).transform("first")
with warnings.catch_warnings():
warnings.simplefilter("ignore")
drop_one = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1"],
)
with_deg = StaggeredTripleDifference(
estimation_method=method, control_group="nevertreated"
).fit(
- data, "outcome", "unit", "period", "first_treat", "eligibility",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ "eligibility",
covariates=["x1", "x2d"],
)
assert np.isfinite(with_deg.overall_se) and with_deg.overall_se > 0
- np.testing.assert_allclose(
- with_deg.overall_se, drop_one.overall_se, rtol=5e-2
- )
+ np.testing.assert_allclose(with_deg.overall_se, drop_one.overall_se, rtol=5e-2)
class TestStaggeredTripleDiffCovariateScaleEquilibration:
diff --git a/tests/test_methodology_sun_abraham.py b/tests/test_methodology_sun_abraham.py
index f3a687ba..f121d43e 100644
--- a/tests/test_methodology_sun_abraham.py
+++ b/tests/test_methodology_sun_abraham.py
@@ -118,6 +118,7 @@ def test_bm_dof_matches_clubsandwich_singleton_and_unit(self, sa_panel, sa_golde
output (p-values, CIs) without changing SEs; this test guards that.
"""
import pandas as pd
+
from diff_diff.linalg import LinearRegression
# Reconstruct the full-dummy design SA's Part G builds (same shape
diff --git a/tests/test_np_npreg_weighted_parity.py b/tests/test_np_npreg_weighted_parity.py
index fdd0091c..2895447d 100644
--- a/tests/test_np_npreg_weighted_parity.py
+++ b/tests/test_np_npreg_weighted_parity.py
@@ -31,10 +31,7 @@
import pytest
GOLDEN_PATH = (
- Path(__file__).resolve().parents[1]
- / "benchmarks"
- / "data"
- / "np_npreg_weighted_golden.json"
+ Path(__file__).resolve().parents[1] / "benchmarks" / "data" / "np_npreg_weighted_golden.json"
)
@@ -72,9 +69,7 @@ def test_intercept_parity(self, weighted_golden, dgp_name):
kernel="epanechnikov",
weights=w,
)
- np.testing.assert_allclose(
- fit.intercept, g["mu_hat"], atol=1e-12, rtol=1e-12
- )
+ np.testing.assert_allclose(fit.intercept, g["mu_hat"], atol=1e-12, rtol=1e-12)
@pytest.mark.parametrize("dgp_name", ["dgp1", "dgp2", "dgp3", "dgp4"])
def test_slope_parity(self, weighted_golden, dgp_name):
@@ -130,7 +125,5 @@ def test_uniform_weights_matches_no_weights(self, weighted_golden):
boundary=float(g["eval_point"]),
kernel="epanechnikov",
)
- np.testing.assert_allclose(
- fit_w.intercept, fit_nw.intercept, atol=1e-14, rtol=1e-14
- )
+ np.testing.assert_allclose(fit_w.intercept, fit_nw.intercept, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(fit_w.slope, fit_nw.slope, atol=1e-14, rtol=1e-14)
diff --git a/tests/test_nprobust_port.py b/tests/test_nprobust_port.py
index f85764f8..efd20a47 100644
--- a/tests/test_nprobust_port.py
+++ b/tests/test_nprobust_port.py
@@ -346,9 +346,7 @@ def test_missing_cluster_id_rejected_none(self):
rng = np.random.default_rng(0)
x = rng.uniform(0.0, 1.0, size=10)
y = rng.normal(size=10)
- cluster = np.array(
- [1, 2, None, 3, 4, 5, 6, 7, 8, 9], dtype=object
- )
+ cluster = np.array([1, 2, None, 3, 4, 5, 6, 7, 8, 9], dtype=object)
with pytest.raises(ValueError, match="missing values"):
lpbwselect_mse_dpi(y, x, cluster=cluster, eval_point=0.0)
@@ -412,8 +410,14 @@ def test_clustered_end_to_end_through_lprobust_bw(self):
# 50 clusters, balanced.
cluster = np.repeat(np.arange(50), G // 50)
res = lpbwselect_mse_dpi(
- y, d, cluster=cluster, eval_point=0.0,
- p=1, deriv=0, kernel="epa", vce="nn",
+ y,
+ d,
+ cluster=cluster,
+ eval_point=0.0,
+ p=1,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
)
assert np.isfinite(res.h_mse_dpi)
assert res.h_mse_dpi > 0.0
@@ -436,7 +440,12 @@ def _load_lprobust_golden():
import json
from pathlib import Path
- path = Path(__file__).resolve().parents[1] / "benchmarks" / "data" / "nprobust_lprobust_golden.json"
+ path = (
+ Path(__file__).resolve().parents[1]
+ / "benchmarks"
+ / "data"
+ / "nprobust_lprobust_golden.json"
+ )
if not path.exists():
pytest.skip(
"Golden values file not found; "
@@ -475,8 +484,16 @@ def test_lprobust_dgp1_parity(self, lprobust_golden):
d, y, g = self._dgp(lprobust_golden, "dgp1")
res = lprobust(
- y, d, eval_point=0.0, h=g["h"], b=g["b"], p=1, q=2, deriv=0,
- kernel="epa", vce="nn",
+ y,
+ d,
+ eval_point=0.0,
+ h=g["h"],
+ b=g["b"],
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
)
# Tiered tolerances per plan commit criterion #1.
np.testing.assert_allclose(res.tau_cl, g["tau_cl"], atol=1e-14, rtol=1e-14)
@@ -490,8 +507,16 @@ def test_lprobust_dgp2_parity(self, lprobust_golden):
d, y, g = self._dgp(lprobust_golden, "dgp2")
res = lprobust(
- y, d, eval_point=0.0, h=g["h"], b=g["b"], p=1, q=2, deriv=0,
- kernel="epa", vce="nn",
+ y,
+ d,
+ eval_point=0.0,
+ h=g["h"],
+ b=g["b"],
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
)
np.testing.assert_allclose(res.tau_cl, g["tau_cl"], atol=1e-12, rtol=1e-12)
np.testing.assert_allclose(res.se_cl, g["se_cl"], atol=1e-12, rtol=1e-12)
@@ -503,8 +528,16 @@ def test_lprobust_dgp3_parity(self, lprobust_golden):
d, y, g = self._dgp(lprobust_golden, "dgp3")
res = lprobust(
- y, d, eval_point=0.0, h=g["h"], b=g["b"], p=1, q=2, deriv=0,
- kernel="epa", vce="nn",
+ y,
+ d,
+ eval_point=0.0,
+ h=g["h"],
+ b=g["b"],
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
)
np.testing.assert_allclose(res.tau_cl, g["tau_cl"], atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(res.se_cl, g["se_cl"], atol=1e-14, rtol=1e-14)
@@ -518,9 +551,10 @@ def test_lprobust_manual_bandwidths_rho_1(self):
rng = np.random.default_rng(20260420)
G = 1000
d = rng.uniform(0.0, 1.0, G)
- y = d + d ** 2 + rng.normal(0, 0.5, G)
- res = lprobust(y, d, eval_point=0.0, h=0.3, b=0.3, p=1, q=2, deriv=0,
- kernel="epa", vce="nn")
+ y = d + d**2 + rng.normal(0, 0.5, G)
+ res = lprobust(
+ y, d, eval_point=0.0, h=0.3, b=0.3, p=1, q=2, deriv=0, kernel="epa", vce="nn"
+ )
assert np.isfinite(res.tau_cl)
assert np.isfinite(res.tau_bc)
assert res.se_cl > 0
@@ -539,8 +573,17 @@ def test_lprobust_cluster_parity(self, lprobust_golden):
y = np.asarray(g["y"], dtype=np.float64)
cluster = np.asarray(g["cluster"])
res = lprobust(
- y, d, eval_point=0.0, h=g["h"], b=g["b"], p=1, q=2, deriv=0,
- kernel="epa", vce="nn", cluster=cluster,
+ y,
+ d,
+ eval_point=0.0,
+ h=g["h"],
+ b=g["b"],
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
+ cluster=cluster,
)
# Cluster IDs happen to match R's first-appearance order here (IDs
# are already sequential 1..50), so bit-parity is achievable.
@@ -560,8 +603,17 @@ def test_lprobust_bwcheck_clip_floor(self):
y = d + rng.normal(0, 0.5, G)
# Very small h; bwcheck should float it up.
res = lprobust(
- y, d, eval_point=0.0, h=0.001, b=0.001, p=1, q=2, deriv=0,
- kernel="epa", vce="nn", bwcheck=21,
+ y,
+ d,
+ eval_point=0.0,
+ h=0.001,
+ b=0.001,
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
+ bwcheck=21,
)
# Bandwidth must have been floored — original 0.001 will not hold.
assert res.h > 0.001
@@ -572,16 +624,23 @@ def test_lprobust_bwcheck_clip_floor(self):
assert res.h == pytest.approx(expected_floor, rel=1e-14)
def test_lprobust_shifted_boundary_dgp5(self, lprobust_golden):
- """Parity at a non-zero boundary (Design 1 continuous-near-d_lower).
- """
+ """Parity at a non-zero boundary (Design 1 continuous-near-d_lower)."""
from diff_diff._nprobust_port import lprobust
g = lprobust_golden["dgp5"]
d = np.asarray(g["d"], dtype=np.float64)
y = np.asarray(g["y"], dtype=np.float64)
res = lprobust(
- y, d, eval_point=g["eval_point_override"], h=g["h"], b=g["b"],
- p=1, q=2, deriv=0, kernel="epa", vce="nn",
+ y,
+ d,
+ eval_point=g["eval_point_override"],
+ h=g["h"],
+ b=g["b"],
+ p=1,
+ q=2,
+ deriv=0,
+ kernel="epa",
+ vce="nn",
)
np.testing.assert_allclose(res.tau_cl, g["tau_cl"], atol=1e-12, rtol=1e-12)
np.testing.assert_allclose(res.tau_bc, g["tau_bc"], atol=1e-12, rtol=1e-12)
@@ -599,9 +658,7 @@ def test_lprobust_cluster_object_none_raises(self):
G = 200
d = rng.uniform(0.0, 1.0, G)
y = d + rng.normal(0, 0.3, G)
- cluster = np.array(
- [i // 10 if i != 5 else None for i in range(G)], dtype=object
- )
+ cluster = np.array([i // 10 if i != 5 else None for i in range(G)], dtype=object)
with pytest.raises(ValueError, match="cluster contains missing"):
lprobust(y, d, eval_point=0.0, h=0.3, b=0.3, cluster=cluster)
@@ -613,9 +670,7 @@ def test_lprobust_cluster_object_nan_raises(self):
G = 200
d = rng.uniform(0.0, 1.0, G)
y = d + rng.normal(0, 0.3, G)
- cluster = np.array(
- [i // 10 if i != 5 else np.nan for i in range(G)], dtype=object
- )
+ cluster = np.array([i // 10 if i != 5 else np.nan for i in range(G)], dtype=object)
with pytest.raises(ValueError, match="cluster contains missing"):
lprobust(y, d, eval_point=0.0, h=0.3, b=0.3, cluster=cluster)
@@ -633,8 +688,9 @@ def test_lprobust_h_gt_b_selects_h_window(self):
# which for epa is 0 <= (d - 0)/h <= 1 => 0 <= d <= h.
h = 0.4
b = 0.2
- res = lprobust(y, d, eval_point=0.0, h=h, b=b, p=1, q=2, deriv=0,
- kernel="epa", vce="nn", bwcheck=None)
+ res = lprobust(
+ y, d, eval_point=0.0, h=h, b=b, p=1, q=2, deriv=0, kernel="epa", vce="nn", bwcheck=None
+ )
# Post-sort d; ind.h = (|d|/h <= 1) & (d <= h) => d in [0, h].
sorted_d = np.sort(d)
# Relaxed: assert n_used > sum of b-window (which is smaller).
@@ -687,7 +743,12 @@ def test_uniform_weights_bit_parity_hc1(self):
x, y = self._panel()
base = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, vce="hc1")
w1 = lprobust(
- y, x, eval_point=0.0, h=0.4, b=0.4, vce="hc1",
+ y,
+ x,
+ eval_point=0.0,
+ h=0.4,
+ b=0.4,
+ vce="hc1",
weights=np.ones(x.size),
)
np.testing.assert_allclose(w1.tau_bc, base.tau_bc, atol=1e-14, rtol=1e-14)
@@ -736,16 +797,14 @@ def test_zero_sum_weights_reject(self):
x, y = self._panel()
with pytest.raises(ValueError, match="sum to zero"):
- lprobust(y, x, eval_point=0.0, h=0.4, b=0.4,
- weights=np.zeros(x.size))
+ lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, weights=np.zeros(x.size))
def test_weights_length_mismatch(self):
from diff_diff._nprobust_port import lprobust
x, y = self._panel()
with pytest.raises(ValueError, match="weights length"):
- lprobust(y, x, eval_point=0.0, h=0.4, b=0.4,
- weights=np.ones(x.size - 1))
+ lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, weights=np.ones(x.size - 1))
def test_zero_weight_observations_drop_from_window(self):
"""Observations with weights[i]=0 filter out via combined ``w>0``
@@ -772,8 +831,7 @@ def test_weights_with_cluster(self):
rng = np.random.default_rng(7)
cluster = rng.integers(0, 20, x.size)
w = rng.uniform(0.5, 1.5, x.size)
- r = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4,
- cluster=cluster, weights=w)
+ r = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, cluster=cluster, weights=w)
assert np.isfinite(r.tau_bc) and np.isfinite(r.se_rb) and r.se_rb > 0
def test_uniform_weights_bit_parity_with_cluster(self):
@@ -783,7 +841,6 @@ def test_uniform_weights_bit_parity_with_cluster(self):
rng = np.random.default_rng(7)
cluster = rng.integers(0, 20, x.size)
base = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, cluster=cluster)
- w1 = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4,
- cluster=cluster, weights=np.ones(x.size))
+ w1 = lprobust(y, x, eval_point=0.0, h=0.4, b=0.4, cluster=cluster, weights=np.ones(x.size))
np.testing.assert_allclose(w1.tau_bc, base.tau_bc, atol=1e-14, rtol=1e-14)
np.testing.assert_allclose(w1.se_rb, base.se_rb, atol=1e-14, rtol=1e-14)
diff --git a/tests/test_power.py b/tests/test_power.py
index 913f913c..f68c148c 100644
--- a/tests/test_power.py
+++ b/tests/test_power.py
@@ -3036,9 +3036,7 @@ def test_clearing_survey_design_falls_back_to_default(self):
initialization must fall back to the default construction."""
from diff_diff.survey import SurveyDesign
- initial = SurveyDesign(
- weights="w0", strata="s0", psu="p0", fpc="f0"
- )
+ initial = SurveyDesign(weights="w0", strata="s0", psu="p0", fpc="f0")
cfg = SurveyPowerConfig(survey_design=initial)
first = cfg._build_survey_design()
assert first is initial
diff --git a/tests/test_prep.py b/tests/test_prep.py
index 68ef8a65..fb452f71 100644
--- a/tests/test_prep.py
+++ b/tests/test_prep.py
@@ -1887,6 +1887,7 @@ def test_top_level_import(self):
def test_jk1_minimum_psu_guard(self):
"""Test that JK1 replicates require at least 2 PSUs."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
# Configured count: 1 PSU total
@@ -1901,6 +1902,7 @@ def test_jk1_minimum_psu_guard(self):
def test_jk1_one_populated_psu_guard(self):
"""Test JK1 guard fires when only one PSU is populated."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
# 2 configured PSUs but only 1 unit -> only 1 populated PSU
@@ -1934,6 +1936,7 @@ def test_repeated_cross_section(self):
def test_invalid_weight_variation(self):
"""Test that invalid weight_variation raises ValueError."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
with pytest.raises(ValueError, match="weight_variation must be"):
@@ -1942,6 +1945,7 @@ def test_invalid_weight_variation(self):
def test_empty_cohort_periods(self):
"""Test that empty cohort_periods raises ValueError."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
with pytest.raises(ValueError, match="cohort_periods must be"):
@@ -1950,6 +1954,7 @@ def test_empty_cohort_periods(self):
def test_cohort_period_out_of_range(self):
"""Test that out-of-range cohort periods raise ValueError."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
# g=1 is invalid: no pre-treatment period (must be >= 2)
@@ -1965,6 +1970,7 @@ def test_cohort_period_out_of_range(self):
def test_cohort_period_non_integer(self):
"""Test that non-integer cohort periods raise ValueError."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
with pytest.raises(ValueError, match="must contain integers"):
@@ -2000,6 +2006,7 @@ def test_default_cohort_periods_small_n_periods(self):
def test_default_cohort_periods_too_small(self):
"""Test that n_periods < 4 with default cohort_periods raises."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
with pytest.raises(ValueError, match="too small"):
@@ -2008,6 +2015,7 @@ def test_default_cohort_periods_too_small(self):
def test_parameter_validation(self):
"""Test upfront validation for invalid parameter values."""
import pytest
+
from diff_diff.prep import generate_survey_did_data
with pytest.raises(ValueError, match="n_units must be positive"):
@@ -2042,7 +2050,6 @@ def test_psu_period_factor_deff_regression(self):
import warnings
from diff_diff import (
- CallawaySantAnna,
DifferenceInDifferences,
SurveyDesign,
)
diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py
index ca4f9a4f..d870749d 100644
--- a/tests/test_replicate_weight_expansion.py
+++ b/tests/test_replicate_weight_expansion.py
@@ -21,6 +21,7 @@
# Fixtures
# ---------------------------------------------------------------------------
+
def _make_simple_panel():
"""2-period panel for DiD/MultiPeriodDiD (treatment/post binary columns)."""
np.random.seed(123)
@@ -34,10 +35,16 @@ def _make_simple_panel():
if treated and t == 1:
y += 2.0 # ATT = 2
y += np.random.normal(0, 0.3)
- rows.append({
- "unit": i, "time": t, "treated": treated, "post": t,
- "outcome": y, "weight": wt,
- })
+ rows.append(
+ {
+ "unit": i,
+ "time": t,
+ "treated": treated,
+ "post": t,
+ "outcome": y,
+ "weight": wt,
+ }
+ )
data = pd.DataFrame(rows)
return data
@@ -60,12 +67,17 @@ def _make_staggered_panel():
if ft > 0 and t >= ft:
y += 2.0
y += np.random.normal(0, 0.4)
- rows.append({
- "unit": i, "time": t, "first_treat": ft,
- "outcome": y, "weight": wt,
- "treated": 1 if ft > 0 else 0,
- "post": 1 if ft > 0 and t >= ft else 0,
- })
+ rows.append(
+ {
+ "unit": i,
+ "time": t,
+ "first_treat": ft,
+ "outcome": y,
+ "weight": wt,
+ "treated": 1 if ft > 0 else 0,
+ "post": 1 if ft > 0 and t >= ft else 0,
+ }
+ )
data = pd.DataFrame(rows)
return data
@@ -112,6 +124,7 @@ def _add_brr_replicates(data, n_rep=16, unit_col="unit"):
# Smoke tests — each estimator × {JK1, BRR}
# ---------------------------------------------------------------------------
+
class TestDiDReplicate:
"""DifferenceInDifferences with replicate weights."""
@@ -120,7 +133,11 @@ def test_did_jk1(self):
rep_cols = _add_jk1_replicates(data, n_rep=10)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = DifferenceInDifferences().fit(
- data, "outcome", "treated", "post", survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ survey_design=sd,
)
assert np.isfinite(result.att)
assert np.isfinite(result.se) and result.se > 0
@@ -132,7 +149,11 @@ def test_did_brr(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = DifferenceInDifferences().fit(
- data, "outcome", "treated", "post", survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ survey_design=sd,
)
assert np.isfinite(result.att)
assert np.isfinite(result.se) and result.se > 0
@@ -144,7 +165,11 @@ def test_did_wild_bootstrap_rejected(self):
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
with pytest.raises((ValueError, NotImplementedError)):
DifferenceInDifferences(inference="wild_bootstrap", cluster="unit").fit(
- data, "outcome", "treated", "post", survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ survey_design=sd,
)
@@ -159,7 +184,12 @@ def test_did_absorb_brr(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = DifferenceInDifferences().fit(
- data, "outcome", "treated", "post", absorb=["group"], survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ absorb=["group"],
+ survey_design=sd,
)
assert np.isfinite(result.att)
assert np.isfinite(result.se) and result.se > 0
@@ -173,7 +203,12 @@ def test_multiperiod_jk1(self):
rep_cols = _add_jk1_replicates(data, n_rep=10)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = MultiPeriodDiD().fit(
- data, "outcome", "treated", "time", post_periods=[1], survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "time",
+ post_periods=[1],
+ survey_design=sd,
)
assert np.isfinite(result.avg_att)
assert np.isfinite(result.avg_se) and result.avg_se > 0
@@ -185,7 +220,12 @@ def test_multiperiod_brr(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = MultiPeriodDiD().fit(
- data, "outcome", "treated", "time", post_periods=[1], survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "time",
+ post_periods=[1],
+ survey_design=sd,
)
assert np.isfinite(result.avg_att)
assert np.isfinite(result.avg_se) and result.avg_se > 0
@@ -208,10 +248,16 @@ def _make_twfe_panel():
if treated and t == 1:
y += 2.0
y += np.random.normal(0, 0.3)
- rows.append({
- "unit": i, "time": t, "treated": treated,
- "post": t, "outcome": y, "weight": wt,
- })
+ rows.append(
+ {
+ "unit": i,
+ "time": t,
+ "treated": treated,
+ "post": t,
+ "outcome": y,
+ "weight": wt,
+ }
+ )
return pd.DataFrame(rows)
def test_twfe_brr(self):
@@ -220,7 +266,12 @@ def test_twfe_brr(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = TwoWayFixedEffects().fit(
- data, "outcome", "treated", "post", "unit", survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ "unit",
+ survey_design=sd,
)
assert np.isfinite(result.att)
assert np.isfinite(result.se) and result.se > 0
@@ -233,7 +284,12 @@ def test_twfe_brr_larger(self):
rep_cols = _add_brr_replicates(data, n_rep=20)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = TwoWayFixedEffects().fit(
- data, "outcome", "treated", "post", "unit", survey_design=sd,
+ data,
+ "outcome",
+ "treated",
+ "post",
+ "unit",
+ survey_design=sd,
)
assert np.isfinite(result.att)
assert np.isfinite(result.se) and result.se > 0
@@ -248,7 +304,12 @@ def test_sun_abraham_brr(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = SunAbraham(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se) and result.overall_se > 0
@@ -261,7 +322,12 @@ def test_sun_abraham_bootstrap_rejected(self):
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
with pytest.raises(ValueError, match="n_bootstrap"):
SunAbraham(n_bootstrap=100).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
@@ -273,7 +339,12 @@ def test_stacked_jk1(self):
rep_cols = _add_jk1_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = StackedDiD().fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se) and result.overall_se > 0
@@ -285,7 +356,12 @@ def test_stacked_brr(self):
rep_cols = _add_brr_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = StackedDiD().fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se) and result.overall_se > 0
@@ -299,7 +375,12 @@ def test_imputation_jk1(self):
rep_cols = _add_jk1_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = ImputationDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se) and result.overall_se > 0
@@ -313,14 +394,20 @@ def test_imputation_event_study_replicate(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = ImputationDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat",
- aggregate="event_study", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ aggregate="event_study",
+ survey_design=sd,
)
assert np.isfinite(result.overall_se) and result.overall_se > 0
assert result.event_study_effects is not None
# At least some identified periods should have finite SE
finite_ses = [
- e for e, eff in result.event_study_effects.items()
+ e
+ for e, eff in result.event_study_effects.items()
if np.isfinite(eff["effect"]) and np.isfinite(eff["se"]) and eff["se"] > 0
]
assert len(finite_ses) > 0, "No event-study periods have finite replicate SE"
@@ -331,8 +418,13 @@ def test_imputation_group_replicate(self):
rep_cols = _add_jk1_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = ImputationDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat",
- aggregate="group", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ aggregate="group",
+ survey_design=sd,
)
assert result.group_effects is not None
for g, eff in result.group_effects.items():
@@ -344,7 +436,12 @@ def test_imputation_bootstrap_rejected(self):
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
with pytest.raises(ValueError, match="n_bootstrap"):
ImputationDiD(n_bootstrap=100).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
@@ -356,7 +453,12 @@ def test_two_stage_jk1(self):
rep_cols = _add_jk1_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = TwoStageDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se) and result.overall_se > 0
@@ -369,8 +471,13 @@ def test_two_stage_event_study_replicate(self):
rep_cols = _add_jk1_replicates(data)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = TwoStageDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat",
- aggregate="event_study", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ aggregate="event_study",
+ survey_design=sd,
)
assert result.event_study_effects is not None
non_ref = {e: eff for e, eff in result.event_study_effects.items() if eff["effect"] != 0.0}
@@ -398,16 +505,35 @@ def test_two_stage_always_treated(self):
# Add always-treated units (first_treat <= min time)
for i in range(50, 55):
for t in range(1, 9):
- data = pd.concat([data, pd.DataFrame([{
- "unit": i, "time": t, "first_treat": 1,
- "outcome": 12.0 + rng.normal(0, 0.3),
- "weight": 1.5, "treated": 1, "post": 1,
- }])], ignore_index=True)
+ data = pd.concat(
+ [
+ data,
+ pd.DataFrame(
+ [
+ {
+ "unit": i,
+ "time": t,
+ "first_treat": 1,
+ "outcome": 12.0 + rng.normal(0, 0.3),
+ "weight": 1.5,
+ "treated": 1,
+ "post": 1,
+ }
+ ]
+ ),
+ ],
+ ignore_index=True,
+ )
rep_cols = _add_jk1_replicates(data, n_rep=10, unit_col="unit")
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
# Should not crash despite always-treated unit exclusion
result = TwoStageDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
# ATT comes from the main fit (always finite once always-treated drop runs)
assert np.isfinite(result.overall_att)
@@ -440,16 +566,35 @@ def test_two_stage_always_treated_event_study_and_group_replicate(self):
data = _make_staggered_panel()
for i in range(50, 55):
for t in range(1, 9):
- data = pd.concat([data, pd.DataFrame([{
- "unit": i, "time": t, "first_treat": 1,
- "outcome": 12.0 + rng.normal(0, 0.3),
- "weight": 1.5, "treated": 1, "post": 1,
- }])], ignore_index=True)
+ data = pd.concat(
+ [
+ data,
+ pd.DataFrame(
+ [
+ {
+ "unit": i,
+ "time": t,
+ "first_treat": 1,
+ "outcome": 12.0 + rng.normal(0, 0.3),
+ "weight": 1.5,
+ "treated": 1,
+ "post": 1,
+ }
+ ]
+ ),
+ ],
+ ignore_index=True,
+ )
rep_cols = _add_jk1_replicates(data, n_rep=10, unit_col="unit")
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
result = TwoStageDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat",
- aggregate="all", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ aggregate="all",
+ survey_design=sd,
)
# Event-study surface: at least one non-reference horizon must have
# replicate-derived finite SE / p_value / conf_int (the replicate
@@ -458,7 +603,8 @@ def test_two_stage_always_treated_event_study_and_group_replicate(self):
# _refit_ts callback, so all four fields must be locked).
assert result.event_study_effects is not None
non_ref_es = {
- e: eff for e, eff in result.event_study_effects.items()
+ e: eff
+ for e, eff in result.event_study_effects.items()
if eff["effect"] != 0.0 and np.isfinite(eff["effect"])
}
assert len(non_ref_es) > 0, "no non-reference event-study effects"
@@ -467,31 +613,27 @@ def test_two_stage_always_treated_event_study_and_group_replicate(self):
f"event-study horizon {e}: replicate SE must be finite "
f"under always-treated drop"
)
- assert np.isfinite(eff["p_value"]), (
- f"event-study horizon {e}: replicate p_value must be finite"
- )
- assert eff["conf_int"] is not None and np.all(np.isfinite(eff["conf_int"])), (
- f"event-study horizon {e}: replicate conf_int bounds must be finite"
- )
+ assert np.isfinite(
+ eff["p_value"]
+ ), f"event-study horizon {e}: replicate p_value must be finite"
+ assert eff["conf_int"] is not None and np.all(
+ np.isfinite(eff["conf_int"])
+ ), f"event-study horizon {e}: replicate conf_int bounds must be finite"
# Group surface: at least one cohort with finite replicate SE /
# p_value / conf_int (same override path applies to group effects).
assert result.group_effects is not None
finite_groups = {
- g: eff for g, eff in result.group_effects.items()
- if np.isfinite(eff["effect"])
+ g: eff for g, eff in result.group_effects.items() if np.isfinite(eff["effect"])
}
assert len(finite_groups) > 0, "no finite group effects"
for g, eff in finite_groups.items():
assert np.isfinite(eff["se"]) and eff["se"] > 0, (
- f"cohort {g}: replicate SE must be finite under "
- f"always-treated drop"
- )
- assert np.isfinite(eff["p_value"]), (
- f"cohort {g}: replicate p_value must be finite"
- )
- assert eff["conf_int"] is not None and np.all(np.isfinite(eff["conf_int"])), (
- f"cohort {g}: replicate conf_int bounds must be finite"
+ f"cohort {g}: replicate SE must be finite under " f"always-treated drop"
)
+ assert np.isfinite(eff["p_value"]), f"cohort {g}: replicate p_value must be finite"
+ assert eff["conf_int"] is not None and np.all(
+ np.isfinite(eff["conf_int"])
+ ), f"cohort {g}: replicate conf_int bounds must be finite"
def test_two_stage_bootstrap_rejected(self):
data = _make_staggered_panel()
@@ -499,7 +641,12 @@ def test_two_stage_bootstrap_rejected(self):
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1")
with pytest.raises(ValueError, match="n_bootstrap"):
TwoStageDiD(n_bootstrap=100).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
@@ -512,7 +659,12 @@ def test_cohort_ses_finite(self):
rep_cols = _add_brr_replicates(data, n_rep=16)
sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR")
result = SunAbraham(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ survey_design=sd,
)
assert result.cohort_effects is not None
for key, eff in result.cohort_effects.items():
diff --git a/tests/test_staggered_triple_diff.py b/tests/test_staggered_triple_diff.py
index cdac60e5..12eab70a 100644
--- a/tests/test_staggered_triple_diff.py
+++ b/tests/test_staggered_triple_diff.py
@@ -641,8 +641,7 @@ def test_collinear_covariates_emit_lstsq_fallback_warning(self):
w for w in caught if "outcome-regression influence-function step" in str(w.message)
]
assert len(or_warnings) == 1, (
- f"Expected exactly one aggregate OR rank-guard warning, "
- f"got {len(or_warnings)}."
+ f"Expected exactly one aggregate OR rank-guard warning, " f"got {len(or_warnings)}."
)
msg = str(or_warnings[0].message)
assert "(g, g_c, t) pair(s)" in msg
diff --git a/tests/test_survey_dcdh.py b/tests/test_survey_dcdh.py
index a44b90da..fa4e825f 100644
--- a/tests/test_survey_dcdh.py
+++ b/tests/test_survey_dcdh.py
@@ -9,7 +9,6 @@
from diff_diff import ChaisemartinDHaultfoeuille, SurveyDesign
from diff_diff.prep_dgp import generate_reversible_did_data
-
# ── Fixtures ────────────────────────────────────────────────────────
@@ -89,9 +88,7 @@ def test_uniform_weights_match_unweighted(self, base_data):
survey_design=sd,
)
# Point estimates should match exactly (uniform weights = unweighted mean)
- assert result_plain.overall_att == pytest.approx(
- result_survey.overall_att, abs=1e-10
- )
+ assert result_plain.overall_att == pytest.approx(result_survey.overall_att, abs=1e-10)
# ── Test: Non-uniform weights change estimate ───────────────────────
@@ -115,7 +112,7 @@ def test_nonuniform_weights_change_att(self, base_data):
multi["pw"] = multi["group"].map(g_weights)
# Make second copy have different weights (still constant within group)
# by giving rows from the second batch higher weight via a multiplier
- multi.loc[len(df):, "pw"] = multi.loc[len(df):, "pw"] * 2.0
+ multi.loc[len(df) :, "pw"] = multi.loc[len(df) :, "pw"] * 2.0
# Since weights now vary within group (first copy vs second copy),
# we need per-observation weights. But dCDH requires group-constant
@@ -147,9 +144,7 @@ def test_nonuniform_weights_change_att(self, base_data):
# weight). The ATTs should match. This confirms the equal-cell
# contract: survey weights don't change the cross-group aggregation.
# The SE will differ because the survey variance accounts for design.
- assert result_plain.overall_att == pytest.approx(
- result_survey.overall_att, abs=1e-8
- )
+ assert result_plain.overall_att == pytest.approx(result_survey.overall_att, abs=1e-8)
# But the SEs should differ when strata/PSU are present
# (tested separately in TestSurveySE)
@@ -194,9 +189,7 @@ class TestSurveySE:
def test_strata_psu_changes_se(self, data_with_survey):
"""Strata + PSU should produce a different SE than weights-only."""
sd_weights_only = SurveyDesign(weights="pw")
- sd_full = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd_full = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r_w = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
@@ -305,19 +298,23 @@ def test_varying_weights_change_att(self, base_data):
multi["pw"] = np.where(np.arange(len(multi)) < len(df), 1.0, 10.0)
sd = SurveyDesign(weights="pw")
result_plain = ChaisemartinDHaultfoeuille(seed=1).fit(
- multi, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ multi,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
)
result_survey = ChaisemartinDHaultfoeuille(seed=1).fit(
- multi, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ multi,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
# Weighted cell means with time-varying noise produce different
# first differences -> different ATT
- assert result_plain.overall_att != pytest.approx(
- result_survey.overall_att, abs=0.01
- )
+ assert result_plain.overall_att != pytest.approx(result_survey.overall_att, abs=0.01)
def test_rejects_replicate_weights_with_bootstrap(self, base_data):
"""Replicate weights combined with n_bootstrap > 0 is rejected.
@@ -392,15 +389,9 @@ def test_bootstrap_survey_auto_inject_no_warning(self, data_with_survey):
with _w.catch_warnings():
_w.simplefilter("error", UserWarning)
# Ignore warnings unrelated to the bootstrap-PSU contract.
- _w.filterwarnings(
- "ignore", message="Single-period placebo SE"
- )
- _w.filterwarnings(
- "ignore", message="pweight weights normalized"
- )
- _w.filterwarnings(
- "ignore", message="Assumption 11"
- )
+ _w.filterwarnings("ignore", message="Single-period placebo SE")
+ _w.filterwarnings("ignore", message="pweight weights normalized")
+ _w.filterwarnings("ignore", message="Assumption 11")
ChaisemartinDHaultfoeuille(n_bootstrap=50, seed=1).fit(
data_with_survey,
outcome="outcome",
@@ -432,21 +423,25 @@ def test_uniform_survey_se_matches_plugin(self, base_data):
sd = SurveyDesign(weights="pw", psu="group")
r_plain = ChaisemartinDHaultfoeuille(seed=1).fit(
- base_data, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ base_data,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
)
r_survey = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
# With PSU=group and uniform weights, survey SE should be
# close to plug-in SE (both assume group-level independence).
# Small-sample corrections (n/(n-1)) cause minor differences.
if np.isfinite(r_plain.overall_se) and np.isfinite(r_survey.overall_se):
- assert r_plain.overall_se == pytest.approx(
- r_survey.overall_se, rel=0.15
- ), (
+ assert r_plain.overall_se == pytest.approx(r_survey.overall_se, rel=0.15), (
f"Survey SE ({r_survey.overall_se:.6f}) should be close to "
f"plug-in SE ({r_plain.overall_se:.6f}) with uniform weights "
f"and PSU=group"
@@ -471,8 +466,11 @@ def test_zero_weight_cell_excluded(self, base_data):
# Should not raise; the zero-weight cell is just absent
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -489,23 +487,22 @@ def test_survey_delta_uses_survey_df(self, data_with_survey):
t-distribution inference with df=df_survey (not z-inference)."""
from scipy import stats
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ survey_design=sd,
)
if not (np.isfinite(r.overall_se) and r.overall_se > 0):
pytest.skip("delta not estimable on this fixture")
assert r.survey_metadata is not None
df_s = r.survey_metadata.df_survey
- assert df_s is not None and df_s > 0, (
- f"expected positive df_survey, got {df_s}"
- )
+ assert df_s is not None and df_s > 0, f"expected positive df_survey, got {df_s}"
t_stat = r.overall_att / r.overall_se
p_t = 2.0 * (1.0 - stats.t.cdf(abs(t_stat), df=df_s))
@@ -526,31 +523,32 @@ def test_survey_delta_t_differs_from_z(self, base_data):
psu_map = {g: i // (n_g // 6) for i, g in enumerate(groups)}
df_["stratum"] = df_["group"].map(strata_map)
df_["cluster"] = df_["group"].map(psu_map)
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ survey_design=sd,
)
if not (np.isfinite(r.overall_se) and r.overall_se > 0):
pytest.skip("delta not estimable on this fixture")
assert r.survey_metadata is not None
df_s = r.survey_metadata.df_survey
- assert df_s is not None and df_s < 30, (
- f"expected small df_survey for t-vs-z gap, got {df_s}"
- )
+ assert (
+ df_s is not None and df_s < 30
+ ), f"expected small df_survey for t-vs-z gap, got {df_s}"
t_stat = r.overall_att / r.overall_se
p_t = 2.0 * (1.0 - stats.t.cdf(abs(t_stat), df=df_s))
p_z = 2.0 * (1.0 - stats.norm.cdf(abs(t_stat)))
# Threaded p-value must match t, not z
assert r.overall_p_value == pytest.approx(p_t, abs=1e-10)
- assert abs(r.overall_p_value - p_z) > 1e-6, (
- "overall_p_value must differ from z-inference when df_survey is small"
- )
+ assert (
+ abs(r.overall_p_value - p_z) > 1e-6
+ ), "overall_p_value must differ from z-inference when df_survey is small"
# ── Test: Survey + controls (DID^X) ─────────────────────────────────
@@ -564,14 +562,16 @@ def test_survey_plus_controls_runs(self, data_with_survey):
rng = np.random.default_rng(7)
df_ = data_with_survey.copy()
df_["x"] = rng.normal(0, 1.0, size=len(df_))
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- controls=["x"], L_max=1, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ controls=["x"],
+ L_max=1,
+ survey_design=sd,
)
assert np.isfinite(r.overall_att)
assert r.survey_metadata is not None
@@ -588,17 +588,19 @@ def test_survey_honest_did_propagates_df(self, data_with_survey):
results.survey_metadata.df_survey (non-None propagation)."""
import warnings
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
with warnings.catch_warnings():
# dCDH HonestDiD emits a methodology-deviation warning
warnings.simplefilter("ignore")
r = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, honest_did=True, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ honest_did=True,
+ survey_design=sd,
)
if r.honest_did_results is None:
pytest.skip("HonestDiD computation returned None on this fixture")
@@ -627,16 +629,23 @@ def test_uniform_weights_het_matches_unweighted(self, base_data):
r_plain = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
)
sd = SurveyDesign(weights="pw")
r_survey = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het", survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
+ survey_design=sd,
)
assert r_plain.heterogeneity_effects is not None
assert r_survey.heterogeneity_effects is not None
@@ -659,15 +668,22 @@ def test_nonuniform_het_changes_beta(self, base_data):
r_plain = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
)
r_survey = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het", survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
+ survey_design=sd,
)
assert r_plain.heterogeneity_effects is not None
assert r_survey.heterogeneity_effects is not None
@@ -675,9 +691,9 @@ def test_nonuniform_het_changes_beta(self, base_data):
b_survey = r_survey.heterogeneity_effects[1]["beta"]
if np.isfinite(b_plain) and np.isfinite(b_survey):
# WLS with non-uniform weights differs from unweighted OLS
- assert b_plain != pytest.approx(b_survey, abs=1e-6), (
- f"plain={b_plain} survey={b_survey} should differ with varying weights"
- )
+ assert b_plain != pytest.approx(
+ b_survey, abs=1e-6
+ ), f"plain={b_plain} survey={b_survey} should differ with varying weights"
def test_survey_het_uses_survey_df(self, data_with_survey):
"""Reported p_value must match t-distribution inference with df_survey."""
@@ -688,14 +704,16 @@ def test_survey_het_uses_survey_df(self, data_with_survey):
groups = sorted(df_["group"].unique())
het_map = {g: rng.uniform(-1, 1) for g in groups}
df_["x_het"] = df_["group"].map(het_map)
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het", survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
+ survey_design=sd,
)
assert r.heterogeneity_effects is not None
entry = r.heterogeneity_effects[1]
@@ -720,28 +738,28 @@ def test_twfe_helper_matches_fit_under_survey(self, data_with_survey):
"""fit and twowayfeweights() produce identical TWFE output under survey."""
from diff_diff.chaisemartin_dhaultfoeuille import twowayfeweights
- sd = SurveyDesign(
- weights="pw", strata="stratum", psu="cluster", nest=True
- )
+ sd = SurveyDesign(weights="pw", strata="stratum", psu="cluster", nest=True)
r = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
helper = twowayfeweights(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
# fit() twfe_* may be None if non-binary treatment; this fixture is binary
assert r.twfe_fraction_negative is not None
assert r.twfe_sigma_fe is not None
assert r.twfe_beta_fe is not None
- assert r.twfe_fraction_negative == pytest.approx(
- helper.fraction_negative, abs=1e-12
- )
+ assert r.twfe_fraction_negative == pytest.approx(helper.fraction_negative, abs=1e-12)
assert r.twfe_sigma_fe == pytest.approx(helper.sigma_fe, abs=1e-12)
assert r.twfe_beta_fe == pytest.approx(helper.beta_fe, abs=1e-12)
@@ -755,8 +773,10 @@ def test_twfe_helper_rejects_non_pweight(self, base_data):
with pytest.raises(ValueError, match="pweight"):
twowayfeweights(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
@@ -778,8 +798,10 @@ def test_twfe_helper_accepts_replicate_weights(self, base_data):
)
result = twowayfeweights(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.beta_fe)
@@ -801,8 +823,10 @@ def test_survey_twfe_matches_obs_level_pweighted_ols(self, data_with_survey):
sd = SurveyDesign(weights="pw")
helper = twowayfeweights(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(helper.beta_fe)
@@ -831,8 +855,10 @@ def test_survey_twfe_matches_obs_level_pweighted_ols(self, data_with_survey):
w_obs = df_["pw"].to_numpy().astype(float)
coef, _, _ = solve_ols(
- X_obs, y_obs,
- weights=w_obs, weight_type="pweight",
+ X_obs,
+ y_obs,
+ weights=w_obs,
+ weight_type="pweight",
return_vcov=False,
)
beta_oracle = float(coef[-1])
@@ -874,8 +900,10 @@ def test_mixed_zero_weight_row_excluded_from_validation(self, base_data):
# Must succeed (not raise fuzzy-DiD ValueError)
result = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -893,8 +921,10 @@ def test_zero_weight_row_with_nan_outcome(self, base_data):
# Must succeed — zero-weight row with NaN outcome is out-of-sample
result = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -913,8 +943,11 @@ def test_zero_weight_row_with_nan_group_id(self, base_data):
sd = SurveyDesign(weights="pw")
# Must succeed — zero-weight row's NaN group id is out-of-sample
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -935,9 +968,13 @@ def test_zero_weight_row_with_nan_control(self, base_data):
sd = SurveyDesign(weights="pw")
result = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, controls=["x"], survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ controls=["x"],
+ survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -959,9 +996,13 @@ def test_zero_weight_row_with_nan_heterogeneity(self, base_data):
# Must succeed — zero-weight row with NaN het is out-of-sample
result = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, heterogeneity="x_het", survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ heterogeneity="x_het",
+ survey_design=sd,
)
assert result.heterogeneity_effects is not None
@@ -976,21 +1017,26 @@ def test_survey_trends_linear_runs(self, data_with_survey):
sd = SurveyDesign(weights="pw")
r = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, trends_linear=True, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ trends_linear=True,
+ survey_design=sd,
)
assert r.survey_metadata is not None
# linear_trends_effects populated per REGISTRY line 614 contract
assert r.linear_trends_effects is not None
# At least one horizon should be estimable with finite value
finite_horizons = [
- h for h, entry in r.linear_trends_effects.items()
+ h
+ for h, entry in r.linear_trends_effects.items()
if np.isfinite(entry.get("effect", np.nan))
]
- assert len(finite_horizons) > 0, (
- "expected at least one horizon with finite linear_trends_effect"
- )
+ assert (
+ len(finite_horizons) > 0
+ ), "expected at least one horizon with finite linear_trends_effect"
# ── Test: Survey + trends_nonparam ──────────────────────────────────
@@ -1004,9 +1050,13 @@ def test_survey_trends_nonparam_runs(self, data_with_survey):
sd = SurveyDesign(weights="pw")
r = ChaisemartinDHaultfoeuille(seed=1).fit(
data_with_survey,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, trends_nonparam="stratum", survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ trends_nonparam="stratum",
+ survey_design=sd,
)
assert r.survey_metadata is not None
assert r.event_study_effects is not None
@@ -1045,22 +1095,22 @@ def _make_join_then_leave_panel(seed=42, n_groups=30, n_periods=8):
else:
d = 0
y = group_fe + 2.0 * t + 5.0 * d + rng.normal(0, 0.3)
- rows.append(
- {"group": g, "period": t, "treatment": d, "outcome": y, "pw": 1.0}
- )
+ rows.append({"group": g, "period": t, "treatment": d, "outcome": y, "pw": 1.0})
return pd.DataFrame(rows)
def test_survey_design2_runs(self):
df_ = self._make_join_then_leave_panel()
sd = SurveyDesign(weights="pw")
# drop_larger_lower=False keeps the 2-switch groups
- r = ChaisemartinDHaultfoeuille(
- seed=1, drop_larger_lower=False
- ).fit(
+ r = ChaisemartinDHaultfoeuille(seed=1, drop_larger_lower=False).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=1, design2=True, survey_design=sd,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=1,
+ design2=True,
+ survey_design=sd,
)
assert r.survey_metadata is not None
assert r.design2_effects is not None
@@ -1099,9 +1149,13 @@ def test_accepts_varying_psu_within_group(self, base_data):
df_["psu"] = df_["period"] % 2
sd = SurveyDesign(weights="pw", strata="stratum", psu="psu")
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=2,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se)
@@ -1113,9 +1167,13 @@ def test_accepts_varying_psu_within_group(self, base_data):
df_const["psu"] = 0 # constant-within-group PSU baseline
sd_const = SurveyDesign(weights="pw", strata="stratum", psu="psu")
r_const = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_const, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd_const, L_max=2,
+ df_const,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd_const,
+ L_max=2,
)
assert result.overall_se != r_const.overall_se, (
"Cell-period allocator must produce a different SE when PSU "
@@ -1136,9 +1194,13 @@ def test_accepts_varying_strata_within_group(self, base_data):
df_["psu"] = df_["group"]
sd = SurveyDesign(weights="pw", strata="stratum", psu="psu", nest=True)
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=2,
)
assert np.isfinite(result.overall_att)
assert np.isfinite(result.overall_se)
@@ -1160,9 +1222,13 @@ def test_heterogeneity_with_varying_psu_succeeds(self, base_data):
df_["psu"] = df_["group"].astype(int) * 2 + (df_["period"].astype(int) % 2)
sd = SurveyDesign(weights="pw", psu="psu")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- heterogeneity="x_het", L_max=1,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ heterogeneity="x_het",
+ L_max=1,
survey_design=sd,
)
assert res.heterogeneity_effects is not None
@@ -1190,9 +1256,13 @@ def test_bootstrap_with_varying_psu_succeeds(self, base_data):
df_["psu"] = df_["group"].astype(int) * 2 + (df_["period"].astype(int) % 2)
sd = SurveyDesign(weights="pw", psu="psu")
res = ChaisemartinDHaultfoeuille(n_bootstrap=200, seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=2,
)
assert res.bootstrap_results is not None
assert np.isfinite(res.bootstrap_results.overall_se)
@@ -1224,30 +1294,36 @@ def test_auto_inject_with_varying_strata_nest_true_succeeds(self, base_data):
df_["stratum"] = df_["period"] % 2
sd_auto = SurveyDesign(weights="pw", strata="stratum", nest=True)
sd_explicit = SurveyDesign(
- weights="pw", strata="stratum", psu="group", nest=True,
+ weights="pw",
+ strata="stratum",
+ psu="group",
+ nest=True,
)
r_auto = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd_auto, L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd_auto,
+ L_max=2,
)
r_explicit = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd_explicit, L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd_explicit,
+ L_max=2,
)
assert np.isfinite(r_auto.overall_att)
assert np.isfinite(r_auto.overall_se)
if np.isfinite(r_auto.overall_se) and np.isfinite(r_explicit.overall_se):
- assert r_auto.overall_se == pytest.approx(
- r_explicit.overall_se, rel=1e-6
- )
+ assert r_auto.overall_se == pytest.approx(r_explicit.overall_se, rel=1e-6)
assert r_auto.survey_metadata is not None
assert r_explicit.survey_metadata is not None
- assert (
- r_auto.survey_metadata.df_survey
- == r_explicit.survey_metadata.df_survey
- )
+ assert r_auto.survey_metadata.df_survey == r_explicit.survey_metadata.df_survey
def test_heterogeneity_auto_inject_with_varying_strata_nest_true_succeeds(self, base_data):
"""PR 3 unblocks heterogeneity + SurveyDesign(strata, nest=True)
@@ -1263,9 +1339,13 @@ def test_heterogeneity_auto_inject_with_varying_strata_nest_true_succeeds(self,
df_["x_het"] = (df_["group"].astype(int) % 2).astype(float)
sd = SurveyDesign(weights="pw", strata="stratum", nest=True)
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- heterogeneity="x_het", L_max=1,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ heterogeneity="x_het",
+ L_max=1,
survey_design=sd,
)
assert res.heterogeneity_effects is not None
@@ -1288,9 +1368,13 @@ def test_heterogeneity_multi_horizon_varying_psu_succeeds(self, base_data):
df_["psu"] = df_["group"].astype(int) * 2 + (df_["period"].astype(int) % 2)
sd = SurveyDesign(weights="pw", psu="psu")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- heterogeneity="x_het", L_max=2,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ heterogeneity="x_het",
+ L_max=2,
survey_design=sd,
)
assert res.heterogeneity_effects is not None
@@ -1316,8 +1400,11 @@ def test_auto_inject_with_varying_strata_raises(self, base_data):
sd = SurveyDesign(weights="pw", strata="stratum")
with pytest.raises(ValueError, match=r"psu="):
ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
@@ -1338,8 +1425,11 @@ def test_within_cell_psu_variation_rejected(self, base_data):
sd = SurveyDesign(weights="pw", strata="stratum", psu="psu")
with pytest.raises(ValueError, match="PSU to be constant within each"):
ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
@@ -1358,8 +1448,11 @@ def test_within_cell_strata_variation_rejected(self, base_data):
sd = SurveyDesign(weights="pw", strata="stratum", psu="psu")
with pytest.raises(ValueError, match="strata to be constant within each"):
ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
@@ -1371,8 +1464,11 @@ def test_accepts_varying_weights_within_group(self, base_data):
df_["pw"] = rng.uniform(0.5, 2.0, size=len(df_))
sd = SurveyDesign(weights="pw")
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -1384,28 +1480,27 @@ def test_auto_inject_psu_matches_explicit_group_psu(self, base_data):
df_ = base_data.copy()
df_["pw"] = 1.0
r_no_psu = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=SurveyDesign(weights="pw"),
)
r_explicit = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=SurveyDesign(weights="pw", psu="group"),
)
- assert r_no_psu.overall_att == pytest.approx(
- r_explicit.overall_att, abs=1e-10
- )
+ assert r_no_psu.overall_att == pytest.approx(r_explicit.overall_att, abs=1e-10)
if np.isfinite(r_no_psu.overall_se) and np.isfinite(r_explicit.overall_se):
- assert r_no_psu.overall_se == pytest.approx(
- r_explicit.overall_se, rel=1e-6
- )
+ assert r_no_psu.overall_se == pytest.approx(r_explicit.overall_se, rel=1e-6)
assert r_no_psu.survey_metadata is not None
assert r_explicit.survey_metadata is not None
- assert (
- r_no_psu.survey_metadata.df_survey
- == r_explicit.survey_metadata.df_survey
- )
+ assert r_no_psu.survey_metadata.df_survey == r_explicit.survey_metadata.df_survey
def test_degenerate_cohort_survey_se_is_nan(self):
"""When every variance-eligible group is its own singleton
@@ -1421,35 +1516,38 @@ def test_degenerate_cohort_survey_se_is_nan(self):
for t in range(5):
d = 1 if t >= f_switch else 0
y = float(g) + 0.5 * t + float(d)
- rows.append({
- "group": g,
- "period": t,
- "treatment": d,
- "outcome": y,
- "pw": 1.0,
- })
+ rows.append(
+ {
+ "group": g,
+ "period": t,
+ "treatment": d,
+ "outcome": y,
+ "pw": 1.0,
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(weights="pw")
import warnings as _warnings
+
with _warnings.catch_warnings(record=True) as w:
_warnings.simplefilter("always")
result = ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
# overall_se must be NaN on degenerate cohorts (not 0.0)
assert np.isnan(result.overall_se), (
- f"Degenerate-cohort survey overall_se must be NaN, "
- f"got {result.overall_se}"
+ f"Degenerate-cohort survey overall_se must be NaN, " f"got {result.overall_se}"
)
# Degenerate-cohort warning must fire
assert any(
- "cohort" in str(wi.message).lower()
- and "identically zero" in str(wi.message).lower()
+ "cohort" in str(wi.message).lower() and "identically zero" in str(wi.message).lower()
for wi in w
), "Expected degenerate-cohort warning to fire under survey path"
@@ -1469,24 +1567,27 @@ def test_subpopulation_preserves_full_design_df_survey(self, base_data):
sd = SurveyDesign(weights="pw")
r_subpop = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
# Reference: explicit psu='group' preserves the full-design
# PSU count because the resolver sees all groups (even those
# entirely zero-weighted). The auto-inject path must match this.
r_explicit = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=SurveyDesign(weights="pw", psu="group"),
)
assert r_subpop.survey_metadata is not None
assert r_explicit.survey_metadata is not None
- assert (
- r_subpop.survey_metadata.df_survey
- == r_explicit.survey_metadata.df_survey
- ), (
+ assert r_subpop.survey_metadata.df_survey == r_explicit.survey_metadata.df_survey, (
f"Auto-inject df_survey={r_subpop.survey_metadata.df_survey} "
f"must match explicit psu='group' df_survey="
f"{r_explicit.survey_metadata.df_survey} "
@@ -1503,9 +1604,13 @@ def test_off_horizon_row_duplication_does_not_change_se(self, base_data):
sd = SurveyDesign(weights="pw")
r_base = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, survey_design=sd,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ survey_design=sd,
)
# Duplicate the first row: cell mean y_gt stays the same
@@ -1515,14 +1620,16 @@ def test_off_horizon_row_duplication_does_not_change_se(self, base_data):
dup = df_.iloc[0].copy()
df_dup = pd.concat([df_, pd.DataFrame([dup])], ignore_index=True)
r_dup = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_dup, outcome="outcome", group="group",
- time="period", treatment="treatment",
- L_max=2, survey_design=sd,
+ df_dup,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ L_max=2,
+ survey_design=sd,
)
if np.isfinite(r_base.overall_se) and np.isfinite(r_dup.overall_se):
- assert r_base.overall_se == pytest.approx(
- r_dup.overall_se, rel=1e-6
- ), (
+ assert r_base.overall_se == pytest.approx(r_dup.overall_se, rel=1e-6), (
f"Row duplication changed SE ({r_base.overall_se} vs "
f"{r_dup.overall_se}) — auto-inject psu=group is not active."
)
@@ -1539,9 +1646,9 @@ def test_cell_allocator_row_sum_identity(self):
at ``t = 2``.
"""
from diff_diff.chaisemartin_dhaultfoeuille import (
- _compute_full_per_group_contributions,
_cohort_recenter,
_cohort_recenter_per_period,
+ _compute_full_per_group_contributions,
)
# D_mat, Y_mat, N_mat shaped (n_groups=4, n_periods=3).
@@ -1575,9 +1682,13 @@ def test_cell_allocator_row_sum_identity(self):
a11_minus_zeroed = np.array([True, True], dtype=bool)
U, U_pp = _compute_full_per_group_contributions(
- D_mat=D_mat, Y_mat=Y_mat, N_mat=N_mat,
- n_10_t_arr=n_10_t_arr, n_00_t_arr=n_00_t_arr,
- n_01_t_arr=n_01_t_arr, n_11_t_arr=n_11_t_arr,
+ D_mat=D_mat,
+ Y_mat=Y_mat,
+ N_mat=N_mat,
+ n_10_t_arr=n_10_t_arr,
+ n_00_t_arr=n_00_t_arr,
+ n_01_t_arr=n_01_t_arr,
+ n_11_t_arr=n_11_t_arr,
a11_plus_zeroed_arr=a11_plus_zeroed,
a11_minus_zeroed_arr=a11_minus_zeroed,
side="overall",
@@ -1628,8 +1739,11 @@ def test_within_cell_check_excludes_zero_weight_rows(self, base_data):
df_ = pd.concat([df_, pd.DataFrame([sample])], ignore_index=True)
sd = SurveyDesign(weights="pw", strata="stratum", psu="psu")
result = ChaisemartinDHaultfoeuille(seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -1664,11 +1778,13 @@ def test_psu_level_byte_identity_under_psu_equals_group(self):
rows = []
for g in range(n_groups):
for t in range(n_periods):
- rows.append({
- "group": int(g),
- "period": int(t),
- "pw": 1.0,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "pw": 1.0,
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(weights="pw", psu="group")
resolved = sd.resolve(df_)
@@ -1749,17 +1865,19 @@ def test_replicate_variance_non_invariance_under_varying_ratios(self):
rows = []
for g in range(2):
for t in range(2):
- rows.append({
- "group": int(g),
- "period": int(t),
- "pw": 1.0,
- # Non-PSU-aligned replicate columns: vary within
- # each group so sum_i ratio_ir * psi_i sees a
- # different weighted average of psi mass under
- # the two allocators.
- "rep0": 0.5 if t == 0 else 1.5,
- "rep1": 1.5 if t == 0 else 0.5,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "pw": 1.0,
+ # Non-PSU-aligned replicate columns: vary within
+ # each group so sum_i ratio_ir * psi_i sees a
+ # different weighted average of psi mass under
+ # the two allocators.
+ "rep0": 0.5 if t == 0 else 1.5,
+ "rep1": 1.5 if t == 0 else 0.5,
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(
weights="pw",
@@ -1781,9 +1899,7 @@ def test_replicate_variance_non_invariance_under_varying_ratios(self):
psi_legacy = np.zeros_like(w)
for i in range(len(w)):
if W_g[obs_group[i]] > 0:
- psi_legacy[i] = psi_g[obs_group[i]] * (
- w[i] / W_g[obs_group[i]]
- )
+ psi_legacy[i] = psi_g[obs_group[i]] * (w[i] / W_g[obs_group[i]])
# New cell-period single-cell expansion (what the Binder TSL
# branch uses). All mass lands on obs at t == 1.
@@ -1848,13 +1964,15 @@ def _make_baseline_fixture() -> pd.DataFrame:
for t in range(n_periods):
d = 1 if (f is not None and t >= f) else 0
y = float(g) + 0.3 * t + 1.5 * d + float(rng.normal(0.0, 0.5))
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": y,
- "pw": 1.0,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": y,
+ "pw": 1.0,
+ }
+ )
return pd.DataFrame(rows)
def test_bootstrap_se_matches_pre_pr4_baseline(self):
@@ -1877,9 +1995,7 @@ def test_bootstrap_se_matches_pre_pr4_baseline(self):
from diff_diff import bootstrap_utils as _bu
from diff_diff import linalg as _la
- bootstrap_rust = bool(
- _bu.HAS_RUST_BACKEND and _bu._rust_bootstrap_weights is not None
- )
+ bootstrap_rust = bool(_bu.HAS_RUST_BACKEND and _bu._rust_bootstrap_weights is not None)
ols_rust = bool(_la.HAS_RUST_BACKEND and _la._rust_solve_ols is not None)
assert bootstrap_rust == ols_rust, (
"Incoherent backend dispatch state between bootstrap_utils "
@@ -1889,22 +2005,25 @@ def test_bootstrap_se_matches_pre_pr4_baseline(self):
)
pure_python = not bootstrap_rust
expected = (
- self._BASELINE_OVERALL_SE_PYTHON
- if pure_python
- else self._BASELINE_OVERALL_SE_RUST
+ self._BASELINE_OVERALL_SE_PYTHON if pure_python else self._BASELINE_OVERALL_SE_RUST
)
df_ = self._make_baseline_fixture()
sd = SurveyDesign(weights="pw", psu="group")
res = ChaisemartinDHaultfoeuille(n_bootstrap=500, seed=42).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert res.bootstrap_results is not None
observed_se = float(res.bootstrap_results.overall_se)
backend_label = "pure-python" if pure_python else "rust"
assert observed_se == pytest.approx(
- expected, rel=0.0, abs=1e-15,
+ expected,
+ rel=0.0,
+ abs=1e-15,
), (
f"Bootstrap SE drifted from pre-PR-4 baseline "
f"(backend={backend_label}). expected={expected!r}, "
@@ -1927,7 +2046,8 @@ def test_bootstrap_cell_level_raises_on_missing_overall_tensor(self):
group_id_to_psu_code = {0: 0, 1: 1}
# Varying PSU: row 0 has two distinct PSU codes.
psu_codes_per_cell = np.array(
- [[0, 1], [0, 0]], dtype=np.int64,
+ [[0, 1], [0, 0]],
+ dtype=np.int64,
)
with pytest.raises(ValueError, match="u_per_period_overall"):
est._compute_dcdh_bootstrap(
@@ -1951,11 +2071,13 @@ def test_bootstrap_cell_level_raises_on_shape_mismatch(self):
eligible_group_ids = np.array([0, 1])
group_id_to_psu_code = {0: 0, 1: 1}
psu_codes_per_cell = np.array(
- [[0, 1], [0, 0]], dtype=np.int64,
+ [[0, 1], [0, 0]],
+ dtype=np.int64,
)
# Wrong shape — 3 columns instead of 2.
u_pp_wrong = np.array(
- [[0.25, 0.25, 0.0], [-0.15, -0.15, 0.0]], dtype=np.float64,
+ [[0.25, 0.25, 0.0], [-0.15, -0.15, 0.0]],
+ dtype=np.float64,
)
with pytest.raises(ValueError, match="shape mismatch"):
est._compute_dcdh_bootstrap(
@@ -1983,7 +2105,8 @@ def test_bootstrap_cell_level_raises_on_sentinel_mass_leak(self):
# whose PSU code is -1 (simulating terminal missingness
# after cohort-recentering leaks mass to a missing cell).
psu_codes_per_cell = np.array(
- [[0, 1, -1], [0, 1, 0]], dtype=np.int64,
+ [[0, 1, -1], [0, 1, 0]],
+ dtype=np.int64,
)
u_pp_overall_with_leak = np.array(
[[0.25, 0.25, -0.15], [-0.15, -0.15, 0.15]],
@@ -2012,12 +2135,14 @@ def test_bootstrap_cell_level_raises_on_missing_horizon_tensor(self):
est = ChaisemartinDHaultfoeuille(n_bootstrap=50, seed=1)
u_overall = np.array([0.5, -0.3], dtype=np.float64)
u_pp_overall = np.array(
- [[0.25, 0.25], [-0.15, -0.15]], dtype=np.float64,
+ [[0.25, 0.25], [-0.15, -0.15]],
+ dtype=np.float64,
)
eligible_group_ids = np.array([0, 1])
group_id_to_psu_code = {0: 0, 1: 1}
psu_codes_per_cell = np.array(
- [[0, 1], [0, 0]], dtype=np.int64,
+ [[0, 1], [0, 0]],
+ dtype=np.int64,
)
# Horizon tuple missing its per-cell tensor (4th slot = None).
u_h = np.array([0.4, -0.2], dtype=np.float64)
@@ -2051,19 +2176,24 @@ def test_bootstrap_cell_level_with_all_zero_weight_group_does_not_crash(self):
pw = 0.0 if g == n_groups - 1 else 1.0
d = 1 if (f is not None and t >= f) else 0
y = float(g) + 0.1 * t + 1.0 * d
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": y,
- "pw": pw,
- "psu": int(g) * 2 + (int(t) % 2), # varying PSU
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": y,
+ "pw": pw,
+ "psu": int(g) * 2 + (int(t) % 2), # varying PSU
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(weights="pw", psu="psu")
res = ChaisemartinDHaultfoeuille(n_bootstrap=50, seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert res.bootstrap_results is not None
@@ -2088,6 +2218,7 @@ def test_bootstrap_zero_weight_group_equivalent_to_removing_it(self):
correct paths collapse to the same identity-draw structure
under PSU=group — we deliberately use varying PSU here.
"""
+
def _make(include_zero_group: bool) -> pd.DataFrame:
rows = []
n_groups = 9 if include_zero_group else 8
@@ -2100,27 +2231,33 @@ def _make(include_zero_group: bool) -> pd.DataFrame:
# Within-group-varying PSU (period parity per
# group) — exercises the cell-level dispatcher.
psu = int(g) * 2 + (int(t) % 2)
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": y,
- "pw": pw,
- "psu": psu,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": y,
+ "pw": pw,
+ "psu": psu,
+ }
+ )
return pd.DataFrame(rows)
sd = SurveyDesign(weights="pw", psu="psu")
res_a = ChaisemartinDHaultfoeuille(n_bootstrap=200, seed=7).fit(
_make(include_zero_group=True),
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
res_b = ChaisemartinDHaultfoeuille(n_bootstrap=200, seed=7).fit(
_make(include_zero_group=False),
- outcome="outcome", group="group",
- time="period", treatment="treatment",
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
assert res_a.bootstrap_results is not None
@@ -2178,19 +2315,22 @@ def test_fit_raises_on_terminal_missingness_with_varying_psu(self):
continue
d = d_pattern[t]
y = float(g) + 0.1 * t + 1.0 * d
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": y,
- "pw": 1.0,
- # Within-group-varying PSU: period parity per group.
- "psu": int(g) * 2 + (int(t) % 2),
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": y,
+ "pw": 1.0,
+ # Within-group-varying PSU: period parity per group.
+ "psu": int(g) * 2 + (int(t) % 2),
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(weights="pw", psu="psu")
import warnings as _w
+
# Analytical path (n_bootstrap=0): the sentinel-mass guard in
# `_survey_se_from_group_if` raises on the same leakage the
# bootstrap guard rejects — both paths use the cell-period
@@ -2199,12 +2339,17 @@ def test_fit_raises_on_terminal_missingness_with_varying_psu(self):
with _w.catch_warnings():
_w.simplefilter("ignore") # terminal-missingness UserWarning
with pytest.raises(
- ValueError, match="no positive-weight observations",
+ ValueError,
+ match="no positive-weight observations",
):
ChaisemartinDHaultfoeuille(n_bootstrap=0, seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
# Bootstrap path (n_bootstrap > 0): same sentinel-mass guard
@@ -2212,12 +2357,17 @@ def test_fit_raises_on_terminal_missingness_with_varying_psu(self):
with _w.catch_warnings():
_w.simplefilter("ignore")
with pytest.raises(
- ValueError, match="no positive-weight observations",
+ ValueError,
+ match="no positive-weight observations",
):
ChaisemartinDHaultfoeuille(n_bootstrap=50, seed=1).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
def test_fit_succeeds_on_terminal_missingness_with_psu_group(self):
@@ -2245,27 +2395,35 @@ def test_fit_succeeds_on_terminal_missingness_with_psu_group(self):
continue
d = d_pattern[t]
y = float(g) + 0.1 * t + 1.0 * d
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": y,
- "pw": 1.0,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": y,
+ "pw": 1.0,
+ }
+ )
df_ = pd.DataFrame(rows)
# Auto-inject: no explicit `psu` → `SurveyDesign` falls back to
# `psu=` at fit() time. Within-group-constant.
sd = SurveyDesign(weights="pw")
import warnings as _w
+
for n_boot in (0, 50):
with _w.catch_warnings():
_w.simplefilter("ignore")
res = ChaisemartinDHaultfoeuille(
- n_bootstrap=n_boot, seed=1,
+ n_bootstrap=n_boot,
+ seed=1,
).fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
assert np.isfinite(res.overall_att), (
f"n_bootstrap={n_boot}: overall_att must be finite "
@@ -2303,14 +2461,16 @@ def test_bootstrap_dense_codes_under_singleton_baseline_excluded_group(self):
d = 1 if t >= 3 else 0 # joiners at period 3
# Groups 1, 2 share PSU=200; groups 3, 4 share PSU=300.
psu = 200 if g in (1, 2) else 300
- rows.append({
- "group": int(g),
- "period": int(t),
- "treatment": int(d),
- "outcome": float(g) + 0.1 * t + 0.5 * d,
- "pw": 1.0,
- "psu": psu,
- })
+ rows.append(
+ {
+ "group": int(g),
+ "period": int(t),
+ "treatment": int(d),
+ "outcome": float(g) + 0.1 * t + 0.5 * d,
+ "pw": 1.0,
+ "psu": psu,
+ }
+ )
df_ = pd.DataFrame(rows)
sd = SurveyDesign(weights="pw", psu="psu")
@@ -2320,19 +2480,21 @@ def test_bootstrap_dense_codes_under_singleton_baseline_excluded_group(self):
original_bootstrap = est._compute_dcdh_bootstrap
def _spy(**kwargs):
- captured["group_id_to_psu_code"] = kwargs.get(
- "group_id_to_psu_code"
- )
+ captured["group_id_to_psu_code"] = kwargs.get("group_id_to_psu_code")
return original_bootstrap(**kwargs)
est._compute_dcdh_bootstrap = _spy # type: ignore[method-assign]
import warnings as _w
+
with _w.catch_warnings():
_w.simplefilter("ignore") # singleton-baseline warning
est.fit(
- df_, outcome="outcome", group="group",
- time="period", treatment="treatment",
+ df_,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
survey_design=sd,
)
@@ -2360,10 +2522,8 @@ def _spy(**kwargs):
"dense code under correct densification."
)
assert dict_passed[3] == dict_passed[4], (
- "Groups 3 and 4 share PSU=300 and must receive the same "
- "dense code."
+ "Groups 3 and 4 share PSU=300 and must receive the same " "dense code."
)
assert dict_passed[1] != dict_passed[3], (
- "Groups in PSU=200 and PSU=300 must receive distinct "
- "dense codes."
+ "Groups in PSU=200 and PSU=300 must receive distinct " "dense codes."
)
diff --git a/tests/test_survey_dcdh_replicate_psu.py b/tests/test_survey_dcdh_replicate_psu.py
index aac02615..a1a64729 100644
--- a/tests/test_survey_dcdh_replicate_psu.py
+++ b/tests/test_survey_dcdh_replicate_psu.py
@@ -19,7 +19,6 @@
from diff_diff import ChaisemartinDHaultfoeuille, SurveyDesign
from diff_diff.chaisemartin_dhaultfoeuille import twowayfeweights
-
# ── Fixtures ────────────────────────────────────────────────────────
@@ -181,6 +180,7 @@ def test_att_raises_on_terminal_missingness_replicate_path(self):
pre-process the panel. Locks this user-facing failure mode.
"""
import warnings as _w
+
rows = []
for g in range(10):
if g < 5:
@@ -216,7 +216,8 @@ def test_att_raises_on_terminal_missingness_replicate_path(self):
with _w.catch_warnings():
_w.simplefilter("ignore") # terminal-missingness UserWarning
with pytest.raises(
- ValueError, match="no positive-weight observations",
+ ValueError,
+ match="no positive-weight observations",
):
ChaisemartinDHaultfoeuille(seed=1).fit(
df_,
@@ -229,7 +230,10 @@ def test_att_raises_on_terminal_missingness_replicate_path(self):
)
def test_att_cell_allocator_with_varying_replicate_ratios(
- self, base_panel, replicate_design, monkeypatch,
+ self,
+ base_panel,
+ replicate_design,
+ monkeypatch,
):
"""Class A ATT replicate contract: `_survey_se_from_group_if`
dispatches replicate fits through the cell-period allocator
@@ -252,6 +256,7 @@ def test_att_cell_allocator_with_varying_replicate_ratios(
the fit's SE.
"""
import numpy as _np
+
import diff_diff.survey as _survey_mod
R = 20
@@ -268,9 +273,7 @@ def _spy(psi, resolved):
captured_psi.append(_np.asarray(psi).copy())
return original_repvar(psi, resolved)
- monkeypatch.setattr(
- _survey_mod, "compute_replicate_if_variance", _spy
- )
+ monkeypatch.setattr(_survey_mod, "compute_replicate_if_variance", _spy)
res = ChaisemartinDHaultfoeuille(seed=1).fit(
df,
@@ -313,20 +316,19 @@ def _spy(psi, resolved):
for i in range(n_obs):
gi = g_to_idx[obs_group[i]]
if W_g[gi] > 0:
- psi_obs_legacy[i] = (
- U_centered_recon[gi] * w[i] / W_g[gi]
- )
+ psi_obs_legacy[i] = U_centered_recon[gi] * w[i] / W_g[gi]
# (a) Fit's overall_se^2 must match the variance recomputed
# from the captured psi_obs. Trivial sanity check — the fit
# computed it from this exact psi_obs via the spy pass-through.
var_cell_recompute, _ = original_repvar(psi_obs_actual, resolved)
assert float(res.overall_se) ** 2 == pytest.approx(
- var_cell_recompute, rel=1e-12,
+ var_cell_recompute,
+ rel=1e-12,
), (
- f"Fit's overall_se^2 does not match the recomputed "
- f"variance on the captured psi_obs — spy / dispatch "
- f"misalignment."
+ "Fit's overall_se^2 does not match the recomputed "
+ "variance on the captured psi_obs — spy / dispatch "
+ "misalignment."
)
# (b) Legacy-reconstructed psi_obs yields a different Rao-Wu
@@ -337,7 +339,9 @@ def _spy(psi, resolved):
# this assertion would fail, surfacing the regression.
var_legacy, _ = original_repvar(psi_obs_legacy, resolved)
assert not _np.isclose(
- var_cell_recompute, var_legacy, rtol=1e-6,
+ var_cell_recompute,
+ var_legacy,
+ rtol=1e-6,
), (
f"Expected cell-allocator variance to differ from legacy "
f"group-level variance on this fixture (per-row "
@@ -420,9 +424,7 @@ def test_jk1_converges_to_tsl(self, replicate_design):
)
@pytest.mark.parametrize("method", ["BRR", "JK1"])
- def test_multi_horizon_under_replicate(
- self, replicate_design, method
- ):
+ def test_multi_horizon_under_replicate(self, replicate_design, method):
"""Multi-horizon DID_l inherits the Class A dispatch — horizon 1
(always identifiable on this panel) should produce a finite
replicate-based SE."""
@@ -464,8 +466,7 @@ def test_placebo_under_replicate(self, replicate_design):
L_max=1,
)
assert res.placebo_event_study, (
- "Expected placebo_event_study to be populated for L_max=1 "
- "reversible panel"
+ "Expected placebo_event_study to be populated for L_max=1 " "reversible panel"
)
# Negative key convention: placebo horizon l → key -l.
assert -1 in res.placebo_event_study, (
@@ -559,9 +560,7 @@ def test_twowayfeweights_accepts_replicate(self, base_panel, replicate_design):
# present but don't change the aggregated diagnostics.
assert res_rep.beta_fe == pytest.approx(res_plain.beta_fe, rel=1e-10)
assert res_rep.sigma_fe == pytest.approx(res_plain.sigma_fe, rel=1e-10)
- assert res_rep.fraction_negative == pytest.approx(
- res_plain.fraction_negative, rel=1e-10
- )
+ assert res_rep.fraction_negative == pytest.approx(res_plain.fraction_negative, rel=1e-10)
# ── 3. PSU-level Hall-Mammen wild bootstrap ─────────────────────────
@@ -632,8 +631,7 @@ def test_no_warning_under_auto_inject(self, base_panel):
survey_design=sd,
)
assert not any(
- "Hall-Mammen wild" in str(w.message)
- or "group-level multiplier" in str(w.message)
+ "Hall-Mammen wild" in str(w.message) or "group-level multiplier" in str(w.message)
for w in caught
), (
"Bootstrap-PSU warning should not fire under auto-inject psu=group, "
@@ -878,9 +876,7 @@ def test_generate_psu_or_group_weights_identity(self):
assert np.array_equal(W_identity, W_plain)
@pytest.mark.parametrize("method", ["BRR", "JK1"])
- def test_honest_did_under_replicate(
- self, base_panel, replicate_design, method
- ):
+ def test_honest_did_under_replicate(self, base_panel, replicate_design, method):
"""HonestDiD bounds should flow through replicate SE and the
reduced df_survey (df = min(n_valid) - 1)."""
R = 20
@@ -928,15 +924,19 @@ def test_rank_deficient_replicate_uses_design_df(self, base_panel):
resolved = sd.resolve(df)
design_df = resolved.df_survey
assert design_df is not None and design_df < R - 1, (
- f"Expected rank deficiency: design df={design_df} "
- f"should be < {R - 1}"
+ f"Expected rank deficiency: design df={design_df} " f"should be < {R - 1}"
)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1, honest_did=True,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
+ honest_did=True,
)
# survey_metadata carries the final (reduced) df — not R - 1.
assert res.survey_metadata.df_survey == design_df, (
@@ -969,15 +969,18 @@ def test_dropped_replicate_reduces_df(self, base_panel):
resolved = sd.resolve(df)
design_df = resolved.df_survey
assert design_df is not None and design_df <= R - 3, (
- f"Expected zero-column rank deficiency: design df={design_df} "
- f"should be <= {R - 3}"
+ f"Expected zero-column rank deficiency: design df={design_df} " f"should be <= {R - 3}"
)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
# Persisted df must be capped by the design df, not R - 1.
assert res.survey_metadata.df_survey is not None
@@ -1012,15 +1015,17 @@ def test_rank_1_replicate_forces_nan_inference(self, base_panel):
df[f"rep{r}"] = rep_col_shared
sd = _build_replicate_design(R, "BRR")
resolved = sd.resolve(df)
- assert resolved.df_survey is None, (
- "Rank-1 replicate matrix must have undefined design df"
- )
+ assert resolved.df_survey is None, "Rank-1 replicate matrix must have undefined design df"
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
# survey_metadata.df_survey stays None — the honest undefined
# signal for downstream consumers.
@@ -1038,7 +1043,8 @@ def test_rank_1_replicate_forces_nan_inference(self, base_panel):
assert np.isnan(info["p_value"])
def test_heterogeneity_replicate_cross_surface_df_consistency(
- self, base_panel,
+ self,
+ base_panel,
):
"""With ``heterogeneity=`` active under a rank-deficient
replicate design, every public surface — top-level dCDH,
@@ -1067,9 +1073,13 @@ def test_heterogeneity_replicate_cross_surface_df_consistency(
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
heterogeneity="x_het",
honest_did=True,
)
@@ -1089,16 +1099,16 @@ def test_heterogeneity_replicate_cross_surface_df_consistency(
if np.isfinite(het_info["se"]):
het_t = het_info["beta"] / het_info["se"]
expected_het_p = 2 * _stats.t.sf(abs(het_t), df=final_df)
- assert het_info["p_value"] == pytest.approx(
- expected_het_p, rel=1e-6
- )
+ assert het_info["p_value"] == pytest.approx(expected_het_p, rel=1e-6)
# HonestDiD bounds were computed from the same survey_metadata;
# asserting the result attribute is populated confirms the df
# flow-through was exercised.
assert res.honest_did_results is not None
def test_heterogeneity_late_nvalid_propagates_to_all_surfaces(
- self, base_panel, monkeypatch,
+ self,
+ base_panel,
+ monkeypatch,
):
"""Regression for PR #311 CI review R4 P2.
@@ -1116,6 +1126,7 @@ def test_heterogeneity_late_nvalid_propagates_to_all_surfaces(
`survey_metadata.df_survey`) must reflect the reduced df.
"""
from scipy import stats as _stats
+
from diff_diff import survey as _survey_mod
R = 20
@@ -1136,15 +1147,17 @@ def count_only(psi, resolved_arg):
main_counter["n"] += 1
return original(psi, resolved_arg)
- monkeypatch.setattr(
- _survey_mod, "compute_replicate_if_variance", count_only
- )
+ monkeypatch.setattr(_survey_mod, "compute_replicate_if_variance", count_only)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
main_call_count = main_counter["n"]
monkeypatch.undo()
@@ -1161,15 +1174,20 @@ def reduce_after_main(psi, resolved_arg):
return var, n_valid
monkeypatch.setattr(
- _survey_mod, "compute_replicate_if_variance",
+ _survey_mod,
+ "compute_replicate_if_variance",
reduce_after_main,
)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
heterogeneity="x_het",
)
@@ -1185,7 +1203,8 @@ def reduce_after_main(psi, resolved_arg):
t_stat = res.overall_att / res.overall_se
expected_p_overall = 2 * _stats.t.sf(abs(t_stat), df=expected_df)
assert res.overall_p_value == pytest.approx(
- expected_p_overall, rel=1e-6,
+ expected_p_overall,
+ rel=1e-6,
), (
"overall_p_value should reflect heterogeneity's late "
"n_valid via the post-recompute block"
@@ -1198,7 +1217,8 @@ def reduce_after_main(psi, resolved_arg):
t_l = info["effect"] / info["se"]
expected_p_l = 2 * _stats.t.sf(abs(t_l), df=expected_df)
assert info["p_value"] == pytest.approx(
- expected_p_l, rel=1e-6,
+ expected_p_l,
+ rel=1e-6,
)
# Placebo event study surface (NEGATIVE keys). This is the
@@ -1210,10 +1230,12 @@ def reduce_after_main(psi, resolved_arg):
if np.isfinite(pl_info["se"]) and pl_info["se"] > 0:
t_pl = pl_info["effect"] / pl_info["se"]
expected_p_pl = 2 * _stats.t.sf(
- abs(t_pl), df=expected_df,
+ abs(t_pl),
+ df=expected_df,
)
assert pl_info["p_value"] == pytest.approx(
- expected_p_pl, rel=1e-6,
+ expected_p_pl,
+ rel=1e-6,
), (
f"placebo_event_study[{_lag_neg}] p-value "
f"must reflect the reduced df "
@@ -1228,10 +1250,12 @@ def reduce_after_main(psi, resolved_arg):
if np.isfinite(het_info["se"]) and het_info["se"] > 0:
het_t = het_info["beta"] / het_info["se"]
expected_het_p = 2 * _stats.t.sf(
- abs(het_t), df=expected_df,
+ abs(het_t),
+ df=expected_df,
)
assert het_info["p_value"] == pytest.approx(
- expected_het_p, rel=1e-6,
+ expected_het_p,
+ rel=1e-6,
)
# Normalized effects surface. This is a DERIVED dict built
@@ -1245,10 +1269,12 @@ def reduce_after_main(psi, resolved_arg):
if np.isfinite(norm_info["se"]) and norm_info["se"] > 0:
norm_t = norm_info["effect"] / norm_info["se"]
expected_norm_p = 2 * _stats.t.sf(
- abs(norm_t), df=expected_df,
+ abs(norm_t),
+ df=expected_df,
)
assert norm_info["p_value"] == pytest.approx(
- expected_norm_p, rel=1e-6,
+ expected_norm_p,
+ rel=1e-6,
), (
f"normalized_effects[{_lag}] p-value must "
f"reflect the reduced df ({expected_df}); got "
@@ -1257,7 +1283,9 @@ def reduce_after_main(psi, resolved_arg):
)
def test_late_nvalid_below_two_forces_all_surfaces_nan(
- self, base_panel, monkeypatch,
+ self,
+ base_panel,
+ monkeypatch,
):
"""Regression for R4 P2 follow-up: when heterogeneity's late
replicate failures drive the final `n_valid` to 1 (reduced df
@@ -1285,15 +1313,17 @@ def count_only(psi, resolved_arg):
main_counter["n"] += 1
return original(psi, resolved_arg)
- monkeypatch.setattr(
- _survey_mod, "compute_replicate_if_variance", count_only
- )
+ monkeypatch.setattr(_survey_mod, "compute_replicate_if_variance", count_only)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
)
main_call_count = main_counter["n"]
monkeypatch.undo()
@@ -1309,15 +1339,17 @@ def reduce_to_one(psi, resolved_arg):
return var, 1
return var, n_valid
- monkeypatch.setattr(
- _survey_mod, "compute_replicate_if_variance", reduce_to_one
- )
+ monkeypatch.setattr(_survey_mod, "compute_replicate_if_variance", reduce_to_one)
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
res = ChaisemartinDHaultfoeuille(seed=1).fit(
- df, outcome="outcome", group="group",
- time="period", treatment="treatment",
- survey_design=sd, L_max=1,
+ df,
+ outcome="outcome",
+ group="group",
+ time="period",
+ treatment="treatment",
+ survey_design=sd,
+ L_max=1,
heterogeneity="x_het",
)
# Final effective df is None (reduced_df = 0 < 1 guard).
diff --git a/tests/test_survey_estimator_validation.py b/tests/test_survey_estimator_validation.py
index 70d3ffd1..1c293a84 100644
--- a/tests/test_survey_estimator_validation.py
+++ b/tests/test_survey_estimator_validation.py
@@ -117,25 +117,18 @@ def test_s1_imputation_control_regression(self, golden):
# Build Omega_0 (untreated observations)
# Matches imputation.py line 341: omega_0_mask = ~df["_treated"]
omega_0 = data[
- (np.isinf(data["first_treat"]))
- | (data["period"] < data["first_treat"])
+ (np.isinf(data["first_treat"])) | (data["period"] < data["first_treat"])
].copy()
omega_0 = omega_0.reset_index(drop=True)
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc"
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
resolved = sd.resolve(omega_0)
survey_weights = omega_0["weight"].values.astype(float)
# Build design matrix: unit + time dummies + covariates
# Algebraically equivalent to within-transformation for covariate coefs
- unit_dummies = pd.get_dummies(
- omega_0["unit"].astype(str), prefix="u", drop_first=True
- )
- time_dummies = pd.get_dummies(
- omega_0["period"].astype(str), prefix="t", drop_first=True
- )
+ unit_dummies = pd.get_dummies(omega_0["unit"].astype(str), prefix="u", drop_first=True)
+ time_dummies = pd.get_dummies(omega_0["period"].astype(str), prefix="t", drop_first=True)
X = np.column_stack(
[
unit_dummies.values.astype(float),
@@ -187,9 +180,7 @@ def test_s1_imputation_control_regression(self, golden):
def test_s1_imputation_end_to_end(self, golden):
"""Smoke test: ImputationDiD.fit() with survey weights produces finite ATT."""
data = self._load_staggered_data(golden)
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc"
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = ImputationDiD()
result = est.fit(
data,
@@ -213,9 +204,7 @@ def test_s2_stacked_did(self, golden):
r = golden["s2_stacked_did"]
data = self._load_staggered_data(golden)
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc"
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = StackedDiD(
kappa_pre=1,
kappa_post=1,
@@ -232,12 +221,8 @@ def test_s2_stacked_did(self, golden):
survey_design=sd,
)
- _assert_close(
- result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S2 ATT"
- )
- _assert_close(
- result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S2 SE"
- )
+ _assert_close(result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S2 ATT")
+ _assert_close(result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S2 SE")
# Compare individual event-study coefficients and SEs
if "event_study" in r:
@@ -269,9 +254,7 @@ def test_s3_sun_abraham(self, golden):
r = golden["s3_sun_abraham"]
data = self._load_staggered_data(golden)
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc"
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = SunAbraham()
result = est.fit(
data,
@@ -283,12 +266,8 @@ def test_s3_sun_abraham(self, golden):
)
# Compare overall IW-aggregated ATT
- _assert_close(
- result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S3 ATT"
- )
- _assert_close(
- result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S3 SE"
- )
+ _assert_close(result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S3 ATT")
+ _assert_close(result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S3 SE")
# Compare individual cohort x rel-time effects and SEs (post-treatment only)
if "cohort_effects" in r:
@@ -320,9 +299,7 @@ def test_s4_triple_diff(self, golden):
r = golden["s4_triple_diff"]
data = self._load_ddd_data(golden)
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc"
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = TripleDifference(estimation_method="reg")
result = est.fit(
data,
@@ -333,16 +310,8 @@ def test_s4_triple_diff(self, golden):
survey_design=sd,
)
- _assert_close(
- result.att, r["att"], COEF_RTOL, COEF_ATOL, "S4 DDD coef"
- )
+ _assert_close(result.att, r["att"], COEF_RTOL, COEF_ATOL, "S4 DDD coef")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "S4 DDD SE")
- _assert_close(
- result.t_stat, r["t_stat"], SE_RTOL, SE_ATOL, "S4 t_stat"
- )
- _assert_close(
- result.conf_int[0], r["ci_lower"], SE_RTOL, SE_ATOL, "S4 CI lower"
- )
- _assert_close(
- result.conf_int[1], r["ci_upper"], SE_RTOL, SE_ATOL, "S4 CI upper"
- )
+ _assert_close(result.t_stat, r["t_stat"], SE_RTOL, SE_ATOL, "S4 t_stat")
+ _assert_close(result.conf_int[0], r["ci_lower"], SE_RTOL, SE_ATOL, "S4 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], SE_RTOL, SE_ATOL, "S4 CI upper")
diff --git a/tests/test_survey_phase3.py b/tests/test_survey_phase3.py
index fbac504b..6e7748e4 100644
--- a/tests/test_survey_phase3.py
+++ b/tests/test_survey_phase3.py
@@ -931,20 +931,24 @@ def test_uniform_weight_equivalence(self, cov_survey_data):
result_survey = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
result_nosurv = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
)
assert result_survey.estimation_path == "dr"
assert result_nosurv.estimation_path == "dr"
- np.testing.assert_allclose(
- result_survey.overall_att, result_nosurv.overall_att, atol=1e-8
- )
+ np.testing.assert_allclose(result_survey.overall_att, result_nosurv.overall_att, atol=1e-8)
def test_scale_invariance(self, cov_survey_data):
"""Multiplying all weights by a constant doesn't change ATT."""
@@ -953,7 +957,10 @@ def test_scale_invariance(self, cov_survey_data):
sd1 = SurveyDesign(weights="weight")
result1 = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd1,
)
@@ -962,13 +969,14 @@ def test_scale_invariance(self, cov_survey_data):
sd2 = SurveyDesign(weights="weight_scaled")
result2 = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd2,
)
- np.testing.assert_allclose(
- result1.overall_att, result2.overall_att, atol=1e-8
- )
+ np.testing.assert_allclose(result1.overall_att, result2.overall_att, atol=1e-8)
def test_nontrivial_weight_effect(self, cov_survey_data):
"""Heterogeneous weights produce different ATT from unweighted."""
@@ -977,13 +985,19 @@ def test_nontrivial_weight_effect(self, cov_survey_data):
sd = SurveyDesign(weights="weight")
result_survey = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
result_nosurv = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
)
# Non-uniform weights (1.0 + 0.3*stratum) should produce different ATT
@@ -994,12 +1008,18 @@ def test_full_design_smoke(self, cov_survey_data):
from diff_diff import EfficientDiD
sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc",
+ weights="weight",
+ strata="stratum",
+ psu="psu",
+ fpc="fpc",
nest=True,
)
result = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
@@ -1015,7 +1035,10 @@ def test_aggregation_with_survey(self, cov_survey_data):
sd = SurveyDesign(weights="weight")
result = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
aggregate="all",
survey_design=sd,
@@ -1036,7 +1059,10 @@ def test_bootstrap_covariates_survey(self, cov_survey_data):
sd = SurveyDesign(weights="weight")
result = EfficientDiD(n_bootstrap=30, seed=42).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
@@ -1051,13 +1077,19 @@ def test_analytical_se_differs_from_unweighted(self, cov_survey_data):
sd = SurveyDesign(weights="weight")
result_survey = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
result_nosurv = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
)
# Non-uniform weights (1.0 + 0.3*stratum) should produce different SEs
@@ -1072,20 +1104,24 @@ def test_bootstrap_se_in_ballpark_of_analytical(self, cov_survey_data):
sd = SurveyDesign(weights="weight")
result_analytical = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
result_boot = EfficientDiD(n_bootstrap=199, seed=42).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
ratio = result_boot.overall_se / result_analytical.overall_se
- assert 0.3 < ratio < 3.0, (
- f"Bootstrap/analytical SE ratio {ratio:.2f} outside [0.3, 3.0]"
- )
+ assert 0.3 < ratio < 3.0, f"Bootstrap/analytical SE ratio {ratio:.2f} outside [0.3, 3.0]"
def test_zero_weight_cohort_skipped(self, cov_survey_data):
"""Zero-weight treated cohort should be skipped with a warning."""
@@ -1098,7 +1134,10 @@ def test_zero_weight_cohort_skipped(self, cov_survey_data):
with pytest.warns(UserWarning, match="zero survey weight"):
result = EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
@@ -1115,7 +1154,10 @@ def test_zero_weight_never_treated_raises(self, cov_survey_data):
with pytest.raises(ValueError, match="zero survey weight"):
EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
survey_design=sd,
)
@@ -1130,7 +1172,10 @@ def test_zero_weight_never_treated_nocov_raises(self, cov_survey_data):
with pytest.raises(ValueError, match="zero survey weight"):
EfficientDiD(n_bootstrap=0).fit(
cov_survey_data,
- "outcome", "unit", "time", "first_treat",
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
survey_design=sd,
)
@@ -1146,8 +1191,7 @@ def test_replicate_weight_aggregation(self):
for unit in range(n_units):
wt = 1.0 + 0.3 * (unit % 5)
unit_repwts[unit] = {
- f"repwt_{r}": wt * (0.5 + np.random.random())
- for r in range(n_rep)
+ f"repwt_{r}": wt * (0.5 + np.random.random()) for r in range(n_rep)
}
rows = []
for unit in range(n_units):
@@ -1159,20 +1203,32 @@ def test_replicate_weight_aggregation(self):
if ft > 0 and t >= ft:
y += 2.0
y += np.random.normal(0, 0.5)
- row = {"unit": unit, "time": t, "first_treat": ft,
- "outcome": y, "weight": wt, "x1": x1}
+ row = {
+ "unit": unit,
+ "time": t,
+ "first_treat": ft,
+ "outcome": y,
+ "weight": wt,
+ "x1": x1,
+ }
row.update(unit_repwts[unit])
rows.append(row)
import pandas as pd
+
data = pd.DataFrame(rows)
rep_cols = [f"repwt_{r}" for r in range(n_rep)]
sd = SurveyDesign(
- weights="weight", replicate_weights=rep_cols,
+ weights="weight",
+ replicate_weights=rep_cols,
replicate_method="JK1",
)
result = EfficientDiD(n_bootstrap=0).fit(
- data, "outcome", "unit", "time", "first_treat",
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
covariates=["x1"],
aggregate="event_study",
survey_design=sd,
@@ -1540,11 +1596,18 @@ def test_fpc_census_zero_variance_stacked(self, staggered_survey_data):
data.loc[mask, "fpc"] = float(n_psu_h)
sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc",
+ weights="weight",
+ strata="stratum",
+ psu="psu",
+ fpc="fpc",
nest=True,
)
result = StackedDiD().fit(
- data, "outcome", "unit", "time", "first_treat",
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
survey_design=sd,
)
assert result.overall_se == pytest.approx(0.0, abs=1e-10)
@@ -1558,12 +1621,18 @@ def test_single_psu_lonely_remove_nan_se(self, staggered_survey_data):
data["single_psu"] = data["stratum"]
sd = SurveyDesign(
- weights="weight", strata="stratum", psu="single_psu",
+ weights="weight",
+ strata="stratum",
+ psu="single_psu",
lonely_psu="remove",
)
with pytest.warns(UserWarning, match="only 1 PSU"):
result = SunAbraham().fit(
- data, "outcome", "unit", "time", "first_treat",
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
survey_design=sd,
)
# All strata removed -> NaN SE
@@ -1579,7 +1648,12 @@ def test_full_design_bootstrap_continuous_did(self, staggered_survey_data):
sd = SurveyDesign(weights="weight", strata="stratum")
result = ContinuousDiD(n_bootstrap=30, seed=42).fit(
- data, "outcome", "unit", "time", "first_treat", "dose",
+ data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
+ "dose",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
@@ -1592,7 +1666,11 @@ def test_full_design_bootstrap_efficient_did(self, staggered_survey_data):
sd = SurveyDesign(weights="weight", strata="stratum")
result = EfficientDiD(n_bootstrap=30, seed=42).fit(
- staggered_survey_data, "outcome", "unit", "time", "first_treat",
+ staggered_survey_data,
+ "outcome",
+ "unit",
+ "time",
+ "first_treat",
survey_design=sd,
)
assert np.isfinite(result.overall_att)
diff --git a/tests/test_survey_phase5.py b/tests/test_survey_phase5.py
index 9f601161..afdccf99 100644
--- a/tests/test_survey_phase5.py
+++ b/tests/test_survey_phase5.py
@@ -246,9 +246,7 @@ def test_full_design_placebo_succeeds(self, sdid_survey_data_full_design):
summary = result.summary()
assert "Survey Design" in summary
- def test_full_design_jackknife_succeeds(
- self, sdid_survey_data_jk_well_formed
- ):
+ def test_full_design_jackknife_succeeds(self, sdid_survey_data_jk_well_formed):
"""Jackknife variance with full design now succeeds (restored capability).
PSU-level LOO with stratum aggregation (Rust & Rao 1996):
@@ -279,9 +277,7 @@ def test_full_design_jackknife_succeeds(
summary = result.summary()
assert "Survey Design" in summary
- def test_placebo_with_pweight_only_full_design_stripped_att_match(
- self, sdid_survey_data
- ):
+ def test_placebo_with_pweight_only_full_design_stripped_att_match(self, sdid_survey_data):
"""Placebo ATT with pweight-only is unchanged when stratum/psu
columns are physically dropped from the input DataFrame.
@@ -378,9 +374,7 @@ def test_summary_includes_survey(self, sdid_survey_data, survey_design_weights):
assert "Survey Design" in summary
assert "pweight" in summary
- def test_bootstrap_with_pweight_only_succeeds(
- self, sdid_survey_data, survey_design_weights
- ):
+ def test_bootstrap_with_pweight_only_succeeds(self, sdid_survey_data, survey_design_weights):
"""variance_method='bootstrap' with pweight-only survey succeeds (PR #352).
Restored capability: the bootstrap loop dispatches to the
@@ -404,9 +398,7 @@ def test_bootstrap_with_pweight_only_succeeds(
assert result.se > 0
assert result.variance_method == "bootstrap"
- def test_bootstrap_full_design_se_differs_from_pweight_only(
- self, sdid_survey_data
- ):
+ def test_bootstrap_full_design_se_differs_from_pweight_only(self, sdid_survey_data):
"""Full-design bootstrap SE differs from pweight-only bootstrap SE.
Resurrects the test_full_design_se_differs_from_weights_only
@@ -632,9 +624,7 @@ def sdid_survey_data_full_design():
unit_weight = 1.0 + np.arange(n_units) * 0.05
unit_stratum = np.array([0] * 15 + [1] * 15)
- unit_psu = np.array(
- [0] * 5 + [1] * 5 + [2] * 5 + [3] * 5 + [4] * 5 + [5] * 5
- )
+ unit_psu = np.array([0] * 5 + [1] * 5 + [2] * 5 + [3] * 5 + [4] * 5 + [5] * 5)
unit_map = {u: i for i, u in enumerate(units)}
idx = data["unit"].map(unit_map).values
@@ -702,10 +692,10 @@ def sdid_survey_data_jk_well_formed():
# stratum-1 units across PSU 3/4/5). Treated units 0-4 straddle
# PSU 0 (units 0-1) and PSU 1 (units 2-4).
unit_psu = np.zeros(n_units, dtype=int)
- unit_psu[0:2] = 0 # PSU 0: treated 0, 1
- unit_psu[2:5] = 1 # PSU 1: treated 2, 3, 4
- unit_psu[5:7] = 0 # PSU 0: control 5, 6
- unit_psu[7:9] = 1 # PSU 1: control 7, 8
+ unit_psu[0:2] = 0 # PSU 0: treated 0, 1
+ unit_psu[2:5] = 1 # PSU 1: treated 2, 3, 4
+ unit_psu[5:7] = 0 # PSU 0: control 5, 6
+ unit_psu[7:9] = 1 # PSU 1: control 7, 8
unit_psu[9:13] = 2 # PSU 2: control 9-12
unit_psu[13:18] = 3 # PSU 3: control 13-17
unit_psu[18:23] = 4 # PSU 4: control 18-22
@@ -759,9 +749,7 @@ def test_placebo_full_design_pseudo_treated_stays_within_treated_strata(
# 0-4 + controls 5-14, stratum 1 has controls 15-29.
treated_units = list(range(5))
control_units = list(range(5, 30))
- unit_to_stratum = (
- sdid_survey_data_full_design.groupby("unit")["stratum"].first().to_dict()
- )
+ unit_to_stratum = sdid_survey_data_full_design.groupby("unit")["stratum"].first().to_dict()
strata_control = np.array([unit_to_stratum[u] for u in control_units])
treated_strata_set = set(unit_to_stratum[u] for u in treated_units)
@@ -818,9 +806,7 @@ def fake_default_rng(seed=None):
f"treated_strata_set {treated_strata_set}"
)
- def test_placebo_full_design_raises_on_zero_control_stratum(
- self, sdid_survey_data_full_design
- ):
+ def test_placebo_full_design_raises_on_zero_control_stratum(self, sdid_survey_data_full_design):
"""Case B: stratum with treated units but zero controls → ValueError."""
df = sdid_survey_data_full_design.copy()
# Move all controls out of stratum 0; treated stays in stratum 0.
@@ -828,9 +814,7 @@ def test_placebo_full_design_raises_on_zero_control_stratum(
sd = SurveyDesign(weights="weight", strata="stratum", psu="psu")
est = SyntheticDiD(variance_method="placebo", n_bootstrap=30, seed=7)
- with pytest.raises(
- ValueError, match=r"at least one control per stratum.*has 0 controls"
- ):
+ with pytest.raises(ValueError, match=r"at least one control per stratum.*has 0 controls"):
est.fit(
df,
outcome="outcome",
@@ -966,9 +950,7 @@ def test_placebo_full_design_raises_on_undersupplied_stratum(
# ``nest=True`` so the shifted PSUs stay unique-within-stratum.
df.loc[df["unit"].isin(range(7, 15)), "stratum"] = 1
- sd = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", nest=True
- )
+ sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", nest=True)
est = SyntheticDiD(variance_method="placebo", n_bootstrap=30, seed=7)
with pytest.raises(
ValueError,
@@ -984,9 +966,7 @@ def test_placebo_full_design_raises_on_undersupplied_stratum(
survey_design=sd,
)
- def test_placebo_full_design_se_differs_from_pweight_only(
- self, sdid_survey_data_full_design
- ):
+ def test_placebo_full_design_se_differs_from_pweight_only(self, sdid_survey_data_full_design):
"""Full-design placebo SE differs from pweight-only placebo SE.
Pweight-only path permutes across ALL controls (unstratified);
@@ -1078,9 +1058,7 @@ def test_placebo_fpc_alone_no_op_warns_and_matches_pweight_only(
assert r_fpc.se == pytest.approx(r_pw.se, rel=1e-12)
assert r_fpc.att == pytest.approx(r_pw.att, abs=1e-12)
- def test_placebo_typo_fpc_column_still_raises(
- self, sdid_survey_data_full_design
- ):
+ def test_placebo_typo_fpc_column_still_raises(self, sdid_survey_data_full_design):
"""R13 P3 fix: typoed FPC column name must still raise on placebo.
The FPC pre-resolve drop on placebo (R11 P1) bypasses
@@ -1131,9 +1109,7 @@ def test_placebo_low_fpc_with_explicit_psu_skips_resolve_validator(
# is below the 3-PSU threshold per stratum.
df["fpc_low"] = 2.0
- sd_low_fpc_psu = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc_low"
- )
+ sd_low_fpc_psu = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc_low")
sd_no_fpc = SurveyDesign(weights="weight", strata="stratum", psu="psu")
est_fpc = SyntheticDiD(variance_method="placebo", n_bootstrap=50, seed=42)
@@ -1183,9 +1159,7 @@ def test_placebo_low_fpc_with_explicit_psu_skips_resolve_validator(
survey_design=sd_low_fpc_psu,
)
- def test_placebo_low_fpc_no_psu_warns_no_validator_block(
- self, sdid_survey_data_full_design
- ):
+ def test_placebo_low_fpc_no_psu_warns_no_validator_block(self, sdid_survey_data_full_design):
"""R10 P1 fix: implicit-PSU FPC validator skipped on placebo.
The implicit-PSU FPC validator (PR #355 R8 P1) rejects designs
@@ -1387,9 +1361,7 @@ def test_jackknife_full_design_stratum_aggregation_formula_magnitude(
expected_se = np.sqrt(factor * (ss0 + ss1))
assert result.se == pytest.approx(expected_se, rel=1e-12)
- def test_jackknife_full_design_fpc_reduces_se_magnitude(
- self, sdid_survey_data_jk_well_formed
- ):
+ def test_jackknife_full_design_fpc_reduces_se_magnitude(self, sdid_survey_data_jk_well_formed):
"""With FPC, SE is reduced by the (1-f_h) multiplier per stratum.
Two fits: one without FPC (f_h=0 so (1-f_h)=1); one with FPC set
@@ -1402,9 +1374,7 @@ def test_jackknife_full_design_fpc_reduces_se_magnitude(
df_fpc["fpc_col"] = 6.0 # n_h=3 per stratum, f_h = 3/6 = 0.5
sd_no_fpc = SurveyDesign(weights="weight", strata="stratum", psu="psu")
- sd_fpc = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", fpc="fpc_col"
- )
+ sd_fpc = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc_col")
est1 = SyntheticDiD(variance_method="jackknife", seed=42)
result_no_fpc = est1.fit(
@@ -1427,9 +1397,7 @@ def test_jackknife_full_design_fpc_reduces_se_magnitude(
survey_design=sd_fpc,
)
# Expected magnitude ratio: SE_fpc/SE_no_fpc = sqrt(1 - 0.5) = 1/sqrt(2)
- assert result_fpc.se == pytest.approx(
- result_no_fpc.se / np.sqrt(2), rel=1e-10
- )
+ assert result_fpc.se == pytest.approx(result_no_fpc.se / np.sqrt(2), rel=1e-10)
def test_jackknife_full_design_se_differs_from_pweight_only(
self, sdid_survey_data_jk_well_formed
@@ -1466,9 +1434,7 @@ def test_jackknife_full_design_se_differs_from_pweight_only(
assert result_pw.att == pytest.approx(result_full.att, abs=1e-10)
assert result_pw.se != pytest.approx(result_full.se, abs=1e-6)
- def test_get_loo_effects_df_raises_on_survey_jackknife(
- self, sdid_survey_data_jk_well_formed
- ):
+ def test_get_loo_effects_df_raises_on_survey_jackknife(self, sdid_survey_data_jk_well_formed):
"""R1 P1 fix: get_loo_effects_df blocks only on full-design survey
jackknife (PSU-level replicates), not on pweight-only jackknife.
@@ -1607,9 +1573,7 @@ def test_jackknife_full_design_full_census_short_circuits_undefined_loo(
assert np.isfinite(result.se)
assert result.se == 0.0
- def test_jackknife_full_design_lonely_psu_adjust_raises(
- self, sdid_survey_data_jk_well_formed
- ):
+ def test_jackknife_full_design_lonely_psu_adjust_raises(self, sdid_survey_data_jk_well_formed):
"""R5 P1 fix: ``SurveyDesign(lonely_psu='adjust')`` on the jackknife
survey path raises NotImplementedError rather than silently being
treated as ``"remove"``.
@@ -1693,9 +1657,7 @@ def test_jackknife_full_design_all_certainty_psu_returns_zero_se(
# Under "remove": same design returns SE=NaN with the "every
# stratum was skipped" warning (no contributing stratum).
- sd_remove = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", lonely_psu="remove"
- )
+ sd_remove = SurveyDesign(weights="weight", strata="stratum", psu="psu", lonely_psu="remove")
est_rem = SyntheticDiD(variance_method="jackknife", seed=42)
with pytest.warns(UserWarning, match=r"every stratum was skipped"):
result_rem = est_rem.fit(
@@ -1716,9 +1678,7 @@ def test_jackknife_full_design_lonely_psu_certainty_equivalent_to_remove(
as ``lonely_psu='remove'`` (both contribute 0 for singleton
strata on the jackknife path).
"""
- sd_remove = SurveyDesign(
- weights="weight", strata="stratum", psu="psu", lonely_psu="remove"
- )
+ sd_remove = SurveyDesign(weights="weight", strata="stratum", psu="psu", lonely_psu="remove")
sd_certainty = SurveyDesign(
weights="weight", strata="stratum", psu="psu", lonely_psu="certainty"
)
@@ -1776,9 +1736,7 @@ def test_jackknife_full_design_undefined_replicate_returns_nan(
)
assert np.isnan(result.se)
- def test_jackknife_full_design_single_psu_stratum_skipped(
- self, sdid_survey_data_full_design
- ):
+ def test_jackknife_full_design_single_psu_stratum_skipped(self, sdid_survey_data_full_design):
"""Stratum with only 1 PSU contributes 0 to total variance.
Degenerate stratum: relabel stratum-0 PSU 1+2 to a new stratum 2
@@ -1808,9 +1766,7 @@ def test_jackknife_full_design_single_psu_stratum_skipped(
assert np.isfinite(result.se)
assert result.se > 0
- def test_jackknife_full_design_unstratified_short_circuit(
- self, sdid_survey_data_full_design
- ):
+ def test_jackknife_full_design_unstratified_short_circuit(self, sdid_survey_data_full_design):
"""No strata + single PSU → SE=NaN (unidentified variance)."""
df = sdid_survey_data_full_design.copy()
df["psu"] = 0 # all units in a single PSU
@@ -2321,15 +2277,27 @@ def test_rao_wu_approximates_block_no_strata(self, trop_survey_data):
data["unit_psu"] = data["unit"]
sd_rw = SurveyDesign(
- weights="weight", strata="single_stratum", psu="unit_psu",
+ weights="weight",
+ strata="single_stratum",
+ psu="unit_psu",
)
sd_block = SurveyDesign(weights="weight")
result_rw = TROP(method="local", n_bootstrap=99, seed=42, max_iter=5).fit(
- data, "outcome", "D", "unit", "time", survey_design=sd_rw,
+ data,
+ "outcome",
+ "D",
+ "unit",
+ "time",
+ survey_design=sd_rw,
)
result_block = TROP(method="local", n_bootstrap=99, seed=42, max_iter=5).fit(
- data, "outcome", "D", "unit", "time", survey_design=sd_block,
+ data,
+ "outcome",
+ "D",
+ "unit",
+ "time",
+ survey_design=sd_block,
)
# Point estimates identical (same weights)
diff --git a/tests/test_survey_r_crossvalidation.py b/tests/test_survey_r_crossvalidation.py
index f978c78c..e735faab 100644
--- a/tests/test_survey_r_crossvalidation.py
+++ b/tests/test_survey_r_crossvalidation.py
@@ -61,12 +61,7 @@ def r_results():
def _load_scenario_data(scenario):
"""Extract DataFrame from a scenario's embedded data dict."""
data = scenario["data"]
- df = pd.DataFrame(
- {
- col: vals
- for col, vals in data.items()
- }
- )
+ df = pd.DataFrame({col: vals for col, vals in data.items()})
# Handle R's Inf serialization (jsonlite writes Inf as "Inf" string)
if "first_treat" in df.columns:
raw = df["first_treat"]
@@ -98,7 +93,9 @@ def _build_survey_design(design_type, has_fpc=True):
"""Construct the SurveyDesign for a given design type."""
if design_type == "strata_psu_fpc":
return SurveyDesign(
- weights="weight", strata="stratum", psu="psu",
+ weights="weight",
+ strata="stratum",
+ psu="psu",
fpc="fpc" if has_fpc else None,
)
elif design_type == "strata_psu_nofpc":
@@ -144,7 +141,10 @@ def _fit_scenario(self, r_results, key):
est = DifferenceInDifferences()
result = est.fit(
- data, "outcome", "treated", "post",
+ data,
+ "outcome",
+ "treated",
+ "post",
covariates=covariates,
survey_design=sd,
)
@@ -153,31 +153,25 @@ def _fit_scenario(self, r_results, key):
@pytest.mark.parametrize("key", TIER1_SCENARIOS)
def test_att_matches_r(self, r_results, key):
result, r = self._fit_scenario(r_results, key)
- _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL,
- f"{key} ATT")
+ _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, f"{key} ATT")
@pytest.mark.parametrize("key", TIER1_SCENARIOS)
def test_se_matches_r(self, r_results, key):
result, r = self._fit_scenario(r_results, key)
- _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL,
- f"{key} SE")
+ _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, f"{key} SE")
@pytest.mark.parametrize("key", TIER1_SCENARIOS)
def test_df_matches_r(self, r_results, key):
result, r = self._fit_scenario(r_results, key)
py_df = result.survey_metadata.df_survey
r_df = r["df"]
- assert py_df == r_df, (
- f"{key} df: Python={py_df}, R={r_df}"
- )
+ assert py_df == r_df, f"{key} df: Python={py_df}, R={r_df}"
@pytest.mark.parametrize("key", TIER1_SCENARIOS)
def test_ci_matches_r(self, r_results, key):
result, r = self._fit_scenario(r_results, key)
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- f"{key} CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- f"{key} CI upper")
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, f"{key} CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, f"{key} CI upper")
@pytest.mark.parametrize("key", TIER1_SCENARIOS)
def test_fpc_reduces_se(self, r_results, key):
@@ -186,15 +180,16 @@ def test_fpc_reduces_se(self, r_results, key):
pytest.skip("Only applies to strata_psu_fpc_s42")
r_fpc = r_results["strata_psu_fpc_s42"]
r_nofpc = r_results["strata_psu_nofpc_s42"]
- assert r_fpc["se"] < r_nofpc["se"], (
- f"FPC SE ({r_fpc['se']:.6f}) should be < non-FPC SE ({r_nofpc['se']:.6f})"
- )
+ assert (
+ r_fpc["se"] < r_nofpc["se"]
+ ), f"FPC SE ({r_fpc['se']:.6f}) should be < non-FPC SE ({r_nofpc['se']:.6f})"
# ---------------------------------------------------------------------------
# Tier 1b: TWFE plumbing test
# ---------------------------------------------------------------------------
+
class TestSvyglmTWFE:
"""Verify TWFE survey plumbing produces correct results.
@@ -208,8 +203,7 @@ def test_twfe_survey_produces_finite_results(self, r_results):
sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = TwoWayFixedEffects()
- result = est.fit(data, "outcome", "treated", "period", "unit",
- survey_design=sd)
+ result = est.fit(data, "outcome", "treated", "period", "unit", survey_design=sd)
assert np.isfinite(result.att), "TWFE ATT should be finite"
assert np.isfinite(result.se), "TWFE SE should be finite"
@@ -221,12 +215,11 @@ def test_twfe_df_matches_expected(self, r_results):
sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
est = TwoWayFixedEffects()
- result = est.fit(data, "outcome", "treated", "period", "unit",
- survey_design=sd)
+ result = est.fit(data, "outcome", "treated", "period", "unit", survey_design=sd)
- assert result.survey_metadata.df_survey == r["df"], (
- f"TWFE df: Python={result.survey_metadata.df_survey}, expected={r['df']}"
- )
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"TWFE df: Python={result.survey_metadata.df_survey}, expected={r['df']}"
def test_twfe_att_close_to_pooled(self, r_results):
"""In balanced 2x2, TWFE ATT ≈ pooled DiD ATT (FWL equivalence)."""
@@ -234,13 +227,12 @@ def test_twfe_att_close_to_pooled(self, r_results):
data = _load_scenario_data(r)
sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")
- pooled = DifferenceInDifferences().fit(
- data, "outcome", "treated", "post", survey_design=sd)
+ pooled = DifferenceInDifferences().fit(data, "outcome", "treated", "post", survey_design=sd)
twfe = TwoWayFixedEffects().fit(
- data, "outcome", "treated", "period", "unit", survey_design=sd)
+ data, "outcome", "treated", "period", "unit", survey_design=sd
+ )
- _assert_close(twfe.att, pooled.att, rtol=1e-3, atol=0.01,
- label="TWFE vs pooled ATT")
+ _assert_close(twfe.att, pooled.att, rtol=1e-3, atol=0.01, label="TWFE vs pooled ATT")
# ---------------------------------------------------------------------------
@@ -281,7 +273,11 @@ def _fit_scenario(self, r_results, key):
n_bootstrap=0,
)
result = est.fit(
- data, "outcome", "unit", "period", "first_treat",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
covariates=covariates,
aggregate="simple",
survey_design=sd,
@@ -299,19 +295,20 @@ def test_att_matches_r_did(self, r_results, key):
# Build R (group, time) -> ATT mapping, sorted for stable comparison
r_gt = {
- (int(g), int(t)): att
- for g, t, att in zip(r["gt_groups"], r["gt_periods"], r["gt_att"])
+ (int(g), int(t)): att for g, t, att in zip(r["gt_groups"], r["gt_periods"], r["gt_att"])
}
- assert len(py_gt) == len(r_gt), (
- f"{key}: Python has {len(py_gt)} GT cells, R has {len(r_gt)}"
- )
+ assert len(py_gt) == len(
+ r_gt
+ ), f"{key}: Python has {len(py_gt)} GT cells, R has {len(r_gt)}"
for (g, t), r_att in sorted(r_gt.items()):
assert (g, t) in py_gt, f"{key}: missing Python GT cell ({g},{t})"
_assert_close(
- py_gt[(g, t)]["effect"], r_att,
- rtol=1e-3, atol=0.05,
+ py_gt[(g, t)]["effect"],
+ r_att,
+ rtol=1e-3,
+ atol=0.05,
label=f"{key} ATT(g={g},t={t})",
)
@@ -323,8 +320,10 @@ def test_overall_att_matches(self, r_results, key):
result, r = self._fit_scenario(r_results, key)
_assert_close(
- result.overall_att, r["overall_att"],
- rtol=1e-3, atol=0.05,
+ result.overall_att,
+ r["overall_att"],
+ rtol=1e-3,
+ atol=0.05,
label=f"{key} overall ATT",
)
@@ -351,22 +350,37 @@ def test_design_se_differs_from_naive(self, r_results, key):
# With survey weights
sd = SurveyDesign(weights="weight")
est_survey = CallawaySantAnna(
- estimation_method="reg", control_group="never_treated",
- base_period="varying", n_bootstrap=0,
+ estimation_method="reg",
+ control_group="never_treated",
+ base_period="varying",
+ n_bootstrap=0,
)
result_survey = est_survey.fit(
- data, "outcome", "unit", "period", "first_treat",
- covariates=covariates, aggregate="simple", survey_design=sd,
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ covariates=covariates,
+ aggregate="simple",
+ survey_design=sd,
)
# Without survey weights
est_naive = CallawaySantAnna(
- estimation_method="reg", control_group="never_treated",
- base_period="varying", n_bootstrap=0,
+ estimation_method="reg",
+ control_group="never_treated",
+ base_period="varying",
+ n_bootstrap=0,
)
result_naive = est_naive.fit(
- data, "outcome", "unit", "period", "first_treat",
- covariates=covariates, aggregate="simple",
+ data,
+ "outcome",
+ "unit",
+ "period",
+ "first_treat",
+ covariates=covariates,
+ aggregate="simple",
)
assert result_survey.overall_se != result_naive.overall_se, (
@@ -379,6 +393,7 @@ def test_design_se_differs_from_naive(self, r_results, key):
# Tier 3: BRR replicate weights
# ---------------------------------------------------------------------------
+
class TestReplicateWeights:
"""Cross-validate BRR replicate weight variance vs R svrepdesign."""
@@ -402,11 +417,13 @@ def _fit_brr(self, r_results):
# Build design matrix manually (no intercept — LR adds it)
from diff_diff.linalg import LinearRegression
- X = np.column_stack([
- data["treated"].values.astype(float),
- data["post"].values.astype(float),
- (data["treated"] * data["post"]).values.astype(float),
- ])
+ X = np.column_stack(
+ [
+ data["treated"].values.astype(float),
+ data["post"].values.astype(float),
+ (data["treated"] * data["post"]).values.astype(float),
+ ]
+ )
y = data["outcome"].values.astype(float)
reg = LinearRegression(
@@ -424,8 +441,7 @@ def test_att_matches_svrepglm(self, r_results):
reg, r = self._fit_brr(r_results)
py_att = reg.coefficients_[3]
- _assert_close(py_att, r["att"], ATT_RTOL, ATT_ATOL,
- "BRR ATT")
+ _assert_close(py_att, r["att"], ATT_RTOL, ATT_ATOL, "BRR ATT")
def test_se_matches_svrepglm(self, r_results):
if "brr_replicate" not in r_results:
@@ -434,14 +450,11 @@ def test_se_matches_svrepglm(self, r_results):
py_se = np.sqrt(reg.vcov_[3, 3])
# Wider tolerance for replicate variance (approximate method)
- _assert_close(py_se, r["se"], rtol=0.05, atol=0.1,
- label="BRR SE")
+ _assert_close(py_se, r["se"], rtol=0.05, atol=0.1, label="BRR SE")
def test_df_matches_svrepglm(self, r_results):
if "brr_replicate" not in r_results:
pytest.skip("BRR scenario not in R results")
reg, r = self._fit_brr(r_results)
- assert reg.survey_df_ == r["df"], (
- f"BRR df: Python={reg.survey_df_}, R={r['df']}"
- )
+ assert reg.survey_df_ == r["df"], f"BRR df: Python={reg.survey_df_}, R={r['df']}"
diff --git a/tests/test_survey_real_data.py b/tests/test_survey_real_data.py
index 79d57ebf..ae6d5a67 100644
--- a/tests/test_survey_real_data.py
+++ b/tests/test_survey_real_data.py
@@ -154,13 +154,11 @@ def test_a1_strata_fpc_weights(self, api_results):
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A1 ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A1 SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"A1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "A1 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "A1 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"A1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "A1 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "A1 CI upper")
def test_a2_strata_no_fpc(self, api_results):
"""A2: TSL with strata + weights, no FPC."""
@@ -173,13 +171,11 @@ def test_a2_strata_no_fpc(self, api_results):
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A2 ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A2 SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"A2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "A2 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "A2 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"A2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "A2 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "A2 CI upper")
def test_a3_weights_only(self, api_results):
"""A3: Weighted OLS baseline (no strata or FPC)."""
@@ -192,10 +188,8 @@ def test_a3_weights_only(self, api_results):
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A3 ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A3 SE")
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "A3 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "A3 CI upper")
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "A3 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "A3 CI upper")
def test_a4_twfe(self, api_results):
"""A4: TWFE plumbing with strata + FPC + weights."""
@@ -205,7 +199,11 @@ def test_a4_twfe(self, api_results):
sd = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
est = TwoWayFixedEffects()
result = est.fit(
- data, "outcome", "treated", "period", "school_id",
+ data,
+ "outcome",
+ "treated",
+ "period",
+ "school_id",
survey_design=sd,
)
@@ -224,12 +222,17 @@ def test_a5_subpopulation_elementary(self, api_results):
# Use subpopulation approach: zero out weights for non-elementary
sd_full = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
sd_sub, data_sub = sd_full.subpopulation(
- data, lambda df: df["stype"] == "E",
+ data,
+ lambda df: df["stype"] == "E",
)
est = DifferenceInDifferences()
result = est.fit(
- data_sub, "outcome", "treated", "post", survey_design=sd_sub,
+ data_sub,
+ "outcome",
+ "treated",
+ "post",
+ survey_design=sd_sub,
)
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A5 subpop ATT")
@@ -247,20 +250,21 @@ def test_a6_covariates(self, api_results):
sd = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
est = DifferenceInDifferences()
result = est.fit(
- data, "outcome", "treated", "post",
+ data,
+ "outcome",
+ "treated",
+ "post",
covariates=["meals", "ell"],
survey_design=sd,
)
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A6 cov ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A6 cov SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"A6 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "A6 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "A6 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"A6 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "A6 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "A6 CI upper")
def test_a7_fay_brr_replicates(self, api_results):
"""A7: Fay's BRR replicate weights from real stratified design."""
@@ -285,21 +289,19 @@ def test_a7_fay_brr_replicates(self, api_results):
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A7 Fay ATT")
_assert_close(result.se, r["se"], REP_SE_RTOL, REP_SE_ATOL, "A7 Fay SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"A7 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL,
- "A7 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL,
- "A7 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"A7 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL, "A7 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL, "A7 CI upper")
def test_fpc_reduces_se(self, api_results):
"""FPC should reduce SE vs no-FPC (A1 SE < A2 SE)."""
r_fpc = api_results["a1_strata_fpc_weights"]
r_nofpc = api_results["a2_strata_weights"]
- assert r_fpc["se"] < r_nofpc["se"], (
- f"FPC SE ({r_fpc['se']:.4f}) should be < no-FPC SE ({r_nofpc['se']:.4f})"
- )
+ assert (
+ r_fpc["se"] < r_nofpc["se"]
+ ), f"FPC SE ({r_fpc['se']:.4f}) should be < no-FPC SE ({r_nofpc['se']:.4f})"
# ============================================================================
@@ -340,20 +342,21 @@ def test_b1_strata_psu_weights(self, nhanes_results):
data = self._load_nhanes_data(nhanes_results)
sd = SurveyDesign(
- weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
+ weights="WTMEC2YR",
+ strata="SDMVSTRA",
+ psu="SDMVPSU",
+ nest=True,
)
est = DifferenceInDifferences()
result = est.fit(data, "outcome", "treated", "post", survey_design=sd)
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B1 ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B1 SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"B1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "B1 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "B1 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"B1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "B1 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "B1 CI upper")
def test_b2_covariates(self, nhanes_results):
"""B2: Covariate-adjusted TSL with strata + PSU."""
@@ -362,24 +365,28 @@ def test_b2_covariates(self, nhanes_results):
data = data.dropna(subset=r.get("covariates", [])).reset_index(drop=True)
sd = SurveyDesign(
- weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
+ weights="WTMEC2YR",
+ strata="SDMVSTRA",
+ psu="SDMVPSU",
+ nest=True,
)
est = DifferenceInDifferences()
result = est.fit(
- data, "outcome", "treated", "post",
+ data,
+ "outcome",
+ "treated",
+ "post",
covariates=r.get("covariates", []),
survey_design=sd,
)
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B2 cov ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B2 cov SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"B2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "B2 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "B2 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"B2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "B2 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "B2 CI upper")
def test_b3_weights_only(self, nhanes_results):
"""B3: Weighted OLS (no clustering), baseline comparison."""
@@ -392,13 +399,11 @@ def test_b3_weights_only(self, nhanes_results):
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B3 ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B3 SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"B3 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "B3 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "B3 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"B3 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "B3 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "B3 CI upper")
def test_b4_subpop_female(self, nhanes_results):
"""B4: Subpopulation — female respondents only."""
@@ -406,26 +411,32 @@ def test_b4_subpop_female(self, nhanes_results):
data = self._load_nhanes_data(nhanes_results)
sd_full = SurveyDesign(
- weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
+ weights="WTMEC2YR",
+ strata="SDMVSTRA",
+ psu="SDMVPSU",
+ nest=True,
)
# Subpopulation: female (RIAGENDR == 2)
sd_sub, data_sub = sd_full.subpopulation(
- data, lambda df: df["RIAGENDR"] == 2,
+ data,
+ lambda df: df["RIAGENDR"] == 2,
)
est = DifferenceInDifferences()
result = est.fit(
- data_sub, "outcome", "treated", "post", survey_design=sd_sub,
+ data_sub,
+ "outcome",
+ "treated",
+ "post",
+ survey_design=sd_sub,
)
_assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B4 subpop ATT")
_assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B4 subpop SE")
- assert result.survey_metadata.df_survey == r["df"], (
- f"B4 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
- )
- _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
- "B4 CI lower")
- _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
- "B4 CI upper")
+ assert (
+ result.survey_metadata.df_survey == r["df"]
+ ), f"B4 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
+ _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL, "B4 CI lower")
+ _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL, "B4 CI upper")
# B5 (CallawaySantAnna RC-DiD) was removed: R's did::att_gt cannot produce
# golden values for a 2-period repeated cross-section, and the time-scale
@@ -455,8 +466,7 @@ def _load_recs_data(self, recs_results):
if col in df.columns:
df[col] = pd.to_numeric(df[col], errors="coerce")
# Identify replicate weight columns
- rep_cols = sorted([c for c in df.columns if c.startswith("NWEIGHT")
- and c != "NWEIGHT"])
+ rep_cols = sorted([c for c in df.columns if c.startswith("NWEIGHT") and c != "NWEIGHT"])
for col in rep_cols:
df[col] = pd.to_numeric(df[col], errors="coerce")
return df, rep_cols
@@ -465,9 +475,7 @@ def test_c1_simple_regression(self, recs_results):
"""C1: Simple regression TOTALBTU ~ KOWNRENT with JK1 replicate SEs."""
r = recs_results["c1_simple"]
data, rep_cols = self._load_recs_data(recs_results)
- data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(
- drop=True
- )
+ data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(drop=True)
# Build design matrix: intercept + KOWNRENT
from diff_diff.linalg import LinearRegression as LR
@@ -491,18 +499,15 @@ def test_c1_simple_regression(self, recs_results):
_assert_close(py_coef, r["coef"], ATT_RTOL, ATT_ATOL, "C1 coef")
_assert_close(py_se, r["se"], REP_SE_RTOL, REP_SE_ATOL, "C1 SE")
- assert reg.survey_df_ == r["df"], (
- f"C1 df: Python={reg.survey_df_}, R={r['df']}"
- )
+ assert reg.survey_df_ == r["df"], f"C1 df: Python={reg.survey_df_}, R={r['df']}"
# Compute CI from coef, SE, and survey df
from scipy.stats import t as t_dist
+
t_crit = t_dist.ppf(0.975, reg.survey_df_)
py_ci_lower = py_coef - t_crit * py_se
py_ci_upper = py_coef + t_crit * py_se
- _assert_close(py_ci_lower, r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL,
- "C1 CI lower")
- _assert_close(py_ci_upper, r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL,
- "C1 CI upper")
+ _assert_close(py_ci_lower, r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL, "C1 CI lower")
+ _assert_close(py_ci_upper, r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL, "C1 CI upper")
def test_c2_full_regression(self, recs_results):
"""C2: Full regression TOTALBTU ~ KOWNRENT + TYPEHUQ + REGIONC."""
@@ -540,21 +545,16 @@ def test_c2_full_regression(self, recs_results):
py_coef = reg.coefficients_[1]
py_se = np.sqrt(reg.vcov_[1, 1])
- _assert_close(py_coef, r["coef_kownrent"], ATT_RTOL, ATT_ATOL,
- "C2 KOWNRENT coef")
- _assert_close(py_se, r["se_kownrent"], REP_SE_RTOL, REP_SE_ATOL,
- "C2 KOWNRENT SE")
- assert reg.survey_df_ == r["df"], (
- f"C2 df: Python={reg.survey_df_}, R={r['df']}"
- )
+ _assert_close(py_coef, r["coef_kownrent"], ATT_RTOL, ATT_ATOL, "C2 KOWNRENT coef")
+ _assert_close(py_se, r["se_kownrent"], REP_SE_RTOL, REP_SE_ATOL, "C2 KOWNRENT SE")
+ assert reg.survey_df_ == r["df"], f"C2 df: Python={reg.survey_df_}, R={r['df']}"
from scipy.stats import t as t_dist
+
t_crit = t_dist.ppf(0.975, reg.survey_df_)
py_ci_lower = py_coef - t_crit * py_se
py_ci_upper = py_coef + t_crit * py_se
- _assert_close(py_ci_lower, r["ci_lower_kownrent"], REP_CI_RTOL, REP_CI_ATOL,
- "C2 CI lower")
- _assert_close(py_ci_upper, r["ci_upper_kownrent"], REP_CI_RTOL, REP_CI_ATOL,
- "C2 CI upper")
+ _assert_close(py_ci_lower, r["ci_lower_kownrent"], REP_CI_RTOL, REP_CI_ATOL, "C2 CI lower")
+ _assert_close(py_ci_upper, r["ci_upper_kownrent"], REP_CI_RTOL, REP_CI_ATOL, "C2 CI upper")
def test_c3_deff_diagnostics(self, recs_results):
"""C3: DEFF diagnostics from real JK1 replicate SEs.
@@ -565,9 +565,7 @@ def test_c3_deff_diagnostics(self, recs_results):
finite, positive, and > 1 (expected for a complex survey design).
"""
data, rep_cols = self._load_recs_data(recs_results)
- data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(
- drop=True
- )
+ data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(drop=True)
from diff_diff.linalg import LinearRegression as LR
@@ -601,6 +599,6 @@ def test_c3_deff_diagnostics(self, recs_results):
assert all(deff.effective_n > 0), "Effective n should be positive"
# DEFF > 1 expected for a complex survey with unequal weights
- assert deff.deff[1] > 1.0, (
- f"DEFF(KOWNRENT) = {deff.deff[1]:.4f}, expected > 1.0 for survey data"
- )
+ assert (
+ deff.deff[1] > 1.0
+ ), f"DEFF(KOWNRENT) = {deff.deff[1]:.4f}, expected > 1.0 for survey data"
diff --git a/tests/test_t19_marketing_pulse_drift.py b/tests/test_t19_marketing_pulse_drift.py
index 6c31ebd6..335fdbbd 100644
--- a/tests/test_t19_marketing_pulse_drift.py
+++ b/tests/test_t19_marketing_pulse_drift.py
@@ -92,9 +92,7 @@ def event_study_results(panel):
message=r"Assumption 7 .* is violated: leavers present",
category=UserWarning,
)
- model = DCDH(
- twfe_diagnostic=False, placebo=True, n_bootstrap=199, seed=42
- )
+ model = DCDH(twfe_diagnostic=False, placebo=True, n_bootstrap=199, seed=42)
return model.fit(
panel,
outcome="sessions",
@@ -197,9 +195,7 @@ def test_assumption7_warning_fires_as_expected(panel):
with warnings.catch_warnings(record=True) as ws:
warnings.simplefilter("always")
with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
- model = DCDH(
- twfe_diagnostic=False, placebo=True, n_bootstrap=49, seed=42
- )
+ model = DCDH(twfe_diagnostic=False, placebo=True, n_bootstrap=49, seed=42)
model.fit(
panel,
outcome="sessions",
@@ -236,9 +232,7 @@ def test_event_study_warning_policy_matches_notebook(panel):
message=r".*encountered in matmul",
category=RuntimeWarning,
)
- model = DCDH(
- twfe_diagnostic=False, placebo=True, n_bootstrap=199, seed=42
- )
+ model = DCDH(twfe_diagnostic=False, placebo=True, n_bootstrap=199, seed=42)
model.fit(
panel,
outcome="sessions",
@@ -297,8 +291,6 @@ def test_a11_warning_does_not_fire():
treatment="promo_on",
)
a11_warnings = [
- w
- for w in ws
- if w.category is UserWarning and "Assumption 11" in str(w.message)
+ w for w in ws if w.category is UserWarning and "Assumption 11" in str(w.message)
]
assert len(a11_warnings) == 0, [str(w.message)[:80] for w in a11_warnings]
diff --git a/tests/test_t22_had_survey_design_drift.py b/tests/test_t22_had_survey_design_drift.py
index efbe281d..237ea775 100644
--- a/tests/test_t22_had_survey_design_drift.py
+++ b/tests/test_t22_had_survey_design_drift.py
@@ -409,6 +409,7 @@ def test_event_study_plot_uses_stored_pointwise_ci_endpoints():
``nbformat`` is not in the runtime deps (per
``feedback_golden_file_pytest_skip``)."""
from pathlib import Path
+
nbformat = pytest.importorskip("nbformat")
nb_path = (
diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py
index c9bd1c7d..0e2b6654 100644
--- a/tests/test_two_stage.py
+++ b/tests/test_two_stage.py
@@ -556,7 +556,6 @@ def test_lsmr_fallback_never_densifies(self):
TODO row tracked)."""
import unittest.mock
-
data = generate_test_data()
def _no_lstsq(*a, **k):
@@ -1815,7 +1814,7 @@ def test_bootstrap_dropped_coefficient_propagates_nan_inference(self):
fin_h = [h for h, e in es.items() if np.isfinite(e["se"]) and e["se"] > 0]
assert nan_h, (
"a dropped Stage-2 coordinate should yield a NaN-se bootstrap horizon; "
- f"got {{h: e['se'] for h, e in es.items()}}"
+ "got {h: e['se'] for h, e in es.items()}"
)
assert fin_h, "identified bootstrap horizons should keep finite SE"
diff --git a/tests/test_visualization.py b/tests/test_visualization.py
index 1bdcb1e5..a6ef132a 100644
--- a/tests/test_visualization.py
+++ b/tests/test_visualization.py
@@ -162,7 +162,9 @@ def test_plot_from_dataframe_with_cband(self):
)
# Verify _extract_plot_data returns cband overrides
- result = _extract_plot_data(df, periods=None, pre_periods=None, post_periods=None, reference_period=0)
+ result = _extract_plot_data(
+ df, periods=None, pre_periods=None, post_periods=None, reference_period=0
+ )
ci_lo = result[7]
ci_hi = result[8]
assert ci_lo is not None, "ci_lower_override should not be None with cband columns"
@@ -191,7 +193,9 @@ def test_plot_from_dataframe_with_nan_cband(self):
)
# All-NaN cband columns should not produce overrides
- result = _extract_plot_data(df, periods=None, pre_periods=None, post_periods=None, reference_period=0)
+ result = _extract_plot_data(
+ df, periods=None, pre_periods=None, post_periods=None, reference_period=0
+ )
ci_lo = result[7]
ci_hi = result[8]
assert ci_lo is None, "ci_lower_override should be None when cband is all-NaN"
@@ -314,6 +318,7 @@ def test_plot_cs_universal_base_period(self):
"""
pytest.importorskip("matplotlib")
import matplotlib.pyplot as plt
+
from diff_diff import generate_staggered_data
data = generate_staggered_data(n_units=200, n_periods=10, seed=42)
@@ -344,6 +349,7 @@ def test_plot_cs_with_anticipation(self):
"""
pytest.importorskip("matplotlib")
import matplotlib.pyplot as plt
+
from diff_diff import generate_staggered_data
data = generate_staggered_data(n_units=200, n_periods=10, seed=42)
@@ -657,6 +663,7 @@ def test_plot_uses_cband_cis_by_default(self, cs_cband_results):
"""Test that cband CIs are used by default when available."""
pytest.importorskip("matplotlib")
import matplotlib.pyplot as plt
+
from diff_diff.visualization import _extract_plot_data
# Verify plot succeeds
@@ -664,7 +671,7 @@ def test_plot_uses_cband_cis_by_default(self, cs_cband_results):
assert ax is not None
# Verify cband CIs are extracted
- (_, _, _, _, _, _, _, ci_lower_override, ci_upper_override) = _extract_plot_data(
+ _, _, _, _, _, _, _, ci_lower_override, ci_upper_override = _extract_plot_data(
cs_cband_results, None, None, None, None
)
assert ci_lower_override is not None
diff --git a/tests/test_weighted_fw.py b/tests/test_weighted_fw.py
index 9e102f7e..edc7e651 100644
--- a/tests/test_weighted_fw.py
+++ b/tests/test_weighted_fw.py
@@ -51,13 +51,19 @@ class TestSCWeightFWWeighted:
def test_reg_weights_none_matches_unweighted(self, small_panel):
"""reg_weights=None must be bit-identical to the unweighted call."""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
unweighted = _sc_weight_fw(Y, zeta=0.3, max_iter=10000, min_decrease=1e-7)
weighted_none = _sc_weight_fw(
- Y, zeta=0.3, max_iter=10000, min_decrease=1e-7, reg_weights=None,
+ Y,
+ zeta=0.3,
+ max_iter=10000,
+ min_decrease=1e-7,
+ reg_weights=None,
)
np.testing.assert_allclose(weighted_none, unweighted, rtol=1e-14, atol=0)
@@ -70,14 +76,19 @@ def test_uniform_reg_weights_matches_unweighted(self, small_panel):
uses Σ rw·ω² while the unweighted loop uses np.sum(ω²) — different
reduction orders).
"""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
rw_uniform = np.ones(small_panel["n_control"])
unweighted = _sc_weight_fw(Y, zeta=0.3, max_iter=10000, min_decrease=1e-7)
weighted_uniform = _sc_weight_fw(
- Y, zeta=0.3, max_iter=10000, min_decrease=1e-7,
+ Y,
+ zeta=0.3,
+ max_iter=10000,
+ min_decrease=1e-7,
reg_weights=rw_uniform,
)
np.testing.assert_allclose(weighted_uniform, unweighted, rtol=1e-12, atol=1e-13)
@@ -89,44 +100,65 @@ def test_python_rust_parity_under_weighted_reg(self, small_panel):
weighted objective is strictly convex on the simplex so both
backends converge to the same minimizer.
"""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
rw = np.array([1.5, 0.5, 1.0, 2.0, 0.7, 1.3, 0.9, 1.8, 0.6, 1.1, 1.4, 0.8])
rust_w = _sc_weight_fw(
- Y, zeta=0.3, max_iter=20000, min_decrease=1e-9, reg_weights=rw,
+ Y,
+ zeta=0.3,
+ max_iter=20000,
+ min_decrease=1e-9,
+ reg_weights=rw,
)
numpy_w = _sc_weight_fw_numpy(
- Y, zeta=0.3, max_iter=20000, min_decrease=1e-9, reg_weights=rw,
+ Y,
+ zeta=0.3,
+ max_iter=20000,
+ min_decrease=1e-9,
+ reg_weights=rw,
)
np.testing.assert_allclose(numpy_w, rust_w, rtol=1e-9, atol=1e-10)
def test_simplex_invariants_under_arbitrary_rw(self, small_panel):
"""For any positive rw, ω sums to 1 and is non-negative."""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
rw = np.array([0.3, 0.5, 1.0, 2.0, 1.5, 0.8, 1.2, 0.6, 1.7, 0.9, 1.4, 1.1])
omega = _sc_weight_fw(
- Y, zeta=0.4, max_iter=10000, min_decrease=1e-6, reg_weights=rw,
+ Y,
+ zeta=0.4,
+ max_iter=10000,
+ min_decrease=1e-6,
+ reg_weights=rw,
)
assert omega.shape == (small_panel["n_control"],)
- assert np.isclose(omega.sum(), 1.0, atol=1e-6), \
- f"weights must sum to 1, got {omega.sum()}"
+ assert np.isclose(omega.sum(), 1.0, atol=1e-6), f"weights must sum to 1, got {omega.sum()}"
assert np.all(omega >= -1e-9), "weights must be non-negative"
def test_return_convergence_tuple_shape(self, small_panel):
"""return_convergence=True returns (weights, bool) under weighted reg."""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
rw = np.linspace(0.5, 2.0, small_panel["n_control"])
result = _sc_weight_fw(
- Y, zeta=0.3, max_iter=10000, min_decrease=1e-6,
- return_convergence=True, reg_weights=rw,
+ Y,
+ zeta=0.3,
+ max_iter=10000,
+ min_decrease=1e-6,
+ return_convergence=True,
+ reg_weights=rw,
)
assert isinstance(result, tuple)
assert len(result) == 2
@@ -146,15 +178,20 @@ def test_reg_weights_length_mismatch_raises(self, small_panel):
wrong-shape call must raise ``ValueError`` with a message naming
the expected dimension.
"""
- Y = np.column_stack([
- small_panel["Y_pre_control"],
- small_panel["Y_pre_treated_mean"].reshape(-1, 1),
- ])
+ Y = np.column_stack(
+ [
+ small_panel["Y_pre_control"],
+ small_panel["Y_pre_treated_mean"].reshape(-1, 1),
+ ]
+ )
expected_t0 = Y.shape[1] - 1
bad_rw = np.ones(expected_t0 + 3)
with pytest.raises(ValueError, match=r"reg_weights"):
_sc_weight_fw(
- Y, zeta=0.3, max_iter=100, min_decrease=1e-6,
+ Y,
+ zeta=0.3,
+ max_iter=100,
+ min_decrease=1e-6,
reg_weights=bad_rw,
)
@@ -180,11 +217,18 @@ def test_uniform_rw_matches_unweighted_helper_within_tolerance(self, small_panel
Y_pre_c = small_panel["Y_pre_control"]
Y_pre_t = small_panel["Y_pre_treated_mean"]
unweighted = compute_sdid_unit_weights(
- Y_pre_c, Y_pre_t, zeta_omega=0.3, min_decrease=1e-6,
+ Y_pre_c,
+ Y_pre_t,
+ zeta_omega=0.3,
+ min_decrease=1e-6,
)
rw_uniform = np.ones(small_panel["n_control"])
weighted = compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, rw_uniform, zeta_omega=0.3, min_decrease=1e-6,
+ Y_pre_c,
+ Y_pre_t,
+ rw_uniform,
+ zeta_omega=0.3,
+ min_decrease=1e-6,
)
np.testing.assert_allclose(weighted, unweighted, rtol=1e-6, atol=1e-6)
@@ -194,7 +238,11 @@ def test_simplex_invariants_under_arbitrary_rw(self, small_panel):
rng = np.random.default_rng(7)
rw = rng.uniform(0.3, 2.5, size=small_panel["n_control"])
omega = compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, rw, zeta_omega=0.3, min_decrease=1e-6,
+ Y_pre_c,
+ Y_pre_t,
+ rw,
+ zeta_omega=0.3,
+ min_decrease=1e-6,
)
assert omega.shape == (small_panel["n_control"],)
assert np.isclose(omega.sum(), 1.0, atol=1e-6)
@@ -206,7 +254,10 @@ def test_rw_shape_mismatch_raises(self, small_panel):
wrong_rw = np.ones(small_panel["n_control"] + 1)
with pytest.raises(ValueError, match="rw_control shape"):
compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, wrong_rw, zeta_omega=0.3,
+ Y_pre_c,
+ Y_pre_t,
+ wrong_rw,
+ zeta_omega=0.3,
)
def test_return_convergence_propagates_AND_of_passes(self, small_panel):
@@ -215,8 +266,13 @@ def test_return_convergence_propagates_AND_of_passes(self, small_panel):
rw = np.linspace(0.5, 2.0, small_panel["n_control"])
# Tight tolerance + few iterations to force non-convergence
omega, converged = compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, rw, zeta_omega=0.3,
- min_decrease=1e-15, max_iter_pre_sparsify=5, max_iter=5,
+ Y_pre_c,
+ Y_pre_t,
+ rw,
+ zeta_omega=0.3,
+ min_decrease=1e-15,
+ max_iter_pre_sparsify=5,
+ max_iter=5,
return_convergence=True,
)
assert isinstance(converged, bool)
@@ -236,11 +292,18 @@ def test_uniform_rw_matches_unweighted_helper_within_tolerance(self, small_panel
Y_pre_c = small_panel["Y_pre_control"]
Y_post_c = small_panel["Y_post_control"]
unweighted = compute_time_weights(
- Y_pre_c, Y_post_c, zeta_lambda=0.05, min_decrease=1e-6,
+ Y_pre_c,
+ Y_post_c,
+ zeta_lambda=0.05,
+ min_decrease=1e-6,
)
rw_uniform = np.ones(small_panel["n_control"])
weighted = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw_uniform, zeta_lambda=0.05, min_decrease=1e-6,
+ Y_pre_c,
+ Y_post_c,
+ rw_uniform,
+ zeta_lambda=0.05,
+ min_decrease=1e-6,
)
np.testing.assert_allclose(weighted, unweighted, rtol=1e-6, atol=1e-6)
@@ -250,7 +313,11 @@ def test_simplex_invariants_under_arbitrary_rw(self, small_panel):
rng = np.random.default_rng(11)
rw = rng.uniform(0.3, 2.5, size=small_panel["n_control"])
lam = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw, zeta_lambda=0.05, min_decrease=1e-6,
+ Y_pre_c,
+ Y_post_c,
+ rw,
+ zeta_lambda=0.05,
+ min_decrease=1e-6,
)
assert lam.shape == (small_panel["n_pre"],)
assert np.isclose(lam.sum(), 1.0, atol=1e-6)
@@ -262,7 +329,10 @@ def test_rw_shape_mismatch_raises(self, small_panel):
wrong_rw = np.ones(small_panel["n_control"] + 1)
with pytest.raises(ValueError, match="rw_control shape"):
compute_time_weights_survey(
- Y_pre_c, Y_post_c, wrong_rw, zeta_lambda=0.05,
+ Y_pre_c,
+ Y_post_c,
+ wrong_rw,
+ zeta_lambda=0.05,
)
def test_non_uniform_rw_beats_unweighted_centering_variant(self, small_panel):
@@ -287,7 +357,9 @@ def test_non_uniform_rw_beats_unweighted_centering_variant(self, small_panel):
# Correct path: what compute_time_weights_survey actually does.
lam_correct = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw,
+ Y_pre_c,
+ Y_post_c,
+ rw,
zeta_lambda=0.05,
min_decrease=1e-8,
max_iter=10000,
@@ -300,15 +372,23 @@ def test_non_uniform_rw_beats_unweighted_centering_variant(self, small_panel):
sqrt_rw = np.sqrt(np.maximum(rw, 0.0))
Y_weighted_unweighted_center = Y_time_raw * sqrt_rw[:, None]
lam_buggy = _sc_weight_fw(
- Y_weighted_unweighted_center, zeta=0.05, intercept=True,
- min_decrease=1e-8, max_iter=10000,
+ Y_weighted_unweighted_center,
+ zeta=0.05,
+ intercept=True,
+ min_decrease=1e-8,
+ max_iter=10000,
)
# Sparsify + refit second pass to match the two-pass shape.
from diff_diff.utils import _sparsify
+
lam_buggy = _sparsify(lam_buggy)
lam_buggy = _sc_weight_fw(
- Y_weighted_unweighted_center, zeta=0.05, intercept=True,
- init_weights=lam_buggy, min_decrease=1e-8, max_iter=10000,
+ Y_weighted_unweighted_center,
+ zeta=0.05,
+ intercept=True,
+ init_weights=lam_buggy,
+ min_decrease=1e-8,
+ max_iter=10000,
)
# Compute the canonical (weighted-centered) objective on both.
@@ -319,7 +399,7 @@ def test_non_uniform_rw_beats_unweighted_centering_variant(self, small_panel):
def weighted_ssr(lam_val: np.ndarray) -> float:
resid = A_wc @ lam_val - b_wc
- return float(np.sum(rw * resid ** 2))
+ return float(np.sum(rw * resid**2))
ssr_correct = weighted_ssr(lam_correct)
ssr_buggy = weighted_ssr(lam_buggy)
@@ -339,7 +419,11 @@ def test_zero_rw_subset_handled(self, small_panel):
rw = np.ones(small_panel["n_control"])
rw[:3] = 0.0 # zero out first 3 controls (e.g., undrawn PSU)
lam = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw, zeta_lambda=0.05, min_decrease=1e-6,
+ Y_pre_c,
+ Y_post_c,
+ rw,
+ zeta_lambda=0.05,
+ min_decrease=1e-6,
)
assert lam.shape == (small_panel["n_pre"],)
assert np.isclose(lam.sum(), 1.0, atol=1e-6)
@@ -369,15 +453,25 @@ def test_unit_survey_python_rust_parity(self, small_panel, monkeypatch):
# Rust path (default)
rust_omega = compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, rw, zeta_omega=0.3,
- max_iter_pre_sparsify=200, max_iter=20000, min_decrease=1e-9,
+ Y_pre_c,
+ Y_pre_t,
+ rw,
+ zeta_omega=0.3,
+ max_iter_pre_sparsify=200,
+ max_iter=20000,
+ min_decrease=1e-9,
)
# Force pure-Python path
monkeypatch.setattr(dd_utils, "HAS_RUST_BACKEND", False)
py_omega = compute_sdid_unit_weights_survey(
- Y_pre_c, Y_pre_t, rw, zeta_omega=0.3,
- max_iter_pre_sparsify=200, max_iter=20000, min_decrease=1e-9,
+ Y_pre_c,
+ Y_pre_t,
+ rw,
+ zeta_omega=0.3,
+ max_iter_pre_sparsify=200,
+ max_iter=20000,
+ min_decrease=1e-9,
)
# Sparsify path is identical, weighted FW is strictly convex; the
@@ -394,14 +488,24 @@ def test_time_survey_python_rust_parity(self, small_panel, monkeypatch):
rw = rng.uniform(0.5, 2.0, size=small_panel["n_control"])
rust_lam = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw, zeta_lambda=0.05,
- max_iter_pre_sparsify=200, max_iter=20000, min_decrease=1e-9,
+ Y_pre_c,
+ Y_post_c,
+ rw,
+ zeta_lambda=0.05,
+ max_iter_pre_sparsify=200,
+ max_iter=20000,
+ min_decrease=1e-9,
)
monkeypatch.setattr(dd_utils, "HAS_RUST_BACKEND", False)
py_lam = compute_time_weights_survey(
- Y_pre_c, Y_post_c, rw, zeta_lambda=0.05,
- max_iter_pre_sparsify=200, max_iter=20000, min_decrease=1e-9,
+ Y_pre_c,
+ Y_post_c,
+ rw,
+ zeta_lambda=0.05,
+ max_iter_pre_sparsify=200,
+ max_iter=20000,
+ min_decrease=1e-9,
)
np.testing.assert_allclose(py_lam, rust_lam, rtol=1e-7, atol=1e-7)