"""

Difference-in-Differences estimators with sklearn-like API.

This module contains the core DiD estimators:

- DifferenceInDifferences: Basic 2x2 DiD estimator

- MultiPeriodDiD: Event-study style DiD with period-specific treatment effects

Additional estimators are in separate modules:

- TwoWayFixedEffects: See diff_diff.twfe

- SyntheticDiD: See diff_diff.synthetic_did

- SyntheticControl: See diff_diff.synthetic_control

For backward compatibility, all estimators are re-exported from this module.

"""

import warnings

from typing import Any, Dict, List, Optional, Tuple

import numpy as np

import pandas as pd

from diff_diff._base import BaseEstimator

from diff_diff._deprecation import (

NOT_SUPPLIED,

resolve_renamed_kwarg,

warn_deprecated_kwarg,

)

from diff_diff.linalg import (

LinearRegression,

_absorbed_fe_vcov_scale,

_expand_vcov_with_nan,

compute_r_squared,

solve_ols,

)

from diff_diff.results import DiDResults, MultiPeriodDiDResults, PeriodEffect

from diff_diff.utils import (

WildBootstrapResults,

absorbed_fe_cr1_k_increment,

absorbed_fe_rank,

build_fe_dummy_blocks,

cluster_nested_fe_dims,

demean_by_groups,

fe_dummy_names,

pre_demean_norms,

safe_inference,

snap_absorbed_regressors,

validate_binary,

validate_covariate_names,

validate_design_term_names,

validate_df_convention,

wild_bootstrap_se,

)

class DifferenceInDifferences(BaseEstimator):

