"""

Staggered Triple Difference (DDD) estimator.

Implements Ortiz-Villavicencio & Sant'Anna (2025) for staggered adoption

settings with an eligibility dimension, combining group-time DDD effects

via GMM-optimal weighting.

Core pairwise DiD computation matches R's triplediff::compute_did() exactly

(Riesz/Hajek normalization, separate M1/M3 OR corrections, hessian = (X'WX)^{-1}*n).

"""

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.linalg import (

_check_propensity_diagnostics,

_rank_guarded_inv,

solve_logit,

)

from diff_diff.staggered_aggregation import (

CallawaySantAnnaAggregationMixin,

)

from diff_diff.staggered_bootstrap import (

CallawaySantAnnaBootstrapMixin,

)

from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults

from diff_diff.utils import safe_inference

if TYPE_CHECKING:

from diff_diff.survey import SurveyDesign

__all__ = [

"StaggeredTripleDifference",

"StaggeredTripleDiffResults",

]

# Type alias for pre-computed structures

PrecomputedData = Dict[str, Any]

class StaggeredTripleDifference(

CallawaySantAnnaBootstrapMixin,

CallawaySantAnnaAggregationMixin,

BaseEstimator,

):

"""

Staggered Triple Difference (DDD) estimator.

Computes group-time average treatment effects ATT(g,t) for settings

with staggered adoption and a binary eligibility dimension, using the

three-DiD decomposition of Ortiz-Villavicencio & Sant'Anna (2025).

Multiple comparison groups are combined via GMM-optimal (inverse-variance)

weighting. Event study, group, and overall aggregations are supported.

Parameters

----------

estimation_method : str, default="dr"

Estimation method: "dr" (doubly robust), "ipw" (inverse probability

weighting), or "reg" (regression adjustment).

alpha : float, default=0.05

Significance level.

anticipation : int, default=0

Number of anticipation periods.

base_period : str, default="varying"

Base period selection: "varying" (consecutive comparisons) or

"universal" (always vs g-1-anticipation).

n_bootstrap : int, default=0

Number of multiplier bootstrap repetitions. 0 disables bootstrap.

bootstrap_weights : str, default="rademacher"

Bootstrap weight distribution: "rademacher", "mammen", or "webb".

seed : int or None, default=None

Random seed for reproducibility.

cband : bool, default=True

Whether to compute simultaneous confidence bands.

pscore_trim : float, default=0.01

Propensity score trimming bound.

cluster : str or None, default=None

Column name for cluster-robust standard errors.

rank_deficient_action : str, default="warn"

Action for rank-deficient design matrices: "warn", "error", "silent".

epv_threshold : float, default=10

Minimum events per variable for propensity score logistic regression.

A warning is emitted when EPV falls below this threshold.

pscore_fallback : str, default="error"

Action when propensity score estimation fails: "error" (raise) or

"unconditional" (fall back to unconditional propensity).

References

----------

Ortiz-Villavicencio, M. & Sant'Anna, P.H.C. (2025). "Better Understanding

Triple Differences Estimators." arXiv:2505.09942v3.

"""

def __init__(

self,

estimation_method: str = "dr",

control_group: str = "notyettreated",

alpha: float = 0.05,

anticipation: int = 0,

base_period: str = "varying",

n_bootstrap: int = 0,

bootstrap_weights: str = "rademacher",

seed: Optional[int] = None,

cband: bool = True,

pscore_trim: float = 0.01,

cluster: Optional[str] = None,

rank_deficient_action: str = "warn",

epv_threshold: float = 10,

pscore_fallback: str = "error",

):

if estimation_method not in ["dr", "ipw", "reg"]:

raise ValueError(

f"estimation_method must be 'dr', 'ipw', or 'reg', " f"got '{estimation_method}'"

)

if control_group not in ["nevertreated", "notyettreated"]:

raise ValueError(

f"control_group must be 'nevertreated' or 'notyettreated', "

f"got '{control_group}'"

)

if not (0 < pscore_trim < 0.5):

raise ValueError(f"pscore_trim must be in (0, 0.5), got {pscore_trim}")

if bootstrap_weights not in ["rademacher", "mammen", "webb"]:

raise ValueError(

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

f"got '{bootstrap_weights}'"

)

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 base_period not in ["varying", "universal"]:

raise ValueError(

f"base_period must be 'varying' or 'universal', " f"got '{base_period}'"

)

