"""

def __init__(

self,

anticipation: int = 0,

alpha: float = 0.05,

cluster: Optional[str] = None,

n_bootstrap: int = 0,

bootstrap_weights: str = "rademacher",

seed: Optional[int] = None,

rank_deficient_action: str = "warn",

horizon_max: Optional[int] = None,

aux_partition: str = "cohort_horizon",

pretrends: bool = False,

):

if rank_deficient_action not in ("warn", "error", "silent"):

raise ValueError(

f"rank_deficient_action must be 'warn', 'error', or 'silent', "

f"got '{rank_deficient_action}'"

)

if bootstrap_weights not in ("rademacher", "mammen", "webb"):

raise ValueError(

f"bootstrap_weights must be 'rademacher', 'mammen', or 'webb', "

f"got '{bootstrap_weights}'"

)

if aux_partition not in ("cohort_horizon", "cohort", "horizon"):

raise ValueError(

f"aux_partition must be 'cohort_horizon', 'cohort', or 'horizon', "

f"got '{aux_partition}'"

)

self.anticipation = anticipation

self.alpha = alpha

self.cluster = cluster

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

self.seed = seed

self.rank_deficient_action = rank_deficient_action

self.horizon_max = horizon_max

self.aux_partition = aux_partition

self.pretrends = pretrends

self.is_fitted_ = False

self.results_: Optional[ImputationDiDResults] = None

# Internal state preserved for pretrend_test()

self._fit_data: Optional[Dict[str, Any]] = None

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: object = None,

) -> ImputationDiDResults:

"""

# Validate inputs

required_cols = [outcome, unit, time, first_treat]

if covariates:

required_cols.extend(covariates)

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

if missing:

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

# pretrends + analytical survey is supported (Phase 8e-iii).

# Replicate-weight surveys need per-replicate lead regression refits

# which are not yet implemented — reject that combination.

if (

self.pretrends

and survey_design is not None

and survey_design.replicate_method is not None

and aggregate in ("event_study", "all")

):

raise NotImplementedError(

"pretrends=True is not yet compatible with replicate-weight "

"survey designs. Analytical survey designs (strata/PSU/FPC) "

"are supported. Use pretrends=False with replicate weights."

)

# Create working copy

df = data.copy()

# Resolve survey design if provided

from diff_diff.survey import (

_inject_cluster_as_psu,

_resolve_effective_cluster,

_resolve_survey_for_fit,

_validate_unit_constant_survey,

)

resolved_survey, survey_weights, _, survey_metadata = _resolve_survey_for_fit(

survey_design, data, "analytical"

)

_uses_replicate_imp = (

resolved_survey is not None and resolved_survey.uses_replicate_variance

)

if _uses_replicate_imp and self.n_bootstrap > 0:

raise ValueError(

"Cannot use n_bootstrap > 0 with replicate-weight survey designs. "

"Replicate weights provide their own variance estimation."

)

# Validate within-unit constancy for panel survey designs

if resolved_survey is not None:

_validate_unit_constant_survey(data, unit, survey_design)

if resolved_survey.weight_type != "pweight":

raise ValueError(

f"ImputationDiD survey support requires weight_type='pweight', "

f"got '{resolved_survey.weight_type}'. The survey variance math "

f"assumes probability weights (pweight)."

)

# FPC is supported — threaded through compute_survey_if_variance()

# in _compute_conservative_variance().

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

# Ensure numeric types

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

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

# Validate absorbing treatment: first_treat must be constant within each unit

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

non_constant = ft_nunique[ft_nunique > 1]

if len(non_constant) > 0:

example_unit = non_constant.index[0]

example_vals = sorted(df.loc[df[unit] == example_unit, first_treat].unique())

warnings.warn(

f"{len(non_constant)} unit(s) have non-constant '{first_treat}' "

f"values (e.g., unit '{example_unit}' has values {example_vals}). "

f"ImputationDiD assumes treatment is an absorbing state "

f"(once treated, always treated) with a single treatment onset "

f"time per unit. Non-constant first_treat violates this assumption "

f"and may produce unreliable estimates.",

UserWarning,

stacklevel=2,

)

# Coerce to per-unit value so downstream code

# (_never_treated, _treated, _rel_time) uses a single

# consistent first_treat per unit.