"""

Difference-in-Differences estimator with sklearn-like interface.

Estimates the Average Treatment effect on the Treated (ATT) using

the canonical 2x2 DiD design or panel data with two-way fixed effects.

Parameters

----------

formula : str, optional

R-style formula for the model (e.g., "outcome ~ treated * post").

If provided, overrides column name parameters.

robust : bool, optional

DEPRECATED legacy alias for ``vcov_type`` (row M-045; warns with

``FutureWarning``, removed in 4.0 - use ``vcov_type=``).

``robust=True`` maps to ``vcov_type="hc1"``; ``robust=False`` maps

to ``vcov_type="classical"``. Explicit ``vcov_type`` overrides

``robust`` unless the pair is contradictory (e.g.

``robust=False, vcov_type="hc2"`` raises).

cluster : str, optional

Column name for cluster-robust standard errors. Combined with

``vcov_type``: with ``"hc1"`` dispatches to CR1 (Liang-Zeger); with

``"hc2_bm"`` dispatches to CR2 Bell-McCaffrey (Pustejovsky-Tipton 2018

symmetric-sqrt + Satterthwaite DOF).

vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, optional

Variance-covariance family. Defaults to the ``robust`` alias.

- ``"classical"``: non-robust OLS SEs, ``sigma_hat^2 * (X'X)^{-1}``.

- ``"hc1"``: heteroskedasticity-robust HC1 with ``n/(n-k)`` adjustment

(library default). With ``cluster=``, uses CR1 (Liang-Zeger).

- ``"hc2"``: leverage-corrected meat (one-way only). Errors with

``cluster=``; use ``"hc2_bm"`` for clustered Bell-McCaffrey.

- ``"hc2_bm"``: one-way HC2 + Imbens-Kolesar (2016) Satterthwaite DOF;

with ``cluster=``, Pustejovsky-Tipton (2018) CR2 cluster-robust.

``MultiPeriodDiD(cluster=..., vcov_type="hc2_bm")`` is supported and

uses a cluster-aware Bell-McCaffrey contrast DOF for the

post-period-average ATT (see ``_compute_cr2_bm_contrast_dof`` in

``linalg.py`` and the REGISTRY.md note). Weighted CR2-BM

(``survey_design=`` paths) is a separate gate.

- ``"conley"``: Conley 1999 spatial-HAC sandwich. Pass

``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=<float>``,

and ``conley_lag_cutoff=<int>`` on the constructor; pass

``unit=<col>`` as a fit-time kwarg to :meth:`fit` (NOT on

``__init__``; unused unless Conley is set; not part of

``get_params()`` / ``set_params()``). The block-decomposed panel

sandwich (matches R ``conleyreg`` with ``lag_cutoff > 0``) sums

within-period spatial pairs plus within-unit Bartlett serial

pairs (lag=0 excluded). Explicit ``cluster=<col>`` enables the

combined spatial + cluster product kernel; ``survey_design=``

and ``inference='wild_bootstrap'`` both raise

``NotImplementedError``.

alpha : float, default=0.05

Significance level for confidence intervals.

inference : str, default="analytical"

Inference method: "analytical" for standard asymptotic inference,

or "wild_bootstrap" for wild cluster bootstrap (recommended when

number of clusters is small, <50).

n_bootstrap : int, default=999

Number of bootstrap replications when inference="wild_bootstrap".

bootstrap_weights : str, default="rademacher"

Type of bootstrap weights: "rademacher" (standard), "webb"

(recommended for <10 clusters), or "mammen" (skewness correction).

p_val_type : str, default="two-tailed"

Shape of the wild cluster bootstrap test (mirrors

``fwildclusterboot::boottest``): "two-tailed" (test on ``|t|``,

two-tailed inverted CI — which may be asymmetric) or "equal-tailed"

(each tail at ``alpha/2``, equal-tailed CI). Only used when

``inference="wild_bootstrap"``.

seed : int, optional

Random seed for reproducibility when using bootstrap inference.

If None (default), results will vary between runs.

rank_deficient_action : str, default "warn"

Action when design matrix is rank-deficient (linearly dependent columns):

- "warn": Issue warning and drop linearly dependent columns (default)

- "error": Raise ValueError

- "silent": Drop columns silently without warning

conley_coords, conley_cutoff_km, conley_metric, conley_kernel, conley_lag_cutoff

Conley (1999) spatial-HAC variance configuration. Pass

``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=<float>``,

and ``conley_lag_cutoff=<int>`` on the constructor; the ``unit``

identifier is passed as a fit-time arg to ``fit(...)`` (NOT on

``__init__``) — it is unused unless ``vcov_type="conley"`` and is

therefore not part of ``get_params()`` / ``set_params()`` (which

return constructor-arg dicts). The block-decomposed panel sandwich

(matching R ``conleyreg`` with ``lag_cutoff > 0``) sums within-period

spatial pairs plus within-unit Bartlett serial pairs (lag=0 excluded

to avoid double-counting). Explicit ``cluster=<col>`` + Conley

enables the combined spatial + cluster product kernel; the cluster

must be constant within each unit across periods (validator-enforced).

DiD has no auto-cluster, so cluster is fully opt-in on the Conley

path — absent ``cluster=``, pure Conley spatial HAC applies.

``survey_design=`` + Conley and ``inference='wild_bootstrap'`` +

Conley both raise ``NotImplementedError``.

df_convention : {"residual", "cluster", "normal"}, default "residual"

Degrees-of-freedom convention for analytical t-statistics, p-values,

and CIs. ``"residual"`` (default) uses the fitted residual df

(``n − K_full``); ``"cluster"`` uses the Stata/fixest cluster df

``G − 1`` on clustered fits — it has no effect on unclustered fits

or on ``vcov_type="conley"`` (the combined Conley+cluster product

kernel has no documented ``G − 1`` df reference and keeps the

residual df); ``"normal"`` deliberately uses normal-theory z

inference at the fallback level on every fit, clustered or not.

Applies only at the fallback level of the df resolution under every

value: survey df and per-coefficient Bell-McCaffrey DOF

(``vcov_type="hc2_bm"``) are more refined small-sample corrections

and always take precedence. Point estimates, SEs, and t-statistics

are unaffected — only the reference distribution changes. The

default flips to ``"cluster"`` at v4 (see the REGISTRY clustered-CR1

inference-df deviation note).

Attributes

----------

results_ : DiDResults

Estimation results after calling fit().

is_fitted_ : bool

Whether the model has been fitted.

Examples

--------

Basic usage with a DataFrame:

>>> import pandas as pd

>>> from diff_diff import DifferenceInDifferences

>>>

>>> # Create sample data

>>> 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]

... })

>>>

>>> # Fit the model

>>> did = DifferenceInDifferences()

>>> results = did.fit(data, outcome='outcome', treatment='treated', post='post')

>>>

>>> # View results

>>> print(results.att) # ATT estimate

>>> results.print_summary() # Full summary table

Using formula interface:

>>> did = DifferenceInDifferences()

>>> results = did.fit(data, formula='outcome ~ treated * post')

Notes

-----

The ATT is computed using the standard DiD formula:

ATT = (E[Y|D=1,T=1] - E[Y|D=1,T=0]) - (E[Y|D=0,T=1] - E[Y|D=0,T=0])

Or equivalently via OLS regression:

Y = α + β₁*D + β₂*T + β₃*(D×T) + ε

Where β₃ is the ATT.

"""

def __init__(

self,

robust: Optional[bool] = None,

cluster: Optional[str] = None,

vcov_type: Optional[str] = None,

alpha: float = 0.05,

inference: str = "analytical",

n_bootstrap: int = 999,

bootstrap_weights: str = "rademacher",

p_val_type: str = "two-tailed",

seed: Optional[int] = None,

rank_deficient_action: str = "warn",

conley_coords: Optional[Tuple[str, str]] = None,

conley_cutoff_km: Optional[float] = None,

conley_metric: str = "haversine",

conley_kernel: str = "bartlett",

conley_lag_cutoff: Optional[int] = None,

df_convention: str = "residual",

):

# Resolve vcov_type from the legacy `robust` alias via the shared

# helper so __init__ and set_params use identical validation logic.

from diff_diff.linalg import resolve_vcov_type

validate_df_convention(df_convention)