if epv_threshold <= 0:

raise ValueError(f"epv_threshold must be > 0, got {epv_threshold}")

if pscore_fallback not in ["error", "unconditional"]:

raise ValueError(

f"pscore_fallback must be 'error' or 'unconditional', " f"got '{pscore_fallback}'"

)

self.estimation_method = estimation_method

self.control_group = control_group

self.alpha = alpha

self.anticipation = anticipation

self.base_period = base_period

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

self.seed = seed

self.cband = cband

self.pscore_trim = pscore_trim

self.cluster = cluster

self.rank_deficient_action = rank_deficient_action

self.epv_threshold = epv_threshold

self.pscore_fallback = pscore_fallback

self.is_fitted_ = False

self.results_: Optional[StaggeredTripleDiffResults] = None

# get_params/set_params come from BaseEstimator.

# ------------------------------------------------------------------

# fit()

# ------------------------------------------------------------------

def fit(

self,

data: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

eligibility: str,

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

aggregate: Optional[str] = None,

balance_e: Optional[int] = None,

survey_design: Optional["SurveyDesign"] = None,

) -> StaggeredTripleDiffResults:

"""

Fit the staggered triple difference estimator.

Parameters

----------

data : pd.DataFrame

Panel data.

outcome : str

Outcome variable column name.

unit : str

Unit identifier column name.

time : str

Time period column name.

first_treat : str

Column with the enabling period for each unit's group.

Use 0 or np.inf for never-enabled units.

eligibility : str

Binary eligibility indicator column (0/1, time-invariant).

covariates : list of str, optional

Covariate column names.

aggregate : str, optional

Aggregation method: "event_study", "group", "simple", or "all".

balance_e : int, optional

Event time to balance on for event study.

survey_design : SurveyDesign, optional

Survey design specification for complex survey data. When

provided, uses survey weights for estimation (weighted Riesz

representers, weighted logit, weighted OLS) and design-based

variance for aggregated SEs (overall, event study, group) via

Taylor Series Linearization or replicate weights. Requires

``weight_type='pweight'``.

Returns

-------

StaggeredTripleDiffResults

"""

from diff_diff.survey import (

_resolve_survey_for_fit,

_validate_unit_constant_survey,

compute_survey_metadata,

)

resolved_survey, survey_weights, survey_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, data, "analytical")

)

if resolved_survey is not None:

_validate_unit_constant_survey(data, unit, survey_design)

if resolved_survey.weight_type != "pweight":

raise ValueError(

f"StaggeredTripleDifference survey support requires "

f"weight_type='pweight', got '{resolved_survey.weight_type}'. "

f"The survey variance math assumes probability weights."

)

if aggregate is not None and aggregate not in [

"event_study",

"group",

"simple",

"all",

]:

raise ValueError(

f"aggregate must be 'event_study', 'group', 'simple', or 'all', "

f"got '{aggregate}'"

)

df = data.copy()

self._validate_inputs(df, outcome, unit, time, first_treat, eligibility, covariates)

if self.cluster is not None:

warnings.warn(

"cluster parameter is accepted but cluster-robust analytical SEs "

"are not yet implemented for staggered DDD. Use n_bootstrap > 0 "

"for unit-level clustered inference via multiplier bootstrap.",

UserWarning,

stacklevel=2,

)

if first_treat != "first_treat":

df["first_treat"] = df[first_treat]

# Surface the inf → 0 recategorization the same way StaggeredDiD does

# (see `staggered.py:1508-1519`). Silently recoding inf would shift

# units between treated and never-treated pools with no signal

# (axis-E silent coercion under the Phase 2 audit).

_inf_mask = np.isposinf(df["first_treat"].values)

if _inf_mask.any():

n_inf_rows = int(_inf_mask.sum())

warnings.warn(

f"{n_inf_rows} row(s) have first_treat=inf; recoding to 0 "

f"(never-treated). Use first_treat=0 to suppress this warning.",

UserWarning,

stacklevel=2,

)

df["first_treat"] = df["first_treat"].replace([np.inf, float("inf")], 0)

precomputed = self._precompute_structures(

df,

outcome,

unit,

time,

eligibility,

covariates,

resolved_survey=resolved_survey,

)

# Recompute survey metadata from unit-level resolved survey

if resolved_survey is not None and survey_metadata is not None:

resolved_survey_unit = precomputed.get("resolved_survey_unit")

if resolved_survey_unit is not None:

unit_w = resolved_survey_unit.weights

survey_metadata = compute_survey_metadata(resolved_survey_unit, unit_w)

# Survey df for t-distribution critical values

df_survey = precomputed.get("df_survey")

if (

df_survey is None

and resolved_survey is not None

and hasattr(resolved_survey, "uses_replicate_variance")

and resolved_survey.uses_replicate_variance

):

df_survey = 0 # Forces NaN inference for undefined replicate df

treatment_groups = precomputed["treatment_groups"]

time_periods = precomputed["time_periods"]

all_units = precomputed["all_units"]

time_to_col = precomputed["time_to_col"]

unit_cohorts = precomputed["unit_cohorts"]

eligibility_per_unit = precomputed["eligibility_per_unit"]

n_units = len(all_units)

pscore_cache: Dict = {}

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

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

comparison_group_counts: Dict[Tuple, int] = {}

gmm_weights_store: Dict[Tuple, Dict] = {}

epv_diagnostics: Optional[Dict[Tuple, Dict[str, Any]]] = (

{} if (covariates and self.estimation_method in ("ipw", "dr")) else None

)

# Trackers for rank-deficient linalg solves across all (g, g_c, t)

# cells. `_compute_did_panel` appends to the OR-side tracker;

# `_compute_pscore` appends to the PS-side tracker. Both surface as

# ONE aggregate warning below rather than fanning out per cell.

self._lstsq_fallback_tracker: List[float] = []

self._ps_lstsq_fallback_tracker: List[float] = []

for g in treatment_groups:

# In universal mode, skip the reference period (t == g-1-anticipation)

# so it's omitted from GT estimation. The event-study mixin injects

# a synthetic reference row with effect=0, matching CS behavior.

if self.base_period == "universal":

universal_base = g - 1 - self.anticipation

valid_periods = [t for t in time_periods if t != universal_base]

else:

valid_periods = time_periods

for t in valid_periods:

base_period_val = self._get_base_period(g, t)

if base_period_val is None:

continue

if base_period_val not in time_to_col:

warnings.warn(

f"Base period {base_period_val} for (g={g}, t={t}) is "

"outside the observed panel. Skipping this cell.",

UserWarning,

stacklevel=2,

)

continue

if t not in time_to_col:

continue

has_never_enabled = bool(np.any(unit_cohorts == 0))

if self.control_group == "nevertreated":

# Only use never-enabled cohort as comparison

valid_gc = [0] if has_never_enabled else []

else:

# Use all valid comparison cohorts (not-yet-treated + never)

# Threshold accounts for anticipation: cohorts that start

# treatment within the anticipation window are contaminated.

nyt_threshold = max(t, base_period_val) + self.anticipation

valid_gc = [gc for gc in treatment_groups if gc > nyt_threshold and gc != g]

if has_never_enabled:

valid_gc = [0] + valid_gc

if not valid_gc:

warnings.warn(

f"No valid comparison groups for (g={g}, t={t}), skipping.",

UserWarning,

stacklevel=2,

)

continue

treated_mask = (unit_cohorts == g) & (eligibility_per_unit == 1)

n_treated = int(np.sum(treated_mask))

if n_treated == 0:

continue

att_vec = []

inf_raw = [] # unrescaled IFs

gc_labels = []

gc_cell_sizes = [] # size_gt_ctrl per surviving gc

for gc in valid_gc:

result = self._compute_ddd_gt_gc(

precomputed,

g,

gc,

t,

base_period_val,

covariates,

pscore_cache,

epv_diagnostics=epv_diagnostics,

)

if result is None:

continue

att_gc, inf_gc, size_gt_ctrl = result

if not np.isfinite(att_gc):

continue

att_vec.append(att_gc)

inf_raw.append(inf_gc)

gc_labels.append(gc)

gc_cell_sizes.append(size_gt_ctrl)

if not att_vec:

continue

# Compute size_gt from SURVIVING comparison cohorts only

# (not from all initially valid gc's)

surviving_units = treated_mask.copy()

for gc in gc_labels:

surviving_units |= (unit_cohorts == gc) | (unit_cohorts == g)

survey_w = precomputed.get("survey_weights")

if survey_w is not None:

size_gt = float(np.sum(survey_w[surviving_units]))

else:

size_gt = float(np.sum(surviving_units))

# Apply IF rescaling now that size_gt is known

inf_matrix = []

