"""

Staggered Difference-in-Differences estimators.

Implements modern methods for DiD with variation in treatment timing,

including the Callaway-Sant'Anna (2021) estimator.

"""

import bisect

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

AggregationKit,

)

from diff_diff.linalg import (

_check_propensity_diagnostics,

_detect_rank_deficiency,

_equilibrated_lstsq,

_format_dropped_columns,

_rank_guarded_inv,

solve_logit,

solve_ols,

)

from diff_diff.staggered_aggregation import (

CallawaySantAnnaAggregationMixin,

)

from diff_diff.staggered_bootstrap import (

CallawaySantAnnaBootstrapMixin,

CSBootstrapResults,

)

# Import from split modules

from diff_diff.staggered_results import (

CallawaySantAnnaResults,

GroupTimeEffect,

)

from diff_diff.utils import safe_inference, safe_inference_batch

if TYPE_CHECKING:

from diff_diff.survey import SurveyDesign

# Re-export for backward compatibility

__all__ = [

"CallawaySantAnna",

"CallawaySantAnnaResults",

"CSBootstrapResults",

"GroupTimeEffect",

]

# Type alias for pre-computed structures

PrecomputedData = Dict[str, Any]

class _DeprecatedFitArg:

"""Sentinel for `fit(aggregate=)` / `fit(balance_e=)` (rows M-020 / M-117).

A plain ``None`` default cannot distinguish "not passed" from "passed

None", so a bare ``None`` default would fire the FutureWarning on EVERY

fit. The warning must fire only when the caller actually supplies the

argument.

"""

def __repr__(self) -> str: # pragma: no cover - debugging aid

return "<not supplied>"

_DEPRECATED_FIT_ARG = _DeprecatedFitArg()

