"""

Utility functions for difference-in-differences estimation.

"""

import warnings

from dataclasses import dataclass, field

from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np

import pandas as pd

from scipy import stats

from diff_diff.linalg import compute_robust_vcov as _compute_robust_vcov_linalg

from diff_diff.linalg import solve_ols as _solve_ols_linalg

# Import Rust backend if available (from _backend to avoid circular imports)

from diff_diff._backend import (

HAS_RUST_BACKEND,

_rust_project_simplex,

_rust_sdid_unit_weights,

_rust_compute_time_weights,

_rust_compute_noise_level,

_rust_sc_weight_fw,

_rust_sc_weight_fw_with_convergence,

_rust_sc_weight_fw_weighted,

_rust_sc_weight_fw_weighted_with_convergence,

)

# Numerical constants for optimization algorithms

_OPTIMIZATION_MAX_ITER = 1000 # Maximum iterations for weight optimization

_OPTIMIZATION_TOL = 1e-8 # Convergence tolerance for optimization

_NUMERICAL_EPS = 1e-10 # Small constant to prevent division by zero

# Cache for critical values to avoid repeated scipy calls

_critical_value_cache: Dict[Tuple[float, Optional[int]], float] = {}

def _get_critical_value(alpha: float, df: Optional[int] = None) -> float:

"""Return cached critical value for (alpha, df) pair."""

key = (alpha, df)

if key not in _critical_value_cache:

if df is not None:

_critical_value_cache[key] = float(stats.t.ppf(1 - alpha / 2, df))

else:

_critical_value_cache[key] = float(stats.norm.ppf(1 - alpha / 2))

return _critical_value_cache[key]

def validate_binary(arr: np.ndarray, name: str) -> None:

"""

Validate that an array contains only binary values (0 or 1).

Parameters

----------

arr : np.ndarray

Array to validate.

name : str

Name of the variable (for error messages).

Raises

------

ValueError

If array contains non-binary values.

"""

unique_values = np.unique(arr[~np.isnan(arr)])

if not np.all(np.isin(unique_values, [0, 1])):

raise ValueError(f"{name} must be binary (0 or 1). " f"Found values: {unique_values}")

def validate_covariate_names(

covariates: Optional[List[str]],

reserved_names: Iterable[str],

*,

estimator: str = "estimator",

) -> None:

"""

Validate that covariate column names do not collide with reserved

structural term names (and are not duplicated within ``covariates``).

Fitted coefficients are stored in a ``name -> value`` dict built by zipping

a variable-name list -- structural term names PLUS the user covariate column

names appended verbatim -- with the coefficient vector. A covariate whose

name equals a reserved structural name (the intercept ``const``, the

treatment/time indicators, the interaction term, period dummies,

fixed-effect dummies, or an internal working column) would silently

overwrite the structural coefficient (Python dict last-write-wins),

corrupting the result with no error. Duplicate names within ``covariates``

collapse to a single dict entry the same way.

The comparison is case-sensitive: column names and dict keys are

case-sensitive, so e.g. ``Const`` does not actually collide with ``const``

and is allowed.

Parameters

----------

covariates : list of str or None

User-supplied covariate column names. ``None`` or empty is a no-op.

reserved_names : iterable of str

Reserved structural term names this estimator builds (estimator-specific).

estimator : str

Estimator name, used in the error message.

Raises

------

ValueError

If a covariate name collides with a reserved structural name, or if

``covariates`` contains duplicate names.

"""

if not covariates:

return

reserved = set(reserved_names)

collisions = sorted({c for c in covariates if c in reserved})

if collisions:

raise ValueError(

f"{estimator}: covariate name(s) {collisions} collide with reserved "

f"structural term name(s). These names are used internally for the "

f"intercept, the treatment/time indicators, the interaction term, "

f"period dummies, fixed-effect dummies, or internal working columns, "

f"and a colliding covariate would silently overwrite the structural "

f"coefficient. Rename the covariate column(s). Reserved names for "

f"this fit: {sorted(reserved)}."

)

seen: set = set()

duplicates = []

for c in covariates:

if c in seen:

duplicates.append(c)