for inf_gc, size_gt_ctrl in zip(inf_raw, gc_cell_sizes):

if size_gt_ctrl > 0:

inf_gc = inf_gc * (size_gt / size_gt_ctrl)

inf_matrix.append(inf_gc)

att_gmm, inf_gmm, gmm_w, se_gt = self._combine_gmm(

np.array(att_vec),

np.array(inf_matrix),

n_units,

)

if not np.isfinite(att_gmm):

continue

# R's single-gc SE uses size_gt in denominator, not n_total.

# For multi-gc (GMM), the size_gt factor is already in Omega

# via the per-gc rescaling, so n_total is correct.

if len(gc_labels) == 1:

se_gt = float(np.sqrt(np.sum(inf_gmm**2) / size_gt**2))

if not np.isfinite(se_gt) or se_gt <= 0:

se_gt = np.nan

t_stat, p_value, conf_int = safe_inference(

att_gmm, se_gt, alpha=self.alpha, df=df_survey

)

# Rescale IF for mixin compatibility.

# R stores IF * (n/size_gt) in inf_func_mat, then uses

# SE = sqrt(sum(IF^2)/n^2) = sqrt(sum(psi^2)) with psi = IF/n.

# We need psi = IF_rescaled / n so mixin's sqrt(sum(psi^2)) works.

# IF is already at size_gt/size_gt_ctrl scale from above.

# Apply the final n/size_gt factor, then divide by n for mixin.

inf_gmm_rescaled = inf_gmm * (n_units / size_gt)

inf_gmm_scaled = inf_gmm_rescaled / n_units

# INVARIANT: np.where over boolean masks -> duplicate-free

# index arrays (fancy-+= scatter contract, see

# staggered_aggregation._combined_if_fast).

treated_idx = np.where(treated_mask)[0]

treated_inf = inf_gmm_scaled[treated_idx]

nonzero_mask = (inf_gmm_scaled != 0) & ~treated_mask

control_idx = np.where(nonzero_mask)[0]

control_inf = inf_gmm_scaled[control_idx]

n_control = int(np.sum(nonzero_mask))

group_time_effects[(g, t)] = {

"effect": att_gmm,

"se": se_gt,

"t_stat": t_stat,

"p_value": p_value,

"conf_int": conf_int,

"n_treated": n_treated,

"n_control": n_control,

}

influence_func_info[(g, t)] = {

"treated_idx": treated_idx,

"control_idx": control_idx,

"treated_inf": treated_inf,

"control_inf": control_inf,

}

comparison_group_counts[(g, t)] = len(gc_labels)

gmm_weights_store[(g, t)] = dict(zip(gc_labels, gmm_w.tolist()))

# Consolidated OR influence-function rank-deficiency warning.

# Finding #17 in the Phase 2 silent-failures audit: the per-pair OR

# solve at _compute_did_panel() previously fell back to lstsq with no

# signal, so near/fully singular X'WX in the covariate expansion went

# to the user as a normal result.

if self._lstsq_fallback_tracker and self.rank_deficient_action != "silent":

n_cells = len(self._lstsq_fallback_tracker)

finite_conds = [c for c in self._lstsq_fallback_tracker if np.isfinite(c)]

max_cond = max(finite_conds) if finite_conds else float("inf")

warnings.warn(

f"Rank-deficient X'WX detected in the outcome-regression "

f"influence-function step for {n_cells} (g, g_c, t) pair(s); "

f"dropped redundant direction(s) via a rank-guarded inverse. "

f"Max condition number of affected X'WX: {max_cond:.2e}. "

f"Standard errors use the identified covariate subset; consider "

f"dropping collinear covariates or using "

f"estimation_method='ipw' to avoid the OR projection.",

UserWarning,

stacklevel=2,

)

# Consolidated PS-Hessian rank-deficiency warning (sibling of the

# OR path above). `_compute_pscore` previously fell back from

# `np.linalg.inv(X'WX)` to `np.linalg.lstsq` with no signal, so

# a rank-deficient propensity-score design silently degraded

# IPW/DR influence-function corrections.

if self._ps_lstsq_fallback_tracker and self.rank_deficient_action != "silent":

n_cells = len(self._ps_lstsq_fallback_tracker)

finite_conds = [c for c in self._ps_lstsq_fallback_tracker if np.isfinite(c)]

max_cond = max(finite_conds) if finite_conds else float("inf")

