diff-stream 0.1.0.dev0__py3-none-any.whl

diff_stream/_core/etwfe_kernels.py

@@ -0,0 +1,288 @@
+"""Compiled passes of the extended TWFE regression.
+
+The expanded design has a column for every cohort, period, cohort-by-covariate,
+period-by-covariate and treatment-cell effect, but a row has only a handful of
+nonzero entries: its cohort and period, the covariates interacted with each,
+and the covariates (and a one) in its treatment cell. The NumPy code builds the
+dense design for a batch of units and multiplies it; the kernels here visit
+only the nonzero entries, and remove the unit means (``fixed_effects="unit"``)
+analytically.
+
+Units are processed in fixed-size chunks whose partial sums are added in order,
+so results do not depend on the number of threads. The NumPy code in
+``etwfe.py`` remains the reference; the tests compare the two.
+"""
+
+import numpy as np
+from numba import njit, prange
+
+ENABLED = True  # tests switch it off to compare with the array path
+CHUNK = 64  # units per chunk
+PARTIAL_BYTES = 128e6  # memory of the chunks' partial Gram matrices
+
+
+@njit(cache=True, inline="always")
+def _row(x, i, t, gi, K, G, T, cohort_fe, base_c, base_t, offset, cell_at, cols, vals):
+    """The nonzero entries of one row of the design: fills ``cols`` and ``vals``
+    and returns their number and the row's treatment cell (-1 for none)."""
+    n = 0
+    at = 0
+    if cohort_fe:
+        cols[n] = gi
+        vals[n] = 1.0
+        n += 1
+        at = G
+    if t >= 1:
+        cols[n] = at + t - 1
+        vals[n] = 1.0
+        n += 1
+    cell = cell_at[gi, t]
+    for k in range(K):
+        v = x[i, t, k]
+        if not np.isfinite(v):
+            v = 0.0
+        cols[n] = base_c + gi * K + k
+        vals[n] = v
+        n += 1
+        if t >= 1:
+            cols[n] = base_t + (t - 1) * K + k
+            vals[n] = v
+            n += 1
+    if cell >= 0:
+        cols[n] = offset + cell * (K + 1)
+        vals[n] = 1.0
+        n += 1
+        for k in range(K):
+            v = x[i, t, k]
+            cols[n] = offset + cell * (K + 1) + 1 + k
+            vals[n] = v if np.isfinite(v) else 0.0
+            n += 1
+    return n, cell
+
+
+@njit(parallel=True, cache=True)
+def etwfe_moments(
+    gidx, y, x, w, G, cohort_fe, unit_fe, base_c, base_t, offset, cell_at, chunk, Q, q, norms,
+    means, cdiag, bad,
+):  # fmt: skip
+    """Gram matrix and moments of the regression, the unit-centered column
+    norms, and the weighted covariate sums of each treatment cell."""
+    n, T = y.shape
+    K = x.shape[2]
+    P = Q.shape[1]
+    chunks = Q.shape[0]
+    nz = 3 * K + 3
+    for c in prange(chunks):
+        cols = np.empty((T, nz), dtype=np.int64)
+        vals = np.empty((T, nz))
+        count = np.zeros(T, dtype=np.int64)
+        s = np.zeros(P)
+        wn = np.zeros(P)
+        seen = np.zeros(P, dtype=np.bool_)
+        touched = np.empty(T * nz, dtype=np.int64)
+        for i in range(c * chunk, min(n, (c + 1) * chunk)):
+            mass = 0.0
+            ysum = 0.0
+            ntouch = 0
+            gi = gidx[i]
+            for t in range(T):
+                m, cell = _row(
+                    x, i, t, gi, K, G, T, cohort_fe, base_c, base_t, offset, cell_at,
+                    cols[t], vals[t],
+                )  # fmt: skip
+                count[t] = m
+                wt = w[i, t]
+                yt = y[i, t] if np.isfinite(y[i, t]) else 0.0
+                mass += wt
+                ysum += wt * yt
+                for a in range(m):
+                    ca = cols[t, a]
+                    wa = wt * vals[t, a]
+                    q[c, ca] += wa * yt
+                    if unit_fe:
+                        norms[c, ca] += wa * vals[t, a]
+                        s[ca] += wa
+                        wn[ca] += wt
+                        if not seen[ca]:
+                            seen[ca] = True
+                            touched[ntouch] = ca
+                            ntouch += 1
+                    for b in range(m):
+                        Q[c, ca, cols[t, b]] += wa * vals[t, b]
+                if cell >= 0:
+                    means[c, cell, 0] += wt
+                    for k in range(K):
+                        v = x[i, t, k]
+                        means[c, cell, 1 + k] += wt * (v if np.isfinite(v) else 0.0)
+            if unit_fe:
+                if mass <= 0.0:
+                    bad[c] = 1
+                else:
+                    for a in range(ntouch):
+                        ca = touched[a]
+                        q[c, ca] -= s[ca] * ysum / mass
+                        for b in range(ntouch):
+                            cb = touched[b]
+                            Q[c, ca, cb] -= s[ca] * s[cb] / mass
+                    # the centered sum of squares itself, without the cancellation
+                    # of the subtraction above: a column the unit means remove
+                    # entirely must come out as zero
+                    for t in range(T):
+                        for a in range(count[t]):
+                            ca = cols[t, a]
+                            d = vals[t, a] - s[ca] / mass
+                            cdiag[c, ca] += w[i, t] * d * d
+                    for a in range(ntouch):
+                        ca = touched[a]
+                        dbar = s[ca] / mass
+                        cdiag[c, ca] += dbar * dbar * (mass - wn[ca])
+                for a in range(ntouch):
+                    s[touched[a]] = 0.0
+                    wn[touched[a]] = 0.0
+                    seen[touched[a]] = False
+
+
+@njit(parallel=True, cache=True)
+def etwfe_scores(
+    gidx, y, x, w, G, cohort_fe, unit_fe, base_c, base_t, offset, cell_at, coef, H, out, bad
+):  # fmt: skip
+    """Each unit's score of the regression times H: ``out[unit, :]``."""
+    n, T = y.shape
+    K = x.shape[2]
+    M = H.shape[1]
+    nz = 3 * K + 3
+    chunks = (n + CHUNK - 1) // CHUNK
+    for c in prange(chunks):
+        cols = np.empty((T, nz), dtype=np.int64)
+        vals = np.empty((T, nz))
+        count = np.zeros(T, dtype=np.int64)
+        resid = np.zeros(T)
+        for i in range(c * CHUNK, min(n, (c + 1) * CHUNK)):
+            gi = gidx[i]
+            mass = 0.0
+            rsum = 0.0
+            for t in range(T):
+                m, cell = _row(
+                    x, i, t, gi, K, G, T, cohort_fe, base_c, base_t, offset, cell_at,
+                    cols[t], vals[t],
+                )  # fmt: skip
+                count[t] = m
+                fit = 0.0
+                for a in range(m):
+                    fit += vals[t, a] * coef[cols[t, a]]
+                yt = y[i, t] if np.isfinite(y[i, t]) else 0.0
+                resid[t] = yt - fit
+                mass += w[i, t]
+                rsum += w[i, t] * resid[t]
+            if unit_fe:
+                if mass <= 0.0:
+                    bad[c] = 1
+                    continue
+                rsum /= mass
+            else:
+                rsum = 0.0
+            for t in range(T):
+                we = w[i, t] * (resid[t] - rsum)
+                if we == 0.0:
+                    continue
+                for a in range(count[t]):
+                    wa = we * vals[t, a]
+                    ca = cols[t, a]
+                    for m in range(M):
+                        out[i, m] += wa * H[ca, m]
+
+
+def _arrays(raw):
+    g, y, x, w = raw
+    return (
+        np.ascontiguousarray(y, dtype=np.float64),
+        np.ascontiguousarray(x, dtype=np.float64),
+        np.ascontiguousarray(w, dtype=np.float64),
+    )
+
+
+class ETWFEKernels:
+    """The compiled passes for an ``etwfe.ExtendedDesign``."""
+
+    def __init__(self, model):
+        self.model = model
+        self.P = len(model.names)
+        self.G = len(model.groups)
+        self.cohort_fe = model.fixed_effects == "cohort"
+        T = model.store.periods
+        K = model.K
+        self.base_c = (self.G if self.cohort_fe else 0) + T - 1
+        self.base_t = self.base_c + self.G * K
+        self.offset = model.offset
+        self.cell_at = np.ascontiguousarray(model.cell_at)
+        # chunks per batch: the partial Gram matrices must fit
+        self.max_chunks = int(PARTIAL_BYTES // (8 * self.P * self.P))
+
+    @property
+    def usable(self):
+        return ENABLED and self.max_chunks >= 2
+
+    def _chunking(self, n):
+        chunks = min((n + CHUNK - 1) // CHUNK, self.max_chunks)
+        return chunks, (n + chunks - 1) // chunks
+
+    def group_index(self, g):
+        return np.searchsorted(self.model.group_array, g).astype(np.int64)
+
+    def moments(self, batches):
+        """(Q, q, norms, means): the Gram matrix, moments, column norms and
+        covariate sums of each cell."""
+        model = self.model
+        P, M, K = self.P, len(model.cells), model.K
+        Q, q, norms = np.zeros((P, P)), np.zeros(P), np.zeros(P)
+        means = np.zeros((M, K + 1))
+        cdiag = np.zeros(P)
+        for raw in batches():
+            y, x, w = _arrays(raw)
+            g = raw[0]
+            chunks, size = self._chunking(len(g))
+            parts = (
+                np.zeros((chunks, P, P)),
+                np.zeros((chunks, P)),
+                np.zeros((chunks, P)),
+                np.zeros((chunks, M, K + 1)),
+                np.zeros((chunks, P)),
+                np.zeros(chunks, dtype=np.int8),
+            )
+            etwfe_moments(
+                self.group_index(g), y, x, w, self.G,
+                self.cohort_fe, model.fixed_effects == "unit", self.base_c, self.base_t,
+                self.offset, self.cell_at, size, *parts,
+            )  # fmt: skip
+            if parts[5].any():
+                raise ValueError("each group requires a positive-weight estimation observation")
+            Q += parts[0].sum(axis=0)
+            q += parts[1].sum(axis=0)
+            norms += parts[2].sum(axis=0)
+            means += parts[3].sum(axis=0)
+            cdiag += parts[4].sum(axis=0)
+        if model.fixed_effects == "unit":
+            # columns the unit means remove entirely are exactly zero, as in
+            # the array path; the subtraction leaves rounding noise
+            Q[np.diag_indices(P)] = cdiag
+            dead = cdiag <= norms * 1e-24
+            Q[dead, :] = 0.0
+            Q[:, dead] = 0.0
+            q[dead] = 0.0
+        return Q, q, norms, means
+
+    def scores(self, raw, coef, H_full):
+        """The batch's scores times H (units by cells); ``H_full`` has a row for
+        every column of the design (zero for the dropped ones)."""
+        model = self.model
+        y, x, w = _arrays(raw)
+        out = np.zeros((len(y), H_full.shape[1]))
+        bad = np.zeros((len(y) + CHUNK - 1) // CHUNK, dtype=np.int8)
+        etwfe_scores(
+            self.group_index(raw[0]), y, x, w, self.G, self.cohort_fe,
+            model.fixed_effects == "unit", self.base_c, self.base_t, self.offset, self.cell_at,
+            np.ascontiguousarray(coef), np.ascontiguousarray(H_full), out, bad,
+        )  # fmt: skip
+        if bad.any():
+            raise ValueError("each group requires a positive-weight estimation observation")
+        return out

diff_stream/_core/functionals.py

@@ -0,0 +1,30 @@
+"""Streaming normalized means and their smooth-functional derivatives.
+
+These primitives have no panel, date, treatment, or cohort conventions.
+"""
+
+import numpy as np
+
+
+class NormalizedMeans:
+    def __init__(self, terms, width):
+        self.weight = np.zeros(terms)
+        self.value = np.zeros(terms)
+        self.x = np.zeros((terms, width))
+        self.xvalue = np.zeros((terms, width))
+
+    def update(self, weights, values, X):
+        self.weight += weights.sum(axis=0)
+        self.value += (weights * values).sum(axis=0)
+        self.x += weights.T @ X
+        self.xvalue += (weights * values).T @ X
+
+    def finish(self):
+        if np.any(self.weight <= 0):
+            raise ValueError("a normalized mean has no contributing observations after trimming")
+        self.mean = self.value / self.weight
+        self.mean_x = self.x / self.weight[:, None]
+        self.centered_x = self.xvalue / self.weight[:, None] - self.mean[:, None] * self.mean_x
+
+    def scores(self, weights, values, signs):
+        return (weights * (values - self.mean) / self.weight) @ signs

diff_stream/_core/multiplier.py

@@ -0,0 +1,90 @@
+"""Bounded streaming sums of cluster influence scores times multipliers.
+
+The multiplier samplers are reused from hdfe-stream's wild bootstrap. Its OLS
+restricted-fit/studentization engine is deliberately not used here. Scores may
+belong to any asymptotically linear estimator; no DiD conventions live here.
+"""
+
+from numbers import Integral
+
+import numpy as np
+from hdfe_stream.core.wildboot import WEIGHTS
+from scipy.special import ndtri
+
+MAX_DRAW_ENTRIES = 1 << 18
+
+
+def validate_bootstrap(reps, seed, weights):
+    if (
+        isinstance(reps, (bool, np.bool_))
+        or not isinstance(reps, Integral)
+        or (reps != 0 and reps < 2)
+    ):
+        raise ValueError("multiplier_bootstrap_reps must be 0 (disabled) or an integer >= 2")
+    if seed is not None and (
+        isinstance(seed, (bool, np.bool_)) or not isinstance(seed, Integral) or seed < 0
+    ):
+        raise ValueError("multiplier_bootstrap_seed must be a nonnegative integer or None")
+    if not isinstance(weights, str) or weights not in WEIGHTS:
+        raise ValueError(f"multiplier_bootstrap_weights must be one of {sorted(WEIGHTS)}")
+
+
+class StreamingMultiplier:
+    """Keep repetitions-by-score-width state, never clusters-by-repetitions.
+
+    A unit receives the same multiplier in every score block. RNG draws are
+    unit-major, so changing numerical batch size preserves assignments for a
+    fixed unit traversal. Repartitioning can change traversal and assignments.
+    """
+
+    def __init__(self, width, reps, seed, weights):
+        validate_bootstrap(reps, seed, weights)
+        self.reps = int(reps)
+        self.rng = np.random.default_rng(seed)
+        self.sampler = WEIGHTS[weights]
+        self.draws = np.zeros((self.reps, width))
+        self.n_clusters = 0
+
+    def update(self, *blocks):
+        rows = len(blocks[0])
+        if (
+            any(len(b) != rows or b.ndim != 2 for b in blocks)
+            or sum(b.shape[1] for b in blocks) != self.draws.shape[1]
+        ):
+            raise ValueError("multiplier score blocks have incompatible shapes")
+        step = max(1, MAX_DRAW_ENTRIES // self.reps)
+        for start in range(0, rows, step):
+            stop = min(rows, start + step)
+            xi = np.asarray(
+                self.sampler((stop - start) * self.reps, self.rng), dtype=float
+            ).reshape(stop - start, self.reps)
+            column = 0
+            for block in blocks:
+                width = block.shape[1]
+                self.draws[:, column : column + width] += xi.T @ block[start:stop]
+                column += width
+        self.n_clusters += rows
+
+
+def multiplier_band(draws, alpha, *, joint=True, analytic_se=None):
+    """IQR-studentized bands, using R did's empirical (type 1) quantiles.
+
+    Optional analytic_se uses fixed analytic studentization instead. Draws
+    are centered perturbations on the estimate scale, not resampled estimates.
+    """
+    if isinstance(alpha, (bool, np.bool_)) or not np.isfinite(alpha) or not 0 < alpha < 1:
+        raise ValueError("alpha must be in (0, 1)")
+    if not isinstance(joint, bool):
+        raise ValueError("joint must be boolean")
+    if analytic_se is None:
+        q = np.quantile(draws, [0.25, 0.75], axis=0, method="inverted_cdf")
+        se = (q[1] - q[0]) / (ndtri(0.75) - ndtri(0.25))
+    else:
+        se = np.asarray(analytic_se)
+    if not np.isfinite(draws).all() or not np.isfinite(se).all() or np.any(se <= 0):
+        raise ValueError(
+            "selected coefficients have zero or non-finite bootstrap/analytic scale; select supported coefficients or increase multiplier_bootstrap_reps"
+        )
+    t = np.abs(draws / se)
+    critical = np.quantile(t.max(axis=1) if joint else t, 1 - alpha, axis=0, method="inverted_cdf")
+    return se, critical

diff_stream/_core/nuisance.py

@@ -0,0 +1,255 @@
+"""Shared-pass binary MLE and multi-outcome control-regression systems.
+
+No panel or treatment-cohort conventions live here. A transform selects a
+system's sample and returns (design, binary group, outcome matrix[, weights]) per batch.
+Intended for eventual hdfe_stream.core; current models are unpenalized, no FE.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.special import expit
+
+from .weights import weighted_linear, weighted_system
+
+
+def solve_checked(matrix, rhs, context, *, verify=True):
+    """Solve in diagonally scaled coordinates; never silently drop columns.
+
+    ``verify=False`` skips the eigenvalue check, for a matrix already checked
+    (its scaled form is the same whatever the right-hand side), which for a
+    matrix of hundreds of columns is most of the time of a solve."""
+    scale = np.sqrt(np.maximum(np.diag(matrix), 0))
+    if np.any(scale <= 0) or not np.isfinite(matrix).all():
+        raise ValueError(f"{context}: singular nuisance design")
+    normalized = matrix / np.outer(scale, scale)
+    if verify and np.linalg.eigvalsh(normalized)[0] <= 1e-12:
+        raise ValueError(f"{context}: singular or ill-conditioned nuisance design")
+    vector = np.ndim(rhs) == 1
+    target = np.asarray(rhs) / (scale if vector else scale[:, None])
+    answer = np.linalg.solve(normalized, target)
+    return answer / (scale if vector else scale[:, None])
+
+
+def fit_linear_systems(batches, transforms, width):
+    """Fit weighted OLS on arbitrary row strata together in one source pass.
+
+    Transforms return (X, y[, weights]). Return coefficient and Gram pairs, so callers can
+    propagate estimation uncertainty through arbitrary smooth functionals.
+    """
+    if not transforms:
+        return []
+    grams = [np.zeros((width, width)) for _ in transforms]
+    rhs = [np.zeros(width) for _ in transforms]
+    counts = np.zeros(len(transforms), dtype=np.int64)
+    for raw in batches():
+        for j, transform in enumerate(transforms):
+            X, y, w = weighted_linear(transform(raw))
+            counts[j] += len(y)
+            grams[j] += X.T @ (w[:, None] * X)
+            rhs[j] += X.T @ (w * y)
+    return [
+        (
+            solve_checked(
+                Q,
+                q,
+                f"stratified outcome regression (system {j}, {counts[j]} observations, {width} columns)",
+            ),
+            Q,
+        )
+        for j, (Q, q) in enumerate(zip(grams, rhs, strict=True))
+    ]
+
+
+@dataclass
+class NuisanceSystem:
+    beta: np.ndarray
+    regression: np.ndarray
+    control_gram: np.ndarray
+    n_treated: int = 0
+    n_controls: int = 0
+    iterations: int = 0
+    treated_weight: float = 0.0
+    control_weight: float = 0.0
+    control_weight_squares: float = 0.0
+    treated_weight_squares: float = 0.0
+    design_scale: np.ndarray | None = None
+
+
+def fit_systems(batches, transforms, method, names, *, tol=1e-10, maxiter=100, fast=None):
+    """Fit independent systems together, scanning the source once per Newton step.
+
+    Multi-outcome control OLS systems share a Gram solve. Logistic MLE uses
+    damped Newton steps, scaled solves and an explicit separation check.
+    ``fast`` (``cs_kernels.BalancedCS``) does the passes in compiled code, with the
+    same sums.
+    """
+    K = len(names)
+    systems = [None] * len(transforms)
+    passes = 0
+    if fast is not None:
+        C = len(transforms)
+        acc = dict(
+            counts=np.zeros((C, 2), dtype=np.int64),
+            wsums=np.zeros((C, 4)),
+            scale=np.zeros((C, K)),
+            gram=np.zeros((C, K, K)),
+            reg=np.zeros((C, K, fast.Lmax)),
+        )
+        for raw in batches():
+            fast.initial(raw, K, method != "ipw", acc)
+        for j in range(C):
+            L = int(fast.sizes[j])
+            systems[j] = NuisanceSystem(
+                np.zeros(K),
+                acc["reg"][j, :, :L].copy(),
+                acc["gram"][j].copy(),
+                n_treated=int(acc["counts"][j, 0]),
+                n_controls=int(acc["counts"][j, 1]),
+                treated_weight=float(acc["wsums"][j, 0]),
+                control_weight=float(acc["wsums"][j, 1]),
+                treated_weight_squares=float(acc["wsums"][j, 2]),
+                control_weight_squares=float(acc["wsums"][j, 3]),
+                design_scale=acc["scale"][j].copy(),
+            )
+    else:
+        for raw in batches():
+            for j, transform in enumerate(transforms):
+                X, d, Y, w = weighted_system(transform(raw))
+                if systems[j] is None:
+                    systems[j] = NuisanceSystem(
+                        np.zeros(K),
+                        np.zeros((K, Y.shape[1])),
+                        np.zeros((K, K)),
+                        design_scale=np.zeros(K),
+                    )
+                s = systems[j]
+                s.n_treated += int(d.sum())
+                s.n_controls += int((~d).sum())
+                s.treated_weight += float(w[d].sum())
+                s.control_weight += float(w[~d].sum())
+                s.control_weight_squares += float(w[~d] @ w[~d])
+                s.treated_weight_squares += float(w[d] @ w[d])
+                s.design_scale += np.einsum("i,ij,ij->j", w, X, X)
+                if method != "ipw":
+                    xc = X[~d]
+                    s.control_gram += xc.T @ (w[~d, None] * xc)
+                    s.regression += xc.T @ (w[~d, None] * Y[~d])
+    passes += 1
+    for s in systems:
+        if not s.n_treated or not s.n_controls:
+            raise ValueError("both nuisance groups need observations")
+        s.design_scale = np.sqrt(s.design_scale / (s.treated_weight + s.control_weight))
+        if method != "ipw":
+            s.regression = solve_checked(s.control_gram, s.regression, "outcome regression")
+        if "Intercept" in names:
+            s.beta[names.index("Intercept")] = np.log(s.treated_weight / s.control_weight)
+    if method == "reg":
+        return systems, passes
+
+    def fast_stats(betas, indices):
+        nonlocal passes
+        C = len(transforms)
+        B = np.array(betas)
+        active = np.zeros(C, dtype=bool)
+        active[list(indices)] = True
+        loss, correction = np.zeros(C), np.zeros(C)
+        grad, hess, margin = np.zeros((C, K)), np.zeros((C, K, K)), np.full(C, np.inf)
+        for raw in batches():
+            inc, g_part, h_part, m_part = fast.newton(raw, B, active, K)
+            # compensated accumulation of the objective, as in the array path
+            step = inc - correction
+            total = loss + step
+            correction = (total - loss) - step
+            loss = total
+            grad += g_part
+            hess += h_part
+            margin = np.minimum(margin, m_part)
+        passes += 1
+        return {
+            j: [float(loss[j]), grad[j].copy(), hess[j].copy(), float(margin[j])] for j in indices
+        }
+
+    def stats(betas, indices):
+        nonlocal passes
+        if fast is not None:
+            return fast_stats(betas, indices)
+        out = {j: [0.0, np.zeros(K), np.zeros((K, K)), np.inf] for j in indices}
+        loss_correction = dict.fromkeys(indices, 0.0)
+        for raw in batches():
+            for j in indices:
+                X, d, _, w = weighted_system(transforms[j](raw))
+                if not len(d):
+                    continue
+                eta = X @ betas[j]
+                p = expit(eta)
+                # Compensated accumulation keeps the objective stable across
+                # many small batches without retaining their individual sums.
+                increment = float(w @ (np.logaddexp(0, eta) - d * eta)) - loss_correction[j]
+                total = out[j][0] + increment
+                loss_correction[j] = (total - out[j][0]) - increment
+                out[j][0] = total
+                out[j][1] += X.T @ (w * (d - p))
+                out[j][2] += X.T @ ((w * p * (1 - p))[:, None] * X)
+                out[j][3] = min(out[j][3], float(np.min((2 * d.astype(float) - 1) * eta)))
+        passes += 1
+        return out
+
+    betas = [s.beta.copy() for s in systems]
+    active = set(range(len(systems)))
+    current = stats(betas, active)
+    for _iteration in range(maxiter):
+        steps, done = {}, set()
+        for j in active:
+            loss, gradient, H, margin = current[j]
+            if margin > 1e-8:
+                raise ValueError("propensity logit: perfect separation / no overlap")
+            step = solve_checked(H, gradient, "propensity logit")
+            # Coefficients have inverse covariate units. Test each coordinate
+            # in weighted RMS design units so changing meters to micrometers
+            # cannot turn an unfinished fit into a converged one. The scale is
+            # fixed from the first pass, rather than shrinking with p*(1-p)
+            # near separation. No extra source pass or row storage is needed.
+            scale = systems[j].design_scale
+            if np.all(np.abs(step) * scale <= tol * (1 + np.abs(betas[j]) * scale)):
+                done.add(j)
+            else:
+                steps[j] = step
+        active -= done
+        if not active:
+            for j, s in enumerate(systems):
+                s.beta = betas[j]
+            return systems, passes
+        candidates = list(betas)
+        for j in active:
+            candidates[j] = betas[j] + steps[j]
+        trial = stats(candidates, active)
+        for _backtrack in range(30):
+            # The objective is a SUM, so its floating-point resolution grows
+            # with sample size. Near the optimum, a genuine Newton improvement
+            # can be smaller than summation noise. An absolute-only allowance
+            # causes repeated tiny backtracks/full source scans at large N.
+            # Keep the coefficient-step convergence test unchanged.
+            bad = {
+                j
+                for j in active
+                if not np.isfinite(trial[j][0])
+                or trial[j][0]
+                > current[j][0]
+                + max(1e-10, 64 * np.finfo(float).eps * max(1.0, abs(current[j][0])))
+            }
+            if not bad:
+                break
+            for j in bad:
+                steps[j] *= 0.5
+                candidates[j] = betas[j] + steps[j]
+            trial.update(stats(candidates, bad))
+        else:
+            raise ValueError(
+                "propensity logit line search failed; check overlap and covariate scaling"
+            )
+        for j in active:
+            betas[j] = candidates[j]
+            current[j] = trial[j]
+            systems[j].iterations += 1
+    raise ValueError("propensity logit did not converge; possible quasi-separation / poor overlap")