"""

Local (observation-specific) estimation method for the TROP estimator.

Contains the TROPLocalMixin class with all methods for the local

estimation pathway, including preprocessing, distance computation,

per-observation weight computation, model fitting, LOOCV scoring,

and bootstrap variance estimation.

This module is used via mixin inheritance — see trop.py for the

main TROP class definition.

"""

import logging

import warnings

from typing import List, Optional, Tuple

import numpy as np

import pandas as pd

logger = logging.getLogger(__name__)

from diff_diff._backend import (

HAS_RUST_BACKEND,

_rust_bootstrap_trop_variance,

_rust_unit_distance_matrix,

)

from diff_diff.bootstrap_utils import (

stratified_bootstrap_indices,

warn_bootstrap_failure_rate,

)

from diff_diff.trop_results import _PrecomputedStructures

from diff_diff.utils import warn_if_not_converged

def _validate_and_pivot_treatment(data, time, unit, treatment, all_periods, all_units):

"""Validate treatment column and create D matrix with missing mask.

Rejects observed rows with missing treatment values (data quality error),

then pivots to (time x unit) matrix. Structural gaps from unbalanced panels

are filled with 0 (assumed untreated) and flagged with a warning.

Returns

-------

D : ndarray

Treatment matrix (n_periods x n_units), int.

missing_mask : ndarray

Boolean mask of structurally absent cells (n_periods x n_units).

"""

n_nan_observed = int(data[treatment].isna().sum())

if n_nan_observed > 0:

raise ValueError(

f"{n_nan_observed} observation(s) have missing treatment values. "

f"TROP requires non-missing treatment indicators for all observed "

f"rows. Remove or impute missing values before fitting."

)

D_raw = data.pivot(index=time, columns=unit, values=treatment).reindex(

index=all_periods, columns=all_units

)

missing_mask = pd.isna(D_raw).values

n_missing_structural = int(missing_mask.sum())

if n_missing_structural > 0:

warnings.warn(

f"{n_missing_structural} missing treatment indicator(s) in the "

f"(time x unit) panel matrix filled with 0 (assumed "

f"untreated). This typically occurs in unbalanced panels.",

UserWarning,

stacklevel=3,

)

D = D_raw.fillna(0).astype(int).values

return D, missing_mask

# Module-level convergence tolerance for SVD singular value truncation.

# Singular values below this threshold after soft-thresholding are treated

# as zero to improve numerical stability.

_CONVERGENCE_TOL_SVD: float = 1e-10

def _soft_threshold_svd(

M: np.ndarray,

threshold: float,

convergence_tol: float = _CONVERGENCE_TOL_SVD,

) -> np.ndarray:

"""

Apply soft-thresholding to singular values (proximal operator for nuclear norm).

Parameters

----------

M : np.ndarray

Input matrix.

threshold : float

Soft-thresholding parameter.

convergence_tol : float, default=1e-10

Singular values below this after thresholding are treated as zero.

Returns

-------

np.ndarray

Matrix with soft-thresholded singular values.

"""

if threshold <= 0:

return M

# Handle NaN/Inf values in input

if not np.isfinite(M).all():

M = np.nan_to_num(M, nan=0.0, posinf=0.0, neginf=0.0)

try:

U, s, Vt = np.linalg.svd(M, full_matrices=False)

except np.linalg.LinAlgError:

# SVD failed, return zero matrix

return np.zeros_like(M)

# Check for numerical issues in SVD output

if not (np.isfinite(U).all() and np.isfinite(s).all() and np.isfinite(Vt).all()):

# SVD produced non-finite values, return zero matrix

return np.zeros_like(M)

s_thresh = np.maximum(s - threshold, 0)

# Use truncated reconstruction with only non-zero singular values

nonzero_mask = s_thresh > convergence_tol

if not np.any(nonzero_mask):

return np.zeros_like(M)

# Truncate to non-zero components for numerical stability

U_trunc = U[:, nonzero_mask]

s_trunc = s_thresh[nonzero_mask]

Vt_trunc = Vt[nonzero_mask, :]