seen.add(c)

if duplicates:

raise ValueError(

f"{estimator}: duplicate covariate name(s) {sorted(set(duplicates))} "

f"in `covariates`. Each covariate maps to one coefficient; duplicates "

f"collapse to a single entry. Remove the duplicate(s)."

)

def validate_design_term_names(

var_names: Iterable[str],

*,

estimator: str = "estimator",

) -> None:

"""

Raise if the assembled design term-name list contains duplicates.

Backstop for :func:`validate_covariate_names`: even after the user

covariates are cleared, a fixed-effect dummy name (``{fe}_{value}``) can

still collide with a structural term — most notably a ``MultiPeriodDiD``

``period_{p}`` event-study key when a non-time fixed effect produces matching

dummy names — or with another dummy. Such a duplicate would silently

overwrite a coefficient when ``var_names`` is zipped into the result's

``coefficients`` dict (Python dict last-write-wins). This checks the FINAL

name list (structural terms + covariates + fixed-effect dummies) right

before the dict is built, catching collisions that depend on the data and so

cannot be known up front.

Parameters

----------

var_names : iterable of str

The fully assembled design-matrix column-name list.

estimator : str

Estimator name, used in the error message.

Raises

------

ValueError

If any name appears more than once.

"""

seen: set = set()

duplicates = []

for name in var_names:

if name in seen:

duplicates.append(name)

seen.add(name)

if duplicates:

raise ValueError(

f"{estimator}: the fitted design has duplicate term name(s) "

f"{sorted(set(duplicates))} — a covariate or fixed-effect dummy name "

f"collides with a structural term (intercept, treatment/time "

f"indicators, the interaction, or period dummies) or with another "

f"column. This would silently overwrite a coefficient in the result. "

f"Rename the offending fixed-effect category or covariate column."

)

def fe_dummy_names(col: pd.Series, prefix: str) -> List[str]:

"""

Reserved fixed-effect dummy column names for the collision guard, matching

``pd.get_dummies(col, prefix=prefix, drop_first=True).columns`` WITHOUT

materializing the dense ``(n x G)`` dummy matrix.

The within-transform ``TwoWayFixedEffects`` path is specifically designed to

avoid expanding high-cardinality fixed-effect dummies (that is its scaling

contract), so the collision guard must reserve those names without building

the dummy block. ``pd.get_dummies`` orders categories via

``pd.Categorical(col).categories`` — sorted unique values for a plain column,

the declared category order for a ``Categorical`` — then ``drop_first=True``

drops the first. This derivation reproduces that exactly (including

``Categorical`` columns with a non-default category order) at ``O(G)`` memory.

Parameters

----------

col : pandas.Series

The fixed-effect / unit / time column.

prefix : str

Dummy-name prefix (the project uses ``fe`` for ``fixed_effects`` and

``_fe_{unit}`` / ``_fe_{time}`` for TWFE unit/time dummies).

Returns

-------

list of str

The kept (post ``drop_first``) dummy column names.

"""

if isinstance(col.dtype, pd.CategoricalDtype):

cats = list(col.cat.categories)

else:

cats = list(pd.Categorical(col).categories)

return [f"{prefix}_{c}" for c in cats[1:]]

def warn_if_not_converged(

converged: bool,

method_name: str,

max_iter: int,

tol: Optional[float] = None,

stacklevel: int = 3,

) -> None:

"""Emit a UserWarning when an iterative solver exhausts max_iter without converging.

Shared helper for axis-B silent-failure fixes (iterative loops that otherwise

return the current iterate without signaling non-convergence).

"""

if converged:

return

tol_suffix = f" (tol={tol})" if tol is not None else ""

warnings.warn(

f"{method_name} did not converge in {max_iter} iterations{tol_suffix}. "

"Results may be inaccurate.",

UserWarning,

stacklevel=stacklevel,

)

def compute_robust_se(

X: np.ndarray, residuals: np.ndarray, cluster_ids: Optional[np.ndarray] = None

) -> np.ndarray:

"""

Compute heteroskedasticity-robust (HC1) or cluster-robust standard errors.

This function is a thin wrapper around the optimized implementation in

diff_diff.linalg for backwards compatibility.

Parameters

----------

X : np.ndarray

Design matrix of shape (n, k).

residuals : np.ndarray

Residuals from regression of shape (n,).

cluster_ids : np.ndarray, optional

Cluster identifiers for cluster-robust SEs.

Returns

-------

np.ndarray

Variance-covariance matrix of shape (k, k).

"""

return _compute_robust_vcov_linalg(X, residuals, cluster_ids)

def compute_confidence_interval(

estimate: float, se: float, alpha: float = 0.05, df: Optional[int] = None

) -> Tuple[float, float]:

"""

Compute confidence interval for an estimate.

Parameters

----------

estimate : float

Point estimate.

se : float

Standard error.

alpha : float

Significance level (default 0.05 for 95% CI).

df : int, optional

Degrees of freedom. If None, uses normal distribution.

Returns

-------

tuple

(lower_bound, upper_bound) of confidence interval.

"""

critical_value = _get_critical_value(alpha, df)

lower = estimate - critical_value * se

upper = estimate + critical_value * se

return (lower, upper)

def compute_p_value(t_stat: float, df: Optional[int] = None, two_sided: bool = True) -> float:

"""

Compute p-value for a t-statistic.

Parameters

----------

t_stat : float

T-statistic.

df : int, optional

Degrees of freedom. If None, uses normal distribution.

two_sided : bool

Whether to compute two-sided p-value (default True).

Returns

-------

float

P-value.

"""

if df is not None:

p_value = stats.t.sf(np.abs(t_stat), df)

else:

p_value = stats.norm.sf(np.abs(t_stat))

if two_sided:

p_value *= 2

return float(p_value)

def safe_inference(effect, se, alpha=0.05, df=None):

"""Compute t_stat, p_value, conf_int with NaN-safe gating.

When SE is non-finite, zero, or negative, ALL inference fields

are set to NaN to prevent misleading statistical output.

Accepts scalar inputs only (not numpy arrays). All existing inference

call sites operate on scalars within loops.

Parameters

----------

effect : float

Point estimate (treatment effect or coefficient).

se : float

Standard error of the estimate.

alpha : float, optional

Significance level for confidence interval (default 0.05).

df : int, optional

Degrees of freedom. If None, uses normal distribution.

Returns

-------

tuple

(t_stat, p_value, (ci_lower, ci_upper)). All NaN when SE is

non-finite, zero, or negative.

"""

if not (np.isfinite(se) and se > 0):

return np.nan, np.nan, (np.nan, np.nan)

if df is not None and df <= 0:

# Undefined degrees of freedom (e.g., rank-deficient replicate design)

return np.nan, np.nan, (np.nan, np.nan)

t_stat = effect / se

p_value = compute_p_value(t_stat, df=df)

conf_int = compute_confidence_interval(effect, se, alpha, df=df)

return t_stat, p_value, conf_int

def safe_inference_batch(effects, ses, alpha=0.05, df=None):

"""Vectorized batch inference for arrays of effects and SEs.

Parameters

----------

effects : np.ndarray

Array of point estimates.

ses : np.ndarray

Array of standard errors.

alpha : float, optional

Significance level (default 0.05).

df : int, optional

Degrees of freedom. If None, uses normal distribution.

Returns

-------

t_stats : np.ndarray

p_values : np.ndarray

ci_lowers : np.ndarray

ci_uppers : np.ndarray

"""

effects = np.asarray(effects, dtype=float)

ses = np.asarray(ses, dtype=float)

n = len(effects)

t_stats = np.full(n, np.nan)

p_values = np.full(n, np.nan)

ci_lowers = np.full(n, np.nan)

ci_uppers = np.full(n, np.nan)

# Undefined df (e.g., rank-deficient replicate design) → all NaN

if df is not None and df <= 0:

return t_stats, p_values, ci_lowers, ci_uppers

valid = np.isfinite(ses) & (ses > 0)

if not np.any(valid):

return t_stats, p_values, ci_lowers, ci_uppers

t_stats[valid] = effects[valid] / ses[valid]

if df is not None:

