"""

Efficient Difference-in-Differences estimator.

Implements the ATT estimator from Chen, Sant'Anna & Xie (2025).

Without covariates, achieves the semiparametric efficiency bound via

closed-form within-group covariances. With covariates, uses a doubly

robust path with sieve outcome regressions, sieve propensity ratios, and

kernel-smoothed conditional Omega*(X) (see class docstring for details).

Under PT-All the model is overidentified and EDiD exploits this for

tighter inference; under PT-Post it reduces to the standard

single-baseline estimator (Callaway-Sant'Anna).

The variance machinery is purely influence-function-based: per-unit EIF

values aggregate via ``sqrt(mean(EIF**2)/n)`` (unclustered, HC1-style),

Liang-Zeger CR1 on cluster-aggregated EIF (under ``cluster=``), or

Taylor Series Linearization on the combined IF (under ``survey_design=``).

Because the per-unit EIF aggregation has no equivalent single design

matrix, analytical-sandwich families ``{classical, hc2, hc2_bm}`` cannot

be defined and the ``vcov_type`` input contract is permanently narrow to

``{"hc1"}`` — see ``docs/methodology/REGISTRY.md`` "IF-based variance

estimators vs analytical-sandwich estimators" for the structural rationale.

"""

import warnings

from typing import TYPE_CHECKING, 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

from diff_diff.aggregation import AggregationKit

from diff_diff.efficient_did_aggregation import (

_cluster_aggregate,

_compute_se_from_eif,

_EfficientAggregationMixin,

)

from diff_diff.efficient_did_bootstrap import (

EDiDBootstrapResults,

EfficientDiDBootstrapMixin,

)

if TYPE_CHECKING:

from diff_diff.survey import ResolvedSurveyDesign

from diff_diff.efficient_did_covariates import (

OMEGA_RIDGE_DEFAULT,

_silverman_bandwidth,

compute_conditional_cells_tiled,

compute_eif_cov,

compute_generated_outcomes_cov,

compute_omega_star_conditional,

compute_per_unit_weights,

estimate_inverse_propensity_sieve,

estimate_outcome_regression,

estimate_propensity_ratio_sieve,

)

from diff_diff.efficient_did_results import EfficientDiDResults, HausmanPretestResult

from diff_diff.efficient_did_weights import (

_omega_star_nocov_gram,

compute_efficient_weights,

compute_eif_nocov,

compute_generated_outcomes_nocov,

compute_omega_star_nocov,

enumerate_valid_triples,

)

from diff_diff.utils import safe_inference

# Re-export for convenience

__all__ = ["EfficientDiD", "EfficientDiDResults", "EDiDBootstrapResults"]

def _validate_and_build_cluster_mapping(

df: pd.DataFrame,

unit: str,

cluster: str,

all_units: list,

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

"""Validate cluster column and build unit-to-cluster-index mapping.

Checks: column exists, no NaN, per-unit constancy, >= 2 clusters.

Returns (cluster_indices, n_clusters).

"""

if cluster not in df.columns:

raise ValueError(f"Cluster column '{cluster}' not found in data.")

if df[cluster].isna().any():

raise ValueError(f"Cluster column '{cluster}' contains missing values.")

cluster_by_unit = df.groupby(unit)[cluster]

if (cluster_by_unit.nunique() > 1).any():

raise ValueError(

f"Cluster column '{cluster}' varies within unit. "

"Cluster assignment must be constant per unit."

)

cluster_col = cluster_by_unit.first().reindex(all_units).values

unique_clusters = np.unique(cluster_col)

n_clusters = len(unique_clusters)

if n_clusters < 2:

raise ValueError(f"Need at least 2 clusters for cluster-robust SEs, got {n_clusters}.")

cluster_to_idx = {c: i for i, c in enumerate(unique_clusters)}

indices = np.array([cluster_to_idx[c] for c in cluster_col])

return indices, n_clusters

def _build_edid_aggregation_kit(

eif_by_gt: Dict[Tuple[Any, Any], np.ndarray],

group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]],

treatment_groups: List[Any],

time_periods: List[Any],

pt_assumption: str,

n_units_total: int,

n_units: int,

cohort_fractions: Dict[float, float],

unit_cohorts: np.ndarray,

cluster_indices: Optional[np.ndarray],

n_clusters: Optional[int],

unit_level_weights: Optional[np.ndarray],

resolved_survey_unit: Optional["ResolvedSurveyDesign"],

df_survey: Optional[float],

alpha: float,

anticipation: int,

) -> Optional[AggregationKit]:

"""Bundle the retained EIF payload + bookkeeping for post-fit aggregate().

Reference bundling only — zero array copies, zero extra computation, so

plain-fit timing is unchanged; the dominant memory payload is the

per-(g,t) EIF dict, O(n_units x n_gt) (full enumeration in the

``docs/methodology/REGISTRY.md`` EfficientDiD M-023 Note).

``df_survey`` MUST be the post-overall snapshot captured in ``fit()``

immediately after the overall inference and before the ES/group gates —

NOT a read of ``estimator._survey_df`` at kit-build time, which on a

deprecated ``fit(aggregate="group"/"all")`` replicate fit is the

post-group value (``_compute_survey_eif_se`` can set it to ``None`` on

a degenerate replicate design). The snapshot is the exact value every

fit-time aggregation seeded from, so post-fit recompute replays it.

"""