# Compute result, suppressing expected numerical warnings from

# ill-conditioned matrices during alternating minimization

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

result = (U_trunc * s_trunc) @ Vt_trunc

# Replace any NaN/Inf in result with zeros

if not np.isfinite(result).all():

result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)

return result

class TROPLocalMixin:

"""Mixin providing local (observation-specific) estimation for TROP.

Methods in this mixin access the following attributes from the main

TROP class via ``self``:

- Solver params: ``max_iter``, ``tol``

- Inference params: ``n_bootstrap``, ``seed``

- State: ``_precomputed``

"""

# Type hints for attributes accessed from the main TROP class

max_iter: int

tol: float

n_bootstrap: int

seed: Optional[int]

_precomputed: Optional[_PrecomputedStructures]

# Convergence tolerance for SVD singular value truncation

CONVERGENCE_TOL_SVD: float = 1e-10

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

# Preprocessing and distance computation

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

def _precompute_structures(

self,

Y: np.ndarray,

D: np.ndarray,

control_unit_idx: np.ndarray,

n_units: int,

n_periods: int,

) -> _PrecomputedStructures:

"""

Pre-compute data structures that are reused across LOOCV and estimation.

This method computes once what would otherwise be computed repeatedly:

- Pairwise unit distance matrix

- Time distance vectors

- Masks and indices

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

D : np.ndarray

Treatment indicator matrix (n_periods x n_units).

control_unit_idx : np.ndarray

Indices of control units.

n_units : int

Number of units.

n_periods : int

Number of periods.

Returns

-------

_PrecomputedStructures

Pre-computed structures for efficient reuse.

"""

# Compute pairwise unit distances (for all observation-specific weights)

# Following Equation 3 (page 7): RMSE between units over pre-treatment

if HAS_RUST_BACKEND and _rust_unit_distance_matrix is not None:

# Use Rust backend for parallel distance computation (4-8x speedup)

unit_dist_matrix = _rust_unit_distance_matrix(Y, D.astype(np.float64))

else:

unit_dist_matrix = self._compute_all_unit_distances(Y, D, n_units, n_periods)

# Pre-compute time distance vectors for each target period

# Time distance: |t - s| for all s and each target t

time_dist_matrix = np.abs(

np.arange(n_periods)[:, np.newaxis] - np.arange(n_periods)[np.newaxis, :]

) # (n_periods, n_periods) where [t, s] = |t - s|

# Control and treatment masks

control_mask = D == 0

treated_mask = D == 1

# Identify treated observations

treated_observations = list(zip(*np.where(treated_mask)))

# Control observations for LOOCV

control_obs = [

(t, i)

for t in range(n_periods)

for i in range(n_units)

if control_mask[t, i] and not np.isnan(Y[t, i])

]

return {

"unit_dist_matrix": unit_dist_matrix,

"time_dist_matrix": time_dist_matrix,

"control_mask": control_mask,

"treated_mask": treated_mask,

"treated_observations": treated_observations,

"control_obs": control_obs,

"control_unit_idx": control_unit_idx,

"D": D,

"Y": Y,

"n_units": n_units,

"n_periods": n_periods,

}

def _compute_all_unit_distances(

self,

Y: np.ndarray,

D: np.ndarray,

n_units: int,

n_periods: int,

) -> np.ndarray:

"""

Compute pairwise unit distance matrix using vectorized operations.

Following Equation 3 (page 7):

dist_unit_{-t}(j, i) = sqrt(sum_u (Y_{iu} - Y_{ju})^2 / n_valid)

For efficiency, we compute a base distance matrix excluding all treated

observations, which provides a good approximation. The exact per-observation

distances are refined when needed.

Uses vectorized numpy operations with masked arrays for O(n^2) complexity

but with highly optimized inner loops via numpy/BLAS.

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

D : np.ndarray

Treatment indicator matrix (n_periods x n_units).

n_units : int

Number of units.

n_periods : int

Number of periods.

Returns

-------

np.ndarray

Pairwise distance matrix (n_units x n_units).

"""

# Mask for valid observations: control periods only (D=0), non-NaN