df[first_treat] = df.groupby(unit)[first_treat].transform("first")

# Identify treatment status

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

# Check for always-treated units (treated in all observed periods)

min_time = df[time].min()

always_treated_mask = (~df["_never_treated"]) & (df[first_treat] <= min_time)

n_always_treated = df.loc[always_treated_mask, unit].nunique()

if n_always_treated > 0:

warnings.warn(

f"{n_always_treated} unit(s) are treated in all observed periods "

f"(first_treat <= {min_time}). These units have no untreated "

"observations and cannot contribute to the counterfactual model. "

"Their treatment effects will be imputed but may be unreliable.",

UserWarning,

stacklevel=2,

)

# Create treatment indicator D_it

# D_it = 1 if t >= first_treat and first_treat > 0

# With anticipation: D_it = 1 if t >= first_treat - anticipation

effective_treat = df[first_treat] - self.anticipation

df["_treated"] = (~df["_never_treated"]) & (df[time] >= effective_treat)

# Identify Omega_0 (untreated) and Omega_1 (treated)

omega_0_mask = ~df["_treated"]

omega_1_mask = df["_treated"]

n_omega_0 = int(omega_0_mask.sum())

n_omega_1 = int(omega_1_mask.sum())

if n_omega_0 == 0:

raise ValueError(

"No untreated observations found. Cannot estimate counterfactual model."

)

if n_omega_1 == 0:

raise ValueError("No treated observations found. Nothing to estimate.")

# Identify groups and time periods

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

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

if len(treatment_groups) == 0:

raise ValueError("No treated units found. Check 'first_treat' column.")

# Unit info

unit_info = (

df.groupby(unit).agg({first_treat: "first", "_never_treated": "first"}).reset_index()

)

n_treated_units = int((~unit_info["_never_treated"]).sum())

# Control units = units with at least one untreated observation

units_in_omega_0 = df.loc[omega_0_mask, unit].unique()

n_control_units = len(units_in_omega_0)

# Cluster variable

cluster_var = self.cluster if self.cluster is not None else unit

if self.cluster is not None and self.cluster not in df.columns:

raise ValueError(

f"Cluster column '{self.cluster}' not found in data. "

f"Available columns: {list(df.columns)}"

)

# Resolve effective cluster and inject cluster-as-PSU for survey variance

if resolved_survey is not None:

cluster_ids_raw = df[cluster_var].values if cluster_var in df.columns else None

effective_cluster_ids = _resolve_effective_cluster(

resolved_survey,

cluster_ids_raw,

cluster_var if self.cluster is not None else None,

)

resolved_survey = _inject_cluster_as_psu(resolved_survey, effective_cluster_ids)

# When survey PSU is present, use it as the effective cluster for

# Theorem 3 variance (PSU overrides unit-level clustering)

if resolved_survey.psu is not None:

# Create a temporary column with PSU IDs for cluster_var

df["_survey_cluster"] = resolved_survey.psu

cluster_var = "_survey_cluster"

# Recompute metadata after PSU injection

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

from diff_diff.survey import compute_survey_metadata

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)

# Compute relative time

df["_rel_time"] = np.where(

~df["_never_treated"],

df[time] - df[first_treat],

np.nan,

)

# ---- Step 1: OLS on untreated observations ----

unit_fe, time_fe, grand_mean, delta_hat, kept_cov_mask = self._fit_untreated_model(

df, outcome, unit, time, covariates, omega_0_mask, weights=survey_weights

)

# ---- Rank condition checks ----

# Check: every treated unit should have >= 1 untreated period (for unit FE)

treated_unit_ids = df.loc[omega_1_mask, unit].unique()

units_with_fe = set(unit_fe.keys())

units_missing_fe = set(treated_unit_ids) - units_with_fe

# Check: every post-treatment period should have >= 1 untreated unit (for time FE)

post_period_ids = df.loc[omega_1_mask, time].unique()

periods_with_fe = set(time_fe.keys())

periods_missing_fe = set(post_period_ids) - periods_with_fe

if units_missing_fe or periods_missing_fe:

parts = []

if units_missing_fe:

sorted_missing = sorted(units_missing_fe)

parts.append(

f"{len(units_missing_fe)} treated unit(s) have no untreated "

f"periods (units: {sorted_missing[:5]}"

f"{'...' if len(units_missing_fe) > 5 else ''})"

)

