Instrumental Variables & Difference-in-Differences

"""Two-stage least squares and difference-in-differences - transcribed.

2SLS in its most literal form: regress the treatment on the instruments, then
regress the outcome on the fitted treatment. This gives the correct POINT
ESTIMATE and, as written, the wrong standard errors - which is the single most
common way this method is implemented incorrectly.
"""


def solve_normal_equations(X, y):
    """OLS by Gaussian elimination on X^T X b = X^T y."""
    n = len(X)
    d = len(X[0])

    gram = [[0.0] * d for _ in range(d)]
    rhs = [0.0] * d
    for i in range(n):
        for j in range(d):
            rhs[j] += X[i][j] * y[i]
            for k in range(d):
                gram[j][k] += X[i][j] * X[i][k]

    augmented = [gram[j][:] + [rhs[j]] for j in range(d)]

    for column in range(d):
        pivot = max(range(column, d), key=lambda r: abs(augmented[r][column]))
        if abs(augmented[pivot][column]) < 1e-12:
            raise ValueError("singular system: regressors are collinear")
        augmented[column], augmented[pivot] = augmented[pivot], augmented[column]

        for row in range(column + 1, d):
            factor = augmented[row][column] / augmented[column][column]
            for k in range(column, d + 1):
                augmented[row][k] -= factor * augmented[column][k]

    beta = [0.0] * d
    for row in reversed(range(d)):
        acc = augmented[row][d]
        for k in range(row + 1, d):
            acc -= augmented[row][k] * beta[k]
        beta[row] = acc / augmented[row][row]
    return beta


def two_stage_least_squares(y, X, Z):
    """X holds the endogenous treatment plus controls; Z the instruments plus
    the same controls."""
    # First stage: project every column of X onto the instrument space.
    fitted = [[0.0] * len(X[0]) for _ in range(len(X))]
    for column in range(len(X[0])):
        target = [X[i][column] for i in range(len(X))]
        coefficients = solve_normal_equations(Z, target)
        for i in range(len(X)):
            fitted[i][column] = sum(Z[i][j] * coefficients[j] for j in range(len(Z[0])))

    # Second stage: outcome on the FITTED treatment. Correct estimate.
    beta = solve_normal_equations(fitted, y)

    # The residuals below are wrong for inference. They are computed against
    # the FITTED treatment, but the model is about the ACTUAL one - so any
    # standard error built from them understates the true uncertainty. Every
    # real 2SLS implementation recomputes residuals against the original X,
    # which is exactly why hand-rolling this from two OLS calls is the common
    # way to get a right-looking estimate with wrong inference.
    residuals = []
    for i in range(len(y)):
        prediction = sum(fitted[i][j] * beta[j] for j in range(len(beta)))
        residuals.append(y[i] - prediction)

    return beta, residuals


def difference_in_differences(outcomes, treated, post):
    """The 2x2 estimator: four cell means, and a difference of differences."""
    sums = {(0, 0): 0.0, (0, 1): 0.0, (1, 0): 0.0, (1, 1): 0.0}
    counts = {(0, 0): 0, (0, 1): 0, (1, 0): 0, (1, 1): 0}

    for i in range(len(outcomes)):
        cell = (treated[i], post[i])
        sums[cell] += outcomes[i]
        counts[cell] += 1

    for cell, count in counts.items():
        if count == 0:
            raise ValueError(f"cell {cell} is empty; the design needs all four")

    treated_change = sums[(1, 1)] / counts[(1, 1)] - sums[(1, 0)] / counts[(1, 0)]
    control_change = sums[(0, 1)] / counts[(0, 1)] - sums[(0, 0)] / counts[(0, 0)]

    # The control group's change estimates what would have happened anyway -
    # which is the parallel-trends assumption, and it is untestable.
    return treated_change - control_change

"""2SLS - correct residuals, clustered errors, first-stage diagnostics."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


@dataclass(frozen=True)
class IvEstimate:
    """The estimate, its uncertainty, and the diagnostic that decides whether
    either is interpretable."""

    coefficients: Vector
    standard_errors: Vector
    first_stage_f: float
    n_clusters: int

    @property
    def is_weak(self) -> bool:
        """First-stage F below 10 is the classical warning, and the modern
        view is that 10 is far too lenient. A weak instrument biases 2SLS back
        toward OLS - toward exactly the bias it was meant to remove - while the
        standard errors continue to look precise."""
        return self.first_stage_f < 10.0

    def confidence_interval(self, index: int, z: float = 1.96) -> tuple[float, float]:
        centre = float(self.coefficients[index])
        half = z * float(self.standard_errors[index])
        return centre - half, centre + half


def two_stage_least_squares(
    y: Vector,
    X: Matrix,
    Z: Matrix,
    clusters: NDArray,
) -> IvEstimate:
    """Fit by 2SLS with cluster-robust standard errors.

    Raises ValueError on malformed input or an under-identified system.
    """
    if X.ndim != 2 or Z.ndim != 2:
        raise ValueError(f"X and Z must be 2-D, got {X.shape} and {Z.shape}")
    if not (X.shape[0] == Z.shape[0] == y.shape[0] == clusters.shape[0]):
        raise ValueError("y, X, Z and clusters must have the same number of rows")
    if Z.shape[1] < X.shape[1]:
        raise ValueError(
            f"under-identified: {Z.shape[1]} instruments for {X.shape[1]} regressors"
        )

    n = y.size

    # Project X onto the instrument space. lstsq rather than the normal
    # equations: forming Z^T Z squares the condition number, and instrument
    # matrices are frequently ill-conditioned because the controls appear in
    # both X and Z.
    first_stage, *_ = np.linalg.lstsq(Z, X, rcond=None)
    projected = Z @ first_stage

    coefficients, *_ = np.linalg.lstsq(projected, y, rcond=None)

    # THE CORRECTION. Residuals are computed against the ORIGINAL X, not the
    # projected one. Using the projected regressors gives the right point
    # estimate and understates the variance, because the model is about the
    # actual treatment rather than its fitted value. This one line is the
    # difference between correct and incorrect inference, and it is why 2SLS
    # should never be assembled from two calls to an OLS routine.
    residuals = y - X @ coefficients

    # Cluster-robust variance: observations within a cluster are correlated,
    # so the meat matrix sums cluster-level score vectors rather than
    # individual ones. Clustering at the observation level instead understates
    # the errors by a factor that grows with cluster size.
    bread = np.linalg.inv(projected.T @ projected)
    meat = np.zeros_like(bread)
    unique = np.unique(clusters)
    for cluster in unique:
        rows = clusters == cluster
        score = projected[rows].T @ residuals[rows]
        meat += np.outer(score, score)

    correction = (len(unique) / (len(unique) - 1)) * ((n - 1) / (n - X.shape[1]))
    covariance = correction * bread @ meat @ bread

    # First-stage F on the excluded instruments: the diagnostic that decides
    # whether anything above is interpretable.
    excluded = Z.shape[1] - (X.shape[1] - 1)
    treatment = X[:, 0]
    residual_full = treatment - Z @ np.linalg.lstsq(Z, treatment, rcond=None)[0]
    controls = Z[:, excluded:]
    residual_reduced = treatment - controls @ np.linalg.lstsq(controls, treatment, rcond=None)[0]
    numerator = (residual_reduced @ residual_reduced - residual_full @ residual_full) / excluded
    denominator = (residual_full @ residual_full) / (n - Z.shape[1])

    return IvEstimate(
        coefficients=coefficients,
        standard_errors=np.sqrt(np.diag(covariance)),
        first_stage_f=float(numerator / denominator),
        n_clusters=len(unique),
    )

"""2SLS - QR projections reused, cluster sums without a Python loop."""

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


class TwoStageLeastSquares:
    """Both stages share one factorization, and the clustering is a scatter-add.

    The instrument matrix is factorized ONCE. Every column of X projects
    through the same Q, so the first stage is a single triangular solve rather
    than one regression per endogenous regressor - and the same factor answers
    the first-stage F test afterwards.
    """

    def __init__(self, Z: Matrix) -> None:
        # One contiguous float64 block: the factorization and every subsequent
        # product read it, and a non-contiguous array forces a copy on each.
        self._Z = np.ascontiguousarray(Z, dtype=np.float64)
        # QR of the instruments, not Z^T Z. Forming the cross-product squares
        # the condition number, and instrument matrices are routinely
        # ill-conditioned because the controls appear in both X and Z - which
        # makes the columns nearly collinear by construction.
        self._q, self._r = np.linalg.qr(self._Z, mode="reduced")

    def project(self, X: Matrix) -> Matrix:
        """P_Z X, for every column at once: Q (Q^T X), never forming P_Z."""
        return self._q @ (self._q.T @ X)

    def fit(self, y: Vector, X: Matrix, clusters: NDArray) -> tuple[Vector, Vector, float]:
        design = np.ascontiguousarray(X, dtype=np.float64)
        target = np.ascontiguousarray(y, dtype=np.float64)
        n, d = design.shape

        projected = self.project(design)
        coefficients = np.linalg.lstsq(projected, target, rcond=None)[0]

        # Residuals against the ORIGINAL regressors - the correction that
        # separates valid inference from the naive two-regression version.
        residuals = target - design @ coefficients

        # Cluster scores without iterating over clusters. Each row's score is
        # projected[i] * residual[i]; summing them within a cluster is a
        # scatter-add over a dense index, which np.add.at performs in one pass.
        codes = np.unique(clusters, return_inverse=True)[1]
        n_clusters = int(codes.max()) + 1

        scores = np.empty((n_clusters, d), dtype=np.float64)      # allocated once
        scores.fill(0.0)
        np.add.at(scores, codes, projected * residuals[:, None])

        # meat = sum_g s_g s_g^T over clusters, which is one GEMM once the
        # per-cluster scores exist as a matrix.
        meat = scores.T @ scores
        bread = np.linalg.inv(projected.T @ projected)

        correction = (n_clusters / (n_clusters - 1)) * ((n - 1) / (n - d))
        covariance = correction * bread @ meat @ bread

        # First-stage F reuses the same factorization: the residual sum of
        # squares from the full instrument set is available from Q^T without
        # solving again.
        treatment = design[:, 0]
        fitted = self._q @ (self._q.T @ treatment)
        residual_full = treatment - fitted
        excluded = self._Z.shape[1] - (d - 1)
        controls = self._Z[:, excluded:]
        reduced = controls @ np.linalg.lstsq(controls, treatment, rcond=None)[0]
        residual_reduced = treatment - reduced

        numerator = (residual_reduced @ residual_reduced - residual_full @ residual_full) / excluded
        denominator = (residual_full @ residual_full) / (n - self._Z.shape[1])

        return coefficients, np.sqrt(np.diag(covariance)), float(numerator / denominator)