"""

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,

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}'"

)

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.pretrends = pretrends

self.is_fitted_ = False

self.results_: Optional[TwoStageDiDResults] = 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,

) -> TwoStageDiDResults:

"""

# ---- Data validation ----

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}")

# 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_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, data, "analytical")

)

_uses_replicate_ts = resolved_survey is not None and resolved_survey.uses_replicate_variance

if _uses_replicate_ts 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"TwoStageDiD 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_stratified_meat_from_psu_scores()

# in _compute_gmm_variance().

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

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

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

# Validate absorbing treatment

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"TwoStageDiD 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,

)

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

min_time = df[time].min()

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

always_treated_units = df.loc[always_treated_mask, unit].unique()

n_always_treated = len(always_treated_units)

if n_always_treated > 0:

unit_list = ", ".join(str(u) for u in always_treated_units[:10])

suffix = f" (and {n_always_treated - 10} more)" if n_always_treated > 10 else ""

survey_note = ""

if survey_weights is not None or resolved_survey is not None:

survey_note = " Associated survey weights and design arrays " "adjusted to match."

warnings.warn(

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

f"(first_treat <= {min_time}): [{unit_list}{suffix}]. "

"These units have no untreated observations and cannot contribute "

f"to the counterfactual model. Excluding from estimation.{survey_note}",

UserWarning,

stacklevel=2,

)

df = df[~df[unit].isin(always_treated_units)].copy()

# Subset survey arrays to match filtered df

if survey_weights is not None:

keep_mask = ~data[unit].isin(always_treated_units)

survey_weights = survey_weights[keep_mask.values]

if resolved_survey is not None:

keep_mask = ~data[unit].isin(always_treated_units)

resolved_survey = replace(

resolved_survey,

weights=resolved_survey.weights[keep_mask.values],

strata=(

resolved_survey.strata[keep_mask.values]

if resolved_survey.strata is not None

else None

),

psu=(

resolved_survey.psu[keep_mask.values]

if resolved_survey.psu is not None

else None

),

fpc=(

resolved_survey.fpc[keep_mask.values]

if resolved_survey.fpc is not None

else None

),

replicate_weights=(

resolved_survey.replicate_weights[keep_mask.values]

if resolved_survey.replicate_weights is not None

else None

),

)

# Recompute n_psu/n_strata after subsetting

new_n_psu = (

len(np.unique(resolved_survey.psu)) if resolved_survey.psu is not None else 0

)

new_n_strata = (

len(np.unique(resolved_survey.strata))

if resolved_survey.strata is not None

else 0

)

resolved_survey = replace(resolved_survey, n_psu=new_n_psu, n_strata=new_n_strata)

# Recompute survey_metadata since it depends on these counts

from diff_diff.survey import compute_survey_metadata

raw_w = (

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

if survey_design.weights

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

)

survey_metadata = compute_survey_metadata(resolved_survey, raw_w)

# Treatment indicator with anticipation

effective_treat = df[first_treat] - self.anticipation

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

# Partition into 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.")

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

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

# GMM variance (PSU overrides unit-level clustering)

if resolved_survey.psu is not None:

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 = (

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

if survey_design.weights

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

)

survey_metadata = compute_survey_metadata(resolved_survey, raw_w)

# Relative time

df["_rel_time"] = np.where(

~df["_never_treated"],

df[time] - df[first_treat],

np.nan,

)

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

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

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)

# ---- Residualize ALL observations ----

y_tilde = self._residualize(

df, outcome, unit, time, covariates, unit_fe, time_fe, grand_mean, delta_hat

)

df["_y_tilde"] = y_tilde

# ---- Stage 2: OLS of y_tilde on treatment indicators ----

# Build design matrices and compute effects + GMM variance

ref_period = -1 - self.anticipation

# 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

if _uses_replicate_ts and _survey_df is None:

_survey_df = 0

# Always compute overall ATT (static specification)

overall_att, overall_se = self._stage2_static(

df=df,

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,

survey_weights=survey_weights,

survey_weight_type=survey_weight_type,

resolved_survey=(resolved_survey if not _uses_replicate_ts else None),

)

# Compute overall ATT 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._stage2_event_study(

df=df,

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,

ref_period=ref_period,

balance_e=balance_e,

kept_cov_mask=kept_cov_mask,

survey_weights=survey_weights,

survey_weight_type=survey_weight_type,

survey_df=_survey_df,

resolved_survey=(resolved_survey if not _uses_replicate_ts else None),

)

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

group_effects = self._stage2_group(

df=df,

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_weight_type=survey_weight_type,

survey_df=_survey_df,

resolved_survey=(resolved_survey if not _uses_replicate_ts else None),

)

# Replicate variance override: derive keys from actual outputs, then refit

_n_valid_rep_ts = None

_vcov_rep_ts = None

if _uses_replicate_ts:

from diff_diff.survey import compute_replicate_refit_variance

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