warnings.warn(

f"Rank-deficient X'WX detected in the propensity-score "

f"Hessian for {n_cells} (g, g_c, t) pair(s); dropped redundant "

f"direction(s) via a rank-guarded inverse. Max condition number "

f"of affected X'WX: {max_cond:.2e}. IPW/DR influence-function "

f"corrections use the identified covariate subset; consider "

f"dropping collinear propensity-score covariates or using "

f"estimation_method='reg' to avoid the PS path.",

UserWarning,

stacklevel=2,

)

# Consolidated EPV summary warning

if epv_diagnostics:

low_epv = {k: v for k, v in epv_diagnostics.items() if v.get("is_low")}

if low_epv:

n_affected = len(low_epv)

n_total = len(epv_diagnostics)

min_entry = min(low_epv.values(), key=lambda v: v["epv"])

min_g = min(low_epv.keys(), key=lambda k: low_epv[k]["epv"])

warnings.warn(

f"Low Events Per Variable (EPV) detected in "

f"{n_affected} of {n_total} cohort-time cell(s). "

f"Minimum EPV: {min_entry['epv']:.1f} (cohort g={min_g[0]}). "

f"Consider estimation_method='reg' or fewer covariates. "

f"Call results.epv_summary() for per-cohort details.",

UserWarning,

stacklevel=2,

)

if not group_time_effects:

raise ValueError(

"No valid group-time effects could be computed. "

"Check that the data has sufficient variation in treatment "

"timing and eligibility."

)

# For aggregation: use eligible-treated-only cohort assignments so

# WIF weights match the point estimate weights (n_treated per cohort,

# i.e. P(S=g, Q=1)). This matches the paper's Eq 4.13 which defines

# aggregation weights over the treated population (G_i defined only

# for Q=1 units). Ineligible units get cohort=0 so they don't

# contribute to pg for any treatment group.

# Both precomputed["unit_cohorts"] AND df["first_treat"] must be

# zeroed for ineligible units because the WIF code reads both.

precomputed_agg = dict(precomputed)

cohorts_for_agg = precomputed["unit_cohorts"].copy()

cohorts_for_agg[eligibility_per_unit == 0] = 0

precomputed_agg["unit_cohorts"] = cohorts_for_agg

df_agg = df.copy()

df_agg.loc[df_agg[eligibility] == 0, "first_treat"] = 0

# Overall ATT via aggregation mixin

overall_att, overall_se, overall_effective_df = self._aggregate_simple(

group_time_effects, influence_func_info, df_agg, unit, precomputed_agg

)

# Preserve the ORIGINAL survey df before the simple-overall statistic mutates

# it below. The Eq. 4.14 overall (overall_att_es) must fall back to this

# original df, never to the simple overall's per-statistic replicate df — the

# two statistics can drop different replicate subsets, so reusing the simple

# overall's df would silently give overall_att_es the wrong p-value/CI.

df_survey_original = df_survey

# Use per-statistic effective df from replicate aggregation if available;

# otherwise fall back to the original df from the survey design.

if overall_effective_df is not None:

df_survey = overall_effective_df

if survey_metadata is not None:

survey_metadata.df_survey = df_survey

overall_t_stat, overall_p_value, overall_conf_int = safe_inference(

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

)

# Aggregations

event_study_effects = None

group_effects = None

es_aggregation = None

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

es_aggregation = self._aggregate_event_study(

group_time_effects,

influence_func_info,

treatment_groups,

time_periods,

balance_e,

df_agg,

unit,

precomputed_agg,

)

event_study_effects = es_aggregation.effects

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

group_effects = self._aggregate_by_group(

group_time_effects,

influence_func_info,

treatment_groups,

precomputed_agg,

df_agg,

unit,

)

# Paper Eq. (4.14) overall ATT (event-study average): an opt-in summary

# alongside the default CS-simple ``overall_att``. ``_aggregate_event_study``

# RETURNS it on its aggregation object (it used to stash it on

# ``self._event_study_overall``); populated only when the event-study

# aggregation ran. Analytical inference here; the bootstrap block below

# overrides the SE when ``n_bootstrap > 0`` (mirroring ``overall_se``).

overall_att_es = None

overall_se_es = None

overall_t_stat_es = None

overall_p_value_es = None

overall_conf_int_es = None

# Whether the ANALYTICAL Eq. 4.14 SE was non-finite while its point estimate

# was finite (i.e. a contributing horizon's influence function was non-finite).