p_values[valid] = 2.0 * stats.t.sf(np.abs(t_stats[valid]), df)

else:

p_values[valid] = 2.0 * stats.norm.sf(np.abs(t_stats[valid]))

crit = _get_critical_value(alpha, df)

ci_lowers[valid] = effects[valid] - crit * ses[valid]

ci_uppers[valid] = effects[valid] + crit * ses[valid]

return t_stats, p_values, ci_lowers, ci_uppers

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

# Wild Cluster Bootstrap

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

@dataclass

class WildBootstrapResults:

"""

Results from wild cluster bootstrap inference.

Attributes

----------

se : float

Bootstrap standard error of the coefficient.

p_value : float

Bootstrap p-value (two-sided).

t_stat_original : float

Original t-statistic from the data.

ci_lower : float

Lower bound of the confidence interval.

ci_upper : float

Upper bound of the confidence interval.

n_clusters : int

Number of clusters in the data.

n_bootstrap : int

Number of bootstrap replications.

weight_type : str

Type of bootstrap weights used ("rademacher", "webb", or "mammen").

alpha : float

Significance level used for confidence interval.

bootstrap_distribution : np.ndarray, optional

Full bootstrap distribution of coefficients (if requested).

References

----------

Cameron, A. C., Gelbach, J. B., & Miller, D. L. (2008).

Bootstrap-Based Improvements for Inference with Clustered Errors.

The Review of Economics and Statistics, 90(3), 414-427.

"""

se: float

p_value: float

t_stat_original: float

ci_lower: float

ci_upper: float

n_clusters: int

n_bootstrap: int

weight_type: str

alpha: float = 0.05

bootstrap_distribution: Optional[np.ndarray] = field(default=None, repr=False)

def summary(self) -> str:

"""Generate formatted summary of bootstrap results."""

lines = [

"Wild Cluster Bootstrap Results",

"=" * 40,

f"Bootstrap SE: {self.se:.6f}",

f"Bootstrap p-value: {self.p_value:.4f}",

f"Original t-stat: {self.t_stat_original:.4f}",

f"CI ({int((1-self.alpha)*100)}%): [{self.ci_lower:.6f}, {self.ci_upper:.6f}]",

f"Number of clusters: {self.n_clusters}",

f"Bootstrap reps: {self.n_bootstrap}",

f"Weight type: {self.weight_type}",

]

return "\n".join(lines)

def print_summary(self) -> None:

"""Print formatted summary to stdout."""

print(self.summary())

def _generate_rademacher_weights(n_clusters: int, rng: np.random.Generator) -> np.ndarray:

"""

Generate Rademacher weights: +1 or -1 with probability 0.5.

Parameters

----------

n_clusters : int

Number of clusters.

rng : np.random.Generator

Random number generator.

Returns

-------

np.ndarray

Array of Rademacher weights.

"""

return np.asarray(rng.choice([-1.0, 1.0], size=n_clusters))

def _generate_webb_weights(n_clusters: int, rng: np.random.Generator) -> np.ndarray:

"""

Generate Webb's 6-point distribution weights.

Values: {-sqrt(3/2), -sqrt(2/2), -sqrt(1/2), sqrt(1/2), sqrt(2/2), sqrt(3/2)}

with equal probabilities (1/6 each), giving E[w]=0 and Var(w)=1.0.

This distribution is recommended for very few clusters (G < 10) as it

provides better finite-sample properties than Rademacher weights.

Parameters

----------

n_clusters : int

Number of clusters.

rng : np.random.Generator

Random number generator.

Returns

-------

np.ndarray

Array of Webb weights.

References

----------

Webb, M. D. (2014). Reworking wild bootstrap based inference for

clustered errors. Queen's Economics Department Working Paper No. 1315.

Note: Uses equal probabilities (1/6 each) matching R's `did` package,

which gives unit variance for consistency with other weight distributions.

"""

values = np.array(

[

-np.sqrt(3 / 2),

-np.sqrt(2 / 2),

-np.sqrt(1 / 2),

np.sqrt(1 / 2),

np.sqrt(2 / 2),

np.sqrt(3 / 2),

]

)

# Equal probabilities (1/6 each) matching R's did package, giving Var(w) = 1.0

return np.asarray(rng.choice(values, size=n_clusters))

def _generate_mammen_weights(n_clusters: int, rng: np.random.Generator) -> np.ndarray:

"""

Generate Mammen's two-point distribution weights.

Values: {-(sqrt(5)-1)/2, (sqrt(5)+1)/2}

with probabilities {(sqrt(5)+1)/(2*sqrt(5)), (sqrt(5)-1)/(2*sqrt(5))}.

This distribution satisfies E[v]=0, E[v^2]=1, E[v^3]=1, which provides

asymptotic refinement for skewed error distributions.

Parameters

----------

n_clusters : int

Number of clusters.

rng : np.random.Generator

Random number generator.

Returns

-------

np.ndarray

Array of Mammen weights.

References

----------

Mammen, E. (1993). Bootstrap and Wild Bootstrap for High Dimensional

Linear Models. The Annals of Statistics, 21(1), 255-285.

"""

sqrt5 = np.sqrt(5)

# Values from Mammen (1993)

val1 = -(sqrt5 - 1) / 2 # approximately -0.618

val2 = (sqrt5 + 1) / 2 # approximately 1.618 (golden ratio)

# Probability of val1

p1 = (sqrt5 + 1) / (2 * sqrt5) # approximately 0.724

return np.asarray(rng.choice([val1, val2], size=n_clusters, p=[p1, 1 - p1]))

def wild_bootstrap_se(

X: np.ndarray,

y: np.ndarray,

residuals: np.ndarray,

cluster_ids: np.ndarray,

coefficient_index: int,

n_bootstrap: int = 999,

weight_type: str = "rademacher",

null_hypothesis: float = 0.0,

alpha: float = 0.05,

seed: Optional[int] = None,

return_distribution: bool = False,

) -> WildBootstrapResults:

"""

Compute wild cluster bootstrap standard errors and p-values.

Implements the Wild Cluster Residual (WCR) bootstrap procedure from

Cameron, Gelbach, and Miller (2008). Uses the restricted residuals

approach (imposing H0: coefficient = null_hypothesis) for more accurate

p-value computation.

Parameters

----------

X : np.ndarray

Design matrix of shape (n, k).

y : np.ndarray

Outcome vector of shape (n,).

residuals : np.ndarray

OLS residuals from unrestricted regression, shape (n,).

cluster_ids : np.ndarray

Cluster identifiers of shape (n,).

coefficient_index : int

Index of the coefficient for which to compute bootstrap inference.

For DiD, this is typically 3 (the treatment*post interaction term).

n_bootstrap : int, default=999

Number of bootstrap replications. Odd numbers are recommended for

exact p-value computation.

weight_type : str, default="rademacher"

Type of bootstrap weights:

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

- "webb": 6-point distribution (recommended for <10 clusters)

- "mammen": Two-point distribution with skewness correction

null_hypothesis : float, default=0.0

Value of the null hypothesis for p-value computation.

alpha : float, default=0.05

Significance level for confidence interval.

seed : int, optional

Random seed for reproducibility. If None (default), results

will vary between runs.

return_distribution : bool, default=False

If True, include full bootstrap distribution in results.

Returns

-------

WildBootstrapResults

Dataclass containing bootstrap SE, p-value, confidence interval,

and other inference results.

Raises

------

ValueError

If weight_type is not recognized or if there are fewer than 2 clusters.

Warns

-----

UserWarning

If the number of clusters is less than 5, as bootstrap inference

may be unreliable.

Examples

--------

>>> from diff_diff.utils import wild_bootstrap_se

>>> results = wild_bootstrap_se(

... X, y, residuals, cluster_ids,

... coefficient_index=3, # ATT coefficient

... n_bootstrap=999,

... weight_type="rademacher",

... seed=42

... )

>>> print(f"Bootstrap SE: {results.se:.4f}")

>>> print(f"Bootstrap p-value: {results.p_value:.4f}")

References

----------

Cameron, A. C., Gelbach, J. B., & Miller, D. L. (2008).

Bootstrap-Based Improvements for Inference with Clustered Errors.

The Review of Economics and Statistics, 90(3), 414-427.

MacKinnon, J. G., & Webb, M. D. (2018). The wild bootstrap for

few (treated) clusters. The Econometrics Journal, 21(2), 114-135.

"""

# Validate inputs

valid_weight_types = ["rademacher", "webb", "mammen"]

if weight_type not in valid_weight_types:

raise ValueError(f"weight_type must be one of {valid_weight_types}, got '{weight_type}'")

unique_clusters = np.unique(cluster_ids)

n_clusters = len(unique_clusters)

if n_clusters < 2:

raise ValueError(f"Wild cluster bootstrap requires at least 2 clusters, got {n_clusters}")

if n_clusters < 5:

warnings.warn(

f"Only {n_clusters} clusters detected. Wild bootstrap inference may be "

"unreliable with fewer than 5 clusters. Consider using Webb weights "

"(weight_type='webb') for improved finite-sample properties.",

UserWarning,

)

# Initialize RNG

rng = np.random.default_rng(seed)

# Select weight generator

weight_generators = {

"rademacher": _generate_rademacher_weights,

"webb": _generate_webb_weights,

"mammen": _generate_mammen_weights,

}

generate_weights = weight_generators[weight_type]

n = X.shape[0]

# Step 1: Compute original coefficient and cluster-robust SE

beta_hat, _, vcov_original = _solve_ols_linalg(X, y, cluster_ids=cluster_ids, return_vcov=True)

original_coef = beta_hat[coefficient_index]

assert vcov_original is not None

se_original = np.sqrt(vcov_original[coefficient_index, coefficient_index])

t_stat_original = (original_coef - null_hypothesis) / se_original

# Step 2: Impose null hypothesis (restricted estimation)

# Create restricted y: y_restricted = y - X[:, coef_index] * null_hypothesis

# This imposes the null that the coefficient equals null_hypothesis

y_restricted = y - X[:, coefficient_index] * null_hypothesis

# Fit restricted model (but we need to drop the column for the restricted coef)

# Actually, for WCR bootstrap we keep all columns but impose the null via residuals

# Re-estimate with the restricted dependent variable.

#

# Use return_fitted=True so we get NaN-safe fitted values from the kept

# columns when solve_ols drops rank-deficient nuisance columns. Without

# this, building y_star via `X @ beta_restricted` would propagate NaN

# through every observation whenever a nuisance column was dropped

# (e.g. always-treated unit dummy collinear with treated*post on the

# full-dummy TWFE HC2/HC2-BM path), poisoning the entire bootstrap loop

# despite the ATT being analytically identified.

beta_restricted, residuals_restricted, fitted_restricted, _ = _solve_ols_linalg(

X, y_restricted, return_vcov=False, return_fitted=True

)

# Create cluster-to-observation mapping for efficiency

cluster_map = {c: np.where(cluster_ids == c)[0] for c in unique_clusters}

cluster_indices = [cluster_map[c] for c in unique_clusters]

# Step 3: Bootstrap loop

# Use NaN for invalid draws (singular bootstrap SE) and filter at the

# p-value step, rather than coercing to t*=0 which biases the p-value

# toward small values (since |0| < |t_original| counts as "non-rejection"

# only when the original t is large).

bootstrap_t_stats = np.full(n_bootstrap, np.nan)

bootstrap_coefs = np.full(n_bootstrap, np.nan)

for b in range(n_bootstrap):

# Generate cluster-level weights

cluster_weights = generate_weights(n_clusters, rng)

# Map cluster weights to observations

obs_weights = np.zeros(n)

for g, indices in enumerate(cluster_indices):

obs_weights[indices] = cluster_weights[g]

# Construct bootstrap sample: y* = fitted_restricted + e_restricted * weights

# (fitted_restricted comes from solve_ols's kept-columns reconstruction,

# so it's NaN-safe even when beta_restricted has NaN on dropped columns)

y_star = fitted_restricted + residuals_restricted * obs_weights

# Estimate bootstrap coefficients with cluster-robust SE

beta_star, residuals_star, vcov_star = _solve_ols_linalg(

X, y_star, cluster_ids=cluster_ids, return_vcov=True

)

bootstrap_coefs[b] = beta_star[coefficient_index]

assert vcov_star is not None

se_star = np.sqrt(vcov_star[coefficient_index, coefficient_index])

# Compute bootstrap t-statistic (under null hypothesis); invalid

# draws (singular SE) leave the NaN sentinel for filtering below.

if se_star > 0 and np.isfinite(beta_star[coefficient_index]):

bootstrap_t_stats[b] = (beta_star[coefficient_index] - null_hypothesis) / se_star

# Step 4: Compute bootstrap inference from VALID (finite) draws only.

#

# All-or-nothing NaN contract (per feedback_bootstrap_nan_on_invalid_contract):

# when bootstrap output is degenerate (fewer than 2 finite t-stats or

# 2 finite coefs), return NaN across the full inference surface (se,

# p_value, both CI endpoints, AND the surfaced t_stat_original). The

# original analytical t_stat is still computed in step 1 for diagnostic

# use but is NOT propagated to the user-facing result when bootstrap

# is degenerate — surfacing it alongside NaN se/p/CI would mix

# analytical and bootstrap inference families on the same coefficient.

finite_mask = np.isfinite(bootstrap_t_stats)

n_valid = int(finite_mask.sum())

valid_coefs = bootstrap_coefs[np.isfinite(bootstrap_coefs)]

lower_percentile = alpha / 2 * 100

upper_percentile = (1 - alpha / 2) * 100

if n_valid >= 2 and valid_coefs.size >= 2:

p_value = float(np.mean(np.abs(bootstrap_t_stats[finite_mask]) >= np.abs(t_stat_original)))

# Ensure p-value is at least 1/(n_valid+1) to avoid exact zero.

p_value = float(max(p_value, 1 / (n_valid + 1)))

se_bootstrap = float(np.std(valid_coefs, ddof=1))

ci_lower = float(np.percentile(valid_coefs, lower_percentile))

ci_upper = float(np.percentile(valid_coefs, upper_percentile))

surfaced_t_stat = t_stat_original

else:

# Degenerate bootstrap (insufficient valid draws): NaN-out the

# entire inference tuple. Downstream consumers (estimator-level

# `_run_wild_bootstrap_inference`) map these fields directly onto

# the result object; this guarantees the (se, t_stat, p_value, ci)

# quadruple moves together rather than reporting analytical t_stat

# with NaN se.

p_value = float("nan")

se_bootstrap = float("nan")

ci_lower = float("nan")

ci_upper = float("nan")

surfaced_t_stat = float("nan")

return WildBootstrapResults(

se=se_bootstrap,

p_value=p_value,

t_stat_original=surfaced_t_stat,

ci_lower=ci_lower,

ci_upper=ci_upper,

n_clusters=n_clusters,

n_bootstrap=n_bootstrap,

weight_type=weight_type,

alpha=alpha,

bootstrap_distribution=bootstrap_coefs if return_distribution else None,

)

def check_parallel_trends(

data: pd.DataFrame,

outcome: str,

time: str,

treatment_group: str,

pre_periods: Optional[List[Any]] = None,

) -> Dict[str, Any]:

"""

Perform a simple check for parallel trends assumption.

This computes the trend (slope) in the outcome variable for both

treatment and control groups during pre-treatment periods.

Parameters

----------

data : pd.DataFrame

Panel data.

outcome : str

Name of outcome variable column.

time : str

Name of time period column.

treatment_group : str

Name of treatment group indicator column.

pre_periods : list, optional

List of pre-treatment time periods. If None, infers from data.

Returns

-------

dict

Dictionary with trend statistics and test results.

"""

if pre_periods is None:

# Assume treatment happens at median time period

all_periods = sorted(data[time].unique())

mid_point = len(all_periods) // 2

pre_periods = all_periods[:mid_point]

pre_data = data[data[time].isin(pre_periods)]

# Compute trends for each group

treated_data = pre_data[pre_data[treatment_group] == 1]

control_data = pre_data[pre_data[treatment_group] == 0]

# Simple linear regression for trends

def compute_trend(group_data: pd.DataFrame) -> Tuple[float, float]:

time_values = group_data[time].values

outcome_values = group_data[outcome].values

# Normalize time to start at 0

time_norm = time_values - time_values.min()

# Compute slope using least squares

n = len(time_norm)

if n < 2:

return np.nan, np.nan

mean_t = np.mean(time_norm)

mean_y = np.mean(outcome_values)

# Check for zero variance in time (all same time period)

time_var = np.sum((time_norm - mean_t) ** 2)

if time_var == 0:

return np.nan, np.nan

slope = np.sum((time_norm - mean_t) * (outcome_values - mean_y)) / time_var

# Compute standard error of slope

y_hat = mean_y + slope * (time_norm - mean_t)

residuals = outcome_values - y_hat

mse = np.sum(residuals**2) / (n - 2)

se_slope = np.sqrt(mse / time_var)

return slope, se_slope

treated_slope, treated_se = compute_trend(treated_data)

control_slope, control_se = compute_trend(control_data)

# Test for difference in trends

slope_diff = treated_slope - control_slope

se_diff = np.sqrt(treated_se**2 + control_se**2)

t_stat, p_value, _ = safe_inference(slope_diff, se_diff)

return {

"treated_trend": treated_slope,

"treated_trend_se": treated_se,

"control_trend": control_slope,

"control_trend_se": control_se,

"trend_difference": slope_diff,

"trend_difference_se": se_diff,

"t_statistic": t_stat,

"p_value": p_value,

"parallel_trends_plausible": p_value > 0.05 if not np.isnan(p_value) else None,

}

def check_parallel_trends_robust(

data: pd.DataFrame,

outcome: str,

time: str,

treatment_group: str,

unit: Optional[str] = None,

pre_periods: Optional[List[Any]] = None,

n_permutations: int = 1000,

seed: Optional[int] = None,

wasserstein_threshold: float = 0.2,

) -> Dict[str, Any]:

"""

Perform robust parallel trends testing using distributional comparisons.

Uses the Wasserstein (Earth Mover's) distance to compare the full

distribution of outcome changes between treated and control groups,

with permutation-based inference.

Parameters

----------

data : pd.DataFrame

Panel data with repeated observations over time.

outcome : str

Name of outcome variable column.

time : str

Name of time period column.

treatment_group : str

Name of treatment group indicator column (0/1).

unit : str, optional

Name of unit identifier column. If provided, computes unit-level

changes. Otherwise uses observation-level data.

pre_periods : list, optional

List of pre-treatment time periods. If None, uses first half of periods.

n_permutations : int, default=1000

Number of permutations for computing p-value.

seed : int, optional

Random seed for reproducibility.

wasserstein_threshold : float, default=0.2

Threshold for normalized Wasserstein distance. Values below this

threshold (combined with p > 0.05) suggest parallel trends are plausible.

Returns

-------

dict

Dictionary containing:

- wasserstein_distance: Wasserstein distance between group distributions

- wasserstein_p_value: Permutation-based p-value

- ks_statistic: Kolmogorov-Smirnov test statistic

- ks_p_value: KS test p-value

- mean_difference: Difference in mean changes

- variance_ratio: Ratio of variances in changes

- treated_changes: Array of outcome changes for treated

- control_changes: Array of outcome changes for control

- parallel_trends_plausible: Boolean assessment

Examples

--------

>>> results = check_parallel_trends_robust(

... data, outcome='sales', time='year',

... treatment_group='treated', unit='firm_id'

... )

>>> print(f"Wasserstein distance: {results['wasserstein_distance']:.4f}")

>>> print(f"P-value: {results['wasserstein_p_value']:.4f}")

Notes

-----

The Wasserstein distance (Earth Mover's Distance) measures the minimum

"cost" of transforming one distribution into another. Unlike simple

mean comparisons, it captures differences in the entire distribution

shape, making it more robust to non-normal data and heterogeneous effects.

A small Wasserstein distance and high p-value suggest the distributions

of pre-treatment changes are similar, supporting the parallel trends

assumption.

"""

# Use local RNG to avoid affecting global random state

rng = np.random.default_rng(seed)