# `robust` is deprecated (rows M-045..M-047; removed in 4.0). None is

# the not-supplied sentinel: default constructions and get_params

# round-trips stay silent, only an explicit robust= warns. The raw

# arg lives at `_robust_arg` (what get_params returns); the PUBLIC

# `self.robust` keeps the RESOLVED legacy bool so pre-3.9 attribute

# readers keep seeing True/False until the 4.0 removal.

if robust is not None:

warn_deprecated_kwarg(type(self).__name__, "robust", "use vcov_type= instead")

self._robust_arg = robust

self.robust = robust if robust is not None else True

self.cluster = cluster

self.vcov_type = resolve_vcov_type(robust, vcov_type)

# Preserve the raw constructor arg (possibly None) alongside the

# resolved `vcov_type`. `get_params()` returns the raw arg so

# sklearn clones preserve the implicit-vs-explicit distinction

# (and therefore the backward-compat remap). Set only in __init__

# and updated in ``set_params`` so the flag transitions match the

# user-visible parameter state.

self._vcov_type_arg = vcov_type

self._vcov_type_explicit = vcov_type is not None

self.alpha = alpha

self.inference = inference

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

# Test shape for wild cluster bootstrap (mirrors fwildclusterboot's

# p_val_type): "two-tailed" (default) or "equal-tailed".

self.p_val_type = p_val_type

self.seed = seed

self.rank_deficient_action = rank_deficient_action

# Conley spatial-HAC parameters; column names (NOT array values) for

# the coords. Validation happens at fit() when `data` is in scope.

self.conley_coords = conley_coords

self.conley_cutoff_km = conley_cutoff_km

self.conley_metric = conley_metric

self.conley_kernel = conley_kernel

# Phase 2 panel block-decomposed kwarg. The conley_time + conley_unit

# arrays are auto-derived from data[time].values + data[unit].values

# at fit-time (panel estimators already take time/unit as column names).

self.conley_lag_cutoff = conley_lag_cutoff

# Inference df convention for clustered analytical fits: "residual"

# (default; t/p/CI at the fitted residual df) or "cluster" (the

# Stata/fixest G-1 convention). Survey df and per-coefficient

# Bell-McCaffrey DOF always take precedence over either. The default

# flips to "cluster" at v4 (REGISTRY clustered-CR1 inference-df

# deviation note).

self.df_convention = df_convention

self.is_fitted_ = False

self.results_ = None

self._coefficients = None

self._vcov = None

self._bootstrap_results = None # Store WildBootstrapResults if used

def fit(

self,

data: pd.DataFrame,

outcome: Optional[str] = None,

treatment: Optional[str] = None,

post: Any = NOT_SUPPLIED,

formula: Optional[str] = None,

covariates: Optional[List[str]] = None,

fixed_effects: Optional[List[str]] = None,

absorb: Optional[List[str]] = None,

survey_design=None,

unit: Optional[str] = None,

time: Any = NOT_SUPPLIED,

) -> DiDResults:

"""

Fit the Difference-in-Differences model.

Parameters

----------

data : pd.DataFrame

DataFrame containing the outcome, treatment, and time variables.

outcome : str

Name of the outcome variable column.

treatment : str

Name of the treatment group indicator column (0/1).

post : str

Name of the post-treatment period indicator column (0/1).

formula : str, optional

R-style formula (e.g., "outcome ~ treated * post").

If provided, overrides outcome, treatment, and post parameters.

covariates : list, optional

List of covariate column names to include as linear controls.

Names must not collide with reserved structural terms (``const``,

the treatment/time column names, the ``{treatment}:{time}``

interaction, fixed-effect dummy names, or internal working columns)

and must be unique; a collision or duplicate raises ``ValueError``

(it would otherwise silently overwrite a structural coefficient).

fixed_effects : list, optional

List of categorical column names to include as fixed effects.

Creates dummy variables for each category (drops first level).

Use for low-dimensional fixed effects (e.g., industry, region).

absorb : list, optional

List of categorical column names for high-dimensional fixed effects.

Uses within-transformation (demeaning) instead of dummy variables.

More efficient for large numbers of categories (e.g., firm, individual).

survey_design : SurveyDesign, optional

Survey design specification for design-based inference. When provided,

uses Taylor Series Linearization for variance estimation and

applies sampling weights to the regression.

unit : str, optional

Name of the unit identifier column. Required ONLY when

``vcov_type="conley"`` — the panel block-decomposed Conley

sandwich (matching R ``conleyreg`` with ``lag_cutoff > 0``)

needs the unit identifier to compute the per-unit serial sum.

Mirrors :meth:`MultiPeriodDiD.fit(unit=...)` and

:meth:`TwoWayFixedEffects.fit(unit=...)`. Fit-time only — NOT

a constructor kwarg, so it is not part of ``get_params()`` /

``set_params()`` (which return constructor-arg dicts).

Ignored when ``vcov_type`` is not ``"conley"``.

Returns

-------

DiDResults

Object containing estimation results.

Raises

------

ValueError

If required parameters are missing or data validation fails, or if

a covariate name collides with a reserved structural term name or

duplicates another covariate.

Examples

--------

Using fixed effects (dummy variables):

>>> did.fit(data, outcome='sales', treatment='treated', post='post',

... fixed_effects=['state', 'industry'])

Using absorbed fixed effects (within-transformation):

>>> did.fit(data, outcome='sales', treatment='treated', post='post',

... absorb=['firm_id'])

The keyword-only ``time`` parameter is a deprecated alias for

``post`` (row M-030); it warns with ``FutureWarning`` and will be

removed in 4.0.

"""