# Captured before any bootstrap override so the terminal warning below does not

# misdiagnose a bootstrap-side NaN SE (e.g. cluster-unidentified) as an

# analytical-IF failure (the bootstrap path emits its own warning).

analytical_overall_es_se_nonfinite = False

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

es_overall = es_aggregation.overall if es_aggregation is not None else None

if es_overall is not None:

overall_att_es = es_overall["att"]

overall_se_es = es_overall["se"]

analytical_overall_es_se_nonfinite = bool(

np.isfinite(overall_att_es) and not np.isfinite(overall_se_es)

)

es_eff_df = es_overall.get("effective_df")

# Fall back to the ORIGINAL survey df, not the simple-overall's mutated

# per-statistic df (P1 fix): overall_att_es has its own replicate df.

df_for_es = es_eff_df if es_eff_df is not None else df_survey_original

overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(

overall_att_es, overall_se_es, alpha=self.alpha, df=df_for_es

)

else:

# Event-study aggregation was requested but yielded no post-treatment

# horizon: this is "requested but undefined" -> NaN + warning (the

# library's overall-aggregation contract, matching _aggregate_simple),

# distinct from "not requested" which leaves the fields None.

warnings.warn(

"Event-study aggregation was requested but no post-treatment "

"horizons are available for the Eq. 4.14 overall (overall_att_es); "

"returning NaN.",

UserWarning,

stacklevel=2,

)

overall_att_es = np.nan

overall_se_es = np.nan

overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(

np.nan, np.nan, alpha=self.alpha, df=df_survey

)

# Reject replicate-weight designs for bootstrap — replicate variance

# is an analytical alternative, not compatible with bootstrap

if (

self.n_bootstrap > 0

and resolved_survey is not None

and hasattr(resolved_survey, "uses_replicate_variance")

and resolved_survey.uses_replicate_variance

):

raise NotImplementedError(

"StaggeredTripleDifference bootstrap (n_bootstrap > 0) is not "

"supported with replicate-weight survey designs. Replicate "

"weights provide analytical variance; use n_bootstrap=0 instead."

)

# Bootstrap

bootstrap_results = None

cband_crit_value = None

if self.n_bootstrap > 0:

bootstrap_results = self._run_multiplier_bootstrap(

group_time_effects,

influence_func_info,

aggregate,

balance_e,

treatment_groups,

time_periods,

df_agg,

unit,

precomputed_agg,

self.cband,

)

if bootstrap_results is not None:

overall_se = bootstrap_results.overall_att_se

overall_t_stat, overall_p_value, overall_conf_int = safe_inference(

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

)

overall_conf_int = bootstrap_results.overall_att_ci

overall_p_value = bootstrap_results.overall_att_p_value

# Mirror the override for the Eq. (4.14) event-study-average overall

# (only when event-study aggregation produced it). A NaN bootstrap SE

# (e.g. cluster-unidentified) correctly NaNs the inference.

if overall_att_es is not None and bootstrap_results.overall_att_es_se is not None:

overall_se_es = bootstrap_results.overall_att_es_se

overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(

overall_att_es, overall_se_es, alpha=self.alpha, df=df_survey

)

overall_conf_int_es = bootstrap_results.overall_att_es_ci

overall_p_value_es = bootstrap_results.overall_att_es_p_value

if bootstrap_results.cband_crit_value is not None:

cband_crit_value = bootstrap_results.cband_crit_value

# Update group-time effects with bootstrap SEs

if bootstrap_results.group_time_ses:

for gt_key in group_time_effects:

if gt_key in bootstrap_results.group_time_ses:

group_time_effects[gt_key]["se"] = bootstrap_results.group_time_ses[

gt_key

]

group_time_effects[gt_key]["conf_int"] = (

bootstrap_results.group_time_cis[gt_key]

)

group_time_effects[gt_key]["p_value"] = (

bootstrap_results.group_time_p_values[gt_key]

)

t_val, _, _ = safe_inference(

group_time_effects[gt_key]["effect"],

bootstrap_results.group_time_ses[gt_key],

alpha=self.alpha,

df=df_survey,

)

group_time_effects[gt_key]["t_stat"] = t_val

if event_study_effects and bootstrap_results.event_study_ses:

for e_key in event_study_effects:

if e_key in bootstrap_results.event_study_ses:

# ses/cis/p_values are populated together.

assert (

bootstrap_results.event_study_cis is not None

and bootstrap_results.event_study_p_values is not None

)

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