_sorted_es_periods_ts = sorted(

e

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

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

)

_sorted_groups_ts = sorted(

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

)

_n_es_ts = len(_sorted_es_periods_ts)

_n_grp_ts = len(_sorted_groups_ts)

# Build full-sample estimate from actual outputs

_full_est_ts = [overall_att]

_full_est_ts.extend([event_study_effects[e]["effect"] for e in _sorted_es_periods_ts])

_full_est_ts.extend([group_effects[g]["effect"] for g in _sorted_groups_ts])

def _refit_ts(w_r):

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

df,

outcome,

unit,

time,

covariates,

omega_0_mask,

weights=w_r,

)

y_tilde_r = self._residualize(

df,

outcome,

unit,

time,

covariates,

ufe_r,

tfe_r,

gm_r,

delta_r,

)

df_tmp = df.copy()

df_tmp["_y_tilde"] = y_tilde_r

results = []

att_r, _ = self._stage2_static(

df=df_tmp,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=ufe_r,

time_fe=tfe_r,

grand_mean=gm_r,

delta_hat=delta_r,

cluster_var=cluster_var,

kept_cov_mask=kcm_r,

survey_weights=w_r,

survey_weight_type="pweight",

)

results.append(att_r)

if _sorted_es_periods_ts:

es_r = self._stage2_event_study(

df=df_tmp,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=ufe_r,

time_fe=tfe_r,

grand_mean=gm_r,

delta_hat=delta_r,

cluster_var=cluster_var,

treatment_groups=treatment_groups,

ref_period=ref_period,

balance_e=balance_e,

kept_cov_mask=kcm_r,

survey_weights=w_r,

survey_weight_type="pweight",

survey_df=None,

)

for e in _sorted_es_periods_ts:

results.append(es_r[e]["effect"] if e in es_r else np.nan)

if _sorted_groups_ts:

grp_r = self._stage2_group(

df=df_tmp,

unit=unit,

time=time,

first_treat=first_treat,

covariates=covariates,

omega_0_mask=omega_0_mask,

omega_1_mask=omega_1_mask,

unit_fe=ufe_r,

time_fe=tfe_r,

grand_mean=gm_r,

delta_hat=delta_r,

cluster_var=cluster_var,

treatment_groups=treatment_groups,

kept_cov_mask=kcm_r,

survey_weights=w_r,

survey_weight_type="pweight",

survey_df=None,

)

for g in _sorted_groups_ts:

results.append(grp_r[g]["effect"] if g in grp_r else np.nan)

return np.array(results)

_vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance(

_refit_ts, np.array(_full_est_ts), resolved_survey

)

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

# Override df if replicates were dropped

if _n_valid_rep_ts < resolved_survey.n_replicates:

_survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0

if survey_metadata is not None:

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

# Recompute overall inference with replicate SE/df

overall_t, overall_p, overall_ci = safe_inference(

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

)

# Override event-study SEs (only for identified effects)

for i, e in enumerate(_sorted_es_periods_ts):

if event_study_effects is not None and e in event_study_effects:

se_e = float(np.sqrt(max(_vcov_rep_ts[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 (only for identified effects)

for j, g in enumerate(_sorted_groups_ts):

if group_effects is not None and g in group_effects:

se_g = float(

np.sqrt(max(_vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + 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, "_y_tilde", "_rel_time"]].copy()

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

tau_finite = treated_df["tau_hat"].notna() & np.isfinite(treated_df["tau_hat"].values)

n_valid_te = int(tau_finite.sum())

if n_valid_te > 0:

if survey_weights is not None:

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

# ---- Bootstrap ----

bootstrap_results = None

if self.n_bootstrap > 0:

try:

bootstrap_results = self._run_bootstrap(

df=df,

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,

treatment_groups=treatment_groups,

ref_period=ref_period,

balance_e=balance_e,

original_att=overall_att,

original_event_study=event_study_effects,

original_group=group_effects,

aggregate=aggregate,

resolved_survey=resolved_survey,

)

except NotImplementedError:

raise # Don't swallow explicit rejections (e.g. lonely_psu="adjust")

except Exception as e:

warnings.warn(

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

UserWarning,

stacklevel=2,

)

if bootstrap_results is not None:

# 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_ = TwoStageDiDResults(

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,

survey_metadata=survey_metadata,

)

self.is_fitted_ = True

return self.results_

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

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

beta = np.zeros(n)

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

)

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, "TwoStageDiD 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()

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 = (

pd.Series(result_after_time, index=idx)

.groupby(unit_vals)

.transform("mean")

.values

)

result_new = result_after_time - unit_means

if np.max(np.abs(result_new - result)) < tol:

result = result_new

converged = True

break

result = result_new

warn_if_not_converged(converged, "TwoStageDiD iterative demean", max_iter, tol)

return result

def _fit_untreated_model(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

weights: Optional[np.ndarray] = None,

) -> Tuple[

Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray]

]:

"""