df: pd.DataFrame,

unit: str,

cluster: str,

all_units: list,

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

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

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

eif_mat: np.ndarray,

cluster_indices: np.ndarray,

n_clusters: int,

) -> np.ndarray:

"""Sum EIF values within clusters and center.

if eif_mat.ndim == 1:

sums = np.bincount(cluster_indices, weights=eif_mat, minlength=n_clusters).astype(float)

else:

sums = np.column_stack(

[

np.bincount(cluster_indices, weights=eif_mat[:, j], minlength=n_clusters)

for j in range(eif_mat.shape[1])

]

).astype(float)

eif: np.ndarray,

n_units: int,

cluster_indices: Optional[np.ndarray] = None,

n_clusters: Optional[int] = None,

) -> float:

"""SE from EIF values, optionally with cluster-robust correction.

if cluster_indices is not None and n_clusters is not None:

centered = _cluster_aggregate(eif, cluster_indices, n_clusters)

correction = n_clusters / (n_clusters - 1) if n_clusters > 1 else 1.0

var = correction * np.sum(centered**2) / (n_units**2)

return float(np.sqrt(max(var, 0.0)))

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

def __init__(

self,

pt_assumption: str = "all",

alpha: float = 0.05,

cluster: Optional[str] = None,

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,

):

self.pt_assumption = pt_assumption

self.alpha = alpha

self.cluster = cluster

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.is_fitted_ = False

self.results_: Optional[EfficientDiDResults] = None

self._unit_resolved_survey = 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}"

)

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

def get_params(self) -> Dict[str, Any]:

"""Get estimator parameters (sklearn-compatible)."""

return {

"pt_assumption": self.pt_assumption,

"anticipation": self.anticipation,

"alpha": self.alpha,

"cluster": self.cluster,

"control_group": self.control_group,

"n_bootstrap": self.n_bootstrap,

"bootstrap_weights": self.bootstrap_weights,

"seed": self.seed,

"sieve_k_max": self.sieve_k_max,

"sieve_criterion": self.sieve_criterion,

"ratio_clip": self.ratio_clip,

"kernel_bandwidth": self.kernel_bandwidth,

}

def set_params(self, **params: Any) -> "EfficientDiD":

"""Set estimator parameters (sklearn-compatible)."""

for key, value in params.items():

if hasattr(self, key):

setattr(self, key, value)

else:

raise ValueError(f"Unknown parameter: {key}")

self._validate_params()

return self

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

def fit(

self,

data: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

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

aggregate: Optional[str] = None,

balance_e: Optional[int] = None,

survey_design: Optional[Any] = None,

store_eif: bool = False,

) -> EfficientDiDResults:

"""Fit the Efficient DiD estimator.

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

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.

first_pos: Dict[Any, int] = {}

for i, u in enumerate(df[unit].values):

if u not in first_pos:

first_pos[u] = i

self._unit_first_panel_row = np.array([first_pos[u] for u in 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:

from diff_diff.survey import ResolvedSurveyDesign

row_idx = self._unit_first_panel_row

unit_weights_s = resolved_survey.weights[row_idx]

unit_strata = (

resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None

)

unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None

unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None

n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0

n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0

self._unit_resolved_survey = resolved_survey.subset_to_units(

row_idx, unit_weights_s, unit_strata, unit_psu, unit_fpc,

n_strata_u, n_psu_u,

)

# 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)

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

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] = {}

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

]

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,

"n_treated": n_treated_per_g[g],

"n_control": n_control_count,

}

eif_by_gt[(g, t)] = np.zeros(n_units)

continue

if use_covariates:

assert covariate_matrix is not None

t_col_val = period_to_col[t]

# Lazily populate nuisance caches for this (g, t)

for gp, tpre in pairs:

tpre_col_val = period_to_col[tpre]

# m_{inf, t, tpre}(X)

key_inf_t = (np.inf, t_col_val, tpre_col_val)

if key_inf_t not in m_hat_cache:

m_hat_cache[key_inf_t] = estimate_outcome_regression(

outcome_wide,

covariate_matrix,

never_treated_mask,

t_col_val,

tpre_col_val,

unit_weights=unit_level_weights,

)

# m_{g', tpre, 1}(X)

key_gp_tpre = (gp, tpre_col_val, effective_p1_col)

if key_gp_tpre not in m_hat_cache:

gp_mask_for_reg = (

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

)

m_hat_cache[key_gp_tpre] = estimate_outcome_regression(

outcome_wide,

covariate_matrix,

gp_mask_for_reg,

tpre_col_val,

effective_p1_col,

unit_weights=unit_level_weights,

)

# r_{g, inf}(X) and r_{g, g'}(X) via sieve (Eq 4.1-4.2)

for comp in {np.inf, gp}:

rkey = (g, comp)

if rkey not in r_hat_cache:

comp_mask = (

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

)

r_hat_cache[rkey] = estimate_propensity_ratio_sieve(

covariate_matrix,

cohort_masks[g],

comp_mask,

k_max=self.sieve_k_max,

criterion=self.sieve_criterion,

ratio_clip=self.ratio_clip,

unit_weights=unit_level_weights,

)

# Per-unit DR generated outcomes: shape (n_units, H)

gen_out = compute_generated_outcomes_cov(

target_g=g,

target_t=t,

valid_pairs=pairs,

outcome_wide=outcome_wide,

cohort_masks=cohort_masks,

never_treated_mask=never_treated_mask,

period_to_col=period_to_col,

period_1_col=effective_p1_col,

cohort_fractions=cohort_fractions,

m_hat_cache=m_hat_cache,

r_hat_cache=r_hat_cache,

)

y_hat = np.mean(gen_out, axis=0) # shape (H,)

# Inverse propensity estimation (algorithm step 4)

# s_hat_{g'}(X) = 1/p_{g'}(X) for Eq 3.12 scaling

for group_id in {g, np.inf} | {gp for gp, _ in pairs}:

if group_id not in s_hat_cache:

group_mask_s = (

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

)

s_hat_cache[group_id] = estimate_inverse_propensity_sieve(

covariate_matrix,

group_mask_s,

k_max=self.sieve_k_max,

criterion=self.sieve_criterion,

unit_weights=unit_level_weights,

)

# Conditional Omega*(X) with per-unit propensities (Eq 3.12)

omega_cond = compute_omega_star_conditional(

target_g=g,

target_t=t,

valid_pairs=pairs,

outcome_wide=outcome_wide,

cohort_masks=cohort_masks,

never_treated_mask=never_treated_mask,

period_to_col=period_to_col,

period_1_col=effective_p1_col,

cohort_fractions=cohort_fractions,

covariate_matrix=covariate_matrix,

s_hat_cache=s_hat_cache,

bandwidth=self.kernel_bandwidth,

unit_weights=unit_level_weights,

)

# Per-unit weights: (n_units, H)

per_unit_w = compute_per_unit_weights(omega_cond)

# ATT = (survey-)weighted mean of per-unit DR scores

if per_unit_w.shape[1] > 0:

per_unit_scores = np.sum(per_unit_w * gen_out, axis=1)

if unit_level_weights is not None:

att_gt = float(np.average(per_unit_scores, weights=unit_level_weights))

else:

att_gt = float(np.mean(per_unit_scores))

else:

att_gt = np.nan

# EIF with per-unit weights (Remark 4.2: plug-in valid)

# Center on scalar ATT, not per-pair means (ensures mean(EIF) ≈ 0)

eif_vals = compute_eif_cov(per_unit_w, gen_out, att_gt, n_units)

eif_by_gt[(g, t)] = eif_vals

else:

# No-covariates path (closed-form)

omega = compute_omega_star_nocov(

target_g=g,

target_t=t,

valid_pairs=pairs,

outcome_wide=outcome_wide,

cohort_masks=cohort_masks,

never_treated_mask=never_treated_mask,

period_to_col=period_to_col,

period_1_col=effective_p1_col,

cohort_fractions=cohort_fractions,

unit_weights=unit_level_weights,

)

weights, _, cond_num = compute_efficient_weights(omega)

stored_weights[(g, t)] = weights

if omega.size > 0:

stored_cond[(g, t)] = cond_num

y_hat = compute_generated_outcomes_nocov(

target_g=g,

target_t=t,

valid_pairs=pairs,

outcome_wide=outcome_wide,

cohort_masks=cohort_masks,

never_treated_mask=never_treated_mask,

period_to_col=period_to_col,

period_1_col=effective_p1_col,

unit_weights=unit_level_weights,

)

att_gt = float(weights @ y_hat) if len(weights) > 0 else np.nan

eif_vals = compute_eif_nocov(

target_g=g,

target_t=t,

weights=weights,

valid_pairs=pairs,

outcome_wide=outcome_wide,

cohort_masks=cohort_masks,

never_treated_mask=never_treated_mask,

period_to_col=period_to_col,

period_1_col=effective_p1_col,

cohort_fractions=cohort_fractions,

n_units=n_units,

unit_weights=unit_level_weights,

)

eif_by_gt[(g, t)] = eif_vals

# Analytical SE = sqrt(mean(EIF^2) / n) [paper p.21]

# With survey: use 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

)

group_time_effects[(g, t)] = {

"effect": att_gt,

"se": se_gt,

"t_stat": t_stat,

"p_value": p_val,

"conf_int": ci,

"n_treated": int(np.sum(cohort_masks[g])),

"n_control": int(np.sum(never_treated_mask)),

}

if not group_time_effects:

raise ValueError(

"Could not estimate any group-time effects. "

"Check data has sufficient observations."

)

# ----- Aggregation -----

overall_att, overall_se = self._aggregate_overall(

group_time_effects,

eif_by_gt,

n_units,

cohort_fractions,

unit_cohorts,

cluster_indices=unit_cluster_indices,

n_clusters=n_clusters,

)

overall_t, overall_p, overall_ci = safe_inference(

overall_att, overall_se, alpha=self.alpha, df=self._survey_df

)

event_study_effects = None

group_effects = None

if aggregate in ("event_study", "all"):

event_study_effects = self._aggregate_event_study(

group_time_effects,

eif_by_gt,

n_units,

cohort_fractions,

treatment_groups,

time_periods,

balance_e,

unit_cohorts=unit_cohorts,

cluster_indices=unit_cluster_indices,

n_clusters=n_clusters,

)

if aggregate in ("group", "all"):

group_effects = self._aggregate_by_group(

group_time_effects,

eif_by_gt,

n_units,

cohort_fractions,

treatment_groups,

unit_cohorts=unit_cohorts,

cluster_indices=unit_cluster_indices,

n_clusters=n_clusters,

)

# ----- Bootstrap -----

# Reject replicate-weight designs for bootstrap — replicate variance

# is an analytical alternative, not compatible with bootstrap

if (

self.n_bootstrap > 0

and self._unit_resolved_survey is not None

and self._unit_resolved_survey.uses_replicate_variance

):

raise NotImplementedError(

"EfficientDiD bootstrap (n_bootstrap > 0) is not supported "

"with replicate-weight survey designs. Replicate weights provide "

"analytical variance; use n_bootstrap=0 instead."

)

bootstrap_results = None

if self.n_bootstrap > 0 and eif_by_gt:

bootstrap_results = self._run_multiplier_bootstrap(

group_time_effects=group_time_effects,

eif_by_gt=eif_by_gt,

n_units=n_units,

aggregate=aggregate,