valid_mask = (D == 0) & ~np.isnan(Y)

# Replace invalid values with NaN for masked computation

Y_masked = np.where(valid_mask, Y, np.nan)

# Transpose to (n_units, n_periods) for easier broadcasting

Y_T = Y_masked.T # (n_units, n_periods)

# Compute pairwise squared differences using broadcasting

# Y_T[:, np.newaxis, :] has shape (n_units, 1, n_periods)

# Y_T[np.newaxis, :, :] has shape (1, n_units, n_periods)

# diff has shape (n_units, n_units, n_periods)

diff = Y_T[:, np.newaxis, :] - Y_T[np.newaxis, :, :]

sq_diff = diff**2

# Count valid (non-NaN) observations per pair

# A difference is valid only if both units have valid observations

valid_diff = ~np.isnan(sq_diff)

n_valid = np.sum(valid_diff, axis=2) # (n_units, n_units)

# Compute sum of squared differences (treating NaN as 0)

sq_diff_sum = np.nansum(sq_diff, axis=2) # (n_units, n_units)

# Compute RMSE distance: sqrt(sum / n_valid)

# Avoid division by zero

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

dist_matrix = np.sqrt(sq_diff_sum / n_valid)

# Set pairs with no valid observations to inf

dist_matrix = np.where(n_valid > 0, dist_matrix, np.inf)

# Ensure diagonal is 0 (same unit distance)

np.fill_diagonal(dist_matrix, 0.0)

return dist_matrix

def _compute_unit_distance_for_obs(

self,

Y: np.ndarray,

D: np.ndarray,

j: int,

i: int,

target_period: int,

) -> float:

"""

Compute observation-specific pairwise distance from unit j to unit i.

This is the exact computation from Equation 3, excluding the target period.

Used when the base distance matrix approximation is insufficient.

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

D : np.ndarray

Treatment indicator matrix.

j : int

Control unit index.

i : int

Treated unit index.

target_period : int

Target period to exclude.

Returns

-------

float

Pairwise RMSE distance.

"""

n_periods = Y.shape[0]

# Mask: exclude target period, both units must be untreated, non-NaN

valid = np.ones(n_periods, dtype=bool)

valid[target_period] = False

valid &= (D[:, i] == 0) & (D[:, j] == 0)

valid &= ~np.isnan(Y[:, i]) & ~np.isnan(Y[:, j])

if np.any(valid):

sq_diffs = (Y[valid, i] - Y[valid, j]) ** 2

return np.sqrt(np.mean(sq_diffs))

else:

return np.inf

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

# Observation-specific estimation

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

def _compute_observation_weights(

self,

Y: np.ndarray,

D: np.ndarray,

i: int,

t: int,

lambda_time: float,

lambda_unit: float,

control_unit_idx: np.ndarray,

n_units: int,

n_periods: int,

) -> np.ndarray:

"""

Compute observation-specific weight matrix for treated observation (i, t).

Following the paper's Algorithm 2 (page 27) and Equation 2 (page 7):

- Time weights theta_s^{i,t} = exp(-lambda_time * |t - s|)

- Unit weights omega_j^{i,t} = exp(-lambda_unit * dist_unit_{-t}(j, i))

Weights are assigned for every unit ``j != i`` (distance-based, per

Eq. 2/3). Treated-cell exclusion is handled by the `(1 - W_{js})`

factor applied inside ``_estimate_model`` via the control mask, not

by gating ``ω_j`` on ``D[t, j]``. Same-cohort donors therefore

contribute via their pre-treatment rows, and future-cohort donors

contribute via rows where both units are still untreated.

Always computes from the function-argument ``Y, D``; does not read

``self._precomputed``. Under bootstrap the caller passes resampled

``Y, D``, and a prior version of this method silently fell through to

the original-panel cache via a ``_precomputed`` branch, producing

stale unit distances.

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

D : np.ndarray

Treatment indicator matrix (n_periods x n_units).

i : int

Treated unit index.

t : int

Treatment period index.

lambda_time : float

Time weight decay parameter.

lambda_unit : float

Unit weight decay parameter.

control_unit_idx : np.ndarray

Indices of never-treated units (for backward compatibility, but not

used for weight computation - we use D matrix directly).

n_units : int

Number of units.

n_periods : int

Number of periods.

Returns

-------

np.ndarray

Weight matrix (n_periods x n_units) for observation (i, t).

"""