def _linear_regression(

X: np.ndarray,

y: np.ndarray,

rank_deficient_action: str = "warn",

weights: Optional[np.ndarray] = None,

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

"""

Fit OLS regression.

Parameters

----------

X : np.ndarray

Feature matrix (n_samples, n_features). Intercept added automatically.

y : np.ndarray

Outcome variable.

rank_deficient_action : str, default "warn"

Action when design matrix is rank-deficient:

- "warn": Issue warning and drop linearly dependent columns (default)

- "error": Raise ValueError

- "silent": Drop columns silently without warning

weights : np.ndarray, optional

Observation weights for WLS. When None, OLS is used.

Returns

-------

beta : np.ndarray

Fitted coefficients (including intercept).

residuals : np.ndarray

Residuals from the fit.

"""

n = X.shape[0]

# Add intercept

X_with_intercept = np.column_stack([np.ones(n), X])

# Use unified OLS backend (no vcov needed)

beta, residuals, _ = solve_ols(

X_with_intercept,

y,

return_vcov=False,

rank_deficient_action=rank_deficient_action,

weights=weights,

)

return beta, residuals

def _cluster_robust_se_from_per_gt_if(

inf_info: Dict[str, Any],

resolved_survey: "Any",

) -> Optional[float]:

"""CR1 Liang-Zeger cluster-robust SE for a single (g,t) ATT.

Builds the per-(g,t) per-index IF vector from ``inf_info`` and routes

through ``compute_survey_if_variance`` so that the per-cell variance

inherits the SAME design-based machinery as the aggregate path:

V = sum_h (1 - f_h) * (n_h / (n_h - 1)) * sum_j (psi_hj - psi_h_bar)^2

where ``psi_hj = sum_{i in PSU j, stratum h} psi_i``. This matches

the documented CR1 contract in REGISTRY.md (synthesized

``SurveyDesign(psu=cluster)`` → ``_compute_stratified_psu_meat``)

and applies the G/(G-1) finite-sample correction, PSU centering,

FPC, and lonely-PSU handling uniformly with overall / event-study

inference.

For the panel path, ``resolved_survey`` is ``resolved_survey_unit``

(length n_units) and the IF index space is per-unit. For the RCS

path, ``resolved_survey`` is the per-obs ``resolved_survey`` (length

n_obs). The helper is index-space agnostic — it just requires

``treated_idx`` / ``control_idx`` in ``inf_info`` to be valid

offsets into ``resolved_survey.psu``.

Return contract (callers depend on this distinction):

* **float SE** — finite cluster-robust variance; caller uses it.

* **NaN** — ``compute_survey_if_variance`` returned NaN (clustered

variance unidentified, e.g., G<2 or lonely-PSU removed all strata).

Caller MUST propagate this NaN through to ``safe_inference`` so

the per-cell inference surface (se / t_stat / p_value / conf_int)

is NaN-consistent — NEVER fall back to the unit-level SE. Falling

back would silently report a different estimator's variance under

a clustered request (``feedback_no_silent_failures``).

* **None** — malformed inputs or invariant violations:

``inf_info`` lacks required IF fields, ``resolved_survey.psu`` is

None, index alignment cannot be verified, or

``compute_survey_if_variance`` returned a negative variance. In

these cases the helper cannot evaluate the contract; caller falls

back to the unit-level SE returned by the underlying estimation

method (no PSU is in play, so unit-level is the documented default).

"""

if (

inf_info is None

or "treated_inf" not in inf_info

or "control_inf" not in inf_info

or "treated_idx" not in inf_info

or "control_idx" not in inf_info

):

return None

treated_idx = np.asarray(inf_info["treated_idx"])

control_idx = np.asarray(inf_info["control_idx"])

treated_inf = np.asarray(inf_info["treated_inf"])

control_inf = np.asarray(inf_info["control_inf"])

psu_array = getattr(resolved_survey, "psu", None)

if psu_array is None:

return None

n = len(psu_array)

if (

treated_idx.size > 0

and (treated_idx.max(initial=-1) >= n or treated_idx.min(initial=0) < 0)

) or (

control_idx.size > 0

and (control_idx.max(initial=-1) >= n or control_idx.min(initial=0) < 0)

):

return None

# Index arrays are unique within each cell by construction at every

# producer (np.where on disjoint masks), so fancy += is exact — same

# scatter contract as staggered_aggregation._combined_if_fast, without

# np.add.at's unbuffered-ufunc overhead (runs once per (g,t) cell when

# cluster= is set).

psi_per_index = np.zeros(n)

if treated_idx.size:

psi_per_index[treated_idx] += treated_inf

if control_idx.size:

psi_per_index[control_idx] += control_inf

# Route through the shared survey helper so the per-cell variance

# gets the same G/(G-1) finite-sample correction, PSU centering,

# FPC handling, and lonely-PSU/G<2→NaN behavior as overall +

# event-study inference (per the documented CR1 contract).

from diff_diff.survey import compute_survey_if_variance

var = compute_survey_if_variance(psi_per_index, resolved_survey)

# Return contract:

# float SE → use it (finite cluster-robust variance)

# NaN → propagate NaN so the caller can NaN-out the inference

# surface rather than silently falling back to the

# unit-level SE (per feedback_no_silent_failures: when

# clustered variance is undefined — e.g., G<2, lonely-PSU

# removed all strata — the user-facing per-cell SE must

# reflect that, not silently revert to a different

# estimator).

# None → malformed (negative variance or other invariant

# violation); caller falls back to the unit-level SE.

if np.isnan(var):

return float("nan")

if var < 0:

return None

return float(np.sqrt(var))

def _safe_inv(

A: np.ndarray,

tracker: Optional[list] = None,

) -> np.ndarray:

"""Rank-guarded generalized inverse of a Gram matrix for analytical SE paths.

Parameters

----------

A : np.ndarray

Square matrix to invert.

tracker : list, optional

When provided, one condition-number sample of ``A`` is appended each

time ``A`` is rank-deficient (near-singular). ``CallawaySantAnna.fit()``

initializes a list and emits a single aggregate `UserWarning` after the

fit finishes, rather than surfacing a separate warning per fallback.

Sibling of finding #17 in the Phase 2 silent-failures audit.

Notes

-----

Delegates to :func:`~diff_diff.linalg._rank_guarded_inv`, which is the sole

owner of the ``tracker`` append. The old ``except LinAlgError: lstsq``

fallback only caught *exactly* singular matrices; a *near*-singular Gram

(e.g. a constant/collinear covariate) returned a garbage inverse (~1e13)

that flowed into the SE. The rank-guarded inverse truncates redundant

directions (finite SE on the identified subset) and returns an all-NaN

matrix only on true rank-0.

"""

inv, _, _ = _rank_guarded_inv(A, tracker=tracker)

return inv

def _nan_gt_entry(

n_treated: int = 0,

n_control: int = 0,

skip_reason: Optional[str] = None,

survey_weight_sum: Optional[float] = None,

) -> Dict[str, Any]:

"""Build a materialized NaN group-time entry for a non-estimable (g, t) cell.

Non-estimable cells (missing base/post period, zero treated/control, zero

survey-weight mass, or a non-finite regression solve) are stored as a NaN

entry in ``group_time_effects`` rather than omitted, so the (g, t) grid is

inspectable (``to_dataframe`` / direct dict access) and the reason is

machine-readable via ``skip_reason`` (one of ``"missing_period"``,

``"zero_treated_control"``, ``"zero_weight_mass"``,

``"non_finite_regression"``; estimable cells carry ``None``).

The cell carries NO ``influence_func_info`` entry: every aggregation and

bootstrap consumer finite-masks (``np.isfinite(effect)``) or filters to IF

members before use, so the NaN cell contributes nothing to any aggregate or

SE — aggregates stay numerically identical to the prior omit behavior, which

matches R ``did``'s ``aggte()``. See REGISTRY.md "CallawaySantAnna" edge

cases for the documented contract.

"""

entry: Dict[str, Any] = {

"effect": np.nan,

"se": np.nan,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (np.nan, np.nan),

"n_treated": int(n_treated),

"n_control": int(n_control),

"skip_reason": skip_reason,

}

if survey_weight_sum is not None:

entry["survey_weight_sum"] = survey_weight_sum

return entry

class CallawaySantAnna(

CallawaySantAnnaBootstrapMixin,

CallawaySantAnnaAggregationMixin,

BaseEstimator,

):

"""

Callaway-Sant'Anna (2021) estimator for staggered Difference-in-Differences.

This estimator handles DiD designs with variation in treatment timing

(staggered adoption) and heterogeneous treatment effects. It avoids the

bias of traditional two-way fixed effects (TWFE) estimators by:

1. Computing group-time average treatment effects ATT(g,t) for each

cohort g (units first treated in period g) and time t.

2. Aggregating these to summary measures (overall ATT, event study, etc.)

using appropriate weights.

Parameters

----------

control_group : str, default="never_treated"

Which units to use as controls:

- "never_treated": Use only never-treated units (recommended)

- "not_yet_treated": Use never-treated and not-yet-treated units

anticipation : int, default=0

Number of periods before treatment where effects may occur.

Set to > 0 if treatment effects can begin before the official

treatment date.

estimation_method : str, default="dr"

Estimation method:

- "dr": Doubly robust (recommended)

- "ipw": Inverse probability weighting

- "reg": Outcome regression

alpha : float, default=0.05

Significance level for confidence intervals.

cluster : str, optional

Column name for cluster-robust standard errors. When set, the

influence-function aggregator clusters at the named level via a

synthesized ``SurveyDesign(psu=cluster_col)`` threaded through the

existing PSU-meat machinery (``_compute_stratified_psu_meat``) and

PSU-level multiplier bootstrap. When ``None`` (default), the

aggregator uses per-unit IF variance (Williams 2000 form). When

``survey_design=SurveyDesign(psu=...)`` is also provided, the

explicit PSU takes precedence; a ``UserWarning`` fires if the bare

``cluster=`` partition differs from the explicit PSU partition.

vcov_type : str, default="hc1"

Variance family. CallawaySantAnna accepts ``{"hc1"}`` only —

``hc1`` means per-unit IF variance when ``cluster=None`` and CR1

Liang-Zeger on the IF when ``cluster=X`` is set. The

analytical-sandwich families (``classical``, ``hc2``, ``hc2_bm``)

and spatial-HAC (``conley``) are rejected at ``__init__`` because

CS's per-(g,t) doubly-robust / IPW / outcome-regression structure

has no single design matrix to compute hat-matrix leverage or

Bell-McCaffrey Satterthwaite DOF on. See REGISTRY.md "IF-based

variance estimators vs analytical-sandwich estimators" for the

structural taxonomy.

n_bootstrap : int, default=0

Number of bootstrap iterations for inference.

If 0, uses analytical standard errors.

Recommended: 999 or more for reliable inference.

.. note:: Memory Usage

Bootstrap multiplier weights are generated and consumed one

draw-block at a time (see :mod:`diff_diff.bootstrap_chunking`), so the

full ``(n_bootstrap, n_units)`` weight matrix is never materialized.

The live weight intermediate is bounded by roughly

``max(~256 MB, 8 * n_units)`` bytes -- a block holds at least one full

draw row -- independent of ``n_bootstrap``. Only the small bootstrap

*output* arrays (``(n_bootstrap, n_group_time)`` and ``(n_bootstrap,)``

per aggregation) stay fully in memory. Stratified survey designs are

the current exception (the full PSU-weight matrix is built up front,

but PSUs are few).

bootstrap_weights : str, default="rademacher"

Type of weights for multiplier bootstrap:

- "rademacher": +1/-1 with equal probability (standard choice)

- "mammen": Two-point distribution (asymptotically valid, matches skewness)

- "webb": Six-point distribution (recommended when n_clusters < 20)

seed : int, optional

Random seed for reproducibility.

rank_deficient_action : str, default="warn"

Action when design matrix is rank-deficient (linearly dependent columns):

- "warn": Issue warning and drop linearly dependent columns (default)

- "error": Raise ValueError

- "silent": Drop columns silently without warning

base_period : str, default="varying"

Method for selecting the base (reference) period for computing

ATT(g,t). Base periods are selected *positionally* (by the nearest

observed period in the sorted panel), matching R ``did::att_gt`` -- so

on gapped (non-consecutive) grids the base is the nearest observed

period, not literal ``t-1`` / ``g-1``. The pre/post split is on the

current period vs the cohort (``t < g`` -> pre), independent of

anticipation; anticipation only shifts the post/universal base. Options:

- "varying": pre-treatment (``t < g``) uses the immediately-preceding

observed period as base; post-treatment uses the last observed

pre-treatment period (largest observed ``p`` with

``p + anticipation < g``).

- "universal": always uses that last observed pre-treatment period as

base.

On consecutive grids these reduce to ``t-1`` / ``g-1-anticipation``.

Both produce identical post-treatment effects. Matches R's

``did::att_gt()`` on gapped panels (base selection, estimable ATT/SE

cells, the ``"universal"`` zero reference cells, and all aggregations).

See :func:`_select_base_period`.

cband : bool, default=True

Whether to compute simultaneous confidence bands (sup-t) for

event study aggregation. Requires ``n_bootstrap > 0``.

When True, results include ``cband_crit_value`` and per-event-time

``cband_conf_int`` entries controlling family-wise error rate.

pscore_trim : float, default=0.01

Trimming bound for propensity scores. Scores are clipped to

``[pscore_trim, 1 - pscore_trim]`` before weight computation

in IPW and DR estimation. Must be in ``(0, 0.5)``.

panel : bool, default=True

Whether the data is a balanced/unbalanced panel (units observed

across multiple time periods). Set to ``False`` for stationary

repeated cross-sections where each observation has a unique unit

ID and units do not repeat across periods. Requires that the

cross-sectional samples are drawn from the same population in

each period (stationarity). Uses cross-sectional DRDID

(Sant'Anna & Zhao 2020, Section 4) with per-observation influence

functions.

allow_unbalanced_panel : bool, default=False

When ``True`` and the input panel is unbalanced (some units are not

observed in every period), route the pooled observations through the

repeated-cross-section levels estimator (matching R

``did::att_gt(allow_unbalanced_panel=TRUE)`` / ``DRDID::reg_did_rc``)

instead of within-cell panel differencing, and cluster the influence

function by unit for the standard error. **Inert on a balanced panel**

(results are byte-identical to the default). When ``False`` (default)

an unbalanced panel is handled by within-cell differencing and a

``UserWarning`` is emitted. ATT matches R bit-for-bit; the SE matches

up to the documented CR1 ``sqrt(G/(G-1))`` finite-sample factor.

``survey_design=`` combined with this flag raises ``NotImplementedError``.

epv_threshold : float, default=10

Events Per Variable threshold for propensity score logit.

When the ratio of minority-class observations to predictor

variables (excluding intercept) falls below this value, a

warning is emitted (or ``ValueError`` raised if

``rank_deficient_action="error"``). Based on Peduzzi et al.

(1996). Only applies to IPW and DR estimation methods.

Use ``diagnose_propensity()`` for a pre-estimation check across

all cohorts.

pscore_fallback : str, default="error"

Action when propensity score estimation fails entirely

(``LinAlgError`` or ``ValueError`` from IRLS):

- "error": Raise the exception (default). Ensures the user is

aware of estimation failures.

- "unconditional": Fall back to unconditional propensity

with a warning. For IPW, this drops all covariates. For DR,

the propensity model becomes unconditional but outcome

regression still uses covariates.

When ``rank_deficient_action="error"``, errors are always

re-raised regardless of this setting.

Attributes

----------

results_ : CallawaySantAnnaResults

Estimation results after calling fit().

is_fitted_ : bool

Whether the model has been fitted.

Examples

--------

Basic usage:

>>> import pandas as pd

>>> from diff_diff import CallawaySantAnna

>>>

>>> # Panel data with staggered treatment

>>> # 'first_treat' = period when unit was first treated (0 if never treated)

>>> data = pd.DataFrame({

... 'unit': [...],

... 'time': [...],

... 'outcome': [...],

... 'first_treat': [...] # 0 for never-treated, else first treatment period

... })

>>>

>>> cs = CallawaySantAnna()

>>> results = cs.fit(data, outcome='outcome', unit='unit',

... time='time', first_treat='first_treat')

>>>

>>> results.print_summary()

With event study aggregation (post-fit - no refit required):

>>> cs = CallawaySantAnna()

>>> results = cs.fit(data, outcome='outcome', unit='unit',

... time='time', first_treat='first_treat')

>>> event_study = results.aggregate('event_study')

>>> event_study.to_dataframe() # doctest: +SKIP

Plotting and the sensitivity analyses (``plot_event_study``,

``compute_honest_did``, ``compute_pretrends_power``) still read the

fit-time surface, so they take a fit that requested aggregation:

>>> from diff_diff import plot_event_study

>>> plotted = cs.fit(data, outcome='outcome', unit='unit',

... time='time', first_treat='first_treat',

... aggregate='event_study') # doctest: +SKIP

>>> plot_event_study(plotted) # doctest: +SKIP

With covariate adjustment (conditional parallel trends):

>>> # When parallel trends only holds conditional on covariates

>>> cs = CallawaySantAnna(estimation_method='dr') # doubly robust

>>> results = cs.fit(data, outcome='outcome', unit='unit',

... time='time', first_treat='first_treat',

... covariates=['age', 'income'])

>>>

>>> # DR is recommended: consistent if either outcome model

>>> # or propensity model is correctly specified

Notes

-----

The key innovation of Callaway & Sant'Anna (2021) is the disaggregated

approach: instead of estimating a single treatment effect, they estimate

ATT(g,t) for each cohort-time pair. This avoids the "forbidden comparison"

problem where already-treated units act as controls.

The ATT(g,t) is identified under parallel trends conditional on covariates:

E[Y(0)_t - Y(0)_g-1 | G=g] = E[Y(0)_t - Y(0)_g-1 | C=1]

where G=g indicates treatment cohort g and C=1 indicates control units.

This uses g-1 as the base period, which applies to post-treatment (t >= g).

With base_period="varying" (default), pre-treatment uses the immediately-

preceding observed period as base for the consecutive comparisons useful in

parallel trends diagnostics. Base periods are selected positionally (nearest

observed period), matching R did::att_gt on gapped grids (see

``_select_base_period``).

References

----------

Callaway, B., & Sant'Anna, P. H. (2021). Difference-in-Differences with

multiple time periods. Journal of Econometrics, 225(2), 200-230.

"""

def __init__(

self,

control_group: str = "never_treated",

anticipation: int = 0,

estimation_method: str = "dr",

alpha: float = 0.05,

cluster: Optional[str] = None,

n_bootstrap: int = 0,

bootstrap_weights: Optional[str] = None,

seed: Optional[int] = None,

rank_deficient_action: str = "warn",

base_period: str = "varying",

cband: bool = True,

pscore_trim: float = 0.01,

panel: bool = True,

allow_unbalanced_panel: bool = False,

epv_threshold: float = 10,

pscore_fallback: str = "error",

vcov_type: str = "hc1",

):

if control_group not in ["never_treated", "not_yet_treated"]:

raise ValueError(

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

f"got '{control_group}'"

)

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

raise ValueError(

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

)

if not (0 < pscore_trim < 0.5):

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

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

)