if not eif_by_gt:

# Unreachable after fit()'s empty-effects raise; kept for the CS

# guard shape (a kit with nothing to re-aggregate is not attached).

return None

return AggregationKit(

bookkeeping={

# PRIVATE SNAPSHOTS of the aggregation inputs (CI review P0):

# aggregate() must never read the MUTABLE public result fields -

# a user edit of results.group_time_effects/groups would

# otherwise mix altered point estimates with the retained

# fit-time EIF variance, yielding plausible-but-invalid

# inference. Row values are scalars/tuples, so per-row dict

# copies suffice. alpha/anticipation are already kit fields.

"group_time_effects": {gt: dict(row) for gt, row in group_time_effects.items()},

"treatment_groups": list(treatment_groups),

"time_periods": list(time_periods),

"pt_assumption": pt_assumption,

"n_units_total": n_units_total,

"n_units": n_units,

"cohort_fractions": cohort_fractions,

"unit_cohorts": unit_cohorts,

"cluster_indices": cluster_indices,

"n_clusters": n_clusters,

"unit_level_weights": unit_level_weights,

"resolved_survey_unit": resolved_survey_unit,

"df_survey": df_survey,

},

influence=eif_by_gt,

alpha=alpha,

anticipation=anticipation,

cband=False,

bootstrap=None,

)

def _hausman_quadratic_form(

delta: np.ndarray,

cov_post: np.ndarray,

cov_all: np.ndarray,

) -> Tuple[float, int, float, int, bool]:

"""Hausman statistic from the event-study delta and the two ES covariances.

Implements the Theorem A.1 test statistic of Chen, Sant'Anna & Xie (2025,

arXiv:2506.17729v1). The variance-difference matrix is

V = aCov(ES_post) - aCov(ES_all) = cov_post - cov_all

(restricted minus efficient, PSD under H0 because the efficient estimator has

the smaller variance), and the statistic is ``H = delta' V^+ delta`` with

``delta = ES_post - ES_all``. ``V`` is inverted by Moore-Penrose pseudoinverse

and the number of strictly positive eigenvalues is used as the chi-square

degrees of freedom -- a finite-sample safeguard for a non-PSD ``V`` that equals

``|E|`` (the number of post-treatment horizons) when ``V`` is well-conditioned.

Parameters

----------

delta : ndarray, shape (|E|,)

Event-study difference ``ES_post - ES_all`` (restricted minus efficient).

cov_post, cov_all : ndarray, shape (|E|, |E|)

Estimator-scale covariances of the restricted (PT-Post) and efficient

(PT-All) event-study vectors.

Returns

-------

H : float

The Hausman statistic (``max(delta' V^+ delta, 0)``); NaN if ``V`` is

non-finite or has no positive eigenvalues.

effective_rank : int

Number of positive eigenvalues of ``V`` (the chi-square degrees of freedom).

p_value : float

Upper-tail ``chi2(effective_rank)`` p-value; NaN when ``H`` is NaN.

n_negative : int

Number of substantially negative eigenvalues of ``V`` (efficiency-reversal

diagnostic).

finite_ok : bool

False when ``V`` contains non-finite entries.

"""

from scipy.stats import chi2

V = cov_post - cov_all

if not np.all(np.isfinite(V)):

return np.nan, 0, np.nan, 0, False

eigvals = np.linalg.eigvalsh(V)

max_eigval = float(np.max(np.abs(eigvals))) if len(eigvals) > 0 else 0.0

tol = max(1e-10 * max_eigval, 1e-15)

n_negative = int(np.sum(eigvals < -tol))

effective_rank = int(np.sum(eigvals > tol))

if effective_rank == 0:

return np.nan, 0, np.nan, n_negative, True

V_pinv = np.linalg.pinv(V, rcond=tol / max_eigval if max_eigval > 0 else 1e-10)

H = max(float(delta @ V_pinv @ delta), 0.0)

p_value = float(chi2.sf(H, df=effective_rank))

return H, effective_rank, p_value, n_negative, True

class EfficientDiD(EfficientDiDBootstrapMixin, _EfficientAggregationMixin, BaseEstimator):