if periods_missing_fe:

sorted_missing = sorted(periods_missing_fe)

parts.append(

f"{len(periods_missing_fe)} post-treatment period(s) have no "

f"untreated units (periods: {sorted_missing[:5]}"

f"{'...' if len(periods_missing_fe) > 5 else ''})"

)

msg = (

"Rank condition violated: "

+ "; ".join(parts)

+ ". Affected treatment effects will be NaN."

)

if self.rank_deficient_action == "error":

raise ValueError(msg)

elif self.rank_deficient_action == "warn":

warnings.warn(msg, UserWarning, stacklevel=2)

# "silent": continue without warning

# ---- Step 2: Impute treatment effects ----

tau_hat, y_hat_0 = self._impute_treatment_effects(

df,

outcome,

unit,

time,

covariates,

omega_1_mask,

unit_fe,

time_fe,

grand_mean,

delta_hat,

)

# Store tau_hat in dataframe

df["_tau_hat"] = np.nan

df.loc[omega_1_mask, "_tau_hat"] = tau_hat

# ---- Step 3: Aggregate ----

# Always compute overall ATT (simple aggregation)

finite_mask = np.isfinite(tau_hat)

valid_tau = tau_hat[finite_mask]

if len(valid_tau) == 0:

overall_att = np.nan

elif survey_weights is not None:

# Survey-weighted ATT: use treated obs' survey weights

treated_survey_w = survey_weights[omega_1_mask.values]

w_finite = treated_survey_w[finite_mask]

overall_att = float(np.average(valid_tau, weights=w_finite))

else:

overall_att = float(np.mean(valid_tau))

# ---- Variance ----

_n_valid_rep_imp = None

_vcov_rep_imp = None

overall_se = np.nan # placeholder; overridden by replicate or conservative path

if not _uses_replicate_imp:

# Conservative variance (Theorem 3)

overall_weights = np.zeros(n_omega_1)

n_valid = int(finite_mask.sum())

if n_valid > 0:

if survey_weights is not None:

treated_sw = survey_weights[omega_1_mask.values]

sw_finite = treated_sw[finite_mask]

overall_weights[finite_mask] = sw_finite / sw_finite.sum()

else:

overall_weights[finite_mask] = 1.0 / n_valid

if n_valid == 0:

overall_se = np.nan

else:

overall_se = self._compute_conservative_variance(

df=df,

outcome=outcome,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=unit_fe,

time_fe=time_fe,

grand_mean=grand_mean,

delta_hat=delta_hat,

weights=overall_weights,

cluster_var=cluster_var,

kept_cov_mask=kept_cov_mask,

survey_weights=survey_weights,

resolved_survey=(resolved_survey if not _uses_replicate_imp else None),

)

# Survey degrees of freedom for t-distribution inference

_survey_df = resolved_survey.df_survey if resolved_survey is not None else None

# Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1

if _uses_replicate_imp and _survey_df is None:

_survey_df = 0 # rank-deficient replicate → NaN inference

# Compute overall inference (may be overridden by replicate below)

overall_t, overall_p, overall_ci = safe_inference(

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

)

# Event study and group aggregation (full-sample, for point estimates)

event_study_effects = None

group_effects = None

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

event_study_effects = self._aggregate_event_study(

df=df,

outcome=outcome,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=unit_fe,

time_fe=time_fe,

grand_mean=grand_mean,

delta_hat=delta_hat,

cluster_var=cluster_var,

treatment_groups=treatment_groups,

balance_e=balance_e,

kept_cov_mask=kept_cov_mask,

survey_weights=survey_weights,

survey_df=_survey_df,

resolved_survey=(resolved_survey if not _uses_replicate_imp else None),

)

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

group_effects = self._aggregate_group(

df=df,

outcome=outcome,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=unit_fe,

time_fe=time_fe,

grand_mean=grand_mean,

delta_hat=delta_hat,

cluster_var=cluster_var,

treatment_groups=treatment_groups,

kept_cov_mask=kept_cov_mask,

survey_weights=survey_weights,

survey_df=_survey_df,

resolved_survey=(resolved_survey if not _uses_replicate_imp else None),

)

# Replicate variance: derive keys from actual outputs (after filtering)