# Time distance: |t - s| following paper's Equation 3 (page 7)

dist_time = np.abs(np.arange(n_periods) - t)

time_weights = np.exp(-lambda_time * dist_time)

# Unit weights ω_j = exp(-λ_unit × dist(j, i)) for all j ≠ i per Eq. 2/3.

# No target-period gate: same-cohort donors enter with distance-based

# weight (their pre-treatment rows contribute via theta_s * omega_j;

# their post-treatment cells are zeroed by the control mask (1-D_{js})

# applied inside ``_estimate_model``). Matches Rust's compute_weight_matrix.

unit_weights = np.zeros(n_units)

if lambda_unit == 0:

# Uniform weights when lambda_unit = 0 — all units get 1.

# Control masking in _estimate_model handles treated-cell exclusion.

unit_weights[:] = 1.0

else:

for j in range(n_units):

if j != i:

# Compute distance excluding target period t (Issue B fix)

dist = self._compute_unit_distance_for_obs(Y, D, j, i, t)

if np.isinf(dist):

unit_weights[j] = 0.0

else:

unit_weights[j] = np.exp(-lambda_unit * dist)

# Target unit gets weight 1 (will be masked out in estimation via

# the control mask, matching the paper's (1-W_{js}) factor).

unit_weights[i] = 1.0

# Weight matrix: outer product (n_periods x n_units)

W = np.outer(time_weights, unit_weights)

return W

def _soft_threshold_svd(

self,

M: np.ndarray,

threshold: float,

) -> np.ndarray:

"""Delegate to module-level ``_soft_threshold_svd``."""

return _soft_threshold_svd(M, threshold, self.CONVERGENCE_TOL_SVD)

def _weighted_nuclear_norm_solve(

self,

Y: np.ndarray,

W: np.ndarray,

L_init: np.ndarray,

alpha: np.ndarray,

beta: np.ndarray,

lambda_nn: float,

max_inner_iter: int = 20,

) -> np.ndarray:

"""

Solve weighted nuclear norm problem using iterative weighted soft-impute.

Issue C fix: Implements the weighted nuclear norm optimization from the

paper's Equation 2 (page 7). The full objective is:

min_L sum W_{ti}(R_{ti} - L_{ti})^2 + lambda_nn||L||_*

This uses proximal gradient descent (Mazumder et al. 2010) with

FISTA/Nesterov acceleration. Lipschitz constant L_f = 2*max(W),

step size eta = 1/(2*max(W)), proximal threshold eta*lambda_nn:

G_k = L_k + (W/max(W)) * (R - L_k)

L_{k+1} = prox_{eta*lambda_nn*||*||_*}(G_k)

IMPORTANT: For observations with W=0 (treated observations), we keep

L values from the previous iteration rather than setting L = R, which

would absorb the treatment effect.

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

W : np.ndarray

Weight matrix (n_periods x n_units), non-negative. W=0 indicates

observations that should not be used for fitting (treated obs).

L_init : np.ndarray

Initial estimate of L matrix.

alpha : np.ndarray

Current unit fixed effects estimate.

beta : np.ndarray

Current time fixed effects estimate.

lambda_nn : float

Nuclear norm regularization parameter.

max_inner_iter : int, default=20

Maximum inner iterations for the proximal algorithm.

Returns

-------

np.ndarray

Updated L matrix estimate.

"""

# Compute target residual R = Y - alpha - beta

R = Y - alpha[np.newaxis, :] - beta[:, np.newaxis]

# Handle invalid values

R = np.nan_to_num(R, nan=0.0, posinf=0.0, neginf=0.0)

# For observations with W=0 (treated obs), keep L_init instead of R

# This prevents L from absorbing the treatment effect