post = resolve_renamed_kwarg(

f"{type(self).__name__}.fit", "time", time, "post", post, default=None

)

# Body-local name; the public parameter is post (M-030).

time = post

# Parse formula if provided

if formula is not None:

outcome, treatment, time, covariates = self._parse_formula(formula, data)

elif outcome is None or treatment is None or time is None:

raise ValueError(

"Must provide either 'formula' or all of 'outcome', 'treatment', and 'post'"

)

# Validate inputs

self._validate_data(data, outcome, treatment, time, covariates)

# Validate binary variables BEFORE any transformations

validate_binary(data[treatment].values, "treatment")

validate_binary(data[time].values, "time")

# Validate fixed effects and absorb columns

if fixed_effects:

for fe in fixed_effects:

if fe not in data.columns:

raise ValueError(f"Fixed effect column '{fe}' not found in data")

if absorb:

for ab in absorb:

if ab not in data.columns:

raise ValueError(f"Absorb column '{ab}' not found in data")

# Resolve survey design if provided

from diff_diff.survey import _resolve_effective_cluster, _resolve_survey_for_fit

resolved_survey, survey_weights, survey_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, data, self.inference)

)

_uses_replicate = resolved_survey is not None and resolved_survey.uses_replicate_variance

if _uses_replicate and self.inference == "wild_bootstrap":

raise ValueError(

"Cannot use inference='wild_bootstrap' with replicate-weight "

"survey designs. Replicate weights provide their own variance "

"estimation."

)

_replicate_vcov_remap = _uses_replicate and self._warn_replicate_vcov_ignored()

# Handle absorbed fixed effects (within-transformation)

working_data = data.copy()

absorbed_vars = []

n_absorbed_effects = 0

# Save raw treatment counts before absorb demeaning

n_treated_raw = int(np.sum(data[treatment].values.astype(float)))

n_control_raw = len(data) - n_treated_raw

# Reject the `absorb + fixed_effects` mutual-exclusion combination

# BEFORE any auto-route. R4 review caught a contract-drift where the

# auto-route silently merged the two arguments on the HC2/HC2-BM

# path — the public API has always treated this combination as

# invalid (different FE-handling paths; mixing them violates the

# FWL theorem on the demeaned half), so keep the explicit rejection

# in front of the auto-route to preserve user-facing behavior.

if absorb and fixed_effects:

raise ValueError(

"Cannot use both absorb and fixed_effects. "

"The absorb within-transformation does not residualize "

"fixed_effects dummies, violating the FWL theorem. "

"Use absorb alone (for high-dimensional FE) "

"or fixed_effects alone (for low-dimensional FE)."

)

# Auto-route absorb → fixed_effects when vcov_type needs the FULL FE

# hat matrix. HC2 leverage and CR2 Bell-McCaffrey DOF both depend on

# the full-design hat; FWL preserves coefficients and residuals but

# not the hat matrix, so the demeaned design's leverage is wrong for

# these vcov families. Building the full-dummy design and routing

# through the existing fixed_effects= branch produces the algebraically

# correct vcov. Empirically matches `lm() + sandwich::vcovHC` and

# `lm() + clubSandwich::vcovCR` (singleton-cluster trick for one-way

# HC2-BM; PT2018 §3.3 unweighted CR2 algebra) at ~1e-14.

# Conley vcov is unaffected: the absorb+Conley path (Wave A) computes

# the panel sandwich on demeaned scores, which is FWL-correct because

# Conley's meat uses only residuals (no leverage term).

# HC1/CR1 paths remain on the demeaned design (no leverage term).

# Note: the user-facing `result.coefficients` under this auto-route

# will include the FE-dummy entries (matching the fixed_effects= path),

# not the slope-only view that a plain `absorb=` returns.

#

# Placement: this auto-route runs BEFORE the legacy multi-absorb +

# survey-weights guard because that guard's rationale ("single-pass

# demeaning is not the correct weighted FWL projection for N > 1

# dimensions") doesn't apply when we're about to swap absorb for

# fixed_effects: the fixed_effects= path builds the full-dummy design

# and solves WLS directly, with no within-transform step. R2 review

# surfaced the scope mismatch (REGISTRY/CHANGELOG said "SUPPORTED" but

# the survey guard fired first on weighted multi-absorb fits).

# Route on the EFFECTIVE vcov family: under a replicate design the