if _uses_replicate_imp:

from diff_diff.survey import compute_replicate_refit_variance

_rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values

_cohorts_treated = df.loc[omega_1_mask, first_treat].values

# Derive keys from actual outputs (excludes filtered/Prop5/ref)

_sorted_rel_times = sorted(

e

for e in (event_study_effects or {}).keys()

if np.isfinite(event_study_effects[e]["effect"])

and event_study_effects[e].get("n_obs", 1) > 0

)

_sorted_groups = sorted(

g for g in (group_effects or {}).keys() if np.isfinite(group_effects[g]["effect"])

)

_n_es = len(_sorted_rel_times)

# Pre-compute balanced cohort mask for balance_e

_balanced_mask_treated = None

if balance_e is not None and _sorted_rel_times:

df_1 = df.loc[omega_1_mask]

rel_times_all = df_1["_rel_time"].values

all_horizons_full = sorted(set(int(h) for h in rel_times_all if np.isfinite(h)))

if self.horizon_max is not None:

all_horizons_full = [h for h in all_horizons_full if abs(h) <= self.horizon_max]

cohort_rel_times = self._build_cohort_rel_times(df, first_treat)

_balanced_mask_treated = self._compute_balanced_cohort_mask(

df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times

)

# Single vectorized refit: [overall, es_e0..., grp_g0...]

def _refit_imp(w_r):

ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model(

df,

outcome,

unit,

time,

covariates,

omega_0_mask,

weights=w_r,

)

tau_r, _ = self._impute_treatment_effects(

df,

outcome,

unit,

time,

covariates,

omega_1_mask,

ufe_r,

tfe_r,

gm_r,

delta_r,

)

fin = np.isfinite(tau_r)

treated_w = w_r[omega_1_mask.values]

results = []

# [0] Overall ATT

tw_fin = treated_w[fin]

tw_sum = np.sum(tw_fin)

results.append(

float(np.sum(tau_r[fin] * tw_fin) / tw_sum) if tw_sum > 0 else np.nan

)

# [1..n_es] Event-study (identified only)

for e in _sorted_rel_times:

mask_e = fin & (_rel_times_treated == e)

if _balanced_mask_treated is not None:

mask_e = mask_e & _balanced_mask_treated

tw_e = treated_w[mask_e]

s = np.sum(tw_e)

results.append(float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan)

# [n_es+1..] Group (identified only)

for g in _sorted_groups:

mask_g = fin & (_cohorts_treated == g)

tw_g = treated_w[mask_g]

s = np.sum(tw_g)

results.append(float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan)

return np.array(results)

# Build full-sample estimate from actual effects

_full_est = [overall_att]

_full_est.extend([event_study_effects[e]["effect"] for e in _sorted_rel_times])

_full_est.extend([group_effects[g]["effect"] for g in _sorted_groups])

_vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance(

_refit_imp, np.array(_full_est), resolved_survey

)

overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0)))

# Override df if replicates were dropped

if _n_valid_rep_imp < resolved_survey.n_replicates:

_survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0

if survey_metadata is not None:

survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None

overall_t, overall_p, overall_ci = safe_inference(

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

)

# Override event-study SEs from vcov diagonal

for i, e in enumerate(_sorted_rel_times):

if event_study_effects is not None and e in event_study_effects:

se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0)))

eff_e = event_study_effects[e]["effect"]

t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df)

event_study_effects[e]["se"] = se_e

event_study_effects[e]["t_stat"] = t_e

event_study_effects[e]["p_value"] = p_e

event_study_effects[e]["conf_int"] = ci_e

# Override group SEs from vcov diagonal

for j, g in enumerate(_sorted_groups):

if group_effects is not None and g in group_effects:

se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0)))

eff_g = group_effects[g]["effect"]

t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df)

group_effects[g]["se"] = se_g

group_effects[g]["t_stat"] = t_g

group_effects[g]["p_value"] = p_g

group_effects[g]["conf_int"] = ci_g

# Build treatment effects dataframe

treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy()

treated_df = treated_df.rename(columns={"_tau_hat": "tau_hat", "_rel_time": "rel_time"})

# Weights consistent with actual ATT: zero for NaN tau_hat

tau_finite = treated_df["tau_hat"].notna()

n_valid_te = int(tau_finite.sum())