"""Efficient DiD estimator (Chen, Sant'Anna & Xie 2025).

Without covariates, achieves the semiparametric efficiency bound for

ATT(g,t) using a closed-form estimator based on within-group sample

means and covariances.

With covariates, uses a doubly robust path: sieve-based propensity

score ratios (Eq 4.1-4.2), sieve outcome regressions (polynomial

basis, AIC/BIC order selection), sieve-estimated inverse propensities

(algorithm step 4), and kernel-smoothed conditional Omega*(X) with

per-unit efficient weights (Eq 3.12). The DR property ensures

consistency if either the outcome regression or the sieve propensity

ratio is correctly specified; because all nuisances are sieves /

kernel smoothers (the paper's flexible-nuisance specification), the

covariate path attains the semiparametric efficiency bound under the

paper's regularity conditions (see REGISTRY.md).

Parameters

----------

pt_assumption : str, default ``"all"``

Parallel trends variant: ``"all"`` (overidentified, uses all

pre-treatment periods and comparison groups) or ``"post"``

(just-identified, single baseline, equivalent to CS).

alpha : float, default 0.05

Significance level.

cluster : str or None

Column name for cluster-robust SEs. When set, analytical SEs

use the Liang-Zeger clustered sandwich estimator on EIF values.

With ``n_bootstrap > 0``, bootstrap weights are generated at the

cluster level (all units in a cluster share the same weight).

vcov_type : str, default ``"hc1"``

Variance-estimator family. Permanently narrow to ``{"hc1"}`` per

the Chen-Sant'Anna-Xie (2025) IF-based variance — analytical-sandwich

families ``{classical, hc2, hc2_bm}`` and ``conley`` are rejected

at ``__init__`` / ``set_params``. See REGISTRY.md for the

methodology rationale (no single design matrix on which hat-matrix

leverage or Bell-McCaffrey Satterthwaite DOF can be defined).

Use ``cluster=<col>`` for Liang-Zeger CR1 on cluster-aggregated EIF;

use ``survey_design=`` for Taylor Series Linearization on the

combined IF.

control_group : str, default ``"never_treated"``

Which units serve as the comparison group:

``"never_treated"`` requires a never-treated cohort (raises if

none exist); ``"last_cohort"`` reclassifies the latest treatment

cohort as pseudo-never-treated and drops periods at

``t >= last_g - anticipation`` so the pseudo-control's

pre-treatment window excludes anticipation-contaminated periods.

Distinct from CallawaySantAnna's ``"not_yet_treated"`` — see

REGISTRY.md for details.

n_bootstrap : int, default 0

Number of multiplier bootstrap iterations (0 = analytical only).

bootstrap_weights : str, default ``"rademacher"``

Bootstrap weight distribution.

seed : int or None

Random seed for reproducibility.

anticipation : int, default 0

Number of anticipation periods (shifts the effective treatment

boundary forward by this amount). When combined with

``control_group="last_cohort"``, also trims the pseudo-control

period set at ``t >= last_g - anticipation`` (see REGISTRY.md).

sieve_k_max : int or None

Maximum polynomial degree for the covariate-path sieves — the

propensity-ratio, inverse-propensity, AND outcome-regression fits all

use it. None = auto (``floor(n_pos^{1/5})`` over each group's

positive-weight support ``n_pos`` — the raw group size when unweighted —

a growing sieve with no fixed ceiling, bounded by ``n_basis < n_pos``;

zero-weight survey rows do not affect order selection). Only

used with covariates. ``sieve_k_max=1`` forces every covariate-path

sieve (outcome regression and both propensity sieves) to degree 1: it

recovers the pre-sieve linear-OLS *outcome regression* but also

degree-1-constrains the propensity sieves, so it does not reproduce the

exact pre-sieve estimator.

sieve_criterion : str, default ``"bic"``

Information criterion (``"aic"`` or ``"bic"``) for the order selection

of all covariate-path sieves (propensity ratio, inverse propensity, and

outcome regression).

ratio_clip : float, default 20.0

Clip sieve propensity ratios to ``[1/ratio_clip, ratio_clip]``.

kernel_bandwidth : float or None

Bandwidth for Gaussian kernel in conditional Omega* estimation.

None = Silverman's rule-of-thumb (automatic).

omega_ridge : float, default ``OMEGA_RIDGE_DEFAULT`` (1e-6)

Relative ridge for the Omega* inversion behind the efficient

weights: solves ``(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1``

instead of inverting the numerically singular Omega* that PT-All's

telescoping overidentified moments produce. Stabilizes per-cell

ATT(g,t) against floating-point-level input/BLAS changes (1-ulp

stability ~1e-9 vs ~1e-4 for the legacy pseudoinverse) without

changing overall-ATT bias/RMSE/coverage (see REGISTRY.md).

``omega_ridge=0`` restores the exact legacy inv/pinv code path

bit-for-bit - including its per-cell condition-number warnings and

the slow O(n^2 H^2) conditional-Omega* loops, so expect the legacy

runtime as well.

Examples

--------

>>> from diff_diff import EfficientDiD

>>> edid = EfficientDiD(pt_assumption="all")

>>> results = edid.fit(data, outcome="y", unit="id", time="t",

... first_treat="first_treat")

>>> results.print_summary()

>>> results.aggregate("event_study").summary() # post-fit aggregation

"""

def __init__(

self,

pt_assumption: str = "all",

alpha: float = 0.05,

cluster: Optional[str] = None,

vcov_type: str = "hc1",

control_group: str = "never_treated",

n_bootstrap: int = 0,

bootstrap_weights: str = "rademacher",

seed: Optional[int] = None,

anticipation: int = 0,

sieve_k_max: Optional[int] = None,

sieve_criterion: str = "bic",

ratio_clip: float = 20.0,

kernel_bandwidth: Optional[float] = None,

omega_ridge: float = OMEGA_RIDGE_DEFAULT,

):