# Default to rademacher if not specified

if bootstrap_weights is None:

bootstrap_weights = "rademacher"

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

)

# vcov_type input contract: CallawaySantAnna is permanently narrow

# to {"hc1"} because the analytical-sandwich families (classical,

# hc2, hc2_bm) require a single regression's hat matrix that CS's

# per-(g,t) doubly-robust / IPW / outcome-regression structure

# doesn't have. See REGISTRY.md "IF-based variance estimators vs

# analytical-sandwich estimators" for the structural taxonomy.

# Factored out so fit() can re-run it: set_params now validates

# eagerly via the BaseEstimator probe re-init, but DIRECT attribute

# mutation (est.vcov_type = ...) still bypasses validation until fit.

self._validate_vcov_type(vcov_type)

self.control_group = control_group

self.anticipation = anticipation

self.estimation_method = estimation_method

self.alpha = alpha

self.cluster = cluster

self.vcov_type = vcov_type

# Track whether vcov_type was explicitly set (for future symmetry

# with SA / StackedDiD / WooldridgeDiD set_params patterns; the

# narrow contract makes the flag a no-op today but consistency

# avoids surprises if the contract ever broadens).

self._vcov_type_explicit = vcov_type != "hc1"

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

self.seed = seed

self.rank_deficient_action = rank_deficient_action

