match-stream 0.1.0__py3-none-any.whl

match_stream/__init__.py

@@ -0,0 +1,45 @@
+"""match-stream: matching on data larger than memory.
+
+Nearest-neighbor matching with or without replacement (on a score or a
+multivariate distance), calipers on the score and on covariates, exact
+strata, kernel and ridge matching, coarsened exact matching, balance
+diagnostics, and treatment effects with Abadie-Imbens or influence-function
+standard errors. Conventions follow R's MatchIt. See
+https://leetucker.github.io/match-stream/.
+
+Layout
+    api.py          match_nn(), match_kernel(), match_cem()
+    bandwidth.py    PairDistance, CrossValidation (kernel bandwidth rules)
+    _spec.py        argument parsing and checks -> Spec
+    _prep.py        pass 0: the score model (feglm_stream), the narrow columns
+    _match.py       nearest neighbor: sort, order, run the kernel, both memory paths
+    _kernel.py      kernel and ridge matching: bandwidths, weights
+    _cem.py         coarsened exact matching (streaming group-by)
+    kernels/        numba: nn.py (score), metric.py (Mahalanobis etc.),
+                    kernel.py (window moments), var.py (conditional variances)
+    matches.py      MatchSet, CEMSet, KernelSet: pairs, weights, counts
+    balance.py      BalanceTable: MatchIt's summary(m, un = TRUE)
+    effects.py      effect(): estimates, AI and influence-function variances; MatchResult
+    simulate.py     simulate_matching()
+    _core/          code bound for hdfe_stream.core (influence functions, IPW,
+                    the first-stage correction, the weighted balance table)
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from .api import match_cem, match_kernel, match_nn
+from .balance import BalanceTable
+from .bandwidth import CrossValidation, PairDistance
+from .effects import MatchResult
+from .matches import CEMSet, KernelSet, MatchSet
+from .simulate import simulate_matching
+
+try:
+    __version__ = _version("match-stream")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.1.0"
+
+__all__ = ["BalanceTable", "CEMSet", "CrossValidation", "KernelSet", "MatchResult", "MatchSet",
+           "PairDistance", "__version__", "match_cem", "match_kernel", "match_nn",
+           "simulate_matching"]

match_stream/_cem.py

@@ -0,0 +1,107 @@
+"""Coarsened exact matching (Iacus, King and Porro), as MatchIt's method="cem".
+
+CEM is a group-by: each matching covariate is cut into bins, a cell is a
+combination of bins and exact-matching values, and a cell is kept when it has
+both treated and control units. Two streaming passes over the source (the
+covariates' ranges, then the cells' counts) and a third that writes the
+matched units' weights; nothing the size of the source is held.
+
+Cutpoints follow MatchIt for a number of bins: equal-width bins over the
+covariate's range among the units used (R's ``cut(x, breaks = k)``; a value on
+a boundary goes to the upper bin here and to the lower one in R, which matters
+only for values exactly on a boundary, of measure zero for continuous data).
+A list gives the cut points themselves (bins are left-closed).
+
+Weights (ATT): a treated unit in a matched cell weighs 1; a control weighs
+(treated in its cell / controls in its cell) x (matched controls / matched
+treated), MatchIt's weights. The ATC swaps the roles.
+"""
+
+from __future__ import annotations
+
+import polars as pl
+from hdfe_stream.core.utils import as_lazy
+
+from ._spec import Spec
+
+ROW = "__ms_row"
+CELL = "__ms_cell"
+
+
+def _bins(x: str, cut, lo: float, hi: float) -> pl.Expr:
+    if isinstance(cut, int):
+        if cut < 1:
+            raise ValueError(f"cutpoints for {x!r} must be a positive number of bins")
+        span = hi - lo
+        if span <= 0:
+            return pl.lit(0, pl.Int32)
+        return ((pl.col(x) - lo) / span * cut).floor().clip(0, cut - 1).cast(pl.Int32)
+    edges = sorted(float(e) for e in cut)
+    return pl.col(x).cut(edges, left_closed=True).to_physical().cast(pl.Int32)
+
+
+def cem(spec: Spec, data, cutpoints) -> tuple[pl.DataFrame, dict, pl.LazyFrame]:
+    """The matched units' weights (row, id, treated, weight, subclass = cell)
+    and the counts."""
+    covs = spec.covariates
+    if not covs and not spec.exact:
+        raise ValueError("CEM needs matching covariates in the formula or exact=")
+    if cutpoints is None:
+        cutpoints = {}
+    elif isinstance(cutpoints, int):
+        cutpoints = {c: cutpoints for c in covs}
+    unknown = set(cutpoints) - set(covs)
+    if unknown:
+        raise ValueError("cutpoints names columns that are not matching covariates: "
+                         f"{sorted(unknown)}")
+    missing_cut = [c for c in covs if c not in cutpoints]
+    if missing_cut:
+        # MatchIt's default is Sturges' rule through R's hist(); not reproduced
+        # exactly, so ask for the number of bins instead of guessing
+        raise ValueError(f"give cutpoints for every covariate (a number of bins or the cut "
+                         f"points): missing {missing_cut}")
+
+    lf = as_lazy(data).with_row_index(ROW).with_columns(pl.col(ROW).cast(pl.Int64))
+    keep = list(dict.fromkeys([ROW, spec.treatment, *spec.exact, *covs,
+                               *([spec.id] if spec.id is not None else [])]))
+    schema = lf.collect_schema()
+    missing = [c for c in keep if c not in schema]
+    if missing:
+        raise ValueError(f"columns not in the data: {missing}")
+    floats = [c for c in keep if schema[c].is_float()]
+    ok = pl.all_horizontal([pl.col(c).is_not_null() for c in keep]
+                           + [pl.col(c).is_not_nan() for c in floats])
+    proj = lf.select(keep)
+    n_source = proj.select(pl.len()).collect(engine="streaming").item()
+    proj = proj.filter(ok).with_columns(pl.col(spec.treatment).cast(pl.Boolean))
+    rng = proj.select([pl.col(c).min().alias(f"lo:{c}") for c in covs]
+                      + [pl.col(c).max().alias(f"hi:{c}") for c in covs]
+                      + [pl.len().alias("n")]).collect(engine="streaming").row(0, named=True)
+
+    bins = [_bins(c, cutpoints[c], rng[f"lo:{c}"], rng[f"hi:{c}"]).alias(f"__b:{c}") for c in covs]
+    keys = [*spec.exact, *[f"__b:{c}" for c in covs]]
+    coarse = proj.with_columns(bins)
+    t = pl.col(spec.treatment)
+    cells = (coarse.group_by(keys)
+             .agg(t.sum().cast(pl.Int64).alias("n_t"), (~t).sum().cast(pl.Int64).alias("n_c"))
+             .filter((pl.col("n_t") > 0) & (pl.col("n_c") > 0))
+             .collect(engine="streaming")
+             .sort(keys).with_row_index(CELL, offset=1)
+             .with_columns(pl.col(CELL).cast(pl.Int64)))
+    mt, mc = int(cells["n_t"].sum()), int(cells["n_c"].sum())
+    focal_t = spec.estimand != "atc"
+    focal_n, other_n = ("n_t", "n_c") if focal_t else ("n_c", "n_t")
+    m_focal, m_other = (mt, mc) if focal_t else (mc, mt)
+    is_focal = t if focal_t else ~t
+    w = (pl.when(is_focal).then(1.0)
+         .otherwise(pl.col(focal_n) / pl.col(other_n) * (m_other / max(m_focal, 1))))
+    out = (coarse.join(cells.lazy(), on=keys, how="inner")
+           .select(pl.col(ROW).alias("row"), *([spec.id] if spec.id is not None else []),
+                   t.alias("treated"), w.alias("weight"), pl.col(CELL).alias("subclass"))
+           .collect(engine="streaming").sort("row"))
+    n_t = int(proj.select(t.cast(pl.Int64).sum()).collect(engine="streaming").item())
+    counts = {"treated": n_t, "control": int(rng["n"]) - n_t, "matched_treated": mt,
+              "matched_control": mc, "unmatched_treated": n_t - mt,
+              "unmatched_control": int(rng["n"]) - n_t - mc, "cells": cells.height,
+              "dropped_missing": n_source - int(rng["n"])}
+    return out, counts, proj

match_stream/_core/__init__.py

@@ -0,0 +1,25 @@
+"""Staging for code bound for ``hdfe_stream.core``.
+
+Each module here has the API it will have in core and no matching-specific
+names, so that other packages (a streaming staggered-DiD package, for one)
+can use it once it moves. Until hdfe-stream releases it, match-stream imports
+it from here.
+
+Modules:
+
+- ``influence``: influence-function variances from streamed per-cluster
+  sums (one-way and multi-way); `weighted_mean_difference`.
+- ``mestimation``: a fitted binary logit/probit first stage (scores,
+  observed information) and the first-stage correction phi + J H^-1 s.
+- ``ipw``: inverse-probability weights (trimming, Hajek scale) and the
+  Hajek IPW difference in means with the estimated-score correction.
+- ``balance``: the weighted two-group balance table.
+
+All four are now in hdfe-stream's ``hdfe_stream.core`` (unreleased; same
+code, docstrings aside). When that hdfe-stream is released, switch the
+imports, pin it, and delete them here; until then edit both copies
+together.
+
+Not here yet: group tallies for ``core.profile.scan_source`` (they extend a
+function that lives in hdfe-stream and serve only ``profile=True``).
+"""

match_stream/_core/balance.py

@@ -0,0 +1,174 @@
+"""A weighted two-group balance table, out of core. Bound for
+``hdfe_stream.core.balance``.
+
+For each term (a numeric column or expression), the weighted means of two
+groups (1 and 0), the standardized mean difference, the variance ratio, and
+the mean and maximum gap between the groups' weighted empirical CDFs. Any
+weighting works: matching weights, inverse-probability weights, entropy
+balancing, or none (weight 1). Everything is streamed sums and one sort per
+term (for the eCDF), so the data are never held in memory.
+
+Conventions (MatchIt's, which cobalt shares; match-stream checks them
+against MatchIt 4.8.2):
+
+- weighted variance: weights normalized to sum to 1 within the group,
+  sum w (x - m)^2 / (1 - sum w^2); for a 0/1 term, m (1 - m);
+- the standardized mean difference is (m1 - m0) / scale, with `scale` the
+  term's standard deviation as the caller defines it (MatchIt: the focal
+  group's, before weighting; cobalt's default for the ATE: the pooled
+  one), else computed here from group 1, group 0 or both pooled (the
+  square root of the mean of the two variances); differences below
+  sqrt(eps) are reported unscaled, as MatchIt;
+- the eCDF gaps are evaluated at each distinct value of the term among all
+  rows given (a row with weight 0 contributes its value but no mass), each
+  group's weights summing to 1; for a 0/1 term both equal |m1 - m0|.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import polars as pl
+
+_G, _W = "__bal_g", "__bal_w"
+EPS_DIFF = np.sqrt(np.finfo(float).eps)
+
+
+def _prepare(lf, group, weight, terms):
+    lf = lf.lazy()
+    g = pl.col(group).cast(pl.Boolean) if isinstance(group, str) else group.cast(pl.Boolean)
+    if weight is None:
+        w = pl.lit(1.0)
+    else:
+        w = pl.col(weight) if isinstance(weight, str) else weight
+    names = [t if isinstance(t, str) else t.meta.output_name() for t in terms]
+    exprs = [(pl.col(t) if isinstance(t, str) else t).cast(pl.Float64).alias(f"__bal_x{i}")
+             for i, t in enumerate(terms)]
+    return lf.select(g.alias(_G), w.cast(pl.Float64).alias(_W), *exprs), names
+
+
+def group_moments(lf, n_terms: int) -> dict:
+    """Per group (1, 0) and term i: sum w, sum w^2, sum w x, sum w x^2, and
+    whether the term only takes 0 and 1 (over all rows given)."""
+    aggs = []
+    for arm, cond in (("1", pl.col(_G)), ("0", ~pl.col(_G))):
+        w = pl.col(_W).filter(cond)
+        aggs += [w.sum().alias(f"sw{arm}"), (w * w).sum().alias(f"sww{arm}")]
+        for i in range(n_terms):
+            xv = pl.col(f"__bal_x{i}").filter(cond)
+            aggs += [(w * xv).sum().alias(f"sx{arm}_{i}"),
+                     (w * xv * xv).sum().alias(f"sxx{arm}_{i}")]
+    for i in range(n_terms):
+        aggs.append(pl.col(f"__bal_x{i}").is_in([0.0, 1.0]).all().alias(f"bin_{i}"))
+    return lf.select(aggs).collect(engine="streaming").row(0, named=True)
+
+
+def mean_var(r: dict, arm: str, i: int, binary: bool) -> tuple[float, float]:
+    """The weighted mean and variance of term i in group `arm` ("1"/"0")."""
+    sw, sww = r[f"sw{arm}"], r[f"sww{arm}"]
+    if not sw:
+        return np.nan, np.nan
+    m = r[f"sx{arm}_{i}"] / sw
+    if binary:
+        return m, m * (1 - m)
+    denom = 1 - sww / sw**2
+    var = (r[f"sxx{arm}_{i}"] / sw - m * m) / denom if denom > 0 else np.nan
+    return m, var
+
+
+def ecdf_gaps(lf, i: int) -> tuple[float, float]:
+    """Mean and maximum |F_1(v) - F_0(v)| over the distinct values v of term i."""
+    norm = pl.col(_W) / pl.col(_W).sum().over(_G)
+    signed = pl.when(pl.col(_G)).then(norm).otherwise(-norm)
+    d = (lf.select(pl.col(f"__bal_x{i}").alias("x"), signed.alias("s"))
+         .sort("x")
+         .with_columns(pl.col("s").cum_sum().abs().alias("gap"))
+         .filter(pl.col("x") != pl.col("x").shift(-1).fill_null(np.inf))
+         .select(pl.col("gap").mean().alias("mean"), pl.col("gap").max().alias("max"))
+         .collect(engine="streaming"))
+    return float(d["mean"][0]), float(d["max"][0])
+
+
+@dataclass
+class GroupSizes:
+    """The rows of groups 1 and 0 (`n1`, `n0`) and their Kish effective
+    sample sizes (`ess1`, `ess0`): (sum w)^2 / sum w^2, the number of
+    equally weighted rows that would give a mean the same variance; equal
+    to the rows without weights, 0 for a group with no weight."""
+
+    n1: int
+    n0: int
+    ess1: float
+    ess0: float
+
+
+def effective_sizes(lf: pl.LazyFrame | pl.DataFrame, group, weight=None) -> GroupSizes:
+    """The rows and Kish effective sample sizes of the two groups in `lf`,
+    `group` and `weight` as in `weighted_balance` (a column or expression;
+    no weight: 1). One streaming pass."""
+    lf, _ = _prepare(lf, group, weight, [])
+    aggs = []
+    for arm, cond in (("1", pl.col(_G)), ("0", ~pl.col(_G))):
+        w = pl.col(_W).filter(cond)
+        aggs += [w.len().alias(f"n{arm}"), w.sum().alias(f"sw{arm}"),
+                 (w * w).sum().alias(f"sww{arm}")]
+    r = lf.select(aggs).collect(engine="streaming").row(0, named=True)
+
+    def ess(arm):
+        return r[f"sw{arm}"] ** 2 / r[f"sww{arm}"] if r[f"sww{arm}"] else 0.0
+    return GroupSizes(n1=int(r["n1"]), n0=int(r["n0"]), ess1=ess("1"), ess0=ess("0"))
+
+
+def weighted_balance(lf: pl.LazyFrame | pl.DataFrame, group, terms, weight=None, *,
+                     scale: str | dict = "pooled", ecdf: bool = True) -> pl.DataFrame:
+    """The balance table: one row per term with ``mean1``, ``mean0``,
+    ``var1``, ``var0``, ``smd``, ``var_ratio`` (var1 / var0; null for a 0/1
+    term or two zero variances; infinite for positive var1 and zero var0),
+    ``ecdf_mean`` and ``ecdf_max`` (|m1 - m0| for a 0/1 term).
+
+    Parameters
+    ----------
+    lf : the rows (a LazyFrame streams)
+    group : column or boolean expression, true for group 1
+    terms : columns or expressions (aliased) to compare
+    weight : column or expression; None weighs every row 1
+    scale : "group1", "group0", "pooled", or a dict term -> standard
+        deviation (e.g. computed before weighting, as MatchIt)
+    ecdf : compute the eCDF gaps (one sort per term)
+    """
+    data, names = _prepare(lf, group, weight, list(terms))
+    r = group_moments(data, len(names))
+    rows = []
+    for i, name in enumerate(names):
+        binary = bool(r[f"bin_{i}"])
+        m1, v1 = mean_var(r, "1", i, binary)
+        m0, v0 = mean_var(r, "0", i, binary)
+        if isinstance(scale, dict):
+            sd = float(scale[name])
+        elif scale == "group1":
+            sd = np.sqrt(v1)
+        elif scale == "group0":
+            sd = np.sqrt(v0)
+        elif scale == "pooled":
+            sd = np.sqrt((v1 + v0) / 2)
+        else:
+            raise ValueError(f"scale must be 'group1', 'group0', 'pooled' or a dict, "
+                             f"not {scale!r}")
+        diff = m1 - m0
+        smd = (
+            (np.copysign(np.inf, diff) if sd == 0 else diff / sd)
+            if abs(diff) > EPS_DIFF else diff
+        )
+        if binary:
+            ratio, emean, emax = None, abs(diff), abs(diff)
+        else:
+            ratio = v1 / v0 if v0 != 0 else (np.inf if v1 > 0 else None)
+            emean, emax = ecdf_gaps(data, i) if ecdf else (np.nan, np.nan)
+        rows.append({"term": name, "mean1": m1, "mean0": m0, "var1": v1, "var0": v0,
+                     "smd": smd, "var_ratio": ratio, "ecdf_mean": emean, "ecdf_max": emax})
+    return pl.DataFrame(rows, schema={"term": pl.String, "mean1": pl.Float64,
+                                      "mean0": pl.Float64, "var1": pl.Float64,
+                                      "var0": pl.Float64, "smd": pl.Float64,
+                                      "var_ratio": pl.Float64, "ecdf_mean": pl.Float64,
+                                      "ecdf_max": pl.Float64})