if n_valid_te > 0:

if survey_weights is not None:

# Survey-weighted: use normalized survey weights for treated obs

treated_sw = survey_weights[omega_1_mask.values]

sw_finite = np.where(tau_finite, treated_sw, 0.0)

sw_sum = sw_finite.sum()

treated_df["weight"] = sw_finite / sw_sum if sw_sum > 0 else 0.0

else:

treated_df["weight"] = np.where(tau_finite, 1.0 / n_valid_te, 0.0)

else:

treated_df["weight"] = 0.0

# Store fit data for pretrend_test

self._fit_data = {

"df": df,

"outcome": outcome,

"unit": unit,

"time": time,

"first_treat": first_treat,

"covariates": covariates,

"omega_0_mask": omega_0_mask,

"omega_1_mask": omega_1_mask,

"cluster_var": cluster_var,

"unit_fe": unit_fe,

"time_fe": time_fe,

"grand_mean": grand_mean,

"delta_hat": delta_hat,

"kept_cov_mask": kept_cov_mask,

"survey_design": survey_design,

"resolved_survey": resolved_survey,

"survey_weights": survey_weights,

}

# Pre-compute cluster psi sums for bootstrap

psi_data = None

if self.n_bootstrap > 0 and n_valid > 0:

try:

# Extract survey weights for untreated obs (same as analytical path)

_sw_0 = survey_weights[omega_0_mask.values] if survey_weights is not None else None

# Extract survey weights for treated obs (event-study/group bootstrap paths)

_sw_1 = survey_weights[omega_1_mask.values] if survey_weights is not None else None

psi_data = self._precompute_bootstrap_psi(

df=df,

outcome=outcome,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=unit_fe,

time_fe=time_fe,

grand_mean=grand_mean,

delta_hat=delta_hat,

cluster_var=cluster_var,

kept_cov_mask=kept_cov_mask,

overall_weights=overall_weights,

event_study_effects=event_study_effects,

group_effects=group_effects,

treatment_groups=treatment_groups,

tau_hat=tau_hat,

balance_e=balance_e,

survey_weights_0=_sw_0,

survey_weights_1=_sw_1,

)

except Exception as e:

warnings.warn(

f"Bootstrap pre-computation failed: {e}. " "Skipping bootstrap inference.",

UserWarning,

stacklevel=2,

)

psi_data = None

# Bootstrap

bootstrap_results = None

if self.n_bootstrap > 0 and psi_data is not None:

bootstrap_results = self._run_bootstrap(

original_att=overall_att,

original_event_study=event_study_effects,

original_group=group_effects,

psi_data=psi_data,

resolved_survey=resolved_survey,

)

# Update inference with bootstrap results

overall_se = bootstrap_results.overall_att_se

overall_t = (

overall_att / overall_se if np.isfinite(overall_se) and overall_se > 0 else np.nan

)

overall_p = bootstrap_results.overall_att_p_value

overall_ci = bootstrap_results.overall_att_ci

# Update event study

if event_study_effects and bootstrap_results.event_study_ses:

for h in event_study_effects:

if (

h in bootstrap_results.event_study_ses

and event_study_effects[h].get("n_obs", 1) > 0

):

event_study_effects[h]["se"] = bootstrap_results.event_study_ses[h]

assert bootstrap_results.event_study_cis is not None

event_study_effects[h]["conf_int"] = bootstrap_results.event_study_cis[h]

assert bootstrap_results.event_study_p_values is not None

event_study_effects[h]["p_value"] = bootstrap_results.event_study_p_values[

h

]

eff_val = event_study_effects[h]["effect"]

se_val = event_study_effects[h]["se"]

event_study_effects[h]["t_stat"] = safe_inference(

eff_val, se_val, alpha=self.alpha

)[0]

# Update group effects

if group_effects and bootstrap_results.group_ses:

for g in group_effects:

if g in bootstrap_results.group_ses:

group_effects[g]["se"] = bootstrap_results.group_ses[g]

assert bootstrap_results.group_cis is not None

group_effects[g]["conf_int"] = bootstrap_results.group_cis[g]

assert bootstrap_results.group_p_values is not None

group_effects[g]["p_value"] = bootstrap_results.group_p_values[g]

eff_val = group_effects[g]["effect"]