# remap to hc1 must also disable this full-dummy swap, or an

# explicit hc2 request would still change the result surface

# (full-dummy coefficients vs absorbed reduced fit) despite the

# "has no effect" warning.

if absorb and not _replicate_vcov_remap and self.vcov_type in ("hc2", "hc2_bm"):

fixed_effects = list(fixed_effects or []) + list(absorb)

absorb = None

absorbed_vars = []

n_absorbed_effects = 0

# Weighted multiple absorbed FE is supported: the absorb path below uses

# iterative alternating projections (demean_by_groups), the exact weighted

# FWL projection for N > 1 dimensions on both balanced and unbalanced panels.

# Validate vcov_type="conley" wire-up. DiD.fit() accepts `unit`

# as a fit-time arg (NOT on __init__) because cluster/unit

# semantics on DiD are opt-in rather than auto-derived (unlike

# MultiPeriodDiD / TwoWayFixedEffects which have a unit declaration

# at fit-time anyway). The panel block-decomposed Conley sandwich

# (matching R conleyreg with lag_cutoff > 0) needs unit/time/coords

# to assemble the within-period spatial and within-unit serial

# sums; we mirror MultiPeriodDiD's reject pattern for missing args

# and the survey/wild-bootstrap incompatibilities.

if self.vcov_type == "conley":

# Shared front-door validation across DiD / MPD / TWFE entry

# points (Wave A holistic fix: replaces the inline drift that

# accumulated across CI R1/R2/R6 — same-class validation gaps

# mirrored across estimator surfaces).

from diff_diff.conley import _validate_conley_estimator_inputs

_validate_conley_estimator_inputs(

estimator_name="DifferenceInDifferences",

data=data,

unit=unit,

conley_coords=self.conley_coords,

conley_cutoff_km=self.conley_cutoff_km,

conley_lag_cutoff=self.conley_lag_cutoff,

survey_design=survey_design,

inference=self.inference,

cluster=self.cluster,

)

if absorb:

# FWL theorem: demean ALL regressors alongside outcome.

# Regressors collinear with absorbed FE (e.g., treatment after

# absorbing unit FE) will zero out and be handled by rank-deficiency.

working_data["_treat_time"] = working_data[treatment].values.astype(

float

) * working_data[time].values.astype(float)

vars_to_demean = [outcome, treatment, time, "_treat_time"] + (covariates or [])

_absorb_regressors = vars_to_demean[1:] # everything except outcome

_pre_norms = pre_demean_norms(working_data, _absorb_regressors, weights=survey_weights)

# Absorbed df MUST be measured before the in-place demean below

# overwrites the group columns with demeaned floats. Equals

# demean_by_groups' historical `sum_d (n_d - 1)` on a connected panel;

# smaller when the incidence graph splits (disconnected/hierarchical).

_absorbed_df = absorbed_fe_rank(

working_data,

list(absorb),

has_intercept_col=True,

weights=survey_weights,

)

# Stash the raw FE columns: the clustered-CR1 K_reference

# increment needs them AFTER the effective cluster resolves,

# but the in-place demean below overwrites them with floats.

_fe_cols_raw = working_data[list(absorb)].copy()

# Method of alternating projections: for N > 1 absorbed dimensions a

# single sequential sweep is only exact on balanced (orthogonal-FE)

# panels; demean_by_groups iterates to the exact (W)LS-FWL residual.

working_data, _ = demean_by_groups( # count superseded by absorbed_fe_rank above

working_data,

vars_to_demean,

list(absorb),

inplace=True,

weights=survey_weights,

)

# FE-spanned regressors demean to numerical junk, not exact zero;

# snap them so rank handling drops them deterministically instead

# of the junk direction perturbing the identified coefficients.

snap_absorbed_regressors(

working_data,

_absorb_regressors,

_pre_norms,

absorbed_desc=f"absorb={list(absorb)}",

group_vars=list(absorb),

rank_deficient_action=self.rank_deficient_action,

display_names={"_treat_time": f"{treatment}:{time}"},

weights=survey_weights,

)

n_absorbed_effects += _absorbed_df

absorbed_vars = list(absorb)

# Extract variables (may be demeaned if absorb was used)

y = working_data[outcome].values.astype(float)

d = working_data[treatment].values.astype(float)

t = working_data[time].values.astype(float)

# Create interaction term

if absorb:

dt = working_data["_treat_time"].values.astype(float)

else:

dt = d * t

# Reject covariate names that collide with reserved structural terms.

# Covariate names are appended verbatim to var_names below and zipped

# into coef_dict, so a covariate named like a structural term would

# silently overwrite that coefficient (dict last-write-wins). The

# reserved set covers the intercept, treatment/time indicators, the

# interaction, the internal _treat_time working column, and any

# fixed-effect dummy names (derived via fe_dummy_names WITHOUT

# materializing the dummy matrix; names match the get_dummies build

# below exactly). validate_design_term_names re-checks the FINAL list.

_reserved = {"const", treatment, time, f"{treatment}:{time}", "_treat_time"}

