Propensity Scores & IPTW

"""Propensity scores and IPTW - the two stages, written out.

Stage 1: a logistic regression predicting WHO GOT TREATED.
Stage 2: a difference of means, weighted by the inverse of that probability.
"""

import math


def fit_propensity(X, treatment, lr=0.05, epochs=2_000):
    """Plain logistic regression by gradient ascent on the log-likelihood."""
    n = len(X)
    d = len(X[0])
    beta = [0.0] * d

    for _ in range(epochs):
        grad = [0.0] * d
        for i in range(n):
            z = sum(beta[j] * X[i][j] for j in range(d))
            p = 1.0 / (1.0 + math.exp(-z))
            residual = treatment[i] - p          # ascent, so t - p not p - t
            for j in range(d):
                grad[j] += residual * X[i][j]
        for j in range(d):
            beta[j] += lr * grad[j] / n

    return beta


def propensity(x, beta):
    z = sum(b * v for b, v in zip(beta, x))
    return 1.0 / (1.0 + math.exp(-z))


def iptw_effect(X, treatment, outcome, beta, clip=0.01):
    """Weighted difference of means.

    Clipping is not optional tidying: a propensity of 0.001 produces a weight
    of 1000, and that single unit would otherwise decide the answer.
    """
    treated_num = treated_den = control_num = control_den = 0.0

    for x, t, y in zip(X, treatment, outcome):
        e = min(max(propensity(x, beta), clip), 1.0 - clip)
        if t == 1:
            w = 1.0 / e
            treated_num += w * y
            treated_den += w
        else:
            w = 1.0 / (1.0 - e)
            control_num += w * y
            control_den += w

    return treated_num / treated_den - control_num / control_den

"""IPTW - typed, with the diagnostics that decide whether to trust the estimate."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray
from sklearn.linear_model import LogisticRegression

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


@dataclass(frozen=True)
class CausalEstimate:
    """An effect is never a bare number - it travels with its diagnostics."""

    ate: float
    std_error: float
    max_weight: float
    effective_sample_size: float
    trimmed_fraction: float

    @property
    def is_trustworthy(self) -> bool:
        """Cheap guard against the two failures that matter in practice."""
        return self.max_weight < 20.0 and self.trimmed_fraction < 0.10


def estimate_ate(
    X: Matrix,
    treatment: Vector,
    outcome: Vector,
    trim: tuple[float, float] = (0.05, 0.95),
) -> CausalEstimate:
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if not (X.shape[0] == treatment.shape[0] == outcome.shape[0]):
        raise ValueError("X, treatment and outcome must describe the same units")
    if not np.isin(treatment, (0, 1)).all():
        raise ValueError("treatment must be binary 0/1")

    model = LogisticRegression(max_iter=1_000)
    model.fit(X, treatment)
    scores = model.predict_proba(X)[:, 1]

    # Common support: a unit with no counterpart in the other arm cannot be
    # compared to anything, and keeping it just manufactures variance.
    low, high = trim
    keep = (scores > low) & (scores < high)
    trimmed_fraction = 1.0 - float(keep.mean())

    scores, t, y = scores[keep], treatment[keep], outcome[keep]

    # Stabilized weights: scaling by the marginal treatment probability keeps
    # the weights centred near 1, which materially reduces their variance.
    marginal = t.mean()
    weights = np.where(t == 1, marginal / scores, (1 - marginal) / (1 - scores))

    treated_mean = np.average(y[t == 1], weights=weights[t == 1])
    control_mean = np.average(y[t == 0], weights=weights[t == 0])
    ate = float(treated_mean - control_mean)

    # Kish effective sample size - how much data the weighting actually left.
    ess = float(weights.sum() ** 2 / np.square(weights).sum())
    std_error = float(np.sqrt(np.var(y) / ess))

    return CausalEstimate(
        ate=ate,
        std_error=std_error,
        max_weight=float(weights.max()),
        effective_sample_size=ess,
        trimmed_fraction=trimmed_fraction,
    )


def standardized_mean_differences(X: Matrix, treatment: Vector, weights: Vector) -> Vector:
    """Balance check. Read this BEFORE the effect - if the covariates are not
    balanced, the effect is measuring the imbalance."""
    t = treatment == 1
    diffs = np.average(X[t], axis=0, weights=weights[t]) - np.average(
        X[~t], axis=0, weights=weights[~t]
    )
    pooled_sd = np.sqrt((X[t].var(axis=0) + X[~t].var(axis=0)) / 2.0)
    return diffs / np.where(pooled_sd == 0, 1.0, pooled_sd)

"""IPTW - cross-fitted propensities and a bootstrap that reuses one draw matrix."""

import numpy as np
from numpy.typing import NDArray
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold

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


def crossfit_propensity(X: Matrix, treatment: Vector, folds: int = 5) -> Vector:
    """Out-of-fold propensity scores.

    A propensity fitted on the same rows it scores is overfitted toward the
    observed assignment, which biases the weights toward 1 and understates the
    correction. Cross-fitting removes that at the cost of one extra pass.
    """
    scores = np.empty(X.shape[0], dtype=np.float64)
    splitter = StratifiedKFold(n_splits=folds, shuffle=True, random_state=0)

    for train_idx, test_idx in splitter.split(X, treatment):
        model = LogisticRegression(max_iter=1_000)
        model.fit(X[train_idx], treatment[train_idx])
        scores[test_idx] = model.predict_proba(X[test_idx])[:, 1]

    return scores


def bootstrap_ate(
    scores: Vector,
    treatment: Vector,
    outcome: Vector,
    draws: int = 1_000,
    rng: np.random.Generator | None = None,
) -> tuple[float, float]:
    """Bootstrap CI over a resample index matrix rather than a Python loop.

    Drawing all `draws` resamples at once turns 1000 sequential O(n) passes
    into a handful of vectorized reductions over one (draws, n) matrix.
    """
    generator = rng if rng is not None else np.random.default_rng(0)
    n = outcome.shape[0]

    marginal = treatment.mean()
    weights = np.where(treatment == 1, marginal / scores, (1 - marginal) / (1 - scores))

    # One (draws, n) index matrix; every resample is a row.
    idx = generator.integers(0, n, size=(draws, n))
    w = weights[idx]
    t = treatment[idx]
    y = outcome[idx]

    # Masked weighted means, computed for all draws in two reductions.
    wt = w * t
    wc = w * (1.0 - t)
    treated = (wt * y).sum(axis=1) / np.maximum(wt.sum(axis=1), 1e-12)
    control = (wc * y).sum(axis=1) / np.maximum(wc.sum(axis=1), 1e-12)

    effects = treated - control
    return float(effects.mean()), float(effects.std(ddof=1))