self.pt_assumption = pt_assumption

self.alpha = alpha

self.cluster = cluster

self.vcov_type = vcov_type

self.control_group = control_group

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

self.seed = seed

self.anticipation = anticipation

self.sieve_k_max = sieve_k_max

self.sieve_criterion = sieve_criterion

self.ratio_clip = ratio_clip

self.kernel_bandwidth = kernel_bandwidth

self.omega_ridge = omega_ridge

self.is_fitted_ = False

self.results_: Optional[EfficientDiDResults] = None

self._unit_resolved_survey: Optional["ResolvedSurveyDesign"] = None

self._validate_params()

def _validate_params(self) -> None:

"""Validate constrained parameters."""

if self.pt_assumption not in ("all", "post"):

raise ValueError(f"pt_assumption must be 'all' or 'post', got '{self.pt_assumption}'")

if self.control_group not in ("never_treated", "last_cohort"):

raise ValueError(

f"control_group must be 'never_treated' or 'last_cohort', "

f"got '{self.control_group}'"

)

valid_weights = ("rademacher", "mammen", "webb")

if self.bootstrap_weights not in valid_weights:

raise ValueError(

f"bootstrap_weights must be one of {valid_weights}, "

f"got '{self.bootstrap_weights}'"

)

if self.sieve_criterion not in ("aic", "bic"):

raise ValueError(

f"sieve_criterion must be 'aic' or 'bic', got '{self.sieve_criterion}'"

)

if not (np.isfinite(self.ratio_clip) and self.ratio_clip > 1.0):

raise ValueError(f"ratio_clip must be finite and > 1.0, got {self.ratio_clip}")

if self.kernel_bandwidth is not None:

if not (np.isfinite(self.kernel_bandwidth) and self.kernel_bandwidth > 0):

raise ValueError(

f"kernel_bandwidth must be finite and > 0 (or None for auto), "

f"got {self.kernel_bandwidth}"

)

if self.sieve_k_max is not None:

if not (isinstance(self.sieve_k_max, (int, np.integer)) and self.sieve_k_max > 0):

raise ValueError(

f"sieve_k_max must be a positive integer (or None for auto), "

f"got {self.sieve_k_max}"

)

if not (np.isfinite(self.omega_ridge) and self.omega_ridge >= 0):

raise ValueError(

f"omega_ridge must be finite and >= 0 (0 = legacy inv/pinv path), "

f"got {self.omega_ridge}"

)

self._validate_vcov_type(self.vcov_type)

@staticmethod

def _validate_vcov_type(vcov_type: str) -> None:

"""Validate ``vcov_type`` against EfficientDiD's narrow IF-based contract.

Permanently accepts ``{"hc1"}`` only — EfficientDiD uses

influence-function-based variance per Chen-Sant'Anna-Xie (2025) achieving

the semiparametric efficiency bound. The per-unit EIF aggregation has no

equivalent single design matrix, so analytical-sandwich families

(``classical``, ``hc2``, ``hc2_bm``) cannot be defined; ``conley`` is

deferred (see DEFERRED.md). Mirrors the narrow-contract pattern in

:class:`ImputationDiD`, :class:`CallawaySantAnna`, and

:class:`TripleDifference`.

"""

_accepted_vcov = {"hc1"}

_if_incompatible_vcov = {"classical", "hc2", "hc2_bm"}

_deferred_vcov = {"conley"}

if vcov_type in _if_incompatible_vcov:

raise ValueError(

f"EfficientDiD(vcov_type={vcov_type!r}) is rejected: "

f"EfficientDiD uses influence-function-based variance per Chen, "

f"Sant'Anna, and Xie (2025) achieving the semiparametric efficiency "

f"bound for ATT(g,t). The per-unit EIF aggregation has no equivalent "

f"single design matrix on which hat matrix leverage or Bell-McCaffrey "

f"Satterthwaite DOF can be defined, so analytical-sandwich families "

f"{{classical, hc2, hc2_bm}} are not paper-prescribed. Use "

f"vcov_type='hc1' (the default) with cluster=<col> for the "

f"Liang-Zeger clustered EIF sandwich estimator."

)

if vcov_type in _deferred_vcov:

raise ValueError(

f"EfficientDiD(vcov_type={vcov_type!r}) is not yet supported: "

f"spatial-HAC composition with EIF aggregation has no reference "

f"implementation today. See DEFERRED.md for the deferred follow-up row. "

f"Use vcov_type='hc1' (the default) with cluster=<col> for "

f"cluster-robust inference."

)

if vcov_type not in _accepted_vcov:

raise ValueError(

f"EfficientDiD(vcov_type={vcov_type!r}) is invalid. "

f"Accepted: {sorted(_accepted_vcov)}."

)

# -- sklearn compatibility ------------------------------------------------

# get_params/set_params come from BaseEstimator.

# -- Main estimation ------------------------------------------------------