if fixed_effects:

for fe in fixed_effects:

_reserved.update(fe_dummy_names(working_data[fe], fe))

validate_covariate_names(covariates, _reserved, estimator="DifferenceInDifferences")

# Build design matrix

X = np.column_stack([np.ones(len(y)), d, t, dt])

var_names = ["const", treatment, time, f"{treatment}:{time}"]

# Add covariates if provided

if covariates:

for cov in covariates:

X = np.column_stack([X, working_data[cov].values.astype(float)])

var_names.append(cov)

# Add fixed effects as dummy variables

if fixed_effects:

# 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

# not wasted and no misleading multicollinearity warning is emitted ahead

# of the intended ValueError.

validate_design_term_names(var_names, estimator="DifferenceInDifferences")

# Extract ATT index (coefficient on interaction term)

att_idx = 3 # Index of interaction term

att_var_name = f"{treatment}:{time}"

assert var_names[att_idx] == att_var_name, (

f"ATT index mismatch: expected '{att_var_name}' at index {att_idx}, "

f"but found '{var_names[att_idx]}'"

)

# Always use LinearRegression for initial fit (unified code path)

# For wild bootstrap, we don't need cluster SEs from the initial fit

cluster_ids = data[self.cluster].values if self.cluster is not None else None

# When survey PSU is present, it overrides cluster for variance estimation

effective_cluster_ids = _resolve_effective_cluster(

resolved_survey, cluster_ids, self.cluster

)

# Inject cluster as effective PSU for survey variance estimation

if resolved_survey is not None and effective_cluster_ids is not None:

from diff_diff.survey import _inject_cluster_as_psu, compute_survey_metadata

resolved_survey = _inject_cluster_as_psu(resolved_survey, effective_cluster_ids)

if resolved_survey.psu is not None and survey_metadata is not None:

raw_w = (

data[survey_design.weights].values.astype(np.float64)

if survey_design.weights

else np.ones(len(data), dtype=np.float64)

)

survey_metadata = compute_survey_metadata(resolved_survey, raw_w)

# When absorb + replicate: pass survey_design=None to prevent

# LinearRegression from computing replicate vcov on already-demeaned

# data (demeaning depends on weights, so replicate refits must re-demean).

_lr_survey = resolved_survey

if _uses_replicate and absorbed_vars:

_lr_survey = None

# Remap implicit "classical" + cluster to CR1 for legacy-alias

# backward compatibility (see `_resolve_effective_vcov_type`).

_fit_vcov_type = (

"hc1"

if _replicate_vcov_remap

else self._resolve_effective_vcov_type(effective_cluster_ids)

)

# Build Conley coord/time/unit arrays when applicable. CRITICAL:

# read from the ORIGINAL `data` frame, NOT `working_data` — `absorb`

# demeans `time` (and any column listed in `absorb`) in working_data,

# so reading `working_data[time]` would silently partition the

# within-period spatial sandwich on residualized floats instead of

# the true pre/post periods (Codex Wave A R1 P0). Coords are likewise

# read from raw `data` for symmetry with TwoWayFixedEffects

# (`twfe.py::TwoWayFixedEffects.fit`) which has the same FWL-

# composability contract: the meat is computed on demeaned scores

# but the kernel grid uses the original space (coords) and time/unit

# indexing. `_compute_conley_vcov` normalizes time labels to dense

# codes 0..T-1 internally, so non-numeric `time` labels (datetime64,

# pd.Period, strings) still work on the MultiPeriodDiD path; DiD's

# binary `time` column is integer 0/1 by convention and is unaffected

# by the normalization.

if _fit_vcov_type == "conley":

# Validated by the conley front-door (_validate_conley_estimator_inputs).

assert self.conley_coords is not None

_conley_coords_arr: Optional[np.ndarray] = np.column_stack(

[

data[self.conley_coords[0]].values.astype(np.float64),

data[self.conley_coords[1]].values.astype(np.float64),

]

)

_conley_time_arr: Optional[np.ndarray] = np.asarray(data[time].values)

_conley_unit_arr: Optional[np.ndarray] = data[unit].values

else:

_conley_coords_arr = None

_conley_time_arr = None

_conley_unit_arr = None

# Clustered-CR1 K_reference adjustment (variance-conventions.md D2/D1):

# absorbed FE not nested in the cluster ADD their conditional rank;

# cluster-nested explicit FE dummies SUBTRACT theirs. Computed only

# for the effective-hc1 clustered analytical lane — under

# wild_bootstrap the analytical fit is deliberately unclustered

# (adjustment travels through the WCB wiring instead), and under a

# survey design the survey variance replaces the CR1 sandwich

# wholesale (moot by design).

_cr1_k_adj = 0

if (

_fit_vcov_type == "hc1"

and self.inference != "wild_bootstrap"

and effective_cluster_ids is not None

and resolved_survey is None

):

if absorbed_vars:

_cr1_k_adj = absorbed_fe_cr1_k_increment(

_fe_cols_raw,

list(absorb),

effective_cluster_ids,

has_intercept_col=True,

weights=survey_weights,

)

elif fixed_effects:

_nested_fe = cluster_nested_fe_dims(

working_data,

list(fixed_effects),

effective_cluster_ids,

weights=survey_weights,

)

if _nested_fe:

_cr1_k_adj = -absorbed_fe_rank(

working_data,

_nested_fe,

has_intercept_col=True,

weights=survey_weights,

)

# Don't forward `robust=self.robust` when the vcov_type has been

# remapped; `robust=False + vcov_type="hc1"` would otherwise trip

# the conflict check inside `LinearRegression.__init__`. The

# remapped vcov_type is the single source of truth for this call.

reg = LinearRegression(

include_intercept=False, # Intercept already in X

cluster_ids=effective_cluster_ids if self.inference != "wild_bootstrap" else None,

alpha=self.alpha,

rank_deficient_action=self.rank_deficient_action,

weights=survey_weights,

weight_type=survey_weight_type,

survey_design=_lr_survey,

vcov_type=_fit_vcov_type,

conley_coords=_conley_coords_arr,

conley_cutoff_km=self.conley_cutoff_km,

conley_metric=self.conley_metric,

conley_kernel=self.conley_kernel,

conley_time=_conley_time_arr,

conley_unit=_conley_unit_arr,

conley_lag_cutoff=self.conley_lag_cutoff,

df_convention=self.df_convention,

).fit(X, y, df_adjustment=n_absorbed_effects, cluster_k_adjustment=_cr1_k_adj)

coefficients = reg.coefficients_

residuals = reg.residuals_

fitted = reg.fitted_values_

assert coefficients is not None

att = coefficients[att_idx]

# Get inference - replicate absorb override, bootstrap, or analytical

if _uses_replicate and absorbed_vars:

# Estimator-level replicate variance: re-demean + re-solve per replicate

from diff_diff.survey import compute_replicate_refit_variance

from diff_diff.utils import safe_inference

_absorb_list = list(absorbed_vars) # capture for closure

# Handle rank-deficient nuisance: refit only identified columns

_id_mask = ~np.isnan(coefficients)

_id_cols = np.where(_id_mask)[0]

_att_idx_reduced = int(np.searchsorted(_id_cols, att_idx))

def _refit_did_absorb(w_r):

nz = w_r > 0

wd = data[nz].copy()

w_nz = w_r[nz]

wd["_treat_time"] = wd[treatment].values.astype(float) * wd[time].values.astype(

float

)

vars_dm = [outcome, treatment, time, "_treat_time"] + (covariates or [])

_rep_norms = pre_demean_norms(wd, vars_dm[1:], weights=w_nz)

wd, _ = demean_by_groups(wd, vars_dm, _absorb_list, inplace=True, weights=w_nz)

# A regressor can become FE-spanned WITHIN a replicate's

# effective sample (half-sample zeroing): snap silently so the

# replicate solve drops it (NaN replicate -> invalid) instead

# of consuming a junk direction.

snap_absorbed_regressors(

wd,

vars_dm[1:],

_rep_norms,

absorbed_desc=f"absorb={_absorb_list}",

group_vars=_absorb_list,

rank_deficient_action="silent",

weights=w_nz,

)

y_r = wd[outcome].values.astype(float)

d_r = wd[treatment].values.astype(float)

t_r = wd[time].values.astype(float)

dt_r = wd["_treat_time"].values.astype(float)

X_r = np.column_stack([np.ones(len(y_r)), d_r, t_r, dt_r])

if covariates:

for cov in covariates:

X_r = np.column_stack([X_r, wd[cov].values.astype(float)])

coef_r, _, _ = solve_ols(

X_r[:, _id_cols],

y_r,

weights=w_nz,

weight_type=survey_weight_type,

rank_deficient_action="silent",

return_vcov=False,

)

return coef_r

vcov_reduced, _n_valid_rep = compute_replicate_refit_variance(

_refit_did_absorb, coefficients[_id_mask], resolved_survey

)

vcov = _expand_vcov_with_nan(vcov_reduced, len(coefficients), _id_cols)

se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0)))

_df_rep = (

survey_metadata.df_survey

if survey_metadata and survey_metadata.df_survey

else 0 # rank-deficient replicate → NaN inference

)

# Replicate-refit path is only reached with a resolved design.

assert resolved_survey is not None

if _n_valid_rep < resolved_survey.n_replicates:

_df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0

if survey_metadata is not None:

survey_metadata.df_survey = _df_rep if _df_rep > 0 else None

t_stat, p_value, conf_int = safe_inference(att, se, alpha=self.alpha, df=_df_rep)

_inference_df_used = float(_df_rep) if _df_rep is not None and _df_rep > 0 else None

elif self.inference == "wild_bootstrap" and self.cluster is not None:

# Override with wild cluster bootstrap inference (bootstrap

# test-inversion based; no reference t-distribution, so no

# effective inference df).