valid_obs_mask = W > 0

R_masked = np.where(valid_obs_mask, R, L_init)

if lambda_nn <= 0:

# No regularization - just return masked residual

# Use soft-thresholding with threshold=0 which returns the input

return R_masked

# Normalize weights so max is 1 (for step size stability)

W_max = np.max(W)

if W_max > 0:

W_norm = W / W_max

else:

W_norm = W

# Initialize L

L = L_init.copy()

L_prev = L.copy()

t_fista = 1.0

# Proximal gradient iteration with FISTA/Nesterov acceleration

# This solves: min_L ||W^{1/2} * (R - L)||_F^2 + lambda||L||_*

# Lipschitz constant L_f = 2*max(W), so eta = 1/(2*max(W))

# Threshold = eta*lambda_nn = lambda_nn/(2*max(W))

for _ in range(max_inner_iter):

L_old = L.copy()

# FISTA momentum

t_fista_new = (1.0 + np.sqrt(1.0 + 4.0 * t_fista**2)) / 2.0

momentum = (t_fista - 1.0) / t_fista_new

L_momentum = L + momentum * (L - L_prev)

# Gradient step from momentum point: L_m + W * (R - L_m)

# For W=0 observations, this keeps L_m unchanged

gradient_step = L_momentum + W_norm * (R_masked - L_momentum)

# Proximal step: soft-threshold singular values

L_prev = L.copy()

threshold = lambda_nn / (2.0 * W_max) if W_max > 0 else lambda_nn / 2.0

L = self._soft_threshold_svd(gradient_step, threshold)

t_fista = t_fista_new

# Check convergence

if np.max(np.abs(L - L_old)) < self.tol:

break

return L