def fit(

self,

data: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

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

aggregate: Any = NOT_SUPPLIED,

balance_e: Any = NOT_SUPPLIED,

survey_design: Optional[Any] = None,

store_eif: bool = False,

) -> EfficientDiDResults:

"""Fit the Efficient DiD estimator.

Parameters

----------

data : DataFrame

Balanced panel data.

outcome : str

Outcome variable column name.

unit : str

Unit identifier column name.

time : str

Time period column name.

first_treat : str

Column indicating first treatment period.

Use 0 or ``np.inf`` for never-treated units.

covariates : list of str, optional

Column names for time-invariant unit-level covariates.

When provided, uses the doubly robust path (outcome regression

+ propensity score ratios).

aggregate : str, optional

DEPRECATED (3.9, removed in 4.0, row M-023): supplying any

value (``None`` included) emits a FutureWarning; the

deprecated path still computes and stores the requested

surfaces exactly as before (``None``, ``"simple"``,

``"event_study"``, ``"group"``, or ``"all"``; unknown strings

act like ``None``, unchanged). Aggregate as a post-fit step

instead: ``results.aggregate('event_study')`` /

``.aggregate('group')`` / ``.aggregate('simple')``. On

bootstrapped fits (``n_bootstrap > 0``) the post-fit

RECOMPUTE levels (``'event_study'``/``'group'``) fail closed

— the deprecated fit-time path remains the supported route

for those bootstrapped aggregated surfaces — while

``aggregate('simple')`` relays the stored bootstrap

inference and stays available (the per-level policy

converged with row M-027).

balance_e : int, optional

DEPRECATED (3.9, removed in 4.0, row M-120): moves onto

post-fit ``aggregate()`` —

``results.aggregate('event_study', balance_e=2)``. EDiD's

balance rule is the ANCHOR-HORIZON rule (keep cohorts with a

finite effect at ``e == balance_e``), the same rule

CallawaySantAnna uses.

survey_design : SurveyDesign, optional

Survey design specification for design-based inference.

Applies survey weights to all means, covariances, and cohort

fractions, and uses Taylor Series Linearization for SE

estimation. Cannot be combined with ``cluster``.

store_eif : bool, default False

Expose per-(g,t) EIF vectors on the PUBLIC

``influence_functions`` results field. Used internally by

:meth:`hausman_pretest`. Since 3.9 (row M-023) the private

aggregation kit ALWAYS retains the per-(g,t) EIF dict to power

post-fit ``results.aggregate()`` — ``store_eif`` governs only

the public field, no longer the retention itself.

Returns

-------

EfficientDiDResults

Raises

------

ValueError

Missing columns, unbalanced panel, non-absorbing treatment,

or PT-Post without a never-treated group.

"""

# M-023/M-120 deprecation shim (CS-style joint warning): a plain

# fit() never warns; supplying EITHER param with ANY value (None

# included) warns once, then the legacy routing below runs

# unchanged - the deprecated path returns exactly the numbers it

# always did (no new value validation; unknown strings still act

# like None). The post-fit successor validates its own vocabulary.

_deprecated_passed = [

n

for n, v in (("aggregate", aggregate), ("balance_e", balance_e))

if v is not NOT_SUPPLIED

]

if _deprecated_passed:

_args = " / ".join(f"{n}=" for n in _deprecated_passed)

warnings.warn(

f"EfficientDiD.fit({_args}) is deprecated and will be "

"removed in 4.0. Fit once, then aggregate as a post-fit "

"step: results = EfficientDiD().fit(...); "

"results.aggregate('event_study') / .aggregate('group') / "

".aggregate('simple'). balance_e moves onto aggregate() "

"alongside it: results.aggregate('event_study', "

"balance_e=2).",

FutureWarning,

stacklevel=2,

)

if aggregate is NOT_SUPPLIED:

aggregate = None

if balance_e is NOT_SUPPLIED:

balance_e = None

self._validate_params()

if self.cluster is not None and survey_design is not None:

raise NotImplementedError(

"cluster and survey_design cannot both be set. "

"Use survey_design with PSU/strata for cluster-robust inference."

)

# Resolve survey design if provided

from diff_diff.survey import _resolve_survey_for_fit, build_unit_first_row_index

resolved_survey, survey_weights, survey_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, data, "analytical")

)

# Validate within-unit constancy for panel survey designs

if resolved_survey is not None:

from diff_diff.survey import _validate_unit_constant_survey

_validate_unit_constant_survey(data, unit, survey_design)

# Store survey df for safe_inference calls (t-distribution with survey df)

self._survey_df = survey_metadata.df_survey if survey_metadata is not None else None

# Guard: replicate design with undefined df → NaN inference

if (

self._survey_df is None

and resolved_survey is not None

and hasattr(resolved_survey, "uses_replicate_variance")

and resolved_survey.uses_replicate_variance

):

self._survey_df = 0

# Bootstrap + survey supported via PSU-level multiplier bootstrap.

# Normalize empty covariates list to None (use nocov path)

if covariates is not None and len(covariates) == 0:

covariates = None