_inference_df_used = None

# K_reference adjustment for the bootstrap's own CR1 factors,

# computed against the RAW cluster ids the bootstrap partitions

# on (NOT effective_cluster_ids — the analytical-lane gate above

# deliberately passed 0 under wild_bootstrap).

_wcb_k_adj = 0

if absorbed_vars:

_wcb_k_adj = absorbed_fe_cr1_k_increment(

_fe_cols_raw,

list(absorb),

cluster_ids,

has_intercept_col=True,

weights=survey_weights,

)

elif fixed_effects:

_nested_wcb = cluster_nested_fe_dims(

working_data,

list(fixed_effects),

cluster_ids,

weights=survey_weights,

)

if _nested_wcb:

_wcb_k_adj = -absorbed_fe_rank(

working_data,

_nested_wcb,

has_intercept_col=True,

weights=survey_weights,

)

se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference(

X, y, residuals, cluster_ids, att_idx, cluster_k_adjustment=_wcb_k_adj

)

else:

# Use analytical inference from LinearRegression

# (handles replicate vcov for no-absorb path automatically)

vcov = reg.vcov_

inference = reg.get_inference(att_idx)

se = inference.se

t_stat = inference.t_stat

p_value = inference.p_value

conf_int = inference.conf_int

_inference_df_used = (

float(inference.df) if inference.df is not None and inference.df > 0 else None

)

r_squared = compute_r_squared(y, residuals)

# Count observations (use raw counts to avoid demeaned values from absorb)

n_treated = n_treated_raw

n_control = n_control_raw

# Create coefficient dictionary

coef_dict = {name: coef for name, coef in zip(var_names, coefficients)}

# Determine inference method and bootstrap info

inference_method = "analytical"

n_bootstrap_used = None

n_clusters_used = None

p_val_type_used = None

if self._bootstrap_results is not None:

inference_method = "wild_bootstrap"

n_bootstrap_used = self._bootstrap_results.n_bootstrap

n_clusters_used = self._bootstrap_results.n_clusters

p_val_type_used = self._bootstrap_results.p_val_type

# Store results

self.results_ = DiDResults(

att=att,

se=se,

t_stat=t_stat,

p_value=p_value,

conf_int=conf_int,

n_obs=len(y),

n_treated=n_treated,

n_control=n_control,

alpha=self.alpha,

coefficients=coef_dict,

vcov=vcov,

residuals=residuals,

fitted_values=fitted,

r_squared=r_squared,

inference_method=inference_method,

n_bootstrap=n_bootstrap_used,

n_clusters=n_clusters_used,

p_val_type=p_val_type_used,

survey_metadata=survey_metadata,

# Report the family that actually produced the SE, which may be

# the remapped "hc1" (CR1) under the legacy alias path, not the

# stored `self.vcov_type`.

vcov_type=_fit_vcov_type,

cluster_name=self.cluster,

conley_lag_cutoff=(self.conley_lag_cutoff if _fit_vcov_type == "conley" else None),

df_convention=self.df_convention,

inference_df=_inference_df_used,

)

self._coefficients = coefficients

self._vcov = vcov

self.is_fitted_ = True

return self.results_

def _fit_ols(

self, X: np.ndarray, y: np.ndarray

) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:

"""

Fit OLS regression.

This method is kept for backwards compatibility. Internally uses the

unified solve_ols from diff_diff.linalg for optimized computation.

Parameters

----------

X : np.ndarray

Design matrix.

y : np.ndarray

Outcome vector.

Returns

-------

tuple

(coefficients, residuals, fitted_values, r_squared)

"""

# Use unified OLS backend

coefficients, residuals, fitted, _ = solve_ols(X, y, return_fitted=True, return_vcov=False)

r_squared = compute_r_squared(y, residuals)

return coefficients, residuals, fitted, r_squared

def _run_wild_bootstrap_inference(

self,

X: np.ndarray,

y: np.ndarray,

residuals: np.ndarray,

cluster_ids: np.ndarray,

coefficient_index: int,

cluster_k_adjustment: int = 0,

) -> Tuple[float, float, Tuple[float, float], float, np.ndarray, WildBootstrapResults]:

"""

Run wild cluster bootstrap inference.

Parameters

----------

X : np.ndarray

Design matrix.

y : np.ndarray

Outcome vector.

residuals : np.ndarray

OLS residuals.

cluster_ids : np.ndarray

Cluster identifiers for each observation.

coefficient_index : int

Index of the coefficient to compute inference for.

cluster_k_adjustment : int, default 0

Signed K_reference adjustment for the bootstrap's CR1 factors

(nestedness computed by the caller against THESE raw cluster

ids). Applied to both the analytical SE inside

``wild_bootstrap_se`` and the stored-vcov recompute below, so

``se == sqrt(vcov[j, j])`` stays exact.

Returns

-------

tuple

(se, p_value, conf_int, t_stat, vcov, bootstrap_results)

"""

bootstrap_results = wild_bootstrap_se(