def _estimate_model(

self,

Y: np.ndarray,

control_mask: np.ndarray,

weight_matrix: np.ndarray,

lambda_nn: float,

n_units: int,

n_periods: int,

exclude_obs: Optional[Tuple[int, int]] = None,

_nonconvergence_tracker: Optional[List[int]] = None,

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

"""

Estimate the model: Y = alpha + beta + L + tau*D + eps with nuclear norm penalty on L.

Uses alternating minimization with vectorized operations:

1. Fix L, solve for alpha, beta via weighted means

2. Fix alpha, beta, solve for L via soft-thresholding

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

control_mask : np.ndarray

Boolean mask for control observations.

weight_matrix : np.ndarray

Pre-computed global weight matrix (n_periods x n_units).

lambda_nn : float

Nuclear norm regularization parameter.

n_units : int

Number of units.

n_periods : int

Number of periods.

exclude_obs : tuple, optional

(t, i) observation to exclude (for LOOCV).

Returns

-------

tuple

(alpha, beta, L) estimated parameters.

"""

W = weight_matrix

# Mask for estimation (control obs only, excluding LOOCV obs if specified)

est_mask = control_mask.copy()

if exclude_obs is not None:

t_ex, i_ex = exclude_obs

est_mask[t_ex, i_ex] = False

# Handle missing values

valid_mask = ~np.isnan(Y) & est_mask

# Initialize

alpha = np.zeros(n_units)

beta = np.zeros(n_periods)

L = np.zeros((n_periods, n_units))

# Pre-compute masked weights for vectorized operations

# Set weights to 0 where not valid

W_masked = W * valid_mask

# Pre-compute weight sums per unit and per time (for denominator)

# shape: (n_units,) and (n_periods,)

weight_sum_per_unit = np.sum(W_masked, axis=0) # sum over periods

weight_sum_per_time = np.sum(W_masked, axis=1) # sum over units

# Handle units/periods with zero weight sum

unit_has_obs = weight_sum_per_unit > 0

time_has_obs = weight_sum_per_time > 0

# Create safe denominators (avoid division by zero)

safe_unit_denom = np.where(unit_has_obs, weight_sum_per_unit, 1.0)

safe_time_denom = np.where(time_has_obs, weight_sum_per_time, 1.0)

# Replace NaN in Y with 0 for computation (mask handles exclusion)

Y_safe = np.where(np.isnan(Y), 0.0, Y)

# Alternating minimization following Algorithm 1 (page 9)

# Minimize: sum W_{ti}(Y_{ti} - alpha_i - beta_t - L_{ti})^2 + lambda_nn||L||_*

converged = False

for _ in range(self.max_iter):

alpha_old = alpha.copy()

beta_old = beta.copy()

L_old = L.copy()

# Step 1: Update alpha and beta (weighted least squares)

# Following Equation 2 (page 7), fix L and solve for alpha, beta

# R = Y - L (residual without fixed effects)

R = Y_safe - L

# Alpha update (unit fixed effects):

# alpha_i = argmin_alpha sum_t W_{ti}(R_{ti} - alpha - beta_t)^2

# Solution: alpha_i = sum_t W_{ti}(R_{ti} - beta_t) / sum_t W_{ti}

R_minus_beta = R - beta[:, np.newaxis] # (n_periods, n_units)

weighted_R_minus_beta = W_masked * R_minus_beta

alpha_numerator = np.sum(weighted_R_minus_beta, axis=0) # (n_units,)

alpha = np.where(unit_has_obs, alpha_numerator / safe_unit_denom, 0.0)

# Beta update (time fixed effects):

# beta_t = argmin_beta sum_i W_{ti}(R_{ti} - alpha_i - beta)^2

# Solution: beta_t = sum_i W_{ti}(R_{ti} - alpha_i) / sum_i W_{ti}

R_minus_alpha = R - alpha[np.newaxis, :] # (n_periods, n_units)

weighted_R_minus_alpha = W_masked * R_minus_alpha

beta_numerator = np.sum(weighted_R_minus_alpha, axis=1) # (n_periods,)

beta = np.where(time_has_obs, beta_numerator / safe_time_denom, 0.0)

# Step 2: Update L with weighted nuclear norm penalty

# Issue C fix: Use weighted soft-impute to properly account for

# observation weights in the nuclear norm optimization.

# Following Equation 2 (page 7): min_L sum W_{ti}(Y - alpha - beta - L)^2 + lambda||L||_*

L = self._weighted_nuclear_norm_solve(

Y_safe, W_masked, L, alpha, beta, lambda_nn, max_inner_iter=10

)

# Check convergence

alpha_diff = np.max(np.abs(alpha - alpha_old))

beta_diff = np.max(np.abs(beta - beta_old))

L_diff = np.max(np.abs(L - L_old))

if max(alpha_diff, beta_diff, L_diff) < self.tol:

converged = True

break

if not converged:

if _nonconvergence_tracker is not None:

_nonconvergence_tracker.append(1)

else:

warn_if_not_converged(

converged,

"TROP local alternating minimization",

self.max_iter,

self.tol,

)

return alpha, beta, L

def _loocv_score_obs_specific(

self,

Y: np.ndarray,

D: np.ndarray,

control_mask: np.ndarray,

control_unit_idx: np.ndarray,

lambda_time: float,

lambda_unit: float,

lambda_nn: float,

n_units: int,

n_periods: int,

) -> float:

"""

Compute leave-one-out cross-validation score with observation-specific weights.

Following the paper's Equation 5 (page 8):

Q(lambda) = sum_{j,s: D_js=0} [tau_js^loocv(lambda)]^2

For each control observation (j, s), treat it as pseudo-treated,

compute observation-specific weights, fit model excluding (j, s),

and sum squared pseudo-treatment effects.

Uses pre-computed structures when available for efficiency.

Parameters

----------

Y : np.ndarray

Outcome matrix (n_periods x n_units).

D : np.ndarray

Treatment indicator matrix (n_periods x n_units).

control_mask : np.ndarray

Boolean mask for control observations.

control_unit_idx : np.ndarray

Indices of control units.

lambda_time : float

Time weight decay parameter.

lambda_unit : float

Unit weight decay parameter.

lambda_nn : float

Nuclear norm regularization parameter.

n_units : int

Number of units.

n_periods : int

Number of periods.

Returns

-------

float

LOOCV score (lower is better).

"""

# Use pre-computed control observations if available

if self._precomputed is not None:

control_obs = self._precomputed["control_obs"]

else:

# Get all control observations

control_obs = [

(t, i)

for t in range(n_periods)

for i in range(n_units)

if control_mask[t, i] and not np.isnan(Y[t, i])

]

# Empty control set check: if no control observations, return infinity

# A score of 0.0 would incorrectly "win" over legitimate parameters

if len(control_obs) == 0:

warnings.warn(

f"LOOCV: No valid control observations for "

f"\u03bb=({lambda_time}, {lambda_unit}, {lambda_nn}). "

"Returning infinite score.",

UserWarning,

)

return np.inf

tau_squared_sum = 0.0

n_valid = 0

nonconverg_tracker: List[int] = []

for t, i in control_obs:

try:

# Compute observation-specific weights for pseudo-treated (i, t)

# Uses pre-computed distance matrices when available

weight_matrix = self._compute_observation_weights(

Y, D, i, t, lambda_time, lambda_unit, control_unit_idx, n_units, n_periods

)

# Estimate model excluding observation (t, i)

alpha, beta, L = self._estimate_model(

Y,

control_mask,

weight_matrix,

lambda_nn,

n_units,

n_periods,

exclude_obs=(t, i),

_nonconvergence_tracker=nonconverg_tracker,

)

# Pseudo treatment effect

tau_ti = Y[t, i] - alpha[i] - beta[t] - L[t, i]

tau_squared_sum += tau_ti**2

n_valid += 1

except (np.linalg.LinAlgError, ValueError):

# Per Equation 5: Q(lambda) must sum over ALL D==0 cells

# Any failure means this lambda cannot produce valid estimates for all cells

warnings.warn(

f"LOOCV: Fit failed for observation ({t}, {i}) with "

f"\u03bb=({lambda_time}, {lambda_unit}, {lambda_nn}). "

"Returning infinite score per Equation 5.",

UserWarning,

)

return np.inf

if nonconverg_tracker:

warn_if_not_converged(

False,

f"TROP local LOOCV: {len(nonconverg_tracker)} of "

f"{len(control_obs)} per-observation fits did not converge "

f"(\u03bb=({lambda_time}, {lambda_unit}, {lambda_nn}))",

self.max_iter,

self.tol,

)

# Return SUM of squared pseudo-treatment effects per Equation 5 (page 8):

# Q(lambda) = sum_{j,s: D_js=0} [tau_js^loocv(lambda)]^2

return tau_squared_sum

def _bootstrap_variance(

self,

data: pd.DataFrame,

outcome: str,

treatment: str,

unit: str,

time: str,

optimal_lambda: Tuple[float, float, float],

Y: Optional[np.ndarray] = None,

D: Optional[np.ndarray] = None,

control_unit_idx: Optional[np.ndarray] = None,

survey_design=None,

unit_weight_arr: Optional[np.ndarray] = None,

resolved_survey=None,

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

"""

Compute bootstrap standard error using unit-level block bootstrap.

When the optional Rust backend is available and the matrix parameters

(Y, D, control_unit_idx) are provided, uses parallelized Rust

implementation for 5-15x speedup. Falls back to Python implementation

if Rust is unavailable or if matrix parameters are not provided.

When a full survey design (strata/PSU/FPC) is present, uses Rao-Wu

rescaled bootstrap instead, which skips the Rust path.

Parameters

----------

data : pd.DataFrame

Original data in long format with unit, time, outcome, and treatment.

outcome : str

Name of the outcome column in data.

treatment : str

Name of the treatment indicator column in data.

unit : str

Name of the unit identifier column in data.

time : str

Name of the time period column in data.

optimal_lambda : tuple of float

Optimal tuning parameters (lambda_time, lambda_unit, lambda_nn)

from cross-validation. Used for model estimation in each bootstrap.

Y : np.ndarray, optional

Outcome matrix of shape (n_periods, n_units). Required for Rust

backend acceleration. If None, falls back to Python implementation.

D : np.ndarray, optional

Treatment indicator matrix of shape (n_periods, n_units) where

D[t,i]=1 indicates unit i is treated at time t. Required for Rust

backend acceleration.

control_unit_idx : np.ndarray, optional

Array of indices for control units (never-treated). Required for

Rust backend acceleration.

survey_design : SurveyDesign, optional

Survey design specification.

unit_weight_arr : np.ndarray, optional

Unit-level survey weights.

resolved_survey : ResolvedSurveyDesign, optional

Resolved survey design (observation-level).

Returns

-------

se : float

Bootstrap standard error of the ATT estimate.

bootstrap_estimates : np.ndarray

Array of ATT estimates from each bootstrap iteration. Length may

be less than n_bootstrap if some iterations failed.

Notes

-----

Uses unit-level block bootstrap where entire unit time series are

resampled with replacement. This preserves within-unit correlation

structure and is appropriate for panel data.

"""

lambda_time, lambda_unit, lambda_nn = optimal_lambda

# Check for full survey design (strata/PSU/FPC present)

_has_full_design = resolved_survey is not None and (

resolved_survey.strata is not None

or resolved_survey.psu is not None

or resolved_survey.fpc is not None

)

# Full survey design: use Python Rao-Wu rescaled bootstrap

if _has_full_design:

return self._bootstrap_rao_wu_local(

data,

outcome,

treatment,

unit,

time,

optimal_lambda,

resolved_survey,

survey_design,

)

# Stratified bootstrap pools (shared by Rust and Python paths).

# Paper's Algorithm 3 (page 27) specifies sampling N_0 control rows

# and N_1 treated rows separately to preserve treatment ratio.

unit_ever_treated = data.groupby(unit)[treatment].max()

treated_units = np.array(unit_ever_treated[unit_ever_treated == 1].index)

control_units = np.array(unit_ever_treated[unit_ever_treated == 0].index)

n_treated_units = len(treated_units)

n_control_units = len(control_units)

# Pre-generate stratified bootstrap indices via numpy (Python-canonical RNG).

# Aligns the RNG layer between backends. Combined with the Rust weight-

# matrix de-normalization and the Python `_compute_observation_weights`

# cache-fallthrough removal (also shipped with this parity work), local-

# method Rust and Python produce matching bootstrap SE up to solver-path

# roundoff (~1e-7); asserted at atol=1e-5 in the parity regression guard.

rng = np.random.default_rng(self.seed)

control_idx, treated_idx = stratified_bootstrap_indices(

rng, n_control_units, n_treated_units, self.n_bootstrap

)

# Try Rust backend for parallel bootstrap (5-15x speedup)

# Only used for pweight-only designs (no strata/PSU/FPC)

if (

HAS_RUST_BACKEND

and _rust_bootstrap_trop_variance is not None

and self._precomputed is not None

and Y is not None

and D is not None

):

try:

control_mask = self._precomputed["control_mask"]

time_dist_matrix = self._precomputed["time_dist_matrix"].astype(np.int64)

bootstrap_estimates, se = _rust_bootstrap_trop_variance(

Y,

D.astype(np.float64),

control_mask.astype(np.uint8),

time_dist_matrix,

lambda_time,

lambda_unit,

lambda_nn,

self.n_bootstrap,

self.max_iter,

self.tol,

control_idx,

treated_idx,

unit_weight_arr,

)

if len(bootstrap_estimates) > 0:

warn_bootstrap_failure_rate(

n_success=len(bootstrap_estimates),

n_attempted=self.n_bootstrap,

context="TROP local bootstrap (Rust)",

)

return float(se), bootstrap_estimates

logger.debug("Rust bootstrap returned 0 samples, falling back to Python")

except Exception as e:

logger.debug("Rust bootstrap variance failed, falling back to Python: %s", e)

warnings.warn(

f"Rust backend failed for bootstrap variance; "

f"falling back to Python. Performance may be reduced. "

f"Error: {e}",

UserWarning,

stacklevel=2,

)

# Python fallback: consume the same indices the Rust branch would have used.

bootstrap_estimates_list = []

nonconverg_tracker: List[int] = []