match_stream/_core/influence.py

@@ -0,0 +1,222 @@
+"""Influence-function variances from streamed sums, with one-way and
+multi-way clustering. Bound for ``hdfe_stream.core.influence``.
+
+An estimator that is (asymptotically) a sum of per-row contributions,
+``theta_hat - theta ~= sum_i phi_i``, has the sandwich variance
+
+    V = c * sum_g (sum_{i in g} phi_i)(sum_{i in g} phi_i)'
+
+over clusters g (each row its own cluster when there is no clustering).
+Only per-cluster sums of phi are needed, so `if_vcov` is one streaming
+``group_by`` per clustering term and holds G x p numbers, never the rows.
+Here phi is the influence function divided by N: it is on the scale of the
+estimate, and V needs no further 1/N^2.
+
+Multi-way clustering is Cameron, Gelbach and Miller's (2011)
+inclusion-exclusion: for dimensions A and B, V_A + V_B - V_(A and B), and so
+on for more. The small-sample factor follows hdfe-stream (pyfixest's
+defaults), so that an estimator that is also a least-squares coefficient
+gets the regression's standard error exactly:
+
+    c = G / (G - 1) * (N - 1) / (N - k)
+
+with G the smallest one-way cluster count (G = N without clustering) and k
+the parameters of that regression (`df_k`; None drops the second factor,
+leaving Stata's ``_robust`` convention G / (G - 1)). Every term of a
+multi-way sum takes the same c. A multi-way variance can have a negative
+diagonal; it is returned as is, with a warning, as in hdfe-stream.
+
+`weighted_mean_difference` is the first consumer: the difference between
+two groups' weighted means with the weights held fixed (matching weights,
+IPW with a known score, ...).
+"""
+
+from __future__ import annotations
+
+import itertools
+import warnings
+from dataclasses import dataclass
+
+import numpy as np
+import polars as pl
+
+
+@dataclass
+class IFVariance:
+    """A variance from `if_vcov`: `vcov` (p x p), the rows `n_obs`, the
+    one-way cluster counts `n_clusters` (by dimension; empty without
+    clustering), the small-sample factor `adj` and the degrees of freedom
+    `df` (smallest G minus 1, or N minus 1)."""
+
+    vcov: np.ndarray
+    n_obs: int
+    n_clusters: dict
+    adj: float
+    df: int
+
+
+def cluster_dims(cluster) -> list[str]:
+    """``None`` -> [], ``"a"`` -> ["a"], ``"a+b"`` or ``["a", "b"]`` -> ["a", "b"]."""
+    if cluster is None:
+        return []
+    if isinstance(cluster, str):
+        dims = [c.strip() for c in cluster.split("+")]
+    else:
+        dims = [str(c).strip() for c in cluster]
+    if not dims or any(not c for c in dims):
+        raise ValueError(f"bad cluster specification {cluster!r}")
+    if len(set(dims)) != len(dims):
+        raise ValueError(f"a cluster dimension is repeated in {cluster!r}")
+    return dims
+
+
+def if_vcov(lf: pl.LazyFrame | pl.DataFrame, phi: list[str], cluster=None, *,
+            df_k: int | None = None) -> IFVariance:
+    """The sandwich variance of an estimator with per-row contributions in
+    the columns `phi` of `lf` (see the module docstring), clustered on
+    `cluster` (a column, ``"a+b"`` or a list for multi-way; None: by row).
+
+    Rows with a null in a cluster column are an error (they belong to no
+    cluster)."""
+    lf = lf.lazy()
+    dims = cluster_dims(cluster)
+    p = len(phi)
+    f = [pl.col(c).cast(pl.Float64) for c in phi]
+    if dims:
+        nulls = (lf.select([pl.col(d).null_count().alias(d) for d in dims])
+                 .collect(engine="streaming").row(0, named=True))
+        bad = {d: n for d, n in nulls.items() if n}
+        if bad:
+            raise ValueError(f"cluster columns have missing values: {bad}")
+    n = int(lf.select(pl.len()).collect(engine="streaming").item())
+
+    def meat(by: list[str] | None) -> tuple[np.ndarray, int]:
+        if by is None:      # each row its own cluster: sum of outer products
+            prods = [(f[a] * f[b]).sum().alias(f"{a}_{b}")
+                     for a in range(p) for b in range(a, p)]
+            vals = lf.select(prods).collect(engine="streaming").row(0)
+            M = np.zeros((p, p))
+            it = iter(vals)
+            for a in range(p):
+                for b in range(a, p):
+                    M[a, b] = M[b, a] = next(it)
+            return M, n
+        s = (lf.group_by(by).agg([x.sum().alias(f"__s{j}") for j, x in enumerate(f)])
+             .collect(engine="streaming"))
+        A = s.select([f"__s{j}" for j in range(p)]).to_numpy()
+        return A.T @ A, s.height
+
+    V = np.zeros((p, p))
+    counts = {}
+    if not dims:
+        V, G = meat(None)
+    else:
+        for r in range(1, len(dims) + 1):
+            sign = 1.0 if r % 2 else -1.0
+            for sub in itertools.combinations(dims, r):
+                M, G_sub = meat(list(sub))
+                V += sign * M
+                if r == 1:
+                    counts[sub[0]] = G_sub
+        G = min(counts.values())
+    adj = G / (G - 1) if G > 1 else np.nan
+    if df_k is not None:
+        adj *= (n - 1) / (n - df_k) if n > df_k else np.nan
+    V = adj * (V + V.T) / 2
+    if np.any(np.diag(V) < 0):
+        warnings.warn("the multi-way clustered variance has negative diagonal entries",
+                      stacklevel=2)
+    return IFVariance(vcov=V, n_obs=n, n_clusters=counts, adj=float(adj), df=G - 1)
+
+
+@dataclass
+class WeightedMeanDifference:
+    """`estimate` = mean1 - mean0 (weighted means of the two groups), its
+    `variance` (an `IFVariance`), and the group totals of the weights."""
+
+    estimate: float
+    mean1: float
+    mean0: float
+    weight1: float
+    weight0: float
+    variance: IFVariance
+
+
+def weighted_mean_difference(lf: pl.LazyFrame | pl.DataFrame, y: str, w: str, group: str,
+                             cluster=None, *, df_k: int | None = 2,
+                             allow_negative: bool = False) -> WeightedMeanDifference:
+    """The difference of the weighted means of `y` between rows with `group`
+    true (1) and false (0), with weights `w` held fixed, and its
+    influence-function variance (`if_vcov`, clustered on `cluster`).
+
+    Each row contributes phi_i = w_i (y_i - mean1) / W1 in group 1 and
+    -w_i (y_i - mean0) / W0 in group 0 (W: the group's total weight). This is
+    the coefficient on `group` in a weighted least-squares regression of `y`
+    on (1, `group`) and, with the default ``df_k=2``, its HC1 / CRV1
+    standard error exactly (as hdfe-stream, pyfixest, or R's
+    ``sandwich::vcovCL``). Rows with zero weight change nothing but N in the
+    small-sample factor; leave them out. Negative weights (ridge or
+    local-linear matching) are an error unless `allow_negative`; each
+    group's total must be positive. Two streaming passes: the group totals,
+    then the clustered sums."""
+    lf = lf.lazy()
+    g = pl.col(group).cast(pl.Boolean)
+    W, Y = pl.col(w).cast(pl.Float64), pl.col(y).cast(pl.Float64)
+    t = (lf.select((W * g).sum().alias("W1"), (W * ~g).sum().alias("W0"),
+                   (W * Y * g).sum().alias("S1"), (W * Y * ~g).sum().alias("S0"),
+                   (W.is_null() | Y.is_null() | pl.col(group).is_null()).sum().alias("nulls"),
+                   (W < 0).sum().alias("neg"))
+         .collect(engine="streaming").row(0, named=True))
+    if t["nulls"]:
+        raise ValueError(f"{t['nulls']:,} rows have a missing {y!r}, {w!r} or {group!r}")
+    if t["neg"] and not allow_negative:
+        raise ValueError(f"{t['neg']:,} rows have a negative weight")
+    if not (t["W1"] > 0 and t["W0"] > 0):
+        raise ValueError("both groups need a positive total weight")
+    m1, m0 = t["S1"] / t["W1"], t["S0"] / t["W0"]
+    phi = (pl.when(g).then(W * (Y - m1) / t["W1"])
+           .otherwise(-W * (Y - m0) / t["W0"]).alias("__phi"))
+    var = if_vcov(lf.with_columns(phi), ["__phi"], cluster, df_k=df_k)
+    return WeightedMeanDifference(estimate=m1 - m0, mean1=m1, mean0=m0, weight1=t["W1"],
+                                  weight0=t["W0"], variance=var)
+
+
+def group_mean_phis(lf: pl.LazyFrame | pl.DataFrame, y: str, w: str, group: str,
+                    wmd: WeightedMeanDifference) -> pl.DataFrame:
+    """The per-row influence contributions behind `weighted_mean_difference`
+    (`wmd` is its result for the same columns): ``phi1`` for the mean of group
+    1 (w_i (y_i - mean1) / W1 on group-1 rows, else 0) and ``phi0`` for the
+    mean of group 0. The difference's contribution is ``phi1 - phi0``."""
+    g = pl.col(group).cast(pl.Boolean)
+    W, Y = pl.col(w).cast(pl.Float64), pl.col(y).cast(pl.Float64)
+    return lf.lazy().select(
+        pl.when(g).then(W * (Y - wmd.mean1) / wmd.weight1).otherwise(0.0).alias("phi1"),
+        pl.when(~g).then(W * (Y - wmd.mean0) / wmd.weight0).otherwise(0.0).alias("phi0"),
+    ).collect(engine="streaming")
+
+
+def joint_if_vcov(parts: list[tuple[pl.DataFrame, dict[int, str]]], p: int, cluster,
+                  factors: np.ndarray) -> IFVariance:
+    """The joint variance of p estimators that live on different tables.
+
+    Each part is ``(table, {j: column})``: a table holding the cluster
+    columns and, for each estimator j it contributes to, its per-row
+    influence contribution in `column`. Estimator j has contributions in one
+    part only, so a row of another part adds zero to it; the tables may share
+    clusters (the same unit, the same source column value), which is what
+    makes the estimators covary. The sums per cluster are those of `if_vcov`
+    over the stacked tables, without its small-sample factor; entry (a, b) is
+    then scaled by sqrt(factors[a] * factors[b]), where `factors[j]` is the
+    factor of estimator j's own fit (`IFVariance.adj`), so that the diagonal
+    equals each estimator's own variance exactly. The returned `adj` is 1."""
+    dims = cluster_dims(cluster)
+    keep = [d for d in dims if d != "__unit"] + (["__unit"] if "__unit" in dims else [])
+    frames = []
+    for tab, cols in parts:
+        frames.append(tab.select(*keep, *(pl.col(c).alias(f"__phi{j}") for j, c in cols.items())))
+    stack = pl.concat(frames, how="diagonal").with_columns(
+        [pl.col(f"__phi{j}").fill_null(0.0) for j in range(p)])
+    raw = if_vcov(stack, [f"__phi{j}" for j in range(p)], dims or None, df_k=None)
+    f = np.sqrt(np.asarray(factors, dtype=float))
+    return IFVariance(vcov=raw.vcov / raw.adj * np.outer(f, f), n_obs=raw.n_obs,
+                      n_clusters=raw.n_clusters, adj=1.0, df=raw.df)