se_val = group_effects[g]["se"]

group_effects[g]["t_stat"] = safe_inference(

eff_val, se_val, alpha=self.alpha

)[0]

# Construct results

self.results_ = ImputationDiDResults(

treatment_effects=treated_df,

overall_att=overall_att,

overall_se=overall_se,

overall_t_stat=overall_t,

overall_p_value=overall_p,

overall_conf_int=overall_ci,

event_study_effects=event_study_effects,

group_effects=group_effects,

groups=treatment_groups,

time_periods=time_periods,

n_obs=len(df),

n_treated_obs=n_omega_1,

n_untreated_obs=n_omega_0,

n_treated_units=n_treated_units,

n_control_units=n_control_units,

alpha=self.alpha,

anticipation=self.anticipation,

bootstrap_results=bootstrap_results,

_estimator_ref=self,

survey_metadata=survey_metadata,

)

self.is_fitted_ = True

return self.results_

# =========================================================================

# Step 1: OLS on untreated observations

# =========================================================================

def _iterative_fe(

self,

y: np.ndarray,

unit_vals: np.ndarray,

time_vals: np.ndarray,

idx: pd.Index,

max_iter: int = 100,

tol: float = 1e-10,

weights: Optional[np.ndarray] = None,

) -> Tuple[Dict[Any, float], Dict[Any, float]]:

"""

n = len(y)

alpha = np.zeros(n) # unit FE broadcast to obs level

beta = np.zeros(n) # time FE broadcast to obs level

# Precompute per-group weight sums (invariant across iterations)

if weights is not None:

w_series = pd.Series(weights, index=idx)

wsum_t = w_series.groupby(time_vals).transform("sum").values

wsum_u = w_series.groupby(unit_vals).transform("sum").values

converged = False

with np.errstate(invalid="ignore", divide="ignore"):

for iteration in range(max_iter):

resid_after_alpha = y - alpha

if weights is not None:

wr_t = pd.Series(resid_after_alpha * weights, index=idx)

beta_new = wr_t.groupby(time_vals).transform("sum").values / wsum_t

else:

beta_new = (

pd.Series(resid_after_alpha, index=idx)

.groupby(time_vals)

.transform("mean")

.values

)

resid_after_beta = y - beta_new

if weights is not None:

wr_u = pd.Series(resid_after_beta * weights, index=idx)

alpha_new = wr_u.groupby(unit_vals).transform("sum").values / wsum_u

else:

alpha_new = (

pd.Series(resid_after_beta, index=idx)

.groupby(unit_vals)

.transform("mean")

.values

)

# Check convergence on FE changes

max_change = max(

np.max(np.abs(alpha_new - alpha)),

np.max(np.abs(beta_new - beta)),

)

alpha = alpha_new

beta = beta_new

if max_change < tol:

converged = True

break

warn_if_not_converged(converged, "ImputationDiD iterative FE solver", max_iter, tol)

unit_fe = pd.Series(alpha, index=idx).groupby(unit_vals).first().to_dict()

time_fe = pd.Series(beta, index=idx).groupby(time_vals).first().to_dict()

return unit_fe, time_fe

@staticmethod

def _iterative_demean(

vals: np.ndarray,

unit_vals: np.ndarray,

time_vals: np.ndarray,

idx: pd.Index,

max_iter: int = 100,

tol: float = 1e-10,

weights: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Demean a vector by iterative alternating projection (unit + time FE removal).

result = vals.copy()

# Precompute per-group weight sums (invariant across iterations)

if weights is not None:

w_series = pd.Series(weights, index=idx)

wsum_t = w_series.groupby(time_vals).transform("sum").values

wsum_u = w_series.groupby(unit_vals).transform("sum").values

converged = False

with np.errstate(invalid="ignore", divide="ignore"):

for _ in range(max_iter):

if weights is not None:

wr_t = pd.Series(result * weights, index=idx)

time_means = wr_t.groupby(time_vals).transform("sum").values / wsum_t

else:

time_means = (

pd.Series(result, index=idx).groupby(time_vals).transform("mean").values

)

result_after_time = result - time_means

if weights is not None:

wr_u = pd.Series(result_after_time * weights, index=idx)

unit_means = wr_u.groupby(unit_vals).transform("sum").values / wsum_u

else:

unit_means = (