use_covariates = covariates is not None

# ----- Validate inputs -----

required_cols = [outcome, unit, time, first_treat]

missing = [c for c in required_cols if c not in data.columns]

if missing:

raise ValueError(f"Missing columns: {missing}")

df = data.copy()

df[time] = pd.to_numeric(df[time])

df[first_treat] = pd.to_numeric(df[first_treat])

# Normalize never-treated: inf -> 0 internally, keep track

df["_never_treated"] = (df[first_treat] == 0) | (df[first_treat] == np.inf)

df.loc[df[first_treat] == np.inf, first_treat] = 0

time_periods = sorted(df[time].unique())

treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0])

# Validate balanced panel

unit_period_counts = df.groupby(unit)[time].nunique()

n_periods = len(time_periods)

if (unit_period_counts != n_periods).any():

raise ValueError(

"Unbalanced panel detected. EfficientDiD requires a balanced "

"panel where every unit is observed in every time period."

)

# Reject non-finite outcomes (NaN/Inf corrupt Omega*/EIF calculations)

non_finite_mask = ~np.isfinite(df[outcome])

if non_finite_mask.any():

n_bad = int(non_finite_mask.sum())

raise ValueError(

f"Found {n_bad} non-finite value(s) in outcome column '{outcome}'. "

"EfficientDiD requires finite outcomes for all unit-period observations."

)

# Reject duplicate (unit, time) rows

dup_mask = df.duplicated(subset=[unit, time], keep=False)

if dup_mask.any():

n_dups = int(dup_mask.sum())

raise ValueError(

f"Found {n_dups} duplicate ({unit}, {time}) rows. "

"EfficientDiD requires exactly one observation per unit-period."

)

# Validate absorbing treatment (vectorized)

ft_nunique = df.groupby(unit)[first_treat].nunique()

bad_units = ft_nunique[ft_nunique > 1]

if len(bad_units) > 0:

uid = bad_units.index[0]

raise ValueError(

f"Non-absorbing treatment detected for unit {uid}: "

"first_treat value changes over time."

)

# Unit info

unit_info = (

df.groupby(unit)

.agg(

{

first_treat: "first",

"_never_treated": "first",

}

)

.reset_index()

)

n_treated_units = int((unit_info[first_treat] > 0).sum())

n_control_units = int(unit_info["_never_treated"].sum())

# Control group logic

if self.control_group == "last_cohort":

# Always reclassify last cohort as pseudo-control when requested

if not treatment_groups:

raise ValueError(

"No treated cohorts found. control_group='last_cohort' requires "

"at least 2 treatment cohorts."

)

last_g = max(treatment_groups)

treatment_groups = [g for g in treatment_groups if g != last_g]

if not treatment_groups:

raise ValueError("Only one treatment cohort; cannot use last_cohort control.")

effective_last = last_g - self.anticipation

time_periods = [t for t in time_periods if t < effective_last]

if len(time_periods) < 2:

raise ValueError(

"Fewer than 2 time periods remain after trimming for last_cohort control."

)

unit_info.loc[unit_info[first_treat] == last_g, first_treat] = 0

unit_info.loc[unit_info[first_treat] == 0, "_never_treated"] = True

n_treated_units = int((unit_info[first_treat] > 0).sum())

n_control_units = int(unit_info["_never_treated"].sum())

elif n_control_units == 0:

raise ValueError(

"No never-treated units found. Use control_group='last_cohort' "

"to use the last treatment cohort as a pseudo-control."

)

# ----- Prepare data -----

all_units = sorted(df[unit].unique())

n_units = len(all_units)

# Build unit-to-first-panel-row index aligned to all_units (sorted)

# order. The previous approach (groupby cumcount == 0) yielded

# first-appearance order which can differ from sorted order when the

# input DataFrame is not pre-sorted by unit.

self._unit_first_panel_row = build_unit_first_row_index(df[unit].values, all_units)

# Build unit-level ResolvedSurveyDesign once (avoids repeated

# construction in _compute_survey_eif_se and ensures consistent

# unit-level df for safe_inference t-distribution).

if resolved_survey is not None:

self._unit_resolved_survey = resolved_survey.subset_to_units_by_row_idx(

self._unit_first_panel_row

)

# Use unit-level df (not panel-level) for t-distribution

self._survey_df = self._unit_resolved_survey.df_survey

# Re-apply replicate guard: undefined df → NaN inference

if self._survey_df is None and self._unit_resolved_survey.uses_replicate_variance:

self._survey_df = 0

else:

self._unit_resolved_survey = None

# Build cluster mapping if cluster-robust SEs requested

if self.cluster is not None:

unit_cluster_indices, n_clusters = _validate_and_build_cluster_mapping(

df, unit, self.cluster, all_units

)

if n_clusters < 50:

warnings.warn(

f"Only {n_clusters} clusters. Analytical clustered SEs may "

"be unreliable. Consider n_bootstrap > 0 for cluster "

"bootstrap inference.",

UserWarning,

stacklevel=2,

)

else:

unit_cluster_indices = None

n_clusters = None

