"""

Mathematical core for the Efficient DiD estimator.

Implements the no-covariates path from Chen, Sant'Anna & Xie (2025):

optimal weighting via the inverse of the conditional covariance matrix Omega*,

generated outcomes from within-group sample means, and the efficient

influence function for analytical standard errors.

All functions are pure (no state), operating on pre-pivoted numpy arrays.

"""

import warnings

from typing import Dict, List, Optional, Tuple

import numpy as np

def enumerate_valid_triples(

target_g: float,

treatment_groups: List[float],

time_periods: List[float],

period_1: float,

pt_assumption: str,

anticipation: int = 0,

never_treated_val: float = np.inf,

) -> List[Tuple[float, float]]:

"""Enumerate valid (g', t_pre) pairs for target (g, t).

Under PT-All, any not-yet-treated cohort g' (including never-treated and

g'=g itself) paired with any baseline t_pre that is pre-treatment for the

*comparison* group g' forms a valid comparison. The target group g appears

only in the first term (Y_t - Y_1), which is independent of t_pre, so

t_pre need not be pre-treatment for g. Under PT-Post, only the

never-treated group with baseline g - 1 - anticipation is valid

(just-identified).

Parameters

----------

target_g : float

Treatment cohort of the target group.

treatment_groups : list of float

All treatment cohort identifiers (finite values only).

time_periods : list of float

All observed time periods, sorted.

period_1 : float

Earliest observed period (universal baseline).

pt_assumption : str

``"all"`` or ``"post"``.

anticipation : int

Number of anticipation periods.

never_treated_val : float

Sentinel for the never-treated group (default ``np.inf``).

Returns

-------

list of (g', t_pre) tuples

Valid comparison pairs. Empty if none exist.

"""

if pt_assumption == "post":

# Just-identified: only (never-treated, g - 1 - anticipation)

baseline = target_g - 1 - anticipation

if baseline >= period_1:

return [(never_treated_val, baseline)]

return []

# PT-All: overidentified

pairs: List[Tuple[float, float]] = []

# Candidate comparison groups: never-treated + all treatment cohorts.

# Including g'=g (same-cohort) is valid under PT-All (Eq 3.9).

# Including g'=∞ (never-treated) produces moments where the second

# and third terms telescope: y_hat = E[Y_t-Y_1|G=g] - E[Y_t-Y_1|G=∞]

# regardless of t_pre. These redundant moments add no information

# beyond the basic 2x2 DiD; the default omega_ridge damps them (the

# legacy omega_ridge=0 pseudoinverse truncates them instead - see the

# Omega* ridge Note in REGISTRY.md). Retained for implementation

# simplicity. The estimator's ridge path additionally drops the

# degenerate PRE-treatment self-pair (g'=g, t_pre=t) at fit level.

candidate_groups: List[float] = [never_treated_val]

for gp in treatment_groups:

candidate_groups.append(gp)

for gp in candidate_groups:

# Determine effective treatment start for comparison group

if np.isinf(gp):

effective_gp = np.inf # never treated

else:

effective_gp = gp - anticipation

for t_pre in time_periods:

if t_pre == period_1:

# period_1 is the universal reference — used as Y_1 in the

# differencing (Eq 3.9 first term). Including t_pre = period_1

# would make the third term Y_1 - Y_1 = 0 (degenerate), so it

# adds no information to Omega* regardless of which g' is used.

continue

# Only require t_pre < g' (pre-treatment for comparison group).

# No constraint on t_pre vs g: the target group appears only in

# the first term (Y_t - Y_1), which is independent of t_pre.

if not np.isinf(effective_gp) and t_pre >= effective_gp:

continue

pairs.append((gp, t_pre))

return pairs

def _sample_cov(

a: np.ndarray,

b: np.ndarray,

w: Optional[np.ndarray] = None,

) -> float:

"""Sample covariance between two 1-D arrays (ddof=1).

Returns 0.0 if fewer than 2 observations.

Parameters

----------

a, b : ndarray, shape (n,)

Data arrays.

w : ndarray, shape (n,), optional

Survey weights. When provided, computes the reliability-weighted

covariance: ``sum(w*(a-a_bar)*(b-b_bar)) / (sum(w) - 1)`` where

``a_bar = average(a, weights=w)``.

"""

n = len(a)

if n < 2:

return 0.0

if w is None:

return float(((a - a.mean()) * (b - b.mean())).sum() / (n - 1))

# Weighted covariance with reliability weights (Bessel-style correction)

a_bar = float(np.average(a, weights=w))

b_bar = float(np.average(b, weights=w))

sum_w = float(np.sum(w))

if sum_w <= 1.0:

return 0.0

return float(np.sum(w * (a - a_bar) * (b - b_bar)) / (sum_w - 1.0))

def compute_omega_star_nocov(

target_g: float,

target_t: float,

valid_pairs: List[Tuple[float, float]],

outcome_wide: np.ndarray,

cohort_masks: Dict[float, np.ndarray],

never_treated_mask: np.ndarray,

period_to_col: Dict[float, int],

period_1_col: int,

cohort_fractions: Dict[float, float],

never_treated_val: float = np.inf,

unit_weights: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Build the |H| x |H| covariance matrix Omega* (Eq 3.12, unconditional).

Each element Omega*[j,k] is the sum of up to five covariance terms

computed from within-group sample covariances scaled by inverse

cohort fractions.

Parameters

----------

target_g : float

Target treatment cohort.

target_t : float

Target time period.

valid_pairs : list of (g', t_pre) tuples

Valid comparison pairs from :func:`enumerate_valid_triples`.

outcome_wide : ndarray, shape (n_units, n_periods)

Pivoted outcome matrix.

cohort_masks : dict

``{cohort: bool_mask}`` over the unit dimension.

never_treated_mask : ndarray of bool

Mask for never-treated units.

period_to_col : dict

``{period: column_index}`` in ``outcome_wide``.

period_1_col : int

Column index of the earliest period (universal baseline Y_1).

cohort_fractions : dict

``{cohort: n_cohort / n}`` for each cohort.

never_treated_val : float

Sentinel for the never-treated group.

unit_weights : ndarray, shape (n_units,), optional

Survey weights at the unit level. When provided, all sample

means and covariances are weighted.

Returns

-------

ndarray, shape (|H|, |H|)

Covariance matrix. Empty (0,0) array if ``valid_pairs`` is empty.

"""

H = len(valid_pairs)

if H == 0:

return np.empty((0, 0))

t_col = period_to_col[target_t]

y1_col = period_1_col

# Pre-extract outcome columns for target group g

g_mask = cohort_masks[target_g]

Y_g = outcome_wide[g_mask] # (n_g, n_periods)

pi_g = cohort_fractions[target_g]

# Extract per-cohort weights (None propagates = unweighted)

w_g = unit_weights[g_mask] if unit_weights is not None else None

w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None

# Y_t - Y_1 for the target group

Yg_t_minus_1 = Y_g[:, t_col] - Y_g[:, y1_col]

# Never-treated outcomes

Y_inf = outcome_wide[never_treated_mask]

pi_inf = cohort_fractions.get(never_treated_val, 0.0)

omega = np.zeros((H, H))

# Hoist Term 1: (1/pi_g) * Var(Y_t - Y_1 | G=g) — same for all (j, k)

term1 = 0.0

if pi_g > 0:

term1 = (1.0 / pi_g) * _sample_cov(Yg_t_minus_1, Yg_t_minus_1, w=w_g)

# Precompute differenced arrays to avoid redundant slicing in the loop

# Never-treated: Y_t - Y_{tpre} and Y_{tpre} - Y_1 for each tpre

inf_t_minus_tpre: Dict[int, np.ndarray] = {}

inf_tpre_minus_1: Dict[int, np.ndarray] = {}

if len(Y_inf) >= 2:

for _, tpre in valid_pairs:

tpre_col = period_to_col[tpre]

if tpre_col not in inf_t_minus_tpre:

inf_t_minus_tpre[tpre_col] = Y_inf[:, t_col] - Y_inf[:, tpre_col]

inf_tpre_minus_1[tpre_col] = Y_inf[:, tpre_col] - Y_inf[:, y1_col]

# Target group: Y_{tpre} - Y_1 for each tpre where g' == target_g

g_tpre_minus_1: Dict[int, np.ndarray] = {}

if pi_g > 0:

for gp, tpre in valid_pairs:

if gp == target_g:

tpre_col = period_to_col[tpre]

if tpre_col not in g_tpre_minus_1:

g_tpre_minus_1[tpre_col] = Y_g[:, tpre_col] - Y_g[:, y1_col]

# Comparison cohort submatrices: cache outcome_wide[cohort_masks[gp]]

gp_outcomes: Dict[float, np.ndarray] = {}

gp_weights: Dict[float, Optional[np.ndarray]] = {}

for gp, _ in valid_pairs:

if not np.isinf(gp) and gp not in gp_outcomes:

if gp in cohort_masks:

gp_outcomes[gp] = outcome_wide[cohort_masks[gp]]

gp_weights[gp] = (

unit_weights[cohort_masks[gp]] if unit_weights is not None else None

)

# Comparison cohort: Y_{tpre} - Y_1 for each (gp, tpre) pair in Term 5

gp_tpre_minus_1: Dict[Tuple[float, int], np.ndarray] = {}

for j in range(H):

gp_j, tpre_j = valid_pairs[j]

tpre_j_col = period_to_col[tpre_j]

for k in range(j, H):

gp_k, tpre_k = valid_pairs[k]

tpre_k_col = period_to_col[tpre_k]

val = term1

# Term 2: (1/pi_inf) * SampleCov(Y_t - Y_{tpre_j}, Y_t - Y_{tpre_k} | G=inf)

if pi_inf > 0 and tpre_j_col in inf_t_minus_tpre:

val += (1.0 / pi_inf) * _sample_cov(

inf_t_minus_tpre[tpre_j_col],

inf_t_minus_tpre[tpre_k_col],

w=w_inf,

)

# Term 3: -1{g == g'_j} / pi_g * SampleCov(Y_t-Y_1, Y_{tpre_j}-Y_1 | G=g)

if gp_j == target_g and tpre_j_col in g_tpre_minus_1:

val -= (1.0 / pi_g) * _sample_cov(

Yg_t_minus_1,

g_tpre_minus_1[tpre_j_col],

w=w_g,

)

# Term 4: -1{g == g'_k} / pi_g * SampleCov(Y_t-Y_1, Y_{tpre_k}-Y_1 | G=g)

if gp_k == target_g and tpre_k_col in g_tpre_minus_1:

val -= (1.0 / pi_g) * _sample_cov(

Yg_t_minus_1,

g_tpre_minus_1[tpre_k_col],

w=w_g,

)

# Term 5: 1{g'_j == g'_k} / pi_{g'_j} * SampleCov(Y_{tpre_j}-Y_1, Y_{tpre_k}-Y_1 | G=g'_j)

if gp_j == gp_k:

if np.isinf(gp_j):

if pi_inf > 0 and tpre_j_col in inf_tpre_minus_1:

val += (1.0 / pi_inf) * _sample_cov(

inf_tpre_minus_1[tpre_j_col],

inf_tpre_minus_1[tpre_k_col],

w=w_inf,

)

else:

pi_gp = cohort_fractions.get(gp_j, 0.0)

if pi_gp > 0 and gp_j in cohort_masks:

Y_gp = gp_outcomes.get(gp_j)

if Y_gp is None:

Y_gp = outcome_wide[cohort_masks[gp_j]]

w_gp = gp_weights.get(gp_j)

if len(Y_gp) >= 2:

# Cache tpre diffs for comparison cohorts

key_j = (gp_j, tpre_j_col)

if key_j not in gp_tpre_minus_1:

gp_tpre_minus_1[key_j] = Y_gp[:, tpre_j_col] - Y_gp[:, y1_col]

key_k = (gp_j, tpre_k_col)

if key_k not in gp_tpre_minus_1:

gp_tpre_minus_1[key_k] = Y_gp[:, tpre_k_col] - Y_gp[:, y1_col]

val += (1.0 / pi_gp) * _sample_cov(

gp_tpre_minus_1[key_j],

gp_tpre_minus_1[key_k],

w=w_gp,

)

omega[j, k] = val

if j != k:

omega[k, j] = val

return omega

def _sample_cov_cols(

A: np.ndarray,

B: np.ndarray,

w: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Column-batched :func:`_sample_cov` (identical semantics per column).

A, B : (n, m) column-matched; returns (m,) with

``out[j] = _sample_cov(A[:, j], B[:, j], w)`` - including the guards:

fewer than 2 rows -> 0.0, and (weighted) ``sum(w) <= 1`` -> 0.0.

"""

n, m = A.shape

if n < 2:

return np.zeros(m)

if w is None:

A0 = A - A.mean(axis=0)

B0 = B - B.mean(axis=0)

return np.einsum("ij,ij->j", A0, B0) / (n - 1)

sum_w = float(np.sum(w))

if sum_w <= 1.0:

return np.zeros(m)

wn = w / sum_w

A0 = A - wn @ A

B0 = B - wn @ B

return np.einsum("i,ij,ij->j", w, A0, B0) / (sum_w - 1.0)

def _omega_star_nocov_gram(

target_g: float,

target_t: float,

valid_pairs: List[Tuple[float, float]],

outcome_wide: np.ndarray,

cohort_masks: Dict[float, np.ndarray],

never_treated_mask: np.ndarray,

period_to_col: Dict[float, int],

period_1_col: int,

cohort_fractions: Dict[float, float],

never_treated_val: float = np.inf,

unit_weights: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Gram/GEMM form of :func:`compute_omega_star_nocov` (v3.7 fast path).

Same five-term structure and guard semantics, but the H(H+1)/2 pair loop

of :func:`_sample_cov` calls collapses to column-batched covariances over

the DISTINCT ``(t_pre_j, t_pre_k)`` combinations (Term 2/5 covariances

depend only on the t_pre columns, not the full pair), assembled by

lookup. Used by the estimator's default ridge path; ``omega_ridge=0``

keeps the legacy loop. Agreement with the loop is at BLAS-reassociation

level (~1e-15 relative).

"""

H = len(valid_pairs)

if H == 0:

return np.empty((0, 0))

t_col = period_to_col[target_t]

y1_col = period_1_col

g_mask = cohort_masks[target_g]

Y_g = outcome_wide[g_mask]

pi_g = cohort_fractions[target_g]

w_g = unit_weights[g_mask] if unit_weights is not None else None

Y_inf = outcome_wide[never_treated_mask]

pi_inf = cohort_fractions.get(never_treated_val, 0.0)

w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None

pcols = sorted({period_to_col[tp] for _, tp in valid_pairs})

pindex = {c: i for i, c in enumerate(pcols)}

n_p = len(pcols)

j_p = [pindex[period_to_col[tp]] for _, tp in valid_pairs]

j_gp = [gp for gp, _ in valid_pairs]

iu, ju = np.triu_indices(n_p)

combo_idx = np.zeros((n_p, n_p), dtype=int)

combo_idx[iu, ju] = np.arange(len(iu))

combo_idx[ju, iu] = combo_idx[iu, ju]

a_g = Y_g[:, t_col] - Y_g[:, y1_col]

term1 = 0.0

if pi_g > 0:

term1 = (1.0 / pi_g) * float(_sample_cov_cols(a_g[:, None], a_g[:, None], w_g)[0])

# Term 2 / Term 5 (never-treated) tables - _sample_cov_cols returns zeros

# for n_inf < 2, matching the legacy table-not-populated skip.

u_inf = Y_inf[:, [t_col]] - Y_inf[:, pcols]

v_inf = Y_inf[:, pcols] - Y_inf[:, [y1_col]]

t2 = _sample_cov_cols(u_inf[:, iu], u_inf[:, ju], w_inf)

t5_inf = _sample_cov_cols(v_inf[:, iu], v_inf[:, ju], w_inf)

# Term 3/4 table (same-cohort pairs only)

t34 = None

if pi_g > 0 and any(gp == target_g for gp in j_gp):

b_g = Y_g[:, pcols] - Y_g[:, [y1_col]]

t34 = _sample_cov_cols(np.repeat(a_g[:, None], n_p, axis=1), b_g, w_g)

# Term 5 tables per finite comparison cohort

t5: Dict[float, np.ndarray] = {}

for gp in set(j_gp):

if np.isinf(gp) or gp in t5:

continue

pi_gp = cohort_fractions.get(gp, 0.0)

if pi_gp <= 0 or gp not in cohort_masks:

continue

Y_c = outcome_wide[cohort_masks[gp]]

w_c = unit_weights[cohort_masks[gp]] if unit_weights is not None else None

v_c = Y_c[:, pcols] - Y_c[:, [y1_col]]

t5[gp] = _sample_cov_cols(v_c[:, iu], v_c[:, ju], w_c)

omega = np.empty((H, H))

for j in range(H):

pj = j_p[j]

gpj = j_gp[j]

for k in range(j, H):

pk = j_p[k]

gpk = j_gp[k]

val = term1

if pi_inf > 0:

val += (1.0 / pi_inf) * t2[combo_idx[pj, pk]]

if gpj == target_g and t34 is not None:

val -= (1.0 / pi_g) * t34[pj]

if gpk == target_g and t34 is not None:

val -= (1.0 / pi_g) * t34[pk]

if gpj == gpk:

if np.isinf(gpj):

if pi_inf > 0:

val += (1.0 / pi_inf) * t5_inf[combo_idx[pj, pk]]

elif gpj in t5:

val += (1.0 / cohort_fractions[gpj]) * t5[gpj][combo_idx[pj, pk]]

omega[j, k] = val

if j != k:

omega[k, j] = val

return omega

def compute_efficient_weights(

omega_star: np.ndarray,

cond_threshold: float = 1e12,

omega_ridge: float = 0.0,

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

"""Compute efficient weights from Omega* inverse (Eq 3.13 / 4.3).

``w = ones @ inv(Omega*) / (ones @ inv(Omega*) @ ones)``

With ``omega_ridge = 0`` (this helper's default), runs the legacy path:

exact inverse, pseudoinverse + per-cell warning when the condition number

exceeds ``cond_threshold``. With ``omega_ridge > 0``, solves the

ridge-regularized system ``(Omega* + lam * max(trace/H, 0) * I) x = 1``

instead - numerically stable on the singular Omega* produced by PT-All's

telescoping moments (see ``OMEGA_RIDGE_DEFAULT`` in

``efficient_did_covariates``) - and emits no per-cell warning (the

estimator consolidates ill-conditioning into one fit-level warning). The

condition number is still computed and returned for diagnostics either

way.

Parameters

----------

omega_star : ndarray, shape (H, H)

Covariance matrix from :func:`compute_omega_star_nocov`.

cond_threshold : float

If condition number exceeds this: legacy path uses pseudoinverse +

warning; ridge path only reports it via the returned cond.

omega_ridge : float

Relative ridge scale; 0 = legacy inv/pinv path.

Returns

-------

weights : ndarray, shape (H,)

Efficient combination weights (sum to 1).

used_pinv : bool

True if pseudoinverse was used (always False on the ridge path).

cond_number : float

Condition number of Omega* (avoids recomputation by caller).

"""

H = omega_star.shape[0]

if H == 0:

return np.array([]), False, 0.0

if H == 1:

return np.array([1.0]), False, 1.0

ones = np.ones(H)

used_pinv = False

# Check for zero matrix

if np.allclose(omega_star, 0.0):

warnings.warn(

"Omega* matrix is all zeros; using uniform weights.",

UserWarning,

stacklevel=2,

)

return ones / H, False, np.inf

cond = float(np.linalg.cond(omega_star))

if omega_ridge > 0:

scale = max(float(np.trace(omega_star)) / H, 0.0)

om_ridged = omega_star + omega_ridge * scale * np.eye(H)

try:

numerator = np.linalg.solve(om_ridged, ones)

except np.linalg.LinAlgError:

# Unreachable for PSD Omega* with lam > 0; minimum-norm backstop.

numerator = np.linalg.pinv(om_ridged) @ ones

denominator = float(numerator @ ones)

if abs(denominator) < 1e-15:

warnings.warn(

"Denominator of efficient weights is near zero; using uniform weights.",

UserWarning,

stacklevel=2,

)

return ones / H, False, cond

return numerator / denominator, False, cond

if cond > cond_threshold:

warnings.warn(

f"Omega* condition number ({cond:.2e}) exceeds threshold "

f"({cond_threshold:.2e}); using pseudoinverse for weights.",

UserWarning,

stacklevel=2,

)

omega_inv = np.linalg.pinv(omega_star)

used_pinv = True

else:

try:

omega_inv = np.linalg.inv(omega_star)

except np.linalg.LinAlgError:

omega_inv = np.linalg.pinv(omega_star)

used_pinv = True

numerator = ones @ omega_inv # shape (H,)

denominator = numerator @ ones # scalar

if abs(denominator) < 1e-15:

warnings.warn(

"Denominator of efficient weights is near zero; using uniform weights.",

UserWarning,

stacklevel=2,

)

return ones / H, used_pinv, cond

weights = numerator / denominator

return weights, used_pinv, cond

def compute_generated_outcomes_nocov(

target_g: float,

target_t: float,

valid_pairs: List[Tuple[float, float]],

outcome_wide: np.ndarray,

cohort_masks: Dict[float, np.ndarray],

never_treated_mask: np.ndarray,

period_to_col: Dict[float, int],

period_1_col: int,

never_treated_val: float = np.inf,

unit_weights: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Compute generated outcome vector (one scalar per valid pair).

In the no-covariates case each generated outcome is a triple-difference

of within-group sample means (Eq 3.9 / 4.4 simplified)::

Y_hat_j = mean(Y_t - Y_1 | G=g)

- mean(Y_t - Y_{t_pre} | G=inf)

- mean(Y_{t_pre} - Y_1 | G=g')

where ``inf`` denotes the never-treated group and ``g'`` is the comparison

cohort for pair *j*.

Parameters

----------

target_g, target_t : float

Target group-time.

valid_pairs : list of (g', t_pre)

Valid comparison pairs.

outcome_wide : ndarray, shape (n_units, n_periods)

cohort_masks, never_treated_mask, period_to_col, period_1_col :

Pre-computed data structures.

never_treated_val : float

Sentinel for never-treated.

unit_weights : ndarray, shape (n_units,), optional

Survey weights at the unit level. When provided, all sample

means become weighted means.

Returns

-------

ndarray, shape (|H|,)

Scalar generated outcome for each pair.

"""

H = len(valid_pairs)

if H == 0:

return np.array([])

t_col = period_to_col[target_t]

y1_col = period_1_col

# Helper: weighted or unweighted mean

def _wmean(vals: np.ndarray, w: Optional[np.ndarray]) -> float:

if w is not None:

return float(np.average(vals, weights=w))

return float(np.mean(vals))

# Per-cohort weights

g_mask = cohort_masks[target_g]

w_g = unit_weights[g_mask] if unit_weights is not None else None

w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None

# Target group mean: mean(Y_t - Y_1 | G = g)

Y_g = outcome_wide[g_mask]

mean_g_t_1 = _wmean(Y_g[:, t_col] - Y_g[:, y1_col], w_g)

# Never-treated outcomes

Y_inf = outcome_wide[never_treated_mask]

y_hat = np.empty(H)

for j, (gp, tpre) in enumerate(valid_pairs):

tpre_col = period_to_col[tpre]

# mean(Y_t - Y_{tpre} | G = inf)

mean_inf_t_tpre = _wmean(Y_inf[:, t_col] - Y_inf[:, tpre_col], w_inf)

# mean(Y_{tpre} - Y_1 | G = g')

if np.isinf(gp):

Y_gp = Y_inf

w_gp = w_inf

else:

Y_gp = outcome_wide[cohort_masks[gp]]

w_gp = unit_weights[cohort_masks[gp]] if unit_weights is not None else None

mean_gp_tpre_1 = _wmean(Y_gp[:, tpre_col] - Y_gp[:, y1_col], w_gp)

y_hat[j] = mean_g_t_1 - mean_inf_t_tpre - mean_gp_tpre_1

return y_hat

def compute_eif_nocov(

target_g: float,

target_t: float,

weights: np.ndarray,

valid_pairs: List[Tuple[float, float]],

outcome_wide: np.ndarray,

cohort_masks: Dict[float, np.ndarray],

never_treated_mask: np.ndarray,

period_to_col: Dict[float, int],

period_1_col: int,

cohort_fractions: Dict[float, float],

n_units: int,

never_treated_val: float = np.inf,

unit_weights: Optional[np.ndarray] = None,

) -> np.ndarray:

"""Compute per-unit efficient influence function values.

For each unit *i* and each valid pair *j*, three terms contribute to

the EIF depending on the unit's cohort membership:

* **Treated term** (unit in cohort g):

``(1/pi_g) * (Y_{i,t} - Y_{i,1} - Y_hat_j) - ATT(g,t)``

* **Never-treated term** (unit in never-treated):

``-(1/pi_g) * (1/pi_inf) * pi_g * (Y_{i,t} - Y_{i,tpre_j} - mean_inf)``

(simplified: contributes the comparison group score for the never-treated)

* **Comparison cohort term** (unit in cohort g'_j):

``-(1/pi_g) * (1/pi_{g'_j}) * pi_g * (Y_{i,tpre_j} - Y_{i,1} - mean_gp)``

These are combined with efficient weights ``w_j``.

The derivation follows Theorem 3.2 and Eq 3.9-3.10, simplified for

the no-covariates case where propensity score ratios equal cohort

fraction ratios.

Parameters

----------

target_g, target_t : float

Target group-time.

weights : ndarray, shape (H,)

Efficient weights.

valid_pairs : list of (g', t_pre)

outcome_wide, cohort_masks, never_treated_mask, period_to_col,

period_1_col, cohort_fractions, n_units, never_treated_val :

Pre-computed data structures.

unit_weights : ndarray, shape (n_units,), optional

Survey weights at the unit level. When provided, within-group

means are weighted means.

Returns

-------

ndarray, shape (n_units,)

EIF value for every unit.

"""

H = len(valid_pairs)

if H == 0:

return np.zeros(n_units)

t_col = period_to_col[target_t]

y1_col = period_1_col

g_mask = cohort_masks[target_g]

Y_g = outcome_wide[g_mask]

pi_g = cohort_fractions[target_g]

Y_inf = outcome_wide[never_treated_mask]

pi_inf = cohort_fractions.get(never_treated_val, 0.0)

# Per-cohort weights

w_g = unit_weights[g_mask] if unit_weights is not None else None

w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None

# Helper for weighted/unweighted mean

def _wmean(vals: np.ndarray, w: Optional[np.ndarray]) -> float:

if w is not None:

return float(np.average(vals, weights=w))

return float(np.mean(vals))

eif = np.zeros(n_units)

# Hoist treated-group computations out of the per-pair loop (j-invariant)

Yg_t_minus_1 = Y_g[:, t_col] - Y_g[:, y1_col]

mean_g_t_1 = _wmean(Yg_t_minus_1, w_g)

treated_demeaned = None

if pi_g > 0:

treated_demeaned = (1.0 / pi_g) * (Yg_t_minus_1 - mean_g_t_1)

# Precompute never-treated diffs per tpre to avoid recomputation

inf_diffs: Dict[int, np.ndarray] = {}

inf_means: Dict[int, float] = {}

for j, (gp, tpre) in enumerate(valid_pairs):

w_j = weights[j]

tpre_col = period_to_col[tpre]

# --- Treated term (units in cohort g) ---

# (1/pi_g) * demeaned(Y_t - Y_1 | G=g) — same for all j

if treated_demeaned is not None:

eif[g_mask] += w_j * treated_demeaned

# --- Never-treated term ---

if tpre_col not in inf_diffs:

inf_diffs[tpre_col] = Y_inf[:, t_col] - Y_inf[:, tpre_col]

inf_means[tpre_col] = _wmean(inf_diffs[tpre_col], w_inf)

if pi_inf > 0:

inf_contrib = -(1.0 / pi_inf) * (inf_diffs[tpre_col] - inf_means[tpre_col])

eif[never_treated_mask] += w_j * inf_contrib

# --- Comparison cohort term ---

# Contribution from units in cohort g'_j for the baseline shift tpre_j - Y_1

if np.isinf(gp):

# Comparison group is never-treated; contribution is folded into

# the never-treated term via Y_{tpre} - Y_1 differencing.

# Additional term: -(1/pi_inf) * demeaned (Y_{tpre} - Y_1 | G=inf)

mean_inf_tpre_1 = _wmean(Y_inf[:, tpre_col] - Y_inf[:, y1_col], w_inf)

if pi_inf > 0:

gp_contrib = -(1.0 / pi_inf) * (

(Y_inf[:, tpre_col] - Y_inf[:, y1_col]) - mean_inf_tpre_1

)

eif[never_treated_mask] += w_j * gp_contrib

else:

gp_mask = cohort_masks[gp]

Y_gp = outcome_wide[gp_mask]

pi_gp = cohort_fractions.get(gp, 0.0)

w_gp = unit_weights[gp_mask] if unit_weights is not None else None

mean_gp_tpre_1 = _wmean(Y_gp[:, tpre_col] - Y_gp[:, y1_col], w_gp)

if pi_gp > 0:

gp_contrib = -(1.0 / pi_gp) * (

(Y_gp[:, tpre_col] - Y_gp[:, y1_col]) - mean_gp_tpre_1

)

eif[gp_mask] += w_j * gp_contrib

return eif