From 9855741a71e89f42ede16967e8794afb99049f60 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 00:52:17 -0400 Subject: [PATCH 1/2] refactor(utils): shared FE-dummy design build across DiD/MPD/TWFE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drop-first pd.get_dummies design construction existed as three inline copies — the DifferenceInDifferences and MultiPeriodDiD fixed_effects= loops and the TwoWayFixedEffects HC2/HC2-BM full-dummy path — whose FE naming / dtype / column-order conventions could drift independently (the drift risk the TODO row flagged). All three now delegate to one diff_diff.utils.build_fe_dummy_blocks helper whose names match fe_dummy_names (the reserved-name collision guard) by construction. Outputs bit-identical (A/B vs pristine origin/main on DiD with non-default-order Categorical FE + covariates, TWFE hc2 + hc2_bm, MultiPeriodDiD multi-FE incl. the fe==time skip); the DiD/MPD paths also drop the per-column np.column_stack accumulation for one block stack. Contract tests lock names==fe_dummy_names (plain, Categorical non-default order, numeric) and values==get_dummies. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 11 +++++++++++ diff_diff/estimators.py | 26 ++++++++++++------------- diff_diff/twfe.py | 19 ++++++++----------- diff_diff/utils.py | 42 +++++++++++++++++++++++++++++++++++++++++ tests/test_utils.py | 37 ++++++++++++++++++++++++++++++++++++ 5 files changed, 110 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a5a82b4..4407ce87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -406,6 +406,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (excluded from the remap). A per-replicate full-dummy HC2 implementation (the TODO row) was investigated and rejected as a costly no-op: it cannot change the replicate variance. Tests lock warn+bit-identity-to-hc1 on all three estimators. +- **Shared FE-dummy design build (`build_fe_dummy_blocks`).** The drop-first + `pd.get_dummies` design construction existed as three inline copies — the + `DifferenceInDifferences` and `MultiPeriodDiD` `fixed_effects=` loops and the + `TwoWayFixedEffects` HC2/HC2-BM full-dummy path — whose FE naming / dtype / + column-order conventions could drift independently (the drift risk flagged in TODO). + All three now delegate to one `diff_diff.utils.build_fe_dummy_blocks` helper whose + names match `fe_dummy_names` (the reserved-name collision guard) by construction. + Outputs are bit-identical (A/B against the previous implementation on DiD with a + non-default-order Categorical FE + covariates, TWFE hc2 + hc2_bm, and MultiPeriodDiD + multi-FE including the `fe == time` skip); the MPD/DiD paths also drop the + per-column `np.column_stack` accumulation (O(k²) copies) for one block stack. - **Per-cell solver fast paths for covariate fits (CallawaySantAnna and every estimator routing through the shared solvers).** Two pure-Python changes in `diff_diff/linalg.py`: (1) `solve_logit`'s IRLS inner step — previously a full diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 0d07a6a4..9b976ea2 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -29,6 +29,7 @@ from diff_diff.results import DiDResults, MultiPeriodDiDResults, PeriodEffect from diff_diff.utils import ( WildBootstrapResults, + build_fe_dummy_blocks, demean_by_groups, fe_dummy_names, pre_demean_norms, @@ -529,13 +530,12 @@ def fit( # Add fixed effects as dummy variables if fixed_effects: - for fe in fixed_effects: - # Create dummies, drop first category to avoid multicollinearity - # Use working_data to be consistent with absorbed FE if both are used - dummies = pd.get_dummies(working_data[fe], prefix=fe, drop_first=True) - for col in dummies.columns: - X = np.column_stack([X, dummies[col].values.astype(float)]) - var_names.append(col) + # Shared drop-first dummy build (names match fe_dummy_names, the + # reserved-name guard above). Use working_data to be consistent + # with absorbed FE if both are used. + _fe_blocks, _fe_names = build_fe_dummy_blocks(working_data, list(fixed_effects)) + X = np.column_stack([X] + _fe_blocks) + var_names.extend(_fe_names) # Reject any duplicate in the FINAL term list (e.g. a fixed-effect dummy # colliding with a structural term) BEFORE the regression — so the fit is @@ -1813,13 +1813,11 @@ def fit( # type: ignore[override] # collapsing the dict and breaking the coefficients-vs-vcov # alignment that downstream consumers rely on). Skip those FEs. if fixed_effects: - for fe in fixed_effects: - if fe == time: - continue - dummies = pd.get_dummies(working_data[fe], prefix=fe, drop_first=True) - for col in dummies.columns: - X = np.column_stack([X, dummies[col].values.astype(float)]) - var_names.append(col) + _mp_fes = [fe for fe in fixed_effects if fe != time] + if _mp_fes: + _fe_blocks, _fe_names = build_fe_dummy_blocks(working_data, _mp_fes) + X = np.column_stack([X] + _fe_blocks) + var_names.extend(_fe_names) # Reject any duplicate in the FINAL term list (e.g. a fixed-effect dummy # colliding with a structural period_{p} key) BEFORE the regression — so diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index fd6b25f2..b98747a2 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -15,6 +15,7 @@ from diff_diff.linalg import LinearRegression from diff_diff.results import DiDResults from diff_diff.utils import ( + build_fe_dummy_blocks, fe_dummy_names, pre_demean_norms, snap_absorbed_regressors, @@ -348,14 +349,13 @@ def fit( # type: ignore[override] ) y = data[outcome].values.astype(np.float64) cov_arrs = [data[c].values.astype(np.float64) for c in (covariates or [])] - unit_dummies_df = pd.get_dummies(data[unit], prefix=f"_fe_{unit}", drop_first=True) - time_dummies_df = pd.get_dummies(data[time], prefix=f"_fe_{time}", drop_first=True) - unit_dummies = unit_dummies_df.values.astype(np.float64) - time_dummies = time_dummies_df.values.astype(np.float64) + # Shared drop-first dummy build (single implementation with the + # DiD/MPD fixed_effects= paths; names match fe_dummy_names). + _fe_blocks, _fe_dummy_names = build_fe_dummy_blocks( + data, [unit, time], prefixes=[f"_fe_{unit}", f"_fe_{time}"] + ) X = np.column_stack( - [np.ones(len(data)), data["_treatment_post"].values] - + cov_arrs - + [unit_dummies, time_dummies] + [np.ones(len(data)), data["_treatment_post"].values] + cov_arrs + _fe_blocks ) # FEs are now in X explicitly; solve_ols's n - k accounting # already subtracts them, so the extra unit + time DOF @@ -367,10 +367,7 @@ def fit( # type: ignore[override] # (matching the MPD invariant # ``len(result.coefficients) == result.vcov.shape[0]``). _twfe_var_names: Optional[List[str]] = ( - ["const", "ATT"] - + list(covariates or []) - + list(unit_dummies_df.columns) - + list(time_dummies_df.columns) + ["const", "ATT"] + list(covariates or []) + _fe_dummy_names ) # Backstop: reject any duplicate in the FINAL term list (e.g. a # unit/time dummy colliding with a structural term or another dummy) diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 4e96e012..04e708aa 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -217,6 +217,48 @@ def fe_dummy_names(col: pd.Series, prefix: str) -> List[str]: return [f"{prefix}_{c}" for c in cats[1:]] +def build_fe_dummy_blocks( + data: pd.DataFrame, + fe_cols: List[str], + prefixes: Optional[List[str]] = None, +) -> Tuple[List[np.ndarray], List[str]]: + """Materialize drop-first fixed-effect dummy blocks and their column names. + + Single shared implementation of the ``pd.get_dummies(col, prefix=..., + drop_first=True)`` design-build used by ``DifferenceInDifferences`` / + ``MultiPeriodDiD`` (``fixed_effects=``) and the ``TwoWayFixedEffects`` + HC2/HC2-BM full-dummy path — previously three inline copies whose FE + naming / dtype / column-order conventions could drift independently. + Names match :func:`fe_dummy_names` (the reserved-name collision guard) + exactly; values are the dense ``float64`` dummy matrix per FE, in + ``get_dummies`` column order. + + Parameters + ---------- + data : pandas.DataFrame + Frame holding the FE columns. + fe_cols : list of str + Fixed-effect column names, in design order. + prefixes : list of str, optional + Dummy-name prefix per FE column (defaults to the column name itself; + TWFE passes ``_fe_{col}`` to keep its internal-name convention). + + Returns + ------- + blocks : list of ndarray + One dense ``(n, G_j - 1)`` float64 dummy block per FE column. + names : list of str + The kept dummy column names across all FEs, in block order. + """ + blocks: List[np.ndarray] = [] + names: List[str] = [] + for fe, prefix in zip(fe_cols, prefixes or fe_cols): + dummies = pd.get_dummies(data[fe], prefix=prefix, drop_first=True) + blocks.append(dummies.values.astype(np.float64)) + names.extend(dummies.columns) + return blocks, names + + def warn_if_not_converged( converged: bool, method_name: str, diff --git a/tests/test_utils.py b/tests/test_utils.py index 6ed57cf9..20da23ed 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2079,3 +2079,40 @@ def test_warns_on_nonconvergence_with_label(self): tol=1e-15, method_name="my solver label", ) + + +class TestBuildFeDummyBlocks: + """Shared FE-dummy design build (DiD/MPD fixed_effects= + TWFE full-dummy + path): names must match fe_dummy_names (the reserved-name collision + guard) and values must match pd.get_dummies exactly.""" + + def test_names_match_fe_dummy_names_contract(self): + from diff_diff.utils import build_fe_dummy_blocks, fe_dummy_names + + df = pd.DataFrame( + { + "plain": ["b", "a", "c", "a"], + "cat": pd.Categorical( + ["x", "z", "y", "z"], categories=["z", "y", "x"] + ), # non-default order + "num": [3, 1, 2, 1], + } + ) + blocks, names = build_fe_dummy_blocks(df, ["plain", "cat", "num"]) + expected = ( + fe_dummy_names(df["plain"], "plain") + + fe_dummy_names(df["cat"], "cat") + + fe_dummy_names(df["num"], "num") + ) + assert names == expected + assert sum(b.shape[1] for b in blocks) == len(names) + + def test_values_match_get_dummies(self): + from diff_diff.utils import build_fe_dummy_blocks + + df = pd.DataFrame({"g": ["b", "a", "c", "a", "b"]}) + blocks, names = build_fe_dummy_blocks(df, ["g"], prefixes=["_fe_g"]) + ref = pd.get_dummies(df["g"], prefix="_fe_g", drop_first=True) + np.testing.assert_array_equal(blocks[0], ref.values.astype(np.float64)) + assert names == list(ref.columns) + assert blocks[0].dtype == np.float64 From c92ba750d8726566f57b36214b61117ff01b9581 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 06:25:33 -0400 Subject: [PATCH 2/2] fix(utils): length-validate prefixes in build_fe_dummy_blocks (review P2) zip(fe_cols, prefixes) would silently skip trailing FE columns if a future caller passed a shorter non-empty prefixes list; now raises ValueError with a mismatched-length test. Co-Authored-By: Claude Fable 5 --- diff_diff/utils.py | 6 ++++++ tests/test_utils.py | 9 +++++++++ 2 files changed, 15 insertions(+) diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 04e708aa..97c7be3d 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -250,6 +250,12 @@ def build_fe_dummy_blocks( names : list of str The kept dummy column names across all FEs, in block order. """ + if prefixes is not None and len(prefixes) != len(fe_cols): + raise ValueError( + f"build_fe_dummy_blocks: prefixes length {len(prefixes)} does not " + f"match fe_cols length {len(fe_cols)}; zip would silently skip " + "trailing FE columns." + ) blocks: List[np.ndarray] = [] names: List[str] = [] for fe, prefix in zip(fe_cols, prefixes or fe_cols): diff --git a/tests/test_utils.py b/tests/test_utils.py index 20da23ed..f2ac8e9f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2116,3 +2116,12 @@ def test_values_match_get_dummies(self): np.testing.assert_array_equal(blocks[0], ref.values.astype(np.float64)) assert names == list(ref.columns) assert blocks[0].dtype == np.float64 + + def test_mismatched_prefixes_length_raises(self): + """Review P2: a shorter non-empty prefixes list must raise, not + silently zip-skip trailing FE columns.""" + from diff_diff.utils import build_fe_dummy_blocks + + df = pd.DataFrame({"a": ["x", "y"], "b": ["u", "v"]}) + with pytest.raises(ValueError, match="prefixes length 1 does not match"): + build_fe_dummy_blocks(df, ["a", "b"], prefixes=["_fe_a"])