period_to_col = {p: i for i, p in enumerate(time_periods)}

period_1 = time_periods[0]

period_1_col = period_to_col[period_1]

# Pivot outcome to wide matrix (n_units, n_periods)

pivot = df.pivot(index=unit, columns=time, values=outcome)

# Reindex to match all_units ordering and time_periods column order

pivot = pivot.reindex(index=all_units, columns=time_periods)

outcome_wide = pivot.values.astype(float)

# Build cohort masks and fractions

unit_info_indexed = unit_info.set_index(unit)

unit_cohorts = unit_info_indexed.reindex(all_units)[first_treat].values.astype(

float

) # 0 = never-treated

cohort_masks: Dict[float, np.ndarray] = {}

for g in treatment_groups:

cohort_masks[g] = unit_cohorts == g

never_treated_mask = unit_cohorts == 0

cohort_masks[np.inf] = never_treated_mask # also keyed by inf sentinel

# ----- Unit-level survey weights -----

# Survey weights in the panel are at obs level (unit x time).

# EfficientDiD works at unit level. Extract one weight per unit

# by taking the first observation per unit (balanced panel, so

# weights should be constant within unit).

unit_level_weights: Optional[np.ndarray] = None

if resolved_survey is not None:

# Use the resolved survey's weights (already normalized per weight_type)

# subset to unit level via _unit_first_panel_row (aligned to all_units)

assert self._unit_resolved_survey is not None

unit_level_weights = self._unit_resolved_survey.weights

self._unit_level_weights = unit_level_weights

cohort_fractions: Dict[float, float] = {}

if unit_level_weights is not None:

# Survey-weighted cohort fractions: sum(w_i for i in cohort) / sum(w_i)

total_w = float(np.sum(unit_level_weights))

for g in treatment_groups:

cohort_fractions[g] = float(np.sum(unit_level_weights[cohort_masks[g]])) / total_w

cohort_fractions[np.inf] = (

float(np.sum(unit_level_weights[never_treated_mask])) / total_w

)

else:

for g in treatment_groups:

cohort_fractions[g] = float(np.sum(cohort_masks[g])) / n_units

cohort_fractions[np.inf] = float(np.sum(never_treated_mask)) / n_units

# ----- Small cohort warnings -----

for g in treatment_groups:

n_g = int(np.sum(cohort_masks[g]))

frac_g = cohort_fractions[g]

if n_g < 2:

warnings.warn(

f"Cohort {g} has only {n_g} unit. Omega* inversion and "

"EIF computation may be numerically unstable.",

UserWarning,

stacklevel=2,

)

elif frac_g < 0.01:

warnings.warn(

f"Cohort {g} represents {frac_g:.1%} of the sample (< 1%). "

"Efficient weights may be imprecise.",

UserWarning,

stacklevel=2,

)

# Guard: never-treated with zero survey weight → no valid comparisons

# Applies to both covariates (DR nuisance) and nocov (weighted means) paths

if cohort_fractions.get(np.inf, 0.0) <= 0 and unit_level_weights is not None:

raise ValueError(

"Never-treated group has zero survey weight. EfficientDiD "

"requires a never-treated control group with positive "

"survey weight for estimation."

)

# ----- Covariate preparation (if provided) -----

covariate_matrix: Optional[np.ndarray] = None

m_hat_cache: Dict[Tuple, np.ndarray] = {}

r_hat_cache: Dict[Tuple[float, float], np.ndarray] = {}

s_hat_cache: Dict[float, np.ndarray] = {} # inverse propensities per group

# Per-fit cache of the polynomial sieve basis, keyed (id(X), degree). The three

# sieve nuisance helpers all build the basis from the same fit-level

# `covariate_matrix`, so this shares each distinct degree's basis across them

# instead of rebuilding it per helper. Lives only for this fit() call.

sieve_basis_cache: Dict[Tuple[int, int], np.ndarray] = {}

if use_covariates:

assert covariates is not None # for type narrowing

# Validate covariate columns exist

missing_cov = [c for c in covariates if c not in data.columns]

if missing_cov:

raise ValueError(f"Missing covariate columns: {missing_cov}")

# Validate no NaN/Inf in covariates

for col_name in covariates:

non_finite_cov = ~np.isfinite(pd.to_numeric(df[col_name], errors="coerce"))

if non_finite_cov.any():

n_bad = int(non_finite_cov.sum())

raise ValueError(

f"Found {n_bad} non-finite value(s) in covariate column "

f"'{col_name}'. Covariates must be finite."

)

# Validate time-invariance: covariates must be constant within each unit

for col_name in covariates:

cov_nunique = df.groupby(unit)[col_name].nunique()

varying = cov_nunique[cov_nunique > 1]

if len(varying) > 0:

uid = varying.index[0]

raise ValueError(

f"Covariate '{col_name}' varies over time for unit {uid}. "

"EfficientDiD requires time-invariant covariates. "

"Extract base-period values before calling fit()."

)

# Extract unit-level covariate matrix from period_1 observations

base_df = df[df[time] == period_1].set_index(unit).reindex(all_units)