e_key

]

event_study_effects[e_key]["conf_int"] = (

bootstrap_results.event_study_cis[e_key]

)

event_study_effects[e_key]["p_value"] = (

bootstrap_results.event_study_p_values[e_key]

)

t_val, _, _ = safe_inference(

event_study_effects[e_key]["effect"],

bootstrap_results.event_study_ses[e_key],

alpha=self.alpha,

df=df_survey,

)

event_study_effects[e_key]["t_stat"] = t_val

if cband_crit_value is not None:

bs_se = bootstrap_results.event_study_ses[e_key]

eff = event_study_effects[e_key]["effect"]

event_study_effects[e_key]["cband_conf_int"] = (

eff - cband_crit_value * bs_se,

eff + cband_crit_value * bs_se,

)

# Update group effects with bootstrap SEs

if (

group_effects

and bootstrap_results.group_effect_ses is not None

and bootstrap_results.group_effect_cis is not None

and bootstrap_results.group_effect_p_values is not None

):

grp_keys = [g for g in group_effects if g in bootstrap_results.group_effect_ses]

for g_key in grp_keys:

group_effects[g_key]["se"] = bootstrap_results.group_effect_ses[g_key]

group_effects[g_key]["conf_int"] = bootstrap_results.group_effect_cis[g_key]

group_effects[g_key]["p_value"] = bootstrap_results.group_effect_p_values[

g_key

]

t_val, _, _ = safe_inference(

group_effects[g_key]["effect"],

bootstrap_results.group_effect_ses[g_key],

alpha=self.alpha,

df=df_survey,

)

group_effects[g_key]["t_stat"] = t_val

# Bootstrap se/p/CI replaced the analytical ones, which

# is what the retained df described - keeping it would

# claim a t-reference that governed nothing.

group_effects[g_key]["df_used"] = None

# Eq. 4.14 overall: an ANALYTICAL non-finite SE under a finite point estimate

# (a contributing horizon's influence function is non-finite, or the variance is

# unidentified — e.g. a single-PSU/cluster design). Surface it — never NaN the SE

# silently. Gated on the analytical-origin flag (captured before the bootstrap

# override) and the final state still being non-finite: a bootstrap that supplies a

# finite SE rescues it (no warning), and a bootstrap that NaNs the SE for unrelated

# reasons (e.g. cluster-unidentified) is reported by the bootstrap's own warning,

# not misdiagnosed here as an analytical-IF failure.

if (

analytical_overall_es_se_nonfinite

and overall_se_es is not None

and not np.isfinite(overall_se_es)

):

warnings.warn(

"Eq. 4.14 overall (overall_att_es) point estimate is defined but its "

"standard error is undefined (NaN): either a contributing post-treatment "

"event-study horizon has a non-finite influence function, or the variance "

"is unidentified (e.g. a single-PSU/cluster survey design). overall_se_es "

"and its inference fields are NaN.",

UserWarning,

stacklevel=2,

)

n_treated_units = int(np.sum((unit_cohorts > 0) & (eligibility_per_unit == 1)))

n_control_units = n_units - n_treated_units

n_never_enabled = int(np.sum(unit_cohorts == 0))

n_eligible = int(np.sum(eligibility_per_unit == 1))

n_ineligible = int(np.sum(eligibility_per_unit == 0))

self.results_ = StaggeredTripleDiffResults(

group_time_effects=group_time_effects,

overall_att=overall_att,

overall_se=overall_se,

overall_t_stat=overall_t_stat,

overall_p_value=overall_p_value,

overall_conf_int=overall_conf_int,

groups=treatment_groups,

time_periods=time_periods,

n_obs=len(df),

n_treated_units=n_treated_units,

n_control_units=n_control_units,

n_never_enabled=n_never_enabled,

n_eligible=n_eligible,

n_ineligible=n_ineligible,

alpha=self.alpha,

control_group=self.control_group,

base_period=self.base_period,

anticipation=self.anticipation,

estimation_method=self.estimation_method,

event_study_effects=event_study_effects,

group_effects=group_effects,

bootstrap_results=bootstrap_results,

cband_crit_value=cband_crit_value,

pscore_trim=self.pscore_trim,

survey_metadata=survey_metadata,

comparison_group_counts=comparison_group_counts,

gmm_weights=gmm_weights_store,

epv_diagnostics=epv_diagnostics if epv_diagnostics else None,

epv_threshold=self.epv_threshold,

pscore_fallback=self.pscore_fallback,

overall_att_es=overall_att_es,

overall_se_es=overall_se_es,

overall_t_stat_es=overall_t_stat_es,

overall_p_value_es=overall_p_value_es,

overall_conf_int_es=overall_conf_int_es,

)