self.base_period = base_period

self.cband = cband

self.pscore_trim = pscore_trim

self.panel = panel

# When True AND the input panel is unbalanced (some units unobserved in

# some periods), route through the repeated-cross-section (RC) levels

# estimator on the pooled observations — matching R

# `did::att_gt(allow_unbalanced_panel=TRUE)` (which sets panel=FALSE ->

# DRDID::reg_did_rc). Inert on a balanced panel (the default within-cell

# differencing path is byte-identical). See fit() for the routing.

self.allow_unbalanced_panel = allow_unbalanced_panel

self.epv_threshold = epv_threshold

self.pscore_fallback = pscore_fallback

self.is_fitted_ = False

self.results_: Optional[CallawaySantAnnaResults] = None

def diagnose_propensity(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

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

) -> pd.DataFrame:

"""

Check Events Per Variable (EPV) across all cohorts without estimation.

Examines the data to identify cohorts where propensity score logit may

be unreliable due to too few events per covariate. Based on Peduzzi

et al. (1996).

This is a raw-count heuristic: it uses total cohort/control unit

counts without filtering for missing outcomes, zero survey weights,

or period-specific validity. The actual fit-time EPV (stored in

``results.epv_diagnostics``) may be lower because ``fit()`` operates

on the valid base/post outcome pair and the positive-weight effective

sample. Use this method as a quick pre-check; rely on

``results.epv_diagnostics`` for authoritative per-cell EPV.

Parameters

----------

df, outcome, unit, time, first_treat, covariates

Same arguments as ``fit()``.

Returns

-------

pd.DataFrame

Per-cohort EPV diagnostics with columns: group, n_treated,

n_control, n_covariates, n_params, epv, status.

"""