covariate_matrix = base_df[list(covariates)].values.astype(float)

# ----- Core estimation: ATT(g, t) for each target -----

# Precompute per-group unit counts (avoid repeated np.sum in loop)

n_treated_per_g = {g: int(np.sum(cohort_masks[g])) for g in treatment_groups}

n_control_count = int(np.sum(never_treated_mask))

group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]] = {}

eif_by_gt: Dict[Tuple[Any, Any], np.ndarray] = {}

stored_weights: Dict[Tuple[Any, Any], np.ndarray] = {}

stored_cond: Dict[Tuple[Any, Any], float] = {}

# Ridge path: ill-conditioned cells consolidate into ONE fit-level

# warning (legacy omega_ridge=0 keeps the per-cell pinv warnings).

_ill_conditioned_cells: List[Tuple[Any, Any, float]] = []

# Ridge-path covariate cells are deferred to a fused unit-tiled pass 2

# (compute_conditional_cells_tiled) after the sieve caches are fully

# populated; each deferred cell holds an order-preserving placeholder

# in group_time_effects until finalization.

_cond_cell_specs: List[Dict[str, Any]] = []

def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any]:

"""Per-cell SE + inference record, shared by the inline and

deferred (pass-2) paths.

Analytical SE = sqrt(mean(EIF^2) / n) [paper p.21]; with survey:

TSL variance via compute_survey_vcov.

"""

if self._unit_resolved_survey is not None:

se_gt = self._compute_survey_eif_se(eif_vals)

else:

se_gt = _compute_se_from_eif(eif_vals, n_units, unit_cluster_indices, n_clusters)

t_stat, p_val, ci = safe_inference(att_gt, se_gt, alpha=self.alpha, df=self._survey_df)

return {

"effect": att_gt,

"se": se_gt,

"t_stat": t_stat,

"p_value": p_val,

"conf_int": ci,

"n_treated": n_treated_per_g[g],

"n_control": n_control_count,

}

for g in treatment_groups:

# Under PT-Post, use per-group baseline Y_{g-1-anticipation}

# instead of the universal Y_1. This implements the weaker

# PT-Post assumption (parallel trends only from g-1 onward),

# matching the Callaway-Sant'Anna estimator exactly.

if self.pt_assumption == "post":

effective_base = g - 1 - self.anticipation

if effective_base not in period_to_col:

warnings.warn(

f"Cohort g={g} dropped: baseline period {effective_base} "

f"(g-1-anticipation) is not in the data.",

UserWarning,

stacklevel=2,

)

continue

effective_p1_col = period_to_col[effective_base]

else:

effective_p1_col = period_1_col

# Guard: skip cohorts with zero survey weight (all units zero-weighted)

if cohort_fractions[g] <= 0:

warnings.warn(

f"Cohort {g} has zero survey weight; skipping.",

UserWarning,

stacklevel=2,

)

continue

# Estimate all (g, t) cells including pre-treatment. Under PT-Post,

# pre-treatment cells serve as placebo/pre-trend diagnostics, matching

# the CallawaySantAnna implementation. Users filter to t >= g for

# post-treatment effects; pre-treatment cells are clearly labeled by

# their (g, t) coordinates in the results object.

for t in time_periods:

# Skip period_1 — it's the universal reference baseline,

# not a target period

if t == period_1:

continue

# Enumerate valid comparison pairs

pairs = enumerate_valid_triples(

target_g=g,

treatment_groups=treatment_groups,

time_periods=time_periods,

period_1=period_1,

pt_assumption=self.pt_assumption,

anticipation=self.anticipation,

)

# Filter out comparison pairs with zero survey weight

if unit_level_weights is not None and pairs:

pairs = [

(gp, tpre)

for gp, tpre in pairs

if np.sum(

unit_level_weights[

never_treated_mask if np.isinf(gp) else cohort_masks[gp]

]

)

> 0

]

# Ridge path: drop the degenerate same-cohort self-pair

# (g' = g, t_pre = t), which arises only for PRE-treatment

# cells (t < g). Its generated outcome telescopes to 0 = 0

# identically (zero-information moment; the exact-null

# direction of Omega*). The legacy pseudoinverse truncates

# that direction, spreading weight over noisy moments; a

# ridge would instead load ~all weight on the zero-variance

# moment, collapsing pre-treatment placebos to a

# deterministic 0 and silently disabling the pre-trend

# diagnostic. Dropping the pair restores data-driven

# placebos. NOT applied at omega_ridge=0 (bit-identical

# legacy behavior).

if self.omega_ridge > 0 and pairs:

pairs = [(gp, tpre) for gp, tpre in pairs if not (gp == g and tpre == t)]

if not pairs:

warnings.warn(

f"No valid comparison pairs for (g={g}, t={t}). " "ATT will be NaN.",

UserWarning,

stacklevel=2,

)

t_stat, p_val, ci = np.nan, np.nan, (np.nan, np.nan)

group_time_effects[(g, t)] = {

"effect": np.nan,

"se": np.nan,

"t_stat": t_stat,

"p_value": p_val,

"conf_int": ci,