self.is_fitted_ = True

return self.results_

# ------------------------------------------------------------------

# Validation

# ------------------------------------------------------------------

def _validate_inputs(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

eligibility: str,

covariates: Optional[List[str]],

) -> None:

"""Validate input data."""

required_cols = [outcome, unit, time, first_treat, eligibility]

if covariates:

required_cols.extend(covariates)

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

if missing:

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

elig_vals = df[eligibility].dropna().unique()

if not set(elig_vals).issubset({0, 1, 0.0, 1.0}):

raise ValueError(

f"Eligibility column '{eligibility}' must be binary (0/1). "

f"Found values: {sorted(elig_vals)}"

)

elig_by_unit = df.groupby(unit)[eligibility].nunique()

varying = elig_by_unit[elig_by_unit > 1]

if len(varying) > 0:

raise ValueError(

f"Eligibility must be time-invariant within units. "

f"Found {len(varying)} units with varying eligibility."

)

for col in [outcome, first_treat, eligibility]:

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

raise ValueError(f"Column '{col}' contains missing values.")

# Reject non-finite outcomes (Inf/-Inf)

if not np.all(np.isfinite(df[outcome])):

raise ValueError(

f"Column '{outcome}' contains non-finite values (Inf/-Inf). "

"All outcome values must be finite."

)

# Reject non-finite covariates

if covariates:

for cov in covariates:

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

raise ValueError(f"Covariate '{cov}' contains missing values.")

if not np.all(np.isfinite(df[cov])):

raise ValueError(f"Covariate '{cov}' contains non-finite values.")

if df[eligibility].nunique() < 2:

raise ValueError(

"Need both eligible (Q=1) and ineligible (Q=0) units. "

f"Only found Q={df[eligibility].unique()[0]}."

)

# Check unique (unit, time) pairs — no duplicate rows

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

if dup.any():

raise ValueError(

f"Duplicate (unit, time) rows found. "

f"{int(dup.sum())} duplicates detected. Panel must have unique rows."

)

# Check balanced panel — every unit observed in exactly the global period set

global_periods = set(df[time].unique())

n_global_periods = len(global_periods)

unit_period_sets = df.groupby(unit)[time].apply(set)

mismatched = unit_period_sets[unit_period_sets != global_periods]

if len(mismatched) > 0:

raise ValueError(

"Unbalanced panel detected. All units must be observed in "

f"all {n_global_periods} periods. "

f"Found {len(mismatched)} units with different period sets."

)

# Check time-invariant first_treat

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

varying_ft = ft_by_unit[ft_by_unit > 1]

if len(varying_ft) > 0:

raise ValueError(

f"first_treat must be time-invariant within units. "

f"Found {len(varying_ft)} units with varying first_treat."

)

# Check time-invariant covariates

if covariates:

for cov in covariates:

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

varying_cov = cov_nunique[cov_nunique > 1]

if len(varying_cov) > 0:

raise ValueError(

f"Covariate '{cov}' must be time-invariant within units. "

f"Found {len(varying_cov)} units with varying values."

)

# ------------------------------------------------------------------

# Precomputation

# ------------------------------------------------------------------

def _precompute_structures(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

eligibility: str,

covariates: Optional[List[str]],

resolved_survey=None,

) -> PrecomputedData:

"""Build precomputed structures for efficient computation."""

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

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

n_units = len(all_units)

n_periods = len(time_periods)

unit_to_idx = {u: i for i, u in enumerate(all_units)}

time_to_col = {t: j for j, t in enumerate(time_periods)}

outcome_matrix = np.full((n_units, n_periods), np.nan)

for _, row in df.iterrows():

u_idx = unit_to_idx[row[unit]]

t_idx = time_to_col[row[time]]

outcome_matrix[u_idx, t_idx] = row[outcome]

unit_df = df.groupby(unit).first().reindex(all_units)

unit_cohorts = unit_df["first_treat"].values.astype(float)

eligibility_per_unit = unit_df[eligibility].values.astype(int)

treatment_groups = sorted([g for g in np.unique(unit_cohorts) if g > 0])