refactor(utils): shared FE-dummy design build across DiD/MPD/TWFE

CHANGELOG.md

@@ -354,6 +354,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

TODO.md

@@ -30,7 +30,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 |-------|----------|--------|--------|----------|
 | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low |
 | `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; (b) discrete-treatment saturated regression (`treatment_type="discrete"`) — **DONE**; (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0` (`control_group="lowest_dose"`) — **DONE** (discrete + continuous mass-point, single-cohort; estimand `ATT(d)−ATT(d_L)`; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low |
-| TWFE's HC2/HC2-BM inline full-dummy build (`twfe.py:280-315`) duplicates the dummy-construction logic in `DifferenceInDifferences(fixed_effects=...)` (`estimators.py:478-486`). Extract a shared helper, or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch (with TWFE-specific cluster-default threading), to reduce drift risk on FE naming / survey behavior / result-surface conventions. Substantive refactor — touches both estimators. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit` | follow-up | Heavy | Low |
 | `HonestDiD` Δ^SD optimal-FLCI center parity — **LANDED (SE-audit B2b)**: replaced the flat Nelder-Mead affine-estimator optimizer with a faithful port of R `HonestDiD::findOptimalFLCI`'s nested convex program (inner min-worst-case-bias at fixed estimator SD `h` via SLSQP QCQP; outer grid-zoom over `h`). Matches R's center + half-length + optimalVec to ~1e-3 (median ~1e-5) across a stress grid; the prior ~9% intermediate-M center drift is removed (widths always matched, coverage unaffected). Analytical folded-normal cv is more accurate than R's MC `.qfoldednormal`. Golden `honest_flci_golden.json` + `TestHonestFLCIParityR`. | `honest_did.py::_flci_solve` | SE-audit | Done | Low |
 
 ### Performance

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_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_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,

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