if not self.panel:

raise NotImplementedError(

"diagnose_propensity() is not yet supported for repeated "

"cross-section data (panel=False). Use fit() with covariates "

"and check results.epv_diagnostics instead."

)

if self.control_group == "not_yet_treated":

raise NotImplementedError(

"diagnose_propensity() is not yet supported for "

"control_group='not_yet_treated' because the control set "

"varies per (g, t) cell. Use fit() with covariates and "

"check results.epv_diagnostics instead."

)

if self.estimation_method == "reg":

return pd.DataFrame(

columns=[

"group",

"n_treated",

"n_control",

"n_covariates",

"n_params",

"epv",

"status",

]

)

if not covariates:

return pd.DataFrame(

columns=[

"group",

"n_treated",

"n_control",

"n_covariates",

"n_params",

"epv",

"status",

]

)

# Normalize np.inf → 0 for never-treated encoding (same as fit())

df = df.copy()

_inf_mask_diag = df[first_treat].isin([np.inf, float("inf")])

if _inf_mask_diag.any():

n_inf_units = df.loc[_inf_mask_diag, unit].nunique()

warnings.warn(

f"{n_inf_units} unit(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)

# Compute time_periods and treatment_groups (same logic as fit())

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

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

precomputed = self._precompute_structures(

df,

outcome,

unit,

time,

first_treat,

covariates,

time_periods=time_periods,

treatment_groups=treatment_groups,

)

cohort_masks = precomputed["cohort_masks"]

never_treated_mask = precomputed["never_treated_mask"]

unit_cohorts = precomputed["unit_cohorts"]

n_covariates = len(covariates)

n_params = n_covariates # predictor count, excluding intercept (Peduzzi convention)

rows = []

for g in sorted(cohort_masks.keys()):

treated_mask = cohort_masks[g]

if self.control_group == "never_treated":

control_mask = never_treated_mask

else:

base_period_val = g - 1 - self.anticipation

nyt_threshold = base_period_val + self.anticipation

control_mask = never_treated_mask | (

(unit_cohorts > nyt_threshold) & (unit_cohorts != g)

)

n_treated = int(np.sum(treated_mask))

n_control = int(np.sum(control_mask))

n_events = min(n_treated, n_control)

epv = n_events / n_params if n_params > 0 else float("inf")

if epv >= self.epv_threshold:

status = "ok"

elif epv >= 2:

status = "low"

else:

status = "critical"

rows.append(

{

"group": g,

"n_treated": n_treated,

"n_control": n_control,

"n_covariates": n_covariates,

"n_params": n_params,

"epv": round(epv, 1),

"status": status,

}

)

return pd.DataFrame(rows)

@staticmethod

def _collapse_survey_to_unit_level(resolved_survey, df, unit_col, all_units):

"""Create unit-level ResolvedSurveyDesign for panel IF-based variance.

Survey design columns are constant within units (validated upstream).

This extracts one row per unit, aligned to ``all_units`` ordering.

"""

from diff_diff.survey import collapse_survey_to_unit_level

return collapse_survey_to_unit_level(resolved_survey, df, unit_col, all_units)

def _precompute_structures(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

covariates: Optional[List[str]],

time_periods: List[Any],

treatment_groups: List[Any],

resolved_survey=None,

) -> PrecomputedData:

"""

Pre-compute data structures for efficient ATT(g,t) computation.

This pivots data to wide format and pre-computes:

- Outcome matrix (units x time periods)

- Covariate matrix (units x covariates) from base period

- Unit cohort membership masks

- Control unit masks

Returns

-------

PrecomputedData

Dictionary with pre-computed structures.

"""

# Get unique units and their cohort assignments

unit_info = df.groupby(unit)[first_treat].first()

all_units = unit_info.index.values

unit_cohorts = unit_info.values

# Create unit index mapping for fast lookups

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

# Pivot outcome to wide format: rows = units, columns = time periods

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

# Reindex to ensure all units are present (handles unbalanced panels)

outcome_wide = outcome_wide.reindex(all_units)

outcome_matrix = outcome_wide.values # Shape: (n_units, n_periods)

period_to_col = {t: i for i, t in enumerate(outcome_wide.columns)}

# Pre-compute cohort masks (boolean arrays)

cohort_masks = {}

for g in treatment_groups:

cohort_masks[g] = unit_cohorts == g

# Never-treated mask

# np.inf was normalized to 0 in fit(), so the np.inf check is defensive only

never_treated_mask = (unit_cohorts == 0) | (unit_cohorts == np.inf)

# Pre-compute covariate matrices by time period if needed

# (covariates are retrieved from the base period of each comparison)

covariate_by_period = None

if covariates:

covariate_by_period = {}

for t in time_periods:

period_data = df[df[time] == t].set_index(unit)

period_cov = period_data.reindex(all_units)[covariates]

covariate_by_period[t] = period_cov.values # Shape: (n_units, n_covariates)

is_balanced = not np.any(np.isnan(outcome_matrix))

# Extract per-unit survey weights (one weight per unit)

if resolved_survey is not None:

sw_by_unit = (

pd.Series(resolved_survey.weights, index=df.index).groupby(df[unit]).first()

)

survey_weights_arr = sw_by_unit.reindex(all_units).values

else:

survey_weights_arr = None

resolved_survey_unit = (

self._collapse_survey_to_unit_level(resolved_survey, df, unit, all_units)

if resolved_survey is not None

else None

)

return {

"all_units": all_units,

"unit_to_idx": unit_to_idx,

"unit_cohorts": unit_cohorts,

"outcome_matrix": outcome_matrix,

"period_to_col": period_to_col,

"observed_sorted": sorted(period_to_col),

"cohort_masks": cohort_masks,

"never_treated_mask": never_treated_mask,

"covariate_by_period": covariate_by_period,

"time_periods": time_periods,

"is_balanced": is_balanced,

"is_panel": True,

"canonical_size": len(all_units),

"survey_weights": survey_weights_arr,

"resolved_survey": resolved_survey,

"resolved_survey_unit": resolved_survey_unit,

"df_survey": (

resolved_survey_unit.df_survey if resolved_survey_unit is not None else None

),

}

def _select_base_period(self, g: Any, t: Any, observed_sorted: List) -> Optional[Any]:

"""Select the base period for cell ``(g, t)``, matching R ``did::att_gt``.

R selects the base period by *position* in the sorted list of observed

periods, not by literal calendar arithmetic: on gapped (non-consecutive)

period grids the base is the nearest observed period, not ``t-1`` /

``g-1``. This reproduces ``did`` 2.5.1 ``compute.att_gt`` exactly and is

byte-identical to the old ``t-1`` / ``g-1-anticipation`` rule on

consecutive grids (where positional == calendar).

Parameters

----------

g, t : cohort and evaluation period (calendar values).

observed_sorted : ascending list of the unique observed periods.

Returns

-------

The base period value, or ``None`` when no valid earlier observed

period exists (a non-estimable cell, materialized by callers as a

``missing_period`` NaN; R cannot estimate it either).

Notes

-----

Following R ``compute.att_gt``, the pre/post split is on the current

period vs the cohort (``t < g`` -> pre), **independent of

anticipation**; anticipation only enters the post/universal base.

- ``universal``, or post-treatment (``t >= g``): base is the last

pre-treatment observed period, i.e. the largest observed ``p`` with

``p + anticipation < g``.

- ``varying`` pre-treatment (``t < g``): base is the immediately

preceding observed period, i.e. the largest observed ``p < t``.

"""

if self.base_period == "universal" or t >= g:

threshold = g - self.anticipation

else: # varying pre-treatment

threshold = t

# Largest observed period strictly below `threshold` (positional

# neighbor). bisect_left gives the insertion index of `threshold`; the

# element just before it is the largest observed value < threshold.

idx = bisect.bisect_left(observed_sorted, threshold)

return observed_sorted[idx - 1] if idx > 0 else None

def _valid_periods_for_group(self, g: Any, time_periods: List, observed_sorted: List) -> List:

"""Evaluation periods ``t`` to attempt as ``ATT(g, t)`` for cohort ``g``.

Centralizes the per-group period filter used by every estimation path

(single source of truth, so the positional-base contract cannot drift):

- ``universal``: all observed periods except the (positional) base

reference period (which is the trivial ``ATT = 0`` cell); the base is

the last observed pre-treatment period, matching R -- NOT literal

``g-1-anticipation`` (which on gapped grids would leave the real base

in and materialize a fake zero cell).

- ``varying``: all periods except the earliest observed one (which can

never be a current period -- it has no earlier base).

Cells that remain non-estimable (``_select_base_period`` returns

``None``) are pruned downstream as ``missing_period`` NaNs.

"""

if self.base_period == "universal":

universal_base = self._select_base_period(g, g, observed_sorted)

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

min_period = observed_sorted[0]

return [t for t in time_periods if t >= g - self.anticipation or t > min_period]

def _compute_att_gt_fast(

self,

precomputed: PrecomputedData,

g: Any,

t: Any,

covariates: Optional[List[str]],

pscore_cache: Optional[Dict] = None,

epv_diagnostics: Optional[Dict] = None,

) -> Tuple[

Optional[float], float, int, int, Optional[Dict[str, Any]], Optional[float], Optional[str]

]:

"""

Compute ATT(g,t) using pre-computed data structures (fast version).

Uses vectorized numpy operations on pre-pivoted outcome matrix

instead of repeated pandas filtering.

Returns

-------

att_gt : float or None

se_gt : float

n_treated : int

n_control : int

inf_func_info : dict or None

survey_weight_sum : float or None

Sum of survey weights for treated units (for aggregation weighting).

skip_reason : str or None

When ``att_gt is None`` (non-estimable cell), the machine-readable

reason (``"missing_period"`` / ``"zero_treated_control"`` /

``"zero_weight_mass"``) so the caller can materialize a NaN cell with

a ``skip_reason``. ``None`` on a successful return.

"""

period_to_col = precomputed["period_to_col"]

outcome_matrix = precomputed["outcome_matrix"]

cohort_masks = precomputed["cohort_masks"]

never_treated_mask = precomputed["never_treated_mask"]

unit_cohorts = precomputed["unit_cohorts"]

covariate_by_period = precomputed["covariate_by_period"]

# Base period selection: positional (sorted-index) neighbor, matching

# R did::att_gt (see _select_base_period). Returns None for a

# non-estimable cell (no earlier observed period).

base_period_val = self._select_base_period(g, t, precomputed["observed_sorted"])

if base_period_val is None or t not in period_to_col: