"""

Unified linear algebra backend for diff-diff.

This module provides optimized OLS and variance estimation with an optional

Rust backend for maximum performance.

The key optimizations are:

1. scipy.linalg.lstsq with 'gelsd' driver (SVD-based, handles rank-deficient matrices)

2. Vectorized cluster-robust SE via groupby (eliminates O(n*clusters) loop)

3. Single interface for all estimators (reduces code duplication)

4. Optional Rust backend for additional speedup (when available)

5. R-style rank deficiency handling: detect, warn, and set NA for dropped columns

The Rust backend is automatically used when available, with transparent

fallback to NumPy/SciPy implementations.

Rank Deficiency Handling

------------------------

When a design matrix is rank-deficient (has linearly dependent columns), the OLS

solution is not unique. This module follows R's `lm()` approach:

1. Detect rank deficiency using pivoted QR decomposition

2. Identify which columns are linearly dependent

3. Drop redundant columns from the solve

4. Set NA (NaN) for coefficients of dropped columns

5. Warn with clear message listing dropped columns

6. Compute valid SEs for remaining (identified) coefficients

This is controlled by the `rank_deficient_action` parameter:

- "warn" (default): Emit warning, set NA for dropped coefficients

- "error": Raise ValueError with dropped column information

- "silent": No warning, but still set NA for dropped coefficients

"""

import warnings

from dataclasses import dataclass

from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload

import numpy as np

import pandas as pd

from scipy import stats

from scipy.linalg import lstsq as scipy_lstsq

from scipy.linalg import qr

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

from diff_diff._backend import (

HAS_RUST_BACKEND,

_rust_compute_robust_vcov,

_rust_solve_ols,

)

# Conley (1999) spatial HAC helpers live in diff_diff.conley to keep this

# module focused on linear-algebra primitives. Imported at the top so the

# `ConleyMetric` type alias is in scope for the public function signatures

# below (which advertise `conley_metric: ConleyMetric`).

from diff_diff.conley import (

ConleyMetric,

_compute_conley_vcov,

_validate_conley_kwargs,

)

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

# Utility Functions

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

def _factorize_cluster_ids(cluster_ids: np.ndarray) -> np.ndarray:

"""

Convert cluster IDs to contiguous integer codes for Rust backend.

Handles string, categorical, or non-contiguous integer cluster IDs by

mapping them to contiguous integers starting from 0.

Parameters

----------

cluster_ids : np.ndarray

Cluster identifiers (can be strings, integers, or categorical).

Returns

-------

np.ndarray

Integer cluster codes (dtype int64) suitable for Rust backend.

"""

# Use pandas factorize for efficient conversion of any dtype

codes, _ = pd.factorize(cluster_ids)

return codes.astype(np.int64)

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

# Rank Deficiency Detection and Handling

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

def _detect_rank_deficiency(

X: np.ndarray,

rcond: Optional[float] = None,

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

"""

Detect rank deficiency using pivoted QR decomposition.

This follows R's lm() approach of using pivoted QR to detect which columns

are linearly dependent. The pivoting ensures we drop the "least important"

columns (those with smallest contribution to the column space).

Rank detection is scale-invariant. A raw pivoted QR runs first; if the design

is genuinely rank-deficient (the deficiency persists under equilibration) its

raw pivot selects which columns to drop, preserving the established drop

order. But if the raw deficiency DISAPPEARS once columns are equilibrated to

unit 2-norm (a scale artifact), the higher equilibrated rank and its pivot

selection are used instead. This repairs the case where a column on a large

raw scale (e.g. an unstandardized covariate ~1e8) inflates the rank threshold

— which is anchored to the largest pivot diagonal — and false-drops

well-scaled columns (intercept / treatment / interaction) on an otherwise

full-rank design. In both branches the retained columns are guaranteed full

rank (pivoted QR has strictly decreasing |R[i,i]|). For a genuinely collinear

design with no scale disparity the dropped column is unchanged.

Parameters

----------

X : ndarray of shape (n, k)

Design matrix.

rcond : float, optional

Relative condition number threshold for determining rank.

Diagonal elements of R smaller than rcond * max(|R_ii|) are treated

as zero. If None, uses 1e-07 to match R's qr() default tolerance.

Returns

-------

rank : int

Numerical rank of the matrix.

dropped_cols : ndarray of int

Indices of columns that are linearly dependent (should be dropped).

Empty if matrix is full rank.

pivot : ndarray of int

Column permutation from QR decomposition.

"""

n, k = X.shape

# R's qr() uses tol = 1e-07 by default (sqrt(eps) ≈ 1.49e-08); we use 1e-07.

if rcond is None:

rcond = 1e-07

def _rank_and_pivot(M: np.ndarray) -> Tuple[int, np.ndarray]:

# Pivoted QR: M @ P = Q @ R. The rank threshold is anchored to the

# largest pivot diagonal |R[0,0]| (decreasing after pivoting).

_Q, R, piv = qr(M, mode="economic", pivoting=True)

r_diag = np.abs(np.diag(R))

if r_diag[0] == 0:

return 0, piv

return int(np.sum(r_diag > rcond * r_diag[0])), piv

# Stage 1 — raw pivoted QR. The common full-rank case exits here with zero

# added cost, and (when the design IS genuinely rank-deficient) this raw

# pivot is what selects which columns to drop, so the established

# drop-column selection for collinear designs is preserved exactly.

rank_raw, pivot_raw = _rank_and_pivot(X)

if rank_raw == k:

return k, np.array([], dtype=int), pivot_raw

# Stage 2 — raw QR reported a deficiency. That can be GENUINE collinearity

# or a SCALE artifact: the threshold is anchored to the largest pivot

# diagonal, so a column on a large raw scale (e.g. an unstandardized

# covariate ~1e8) inflates the threshold and false-drops well-scaled columns

# (intercept / treatment / interaction) on an otherwise full-rank design.

# Equilibrate each column to unit 2-norm and re-detect the rank COUNT only;

# if it is higher, the raw drop was scale-induced. Zero-norm columns keep

# scale 1.0 (no divide-by-zero; still pivot last).

col_norms = np.sqrt(np.einsum("ij,ij->j", X, X))

safe_norms = np.where(col_norms > 0, col_norms, 1.0)

rank_eq, pivot_eq = _rank_and_pivot(X / safe_norms)

if rank_eq > rank_raw:

# SCALE-INDUCED under-count: the raw threshold was inflated by a

# large-scale column. Trust the equilibrated rank AND its pivot — its

# first `rank_eq` columns are independent under the scale-corrected

# criterion (pivoted QR has strictly decreasing |R[i,i]|, so the kept

# set is guaranteed full rank and the retained design is identified).

# The raw pivot's tail is NOT reused here because, with a scale-corrupted

# ordering, its first `rank_eq` columns need not coincide with a

# scale-corrected independent subset.

rank, pivot = rank_eq, pivot_eq

else:

# GENUINE collinearity (no scale disparity): preserve the established

# raw pivot drop selection so the dropped column is unchanged.

rank, pivot = rank_raw, pivot_raw

# Columns after the rank position (in pivot order) are linearly dependent.

# The pivot indexes the ORIGINAL columns, so no remapping is needed.

if rank < k:

dropped_cols = np.sort(pivot[rank:])

else:

dropped_cols = np.array([], dtype=int)

return rank, dropped_cols, pivot

def _format_dropped_columns(

dropped_cols: np.ndarray,

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

) -> str:

"""

Format dropped column information for error/warning messages.

Parameters

----------

dropped_cols : ndarray of int

Indices of dropped columns.

column_names : list of str, optional

Names for the columns. If None, uses indices.

Returns

-------

str

Formatted string describing dropped columns.

"""

if len(dropped_cols) == 0:

return ""

if column_names is not None:

names = [column_names[i] if i < len(column_names) else f"column {i}" for i in dropped_cols]

if len(names) == 1:

return f"'{names[0]}'"

elif len(names) <= 5:

return ", ".join(f"'{n}'" for n in names)

else:

shown = ", ".join(f"'{n}'" for n in names[:5])

return f"{shown}, ... and {len(names) - 5} more"

else:

if len(dropped_cols) == 1:

return f"column {dropped_cols[0]}"

elif len(dropped_cols) <= 5:

return ", ".join(f"column {i}" for i in dropped_cols)

else:

shown = ", ".join(f"column {i}" for i in dropped_cols[:5])

return f"{shown}, ... and {len(dropped_cols) - 5} more"

def _expand_coefficients_with_nan(

coef_reduced: np.ndarray,

k_full: int,

kept_cols: np.ndarray,

) -> np.ndarray:

"""

Expand reduced coefficients to full size, filling dropped columns with NaN.

Parameters

----------

coef_reduced : ndarray of shape (rank,)

Coefficients for kept columns only.

k_full : int

Total number of columns in original design matrix.

kept_cols : ndarray of int

Indices of columns that were kept.

Returns

-------

ndarray of shape (k_full,)

Full coefficient vector with NaN for dropped columns.

"""

coef_full = np.full(k_full, np.nan)

coef_full[kept_cols] = coef_reduced

return coef_full

def _expand_vcov_with_nan(

vcov_reduced: np.ndarray,

k_full: int,

kept_cols: np.ndarray,

) -> np.ndarray:

"""

Expand reduced vcov matrix to full size, filling dropped entries with NaN.

Parameters

----------

vcov_reduced : ndarray of shape (rank, rank)

Variance-covariance matrix for kept columns only.

k_full : int

Total number of columns in original design matrix.

kept_cols : ndarray of int

Indices of columns that were kept.

Returns

-------

ndarray of shape (k_full, k_full)

Full vcov matrix with NaN for dropped rows/columns.

"""

vcov_full = np.full((k_full, k_full), np.nan)

# Use advanced indexing to fill in the kept entries

ix = np.ix_(kept_cols, kept_cols)

vcov_full[ix] = vcov_reduced

return vcov_full

def _equilibrated_lstsq(X: np.ndarray, y: np.ndarray) -> np.ndarray:

"""Solve OLS via column-equilibrated lstsq for scale robustness.

``scipy_lstsq``'s ``cond=1e-7`` truncates singular values relative to the

largest, so a column on a large scale truncates the genuine small-scale

direction and returns finite-but-wrong coefficients (e.g. a covariate ~1e8

alongside a 0/1 dummy). Equilibrating each column to unit 2-norm before the

solve and unscaling the result (``beta = beta_scaled / norm``) is

algebraically exact for the raw system and numerically scale-invariant.

Zero-norm columns use scale 1.0 (their coefficient is unidentified anyway and

is handled by the rank-deficiency path upstream).

"""

col_norms = np.sqrt(np.einsum("ij,ij->j", X, X))

safe_norms = np.where(col_norms > 0, col_norms, 1.0)

coef_scaled = scipy_lstsq(

X / safe_norms, y, lapack_driver="gelsd", check_finite=False, cond=1e-07

)[0]

return coef_scaled / safe_norms

def _rank_guarded_inv(

A: np.ndarray,

*,

rcond: float = 1e-10,

tracker: Optional[list] = None,

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

"""Rank-guarded (generalized) inverse of a symmetric PSD Gram matrix.

Influence-function standard errors invert a covariate Gram matrix

``A = X'WX`` (or a propensity-score Hessian). A constant or collinear

covariate makes ``A`` *near*-singular, but ``np.linalg.solve`` / ``inv``

only raise ``LinAlgError`` on an *exactly* singular matrix, so they return a

garbage inverse (entries ~1e13) that flows straight into the SE. This helper

detects near-singularity in a scale-invariant way and returns a finite

generalized inverse on the identified subspace, truncating redundant

directions. It returns an all-NaN matrix only when the design collapses to

rank 0.

Parameters

----------

A : np.ndarray

Square, symmetric, positive-semidefinite Gram matrix (k x k).

rcond : float, default 1e-10

Relative eigenvalue threshold applied to the *symmetrically

equilibrated* matrix ``D^{-1/2} A D^{-1/2}`` (``D = diag(A)``): a

direction is truncated when its equilibrated eigenvalue is

``<= rcond * max_eigenvalue``. ``1e-10`` (not the design-side ``1e-7``

of :func:`_detect_rank_deficiency`) because a Gram matrix squares the

condition number of ``X``; matches EfficientDiD's ``tol / max_eigval``

relative threshold.

tracker : list, optional

When provided and at least one direction is truncated, ONE

condition-number sample of ``A`` is appended (under ``np.errstate``).

Callers pass the per-fit fallback tracker and must NOT append

themselves: this helper is the sole owner, so the aggregate fallback

warning counts each deficient inversion exactly once.

Returns

-------

(A_ginv, n_dropped, rank) : Tuple[np.ndarray, int, int]

``A_ginv`` is the generalized inverse (all-NaN when ``rank == 0``).

``n_dropped`` is the number of truncated directions, ``rank`` is

``k - n_dropped``.

Notes

-----

Column-drop (not minimum-norm) generalized inverse: when ``A`` is

rank-deficient the guarded path keeps the ``rank`` most-independent columns

(pivoted QR on the equilibrated Gram) and inverts that principal submatrix,

zero-filling the dropped rows/cols. This is a column-drop in the SAME FAMILY

as the point estimate (``_detect_rank_deficiency`` / R's ``lm()``; drop

redundant columns, not a minimum-norm pseudo-inverse), but the column

selection is computed on the equilibrated Gram and is scale-invariant — so it

may drop a different *member* of a collinear set than the point estimate's

raw pivot under mixed-scale *exact* collinearity (a documented deviation that

leaves the SE unchanged, since the identified subspace is the same whichever

redundant member is dropped; see ``docs/methodology/REGISTRY.md``). It is

still a column-drop, so the influence-function SE equals the well-conditioned

(near-collinear) limit: replacing the exactly-collinear covariate with a

near-collinear (full-rank) one yields the same SE to working precision. A

minimum-norm pseudo-inverse would instead diverge from column-drop whenever

the IF multiplier leaves ``range(A)`` — e.g. an outcome-regression bread fit

on the *control* (or a treated sub-cell) sample multiplied by a mean from a

cell where the covariate is NOT collinear — so it is rejected here. With

column-drop there is no such divergence; a covariate that is rank-deficient

only within one cell still legitimately enters the other cells' full-rank

fits, so the ATT and SE reflect that (poor) covariate specification.

The fast (well-conditioned) path returns ``np.linalg.solve(A, I)``

unchanged, so well-conditioned fits are numerically unaffected.

"""

k = A.shape[0]

if k == 0:

return np.zeros((0, 0), dtype=float), 0, 0

# Symmetric equilibration: scale row/col i by sqrt(A[i, i]) so the

# eigenvalue threshold is scale-invariant. Zero/negative diagonal -> 1.0

# (such a direction has no information and pivots into the truncated tail).

diag = np.diag(A).astype(float)

scales = np.sqrt(np.where(diag > 0.0, diag, 1.0))

inv_scales = 1.0 / scales

A_eq = A * inv_scales[:, None] * inv_scales[None, :]

# eigvalsh reads one triangle, so asymmetric round-off in A_eq is ignored.

# Eigenvalues give the scale-invariant rank; pivoted QR (below) selects which

# columns to keep when the design is deficient.

eigvals = np.linalg.eigvalsh(A_eq)

max_eig = float(eigvals[-1]) if eigvals.size else 0.0

thresh = rcond * max_eig if max_eig > 0.0 else 0.0

keep = eigvals > thresh

n_keep = int(np.count_nonzero(keep))

# Fast path: full rank -> exact solve (bit-identical to the prior code).

if max_eig > 0.0 and n_keep == k:

return np.linalg.solve(A, np.eye(k)), 0, k

# Rank-deficient: record one condition-number sample for the aggregate

# fallback warning (the helper is the sole owner of this append).

if tracker is not None:

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

tracker.append(float(np.linalg.cond(A)))

if n_keep == 0:

return np.full((k, k), np.nan), k, 0

# Column-drop generalized inverse: keep the n_keep most-independent columns

# (pivoted QR on the equilibrated Gram), invert that principal submatrix, and

# zero-fill the dropped rows/cols before un-scaling. This is the same

# generalized-inverse family the point estimate uses (drop redundant columns,

# not a minimum-norm pseudo-inverse, which diverges from column-drop when the

# IF multiplier leaves range(A) — e.g. a treated-cell mean). Equilibrating the

# selection (rather than a raw pivot) keeps the cell-only-aliasing SE equal to

# the well-conditioned near-collinear limit (se_ratio ~ 1).

_Q, _R, piv = qr(A_eq, mode="economic", pivoting=True)

kept = np.sort(piv[:n_keep])

A_eq_ginv = np.zeros((k, k), dtype=float)

A_eq_ginv[np.ix_(kept, kept)] = np.linalg.inv(A_eq[np.ix_(kept, kept)])

A_ginv = A_eq_ginv * inv_scales[:, None] * inv_scales[None, :]

n_dropped = k - n_keep

return A_ginv, n_dropped, n_keep

def _solve_ols_rust(

X: np.ndarray,

y: np.ndarray,

*,

cluster_ids: Optional[np.ndarray] = None,

return_vcov: bool = True,

return_fitted: bool = False,

) -> Optional[

Union[

Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],

Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],

]

]:

"""

Rust backend implementation of solve_ols for full-rank matrices.

This is only called when:

1. The Rust backend is available

2. The design matrix is full rank (no rank deficiency handling needed)

For rank-deficient matrices, the Python backend is used instead to

properly handle R-style NA coefficients for dropped columns.

Why the backends differ (by design):

- Rust uses SVD-based solve (minimum-norm solution for rank-deficient)

- Python uses pivoted QR to identify and drop linearly dependent columns

- ndarray-linalg doesn't support QR with pivoting, so Rust can't identify

which specific columns to drop

- For full-rank matrices, both approaches give identical results

- For rank-deficient matrices, only Python can provide R-style NA handling

Parameters

----------

X : np.ndarray

Design matrix of shape (n, k), must be full rank.

y : np.ndarray

Response vector of shape (n,).

cluster_ids : np.ndarray, optional

Cluster identifiers for cluster-robust SEs.

return_vcov : bool

Whether to compute variance-covariance matrix.

return_fitted : bool

Whether to return fitted values.

Returns

-------

coefficients : np.ndarray

OLS coefficients of shape (k,).

residuals : np.ndarray

Residuals of shape (n,).

fitted : np.ndarray, optional

Fitted values if return_fitted=True.

vcov : np.ndarray, optional

Variance-covariance matrix if return_vcov=True.

None

If Rust backend detects numerical instability and caller should

fall back to Python backend.

"""

# Convert cluster_ids to int64 for Rust (handles string/categorical IDs)

if cluster_ids is not None:

cluster_ids = _factorize_cluster_ids(cluster_ids)

# Call Rust backend with fallback on numerical instability

try:

coefficients, residuals, vcov = _rust_solve_ols(

X, y, cluster_ids=cluster_ids, return_vcov=return_vcov

)

except ValueError as e:

error_msg = str(e).lower()

if "numerically unstable" in error_msg or "singular" in error_msg:

warnings.warn(

f"Rust backend detected numerical instability: {e}. "

"Falling back to Python backend.",

UserWarning,

stacklevel=3,

)

return None # Signal caller to use Python fallback

raise

# Convert to numpy arrays

coefficients = np.asarray(coefficients)

residuals = np.asarray(residuals)

if vcov is not None:

vcov = np.asarray(vcov)

# Return with optional fitted values

if return_fitted:

fitted = np.dot(X, coefficients)

return coefficients, residuals, fitted, vcov

else:

return coefficients, residuals, vcov

@overload

def solve_ols(

X: np.ndarray,

y: np.ndarray,

*,

cluster_ids: Optional[np.ndarray] = ...,

return_vcov: bool = ...,

return_fitted: Literal[False] = ...,

check_finite: bool = ...,

rank_deficient_action: str = ...,

column_names: Optional[List[str]] = ...,

skip_rank_check: bool = ...,

weights: Optional[np.ndarray] = ...,

weight_type: str = ...,

vcov_type: str = ...,

conley_coords: Optional[np.ndarray] = ...,

conley_cutoff_km: Optional[float] = ...,

conley_metric: ConleyMetric = ...,

conley_kernel: str = ...,

conley_time: Optional[np.ndarray] = ...,

conley_unit: Optional[np.ndarray] = ...,

conley_lag_cutoff: Optional[int] = ...,

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

@overload

def solve_ols(

X: np.ndarray,

y: np.ndarray,

*,

cluster_ids: Optional[np.ndarray] = ...,

return_vcov: bool = ...,

return_fitted: Literal[True],

check_finite: bool = ...,

rank_deficient_action: str = ...,

column_names: Optional[List[str]] = ...,

skip_rank_check: bool = ...,

weights: Optional[np.ndarray] = ...,

weight_type: str = ...,

vcov_type: str = ...,

conley_coords: Optional[np.ndarray] = ...,

conley_cutoff_km: Optional[float] = ...,

conley_metric: ConleyMetric = ...,

conley_kernel: str = ...,

conley_time: Optional[np.ndarray] = ...,

conley_unit: Optional[np.ndarray] = ...,

conley_lag_cutoff: Optional[int] = ...,

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

@overload

def solve_ols(

X: np.ndarray,

y: np.ndarray,

*,

cluster_ids: Optional[np.ndarray] = ...,

return_vcov: bool = ...,

return_fitted: bool,

check_finite: bool = ...,

rank_deficient_action: str = ...,

column_names: Optional[List[str]] = ...,

skip_rank_check: bool = ...,

weights: Optional[np.ndarray] = ...,

weight_type: str = ...,

vcov_type: str = ...,

conley_coords: Optional[np.ndarray] = ...,

conley_cutoff_km: Optional[float] = ...,

conley_metric: ConleyMetric = ...,

conley_kernel: str = ...,

conley_time: Optional[np.ndarray] = ...,

conley_unit: Optional[np.ndarray] = ...,

conley_lag_cutoff: Optional[int] = ...,

) -> Union[

Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],

Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],

]: ...

_VALID_WEIGHT_TYPES = {"pweight", "fweight", "aweight"}

def _validate_weights(weights, weight_type, n):

"""Validate weights array and weight_type for solve_ols/LinearRegression."""

if weight_type not in _VALID_WEIGHT_TYPES:

raise ValueError(

f"weight_type must be one of {_VALID_WEIGHT_TYPES}, " f"got '{weight_type}'"

)

if weights is not None:

weights = np.asarray(weights, dtype=np.float64)

if weights.shape[0] != n:

raise ValueError(f"weights length ({weights.shape[0]}) must match " f"X rows ({n})")

if np.any(np.isnan(weights)):

raise ValueError("Weights contain NaN values")

if np.any(np.isinf(weights)):

raise ValueError("Weights contain Inf values")

if np.any(weights < 0):

raise ValueError("Weights must be non-negative")

if np.sum(weights) <= 0:

raise ValueError(

"Weights sum to zero — no observations have positive weight. "

"Cannot fit a model on an empty effective sample."

)

if weight_type == "fweight":

fractional = weights - np.round(weights)

if np.any(np.abs(fractional) > 1e-10):

raise ValueError(

"Frequency weights (fweight) must be non-negative integers. "

"Fractional values detected. Use pweight for non-integer weights."

)

return weights

def solve_ols(

X: np.ndarray,

y: np.ndarray,

*,

cluster_ids: Optional[np.ndarray] = None,

return_vcov: bool = True,

return_fitted: bool = False,

check_finite: bool = True,

rank_deficient_action: str = "warn",

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

skip_rank_check: bool = False,

weights: Optional[np.ndarray] = None,

weight_type: str = "pweight",

vcov_type: str = "hc1",

conley_coords: Optional[np.ndarray] = None,

conley_cutoff_km: Optional[float] = None,

conley_metric: ConleyMetric = "haversine",

conley_kernel: str = "bartlett",

conley_time: Optional[np.ndarray] = None,

conley_unit: Optional[np.ndarray] = None,

conley_lag_cutoff: Optional[int] = None,

) -> Union[

Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],

Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],

]:

"""

Solve OLS regression with optional clustered standard errors.

This is the unified OLS solver for all diff-diff estimators. It uses

scipy's optimized LAPACK routines and vectorized variance estimation.

Parameters

----------

X : ndarray of shape (n, k)

Design matrix (should include intercept if desired).

y : ndarray of shape (n,)

Response vector.

cluster_ids : ndarray of shape (n,), optional

Cluster identifiers for cluster-robust standard errors.

If None, HC1 (heteroskedasticity-robust) SEs are computed.

return_vcov : bool, default True

Whether to compute and return the variance-covariance matrix.

Set to False for faster computation when SEs are not needed.

return_fitted : bool, default False

Whether to return fitted values in addition to residuals.

check_finite : bool, default True

Whether to check that X and y contain only finite values (no NaN/Inf).

Set to False for faster computation if you are certain your data is clean.

rank_deficient_action : str, default "warn"

How to handle rank-deficient design matrices:

- "warn": Emit warning and set NaN for dropped coefficients (R-style)

- "error": Raise ValueError with dropped column information

- "silent": No warning, but still set NaN for dropped coefficients

column_names : list of str, optional

Names for the columns (used in warning/error messages).

If None, columns are referred to by their indices.

skip_rank_check : bool, default False

If True, skip the pivoted QR rank check and use Rust backend directly

(when available). This saves O(nk²) computation but will not detect

rank-deficient matrices. Use only when you know the design matrix is

full rank. If the matrix is actually rank-deficient, results may be

incorrect (minimum-norm solution instead of R-style NA handling).

weights : ndarray of shape (n,), optional

Observation weights for Weighted Least Squares. When provided,

minimizes sum(w_i * (y_i - X_i @ beta)^2). Weights should be

pre-normalized (e.g., mean=1 for pweights).

weight_type : str, default "pweight"

Type of weights: "pweight" (inverse selection probability),

"fweight" (frequency), or "aweight" (inverse variance).

Affects variance estimation but not coefficient computation.

vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, default "hc1"

Variance-covariance family forwarded to :func:`compute_robust_vcov`:

- ``"classical"``: non-robust OLS SE, ``sigma_hat^2 * (X'X)^{-1}``.

One-way only; raises if ``cluster_ids`` is also passed.

- ``"hc1"``: heteroskedasticity-robust HC1 with ``n/(n-k)`` adjustment

(default). With ``cluster_ids``, dispatches to CR1 (Liang-Zeger).

- ``"hc2"``: leverage-corrected meat. One-way only; raises with

``cluster_ids`` (use ``"hc2_bm"`` for clustered Bell-McCaffrey).

- ``"hc2_bm"``: HC2 + Imbens-Kolesar (2016) Satterthwaite DOF one-way;

Pustejovsky-Tipton (2018) CR2 Bell-McCaffrey with ``cluster_ids``.

With ``weights``, dispatches to the clubSandwich WLS-CR2 port —

supported for ``weight_type="pweight"`` only. ``aweight`` and

``fweight`` raise ``NotImplementedError`` (port matches the

``pweight`` convention only; aweight/fweight derivations are a

separate methodology task).

- ``"conley"``: Conley (1999) spatial-HAC sandwich. Requires

``conley_coords`` (n × 2 array) and ``conley_cutoff_km`` (positive

bandwidth, no default per Conley 1999 Section 5's sensitivity-grid

recommendation). Two operating modes: cross-sectional (single-period

design or pooled cross-section) and panel block-decomposed (matches

R ``conleyreg`` with ``lag_cutoff > 0``); switch by passing the

three co-required kwargs ``conley_time`` / ``conley_unit`` /

``conley_lag_cutoff``. Combining with ``cluster_ids`` applies the

combined spatial + cluster product kernel (a diff-diff convention;

see :func:`compute_robust_vcov` for details). Combining with

``weights`` raises ``NotImplementedError`` regardless of

``weight_type``: weighted Conley is not implemented on the generic

linalg surface. For probability-sampling weights (``pweight`` /

``survey_design``) the deferral additionally reflects an open

methodological question — no canonical extension of Conley (1999)

exists for weighted spatial-HAC under probability sampling.

conley_coords : ndarray of shape (n, 2), optional

Required when ``vcov_type="conley"``. Two-column array of

``[lat, lon]`` (degrees, for ``conley_metric="haversine"``) or

projected coordinates (for ``conley_metric="euclidean"`` / callable

metric).

conley_cutoff_km : float, optional

Required when ``vcov_type="conley"``. Positive finite bandwidth in

km (haversine) or coord units (euclidean / callable).

conley_metric : {"haversine", "euclidean", callable}, default "haversine"

Distance metric. Haversine uses Earth's mean radius 6371.01 km

(matching R ``conleyreg``). Euclidean treats coords as already

projected. Callable signature ``(coords1, coords2) -> n×n``.

conley_kernel : {"bartlett", "uniform"}, default "bartlett"

Kernel evaluated on pairwise distance ``d_ij/h``. Both kernels emit

a ``UserWarning`` if the resulting meat is materially indefinite;

the radial 1-D Bartlett (matching R ``conleyreg``) is not formally

PSD-guaranteed — see :func:`compute_robust_vcov`.

Returns

-------

coefficients : ndarray of shape (k,)

OLS coefficient estimates. For rank-deficient matrices, coefficients

of linearly dependent columns are set to NaN.

residuals : ndarray of shape (n,)

Residuals (y - fitted). For rank-deficient matrices, uses only

identified coefficients to compute fitted values.

fitted : ndarray of shape (n,), optional

Fitted values. For full-rank matrices, this is X @ coefficients.

For rank-deficient matrices, uses only identified coefficients

(X_reduced @ coefficients_reduced). Only returned if return_fitted=True.

vcov : ndarray of shape (k, k) or None

Variance-covariance matrix (HC1 or cluster-robust).

For rank-deficient matrices, rows/columns for dropped coefficients

are filled with NaN. None if return_vcov=False.

Notes

-----

This function detects rank-deficient matrices using pivoted QR decomposition

and handles them following R's lm() approach:

1. Detect linearly dependent columns via pivoted QR

2. Drop redundant columns and solve the reduced system

3. Set NaN for coefficients of dropped columns

4. Compute valid SEs for identified coefficients only

5. Expand vcov matrix with NaN for dropped rows/columns

The cluster-robust standard errors use the sandwich estimator with the

standard small-sample adjustment: (G/(G-1)) * ((n-1)/(n-k)).

Examples

--------

>>> import numpy as np

>>> from diff_diff.linalg import solve_ols

>>> X = np.column_stack([np.ones(100), np.random.randn(100)])

>>> y = 2 + 3 * X[:, 1] + np.random.randn(100)

>>> coef, resid, vcov = solve_ols(X, y)

>>> print(f"Intercept: {coef[0]:.2f}, Slope: {coef[1]:.2f}")

For rank-deficient matrices with collinear columns:

>>> X = np.random.randn(100, 3)

>>> X[:, 2] = X[:, 0] + X[:, 1] # Perfect collinearity

>>> y = np.random.randn(100)

>>> coef, resid, vcov = solve_ols(X, y) # Emits warning

>>> print(np.isnan(coef[2])) # Dropped column has NaN coefficient

True

"""

# Validate inputs

X = np.asarray(X, dtype=np.float64)

y = np.asarray(y, dtype=np.float64)

if X.ndim != 2:

raise ValueError(f"X must be 2-dimensional, got shape {X.shape}")

if y.ndim != 1:

raise ValueError(f"y must be 1-dimensional, got shape {y.shape}")

if X.shape[0] != y.shape[0]:

raise ValueError(

f"X and y must have same number of observations: " f"{X.shape[0]} vs {y.shape[0]}"

)

n, k = X.shape

if n < k:

raise ValueError(

f"Fewer observations ({n}) than parameters ({k}). "

"Cannot solve underdetermined system."

)

# Validate rank_deficient_action

valid_actions = {"warn", "error", "silent"}

if rank_deficient_action not in valid_actions:

raise ValueError(

f"rank_deficient_action must be one of {valid_actions}, "

f"got '{rank_deficient_action}'"

)

# Check for NaN/Inf values if requested

if check_finite:

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

raise ValueError(

"X contains NaN or Inf values. "

"Clean your data or set check_finite=False to skip this check."

)

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

raise ValueError(

"y contains NaN or Inf values. "

"Clean your data or set check_finite=False to skip this check."

)

# Front-door Conley-specific validation. Runs BEFORE the routing/backend

# branching so `return_vcov=False` cannot bypass the conley+weights /

# conley+cluster_ids guards. Conley needs this because survey-replicate

# paths in MultiPeriodDiD pass `return_vcov=False` + `weights=survey_w`

# and would otherwise silently return survey SEs under a Conley request.

# The full `_validate_vcov_args` is NOT called here because other vcov

# types (e.g. `hc2_bm` + replicate weights) intentionally fall through

# to the survey-vcov path with the analytical validator skipped — that

# legacy contract is preserved.

if vcov_type == "conley":

if cluster_ids is not None or weights is not None:

_validate_vcov_args(vcov_type, cluster_ids, weights)

# WLS transformation: apply sqrt(w) scaling to X and y

# This happens BEFORE routing to Rust or NumPy backends — they receive

# pre-transformed X_w, y_w and solve standard OLS.

# Residuals are back-transformed to original scale afterward.

_original_X = None

_original_y = None

if weights is not None:

weights = _validate_weights(weights, weight_type, n)

_original_X = X

_original_y = y

sqrt_w = np.sqrt(weights)

X = X * sqrt_w[:, np.newaxis]

y = y * sqrt_w

# When weights are present, compute vcov separately on original-scale data

# to avoid double-weighting. The backend only computes point estimates.

_weighted_vcov_external = weights is not None

_backend_return_vcov = return_vcov and not _weighted_vcov_external

# Fast path: skip rank check and use Rust directly when requested

# This saves O(nk²) QR overhead but won't detect rank-deficient matrices

result = None # Will hold the tuple from backend functions

if skip_rank_check:

if (

HAS_RUST_BACKEND

and _rust_solve_ols is not None

and weights is None

and vcov_type == "hc1"

):

result = _solve_ols_rust(

X,

y,

cluster_ids=cluster_ids,

return_vcov=_backend_return_vcov,

return_fitted=return_fitted,

)

# result is None on numerical instability → fall through

if result is None:

result = _solve_ols_numpy(

X,

y,

cluster_ids=cluster_ids,

return_vcov=_backend_return_vcov,

return_fitted=return_fitted,

rank_deficient_action=rank_deficient_action,

column_names=column_names,

_skip_rank_check=True,

vcov_type=vcov_type,

conley_coords=conley_coords,

conley_cutoff_km=conley_cutoff_km,

conley_metric=conley_metric,

conley_kernel=conley_kernel,

conley_time=conley_time,

conley_unit=conley_unit,

conley_lag_cutoff=conley_lag_cutoff,

)

else:

# Check for rank deficiency using fast pivoted QR decomposition.

# Rank detection operates on (possibly weighted) X since collinearity

# depends on the weighted column space.

rank, dropped_cols, pivot = _detect_rank_deficiency(X)

is_rank_deficient = len(dropped_cols) > 0

# Routing strategy:

# - Full-rank + Rust available + no weights + HC1 vcov_type → fast Rust

# - Weighted or rank-deficient or non-HC1 vcov_type → Python backend

# - Rust numerical instability → Python fallback (via None return)

if (

HAS_RUST_BACKEND

and _rust_solve_ols is not None

and not is_rank_deficient

and weights is None

and vcov_type == "hc1"

):

result = _solve_ols_rust(

X,

y,

cluster_ids=cluster_ids,

return_vcov=_backend_return_vcov,

return_fitted=return_fitted,

)

if result is not None:

vcov_check = result[-1]

if _backend_return_vcov and vcov_check is not None and np.any(np.isnan(vcov_check)):

warnings.warn(

"Rust backend detected ill-conditioned matrix (NaN in variance-covariance). "

"Re-running with Python backend for proper rank detection.",

UserWarning,

stacklevel=2,

)

result = None # Force Python fallback below

if result is None:

result = _solve_ols_numpy(

X,

y,

cluster_ids=cluster_ids,

return_vcov=_backend_return_vcov,

return_fitted=return_fitted,

rank_deficient_action=rank_deficient_action,

column_names=column_names,

_precomputed_rank_info=(rank, dropped_cols, pivot),

vcov_type=vcov_type,

conley_coords=conley_coords,

conley_cutoff_km=conley_cutoff_km,

conley_metric=conley_metric,

conley_kernel=conley_kernel,

conley_time=conley_time,

conley_unit=conley_unit,

conley_lag_cutoff=conley_lag_cutoff,

)

# Back-transform residuals and compute weighted vcov on original-scale data.

# The WLS transform (sqrt(w) scaling) is for point estimates only. Vcov must

# be computed on original X and residuals with weights applied exactly once.

if _original_X is not None and _original_y is not None:

if return_fitted:

coefficients, _resid_w, _fitted_w, vcov_out = result

else:

coefficients, _resid_w, vcov_out = result

# Handle rank-deficient case: use only identified columns for fitted values

# to avoid NaN propagation from dropped coefficients

nan_mask = np.isnan(coefficients)

if np.any(nan_mask):

kept_cols = np.where(~nan_mask)[0]

fitted_orig = np.dot(_original_X[:, kept_cols], coefficients[kept_cols])

else:

fitted_orig = np.dot(_original_X, coefficients)

residuals_orig = _original_y - fitted_orig

if return_vcov:

if np.any(nan_mask):

kept_cols = np.where(~nan_mask)[0]