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

diff_stream/cohort_scores.py

@@ -0,0 +1,68 @@
+"""Joint covariance of cell scores and cohort-share scores, by cohort.
+
+The CS covariance sums, over units i, the outer products of [scores_i, share_i]
+where share_i = anchor_i * (e_{g_i} - p): e_g marks the cells of cohort g and p
+is a constant vector. Two facts make this cheaper than one dense product per
+batch (ClusterScoreAccumulator), with the same sums:
+
+- A unit's scores are zero outside the cells its cohort enters (its own cells,
+  and those where it is a control), so each cohort's rows add a dense product
+  only into those columns. Only never-treated units touch every column.
+- The share terms depend on a unit only through its cohort and anchor, so the
+  share and score-share blocks need just the per-cohort sums of anchor**2 and
+  of anchor * scores.
+
+With many cohorts and periods (hundreds of cells) this removes most of the
+work; with few cells it costs about the same as the dense product.
+"""
+
+import numpy as np
+
+
+class CohortScoreCovariance:
+    """Accumulate V = sum s s', cross = sum s h', shares = sum h h' over units.
+
+    ``cell_cohorts`` holds each cell's cohort; ``center`` is p (zeros for an
+    uncentered indicator share).
+    """
+
+    def __init__(self, cell_cohorts, center=None):
+        self.cell_cohorts = np.asarray(cell_cohorts)
+        M = len(self.cell_cohorts)
+        self.center = np.zeros(M) if center is None else np.asarray(center, dtype=float)
+        self.V = np.zeros((M, M))
+        self._weighted = {}  # cohort -> sum of anchor * scores (M,)
+        self._square = {}  # cohort -> sum of anchor**2
+        self.n_clusters = 0
+
+    def add(self, g, scores, anchor=None):
+        """Add one row per unit: cohort labels, cell scores and share anchors."""
+        anchor = np.ones(len(g)) if anchor is None else anchor
+        order = np.argsort(g, kind="stable")
+        labels = g[order]
+        starts = np.r_[0, np.flatnonzero(labels[1:] != labels[:-1]) + 1, len(labels)]
+        M = self.V.shape[0]
+        for a, b in zip(starts[:-1], starts[1:], strict=True):
+            rows = order[a:b]
+            k = labels[a].item()
+            S = scores[rows]
+            used = np.flatnonzero(np.any(S != 0, axis=0))
+            if len(used) == M:
+                self.V += S.T @ S
+            elif len(used):
+                block = S[:, used]
+                self.V[np.ix_(used, used)] += block.T @ block
+            w = anchor[rows]
+            self._weighted[k] = self._weighted.get(k, 0.0) + w @ S
+            self._square[k] = self._square.get(k, 0.0) + float(w @ w)
+        self.n_clusters += len(g)
+
+    def blocks(self):
+        """(V, cross, shares): the blocks of the joint [scores, share] sums."""
+        M = len(self.cell_cohorts)
+        cross, shares = np.zeros((M, M)), np.zeros((M, M))
+        for k, weighted in self._weighted.items():
+            h = (self.cell_cohorts == k) - self.center
+            cross += np.outer(weighted, h)
+            shares += self._square[k] * np.outer(h, h)
+        return self.V, cross, shares

diff_stream/cs.py

@@ -0,0 +1,197 @@
+"""CS comparison populations and earlier-period covariate conventions."""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ._core import cs_kernels
+from ._core.contrasts import fit_contrasts
+from ._core.nuisance import fit_systems
+from .cohort_scores import CohortScoreCovariance
+
+
+@dataclass
+class Comparison:
+    cohort: int
+    controls: tuple
+    xperiod: int
+    cells: list
+    targets: np.ndarray
+    bases: np.ndarray
+    weightperiod: int = 0
+    # members[label] marks the cohort and its controls; cohort labels are
+    # nonnegative integers below len(members).
+    members: np.ndarray = None
+
+    def weights(self, raw, mask):
+        w = raw[3]
+        return w[mask, self.weightperiod] if w.ndim == 2 else w[mask]
+
+    def select(self, raw):
+        g, y, x = raw[:3]
+        if self.members is None:
+            mask = (g == self.cohort) | np.isin(g, self.controls)
+        else:
+            mask = self.members[g]
+        rows = np.flatnonzero(mask)[:, None]
+        return (
+            mask,
+            x[rows[:, 0], self.xperiod, :],
+            g[rows[:, 0]] == self.cohort,
+            y[rows, self.targets] - y[rows, self.bases],
+        )
+
+    def __call__(self, raw):
+        mask, X, d, Y = self.select(raw)
+        return (X, d, Y, self.weights(raw, mask)) if len(raw) == 4 else (X, d, Y)
+
+
+# The fit makes about ten passes over the staged units. When the decoded batches
+# take at most this share of the memory left they are kept after the first pass.
+RETAIN_SHARE = 0.1
+
+
+def balanced_bytes(store, width):
+    """Memory of all decoded batches of a balanced panel: outcomes, covariates, weights."""
+    per_period = 8 * (1 + width + (0 if store.weights_time_invariant else 1))
+    return store.n_units * store.periods * per_period
+
+
+class RetainedBatches:
+    """Calls to ``make()`` that keep their batches after the first complete pass."""
+
+    def __init__(self, make, retain):
+        self.make, self.retain, self.items = make, retain, None
+
+    def __call__(self):
+        if self.items is not None:
+            return iter(self.items)
+        return self._first() if self.retain else self.make()
+
+    def _first(self):
+        items = []
+        for raw in self.make():
+            items.append(raw)
+            yield raw
+        self.items = items
+
+
+def fit_cs(
+    store,
+    groups,
+    cells,
+    bases,
+    controls,
+    method,
+    trim_level,
+    multiplier=None,
+    *,
+    balance_collector=None,
+):
+    """Fit cached comparison systems, then assemble cross-cell unit covariance."""
+    comparisons = {}
+    for j, (g, t) in enumerate(cells):
+        controls_g = tuple(groups[k] for k in np.flatnonzero(controls[j]))
+        tindex = t - store.times[0]
+        xb = 0 if store.x_time_invariant else min(tindex, int(bases[j]))
+        wp = min(tindex, int(bases[j])) if not store.weights_time_invariant else 0
+        key = (g, controls_g, xb, wp)
+        if key not in comparisons:
+            comparisons[key] = [[], [], []]
+        comparisons[key][0].append(j)
+        comparisons[key][1].append(tindex)
+        comparisons[key][2].append(bases[j])
+    size = max(groups) + 1
+
+    def members(g, c):
+        table = np.zeros(size, dtype=bool)
+        table[[g, *c]] = True
+        return table
+
+    transforms = [
+        Comparison(g, c, xb, ix, np.array(t, dtype=int), np.array(b, dtype=int), wp, members(g, c))
+        for (g, c, xb, wp), (ix, t, b) in comparisons.items()
+    ]
+
+    names = store.design.names if store.design is not None else ["Intercept"]
+    batches = RetainedBatches(
+        lambda: store.batches(
+            include_x=True, include_weights=True, row_weights=not store.weights_time_invariant
+        ),
+        retain=store.is_balanced
+        and balanced_bytes(store, len(names)) <= RETAIN_SHARE * store.budget.headroom_gb() * 1e9,
+    )
+    M = len(cells)
+    fast = (
+        cs_kernels.BalancedCS(transforms, M) if cs_kernels.ENABLED and store.is_balanced else None
+    )
+    try:
+        systems, passes = fit_systems(batches, transforms, method, names, fast=fast)
+    except ValueError:
+        # A nearly collinear design can converge in one summation order and not
+        # in another; the array path is the reference.
+        if fast is None:
+            raise
+        systems, passes = fit_systems(batches, transforms, method, names)
+    fitted = fit_contrasts(batches, transforms, systems, method, trim_level, fast=fast)
+    beta = np.zeros(M)
+    for c, f in zip(transforms, fitted, strict=True):
+        beta[c.cells] = f.estimate
+    cohorts = np.array([g for g, t in cells])
+    counts = np.array([store.gw[g] for g, t in cells], dtype=float)
+    # Unadjusted sums of [scores, share] outer products: api.py owns the
+    # explicit CRV0/CRV1 factor for all covariance blocks.
+    accumulator = CohortScoreCovariance(cohorts, counts / store.n_units)
+    packed = fast.pack_fitted(fitted, len(names)) if fast is not None else None
+    for raw in batches():
+        g, _, _, w = raw
+        anchor = w[:, 0] if w.ndim == 2 else w
+        if fast is not None:
+            scores = fast.scores(raw, method, packed, trim_level)
+        else:
+            scores = np.zeros((len(g), M))
+            for c, f in zip(transforms, fitted, strict=True):
+                mask, X, d, Y = c.select(raw)
+                scores[np.ix_(np.flatnonzero(mask), c.cells)] = f.scores(
+                    X, d, Y, c.weights(raw, mask)
+                )
+        accumulator.add(g, scores, anchor)
+        if multiplier is not None:
+            share = anchor[:, None] * ((g[:, None] == cohorts[None, :]) - counts / store.n_units)
+            multiplier.update(scores, share)
+    V, cross, share_cov = accumulator.blocks()
+    if balance_collector is not None:
+        balance_collector.balanced(transforms, fitted, cells, batches)
+    metadata = dict(
+        est_method=method,
+        covariate_names=names,
+        covariates_time_invariant=store.x_time_invariant,
+        nuisance_systems=len(transforms),
+        propensity_models=0 if method == "reg" else len(transforms),
+        nuisance_source_passes=passes + 2,
+        trim_level=trim_level,
+        probability_cap=1 - 1e-6 if method != "reg" else None,
+        comparisons=[
+            dict(
+                cohort=c.cohort,
+                controls=list(c.controls),
+                covariate_time=int(store.times[c.xperiod]),
+                weight_time=int(store.times[c.weightperiod])
+                if not store.weights_time_invariant
+                else None,
+                cells=c.cells,
+                **f.diagnostics,
+            )
+            for c, f in zip(transforms, fitted, strict=True)
+        ],
+    )
+    metadata["warnings"] = []
+    if any(f.diagnostics["trimmed_controls"] for f in fitted):
+        metadata["warnings"].append(
+            "Propensity trimming removed control contributions in some comparisons; inspect diagnostics['comparisons']."
+        )
+    if any(f.diagnostics["capped_probabilities"] for f in fitted):
+        metadata["warnings"].append(
+            "Some propensity probabilities were capped at 1-1e-6 under the R DRDID convention."
+        )
+    return beta, counts, V, cross, share_cov, metadata

diff_stream/cs_unbalanced.py

@@ -0,0 +1,336 @@
+"""R did-compatible observed-row CS functionals, clustered by original unit.
+
+Each date pair uses repeated-cross-section mean functionals on its observed
+rows. Units seen at only one date contribute; scores are folded by unit.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.special import expit
+
+from . import cs as cs_module
+from ._core import cs_row_kernels
+from ._core.functionals import NormalizedMeans
+from ._core.nuisance import fit_linear_systems, fit_systems, solve_checked
+from .cohort_scores import CohortScoreCovariance
+from .cs import RetainedBatches, balanced_bytes
+
+
+@dataclass
+class RowComparison:
+    cohort: int
+    controls: tuple
+    target: int
+    base: int
+
+    def weights(self, raw, unit, post):
+        w = raw[3]
+        return w[unit, np.where(post, self.target, self.base)] if w.ndim == 2 else w[unit]
+
+    def select(self, raw):
+        g, y, x = raw[:3]
+        mask = (g == self.cohort) | np.isin(g, self.controls)
+        dates = np.array([self.base, self.target])
+        unit, period = np.nonzero(mask[:, None] & np.isfinite(y[:, dates]))
+        ti = dates[period]
+        return unit, x[unit, ti], g[unit] == self.cohort, y[unit, ti], period == 1
+
+    def __call__(self, raw):
+        unit, X, d, y, post = self.select(raw)
+        return (
+            (X, d, y[:, None], self.weights(raw, unit, post))
+            if len(raw) == 4
+            else (X, d, y[:, None])
+        )
+
+
+class RowFunctional:
+    def __init__(self, comparison, method, trim, system, models, width):
+        self.c, self.method, self.trim, self.system = comparison, method, trim, system
+        # Model order is control pre/post, treated pre/post.
+        self.models = models
+        if method == "ipw":
+            self.signs = np.array([1, -1, -1, 1.0])
+            self.value_derivative = np.zeros((4, 0))
+        elif method == "reg":
+            self.signs = np.array([1, -1, -1.0])
+            self.value_derivative = np.array([[0, 0], [0, 0], [-1, 1.0]])
+        else:
+            self.signs = np.array([1, -1, -1, 1, 1, -1, -1, 1.0])
+            self.value_derivative = np.array(
+                [
+                    [0, -1, 0, 0],
+                    [-1, 0, 0, 0],
+                    [0, -1, 0, 0],
+                    [-1, 0, 0, 0],
+                    [0, -1, 0, 1],
+                    [0, -1, 0, 1],
+                    [-1, 0, 1, 0],
+                    [-1, 0, 1, 0.0],
+                ]
+            )
+        self.means = NormalizedMeans(len(self.signs), width)
+        self.H = np.zeros((width, width))
+        self.trimmed = 0
+        self.capped = 0
+        self.stratum_counts = np.zeros(4, dtype=int)
+        self.probability_min, self.probability_max = np.inf, -np.inf
+        self.control_weight_squares = np.zeros(2)
+        self.sample_weight = np.zeros(4)
+        self.sample_weight_squares = np.zeros(4)
+
+    def rows(self, raw):
+        unit, X, d, y, post = self.c.select(raw)
+        tp, tm = d & post, d & ~post
+        cp, cm = ~d & post, ~d & ~post
+        p = np.minimum(expit(X @ self.system.beta), 1 - 1e-6) if self.system is not None else None
+        odds = (p / (1 - p)) * (p < self.trim) if p is not None else np.ones(len(d))
+        if self.method == "ipw":
+            weights = np.column_stack([tp, tm, cp * odds, cm * odds])
+            values = np.broadcast_to(y[:, None], weights.shape)
+        elif self.method == "reg":
+            weights = np.column_stack([tp, tm, d]).astype(float)
+            values = np.column_stack([y, y, X @ (self.models[1][0] - self.models[0][0])])
+        else:
+            m = [X @ b for b, _ in self.models]
+            r = y - np.where(post, m[1], m[0])
+            weights = np.column_stack([tp, tm, cp * odds, cm * odds, d, tp, d, tm])
+            values = np.column_stack(
+                [r, r, r, r, m[3] - m[1], m[3] - m[1], m[2] - m[0], m[2] - m[0]]
+            )
+        if len(raw) == 4:
+            weights = weights * self.c.weights(raw, unit, post)[:, None]
+        return unit, X, d, y, post, p, weights, values
+
+    def update(self, raw):
+        unit, X, d, _, post, p, w, v = self.rows(raw)
+        sw = self.c.weights(raw, unit, post) if len(raw) == 4 else np.ones(len(d))
+        for j, mask in enumerate((~d & ~post, ~d & post, d & ~post, d & post)):
+            self.sample_weight[j] += sw[mask].sum()
+            self.sample_weight_squares[j] += sw[mask] @ sw[mask]
+        self.means.update(w, v, X)
+        self.stratum_counts += [
+            np.sum(~d & ~post),
+            np.sum(~d & post),
+            np.sum(d & ~post),
+            np.sum(d & post),
+        ]
+        if p is not None:
+            self.H += X.T @ ((sw * p * (1 - p))[:, None] * X)
+            self.trimmed += int(((~d) & (p >= self.trim)).sum())
+            self.capped += int((expit(X @ self.system.beta) > 1 - 1e-6).sum())
+            self.control_weight_squares += (w[:, [3, 2]] ** 2).sum(axis=0)
+            if len(p):
+                self.probability_min = min(self.probability_min, float(p.min()))
+                self.probability_max = max(self.probability_max, float(p.max()))
+
+    def load(self, acc, j):
+        """Take the sums of the compiled pass (``cs_row_kernels``) for comparison ``j``."""
+        n = len(self.signs)
+        self.means.weight = acc["weight"][j, :n].copy()
+        self.means.value = acc["value"][j, :n].copy()
+        self.means.x = acc["xw"][j, :n].copy()
+        self.means.xvalue = acc["xv"][j, :n].copy()
+        self.stratum_counts = acc["strata"][j].copy()
+        self.sample_weight = acc["sw"][j].copy()
+        self.sample_weight_squares = acc["sw2"][j].copy()
+        if self.system is not None:
+            self.H = acc["hess"][j].copy()
+            self.trimmed = int(acc["trimmed"][j])
+            self.capped = int(acc["capped"][j])
+            self.control_weight_squares = acc["cws"][j].copy()
+            if acc["pmax"][j] >= acc["pmin"][j]:
+                self.probability_min = float(acc["pmin"][j])
+                self.probability_max = float(acc["pmax"][j])
+
+    def finish(self):
+        self.means.finish()
+        self.estimate = self.means.mean @ self.signs
+        J = (self.value_derivative.T * self.signs) @ self.means.mean_x
+        self.reg_correction = [
+            solve_checked(Q, j, "outcome score correction")
+            for (_, Q), j in zip(self.models, J, strict=True)
+        ]
+        if self.system is not None:
+            Jps = -self.means.centered_x[2] + self.means.centered_x[3]
+            self.ps_correction = solve_checked(self.H, Jps, "propensity score correction")
+
+    def scores(self, raw):
+        unit, X, d, y, post, p, w, v = self.rows(raw)
+        phi = self.means.scores(w, v, self.signs)
+        sw = self.c.weights(raw, unit, post) if len(raw) == 4 else np.ones(len(d))
+        strata = [~d & ~post, ~d & post, d & ~post, d & post]
+        for (b, _), correction, mask in zip(
+            self.models, self.reg_correction, strata[: len(self.models)], strict=True
+        ):
+            phi += sw * mask * (y - X @ b) * (X @ correction)
+        if self.system is not None:
+            phi += sw * (d - p) * (X @ self.ps_correction)
+        return unit, phi
+
+
+def fit_cs_unbalanced(
+    store, groups, cells, bases, controls, method, trim, multiplier=None, *, balance_collector=None
+):
+    K = len(store.design.names) if store.design is not None else 1
+    names = store.design.names if store.design is not None else ["Intercept"]
+    comparisons = [
+        RowComparison(g, tuple(np.asarray(groups)[c]), t - store.times[0], int(b))
+        for (g, t), b, c in zip(cells, bases, controls, strict=True)
+    ]
+
+    fast = cs_row_kernels.ENABLED
+    retain = (
+        fast
+        and balanced_bytes(store, K) <= cs_module.RETAIN_SHARE * store.budget.headroom_gb() * 1e9
+    )
+    batches = RetainedBatches(
+        lambda: store.batches(
+            include_x=True, include_weights=True, row_weights=not store.weights_time_invariant
+        ),
+        retain=retain,
+    )
+    size = int(max(np.max(groups), *(c.cohort for c in comparisons))) + 1
+
+    linear_keys = {}
+    model_indices = []
+    for c in comparisons:
+        keys = [] if method == "ipw" else [(c.controls, c.base), (c.controls, c.target)]
+        if method == "dr":
+            keys += [((c.cohort,), c.base), ((c.cohort,), c.target)]
+        model_indices.append([linear_keys.setdefault(key, len(linear_keys)) for key in keys])
+
+    def linear_transform(cohorts, ti):
+        def select(raw):
+            g, y, x, w = raw
+            mask = np.isin(g, cohorts) & np.isfinite(y[:, ti])
+            return x[mask, ti], y[mask, ti], w[mask, ti] if w.ndim == 2 else w[mask]
+
+        return select
+
+    if fast and linear_keys:
+        gram, moments = cs_row_kernels.fit_linear_rows(batches, list(linear_keys), size, K)
+        models = [
+            (solve_checked(Q, q, f"stratified outcome regression (system {j}, {K} columns)"), Q)
+            for j, (Q, q) in enumerate(zip(gram, moments, strict=True))
+        ]
+    else:
+        models = fit_linear_systems(batches, [linear_transform(*key) for key in linear_keys], K)
+    ps_keys, ps_indices, ps_transforms = {}, [], []
+    for c in comparisons:
+        key = (c.cohort, c.controls, *sorted([c.base, c.target]))
+        if key not in ps_keys:
+            ps_keys[key] = len(ps_keys)
+            ps_transforms.append(c)
+        ps_indices.append(ps_keys[key])
+    kernels = cs_row_kernels.RowKernels(comparisons, size) if fast else None
+
+    def propensity_fits(use_kernels):
+        return fit_systems(
+            batches,
+            ps_transforms,
+            "ipw",
+            names,
+            fast=cs_row_kernels.RowKernels(ps_transforms, size) if use_kernels else None,
+        )
+
+    if method == "reg":
+        systems, passes = [], 0
+    else:
+        try:
+            systems, passes = propensity_fits(fast)
+        except ValueError:
+            # A nearly collinear design can converge in one summation order and
+            # not in another; the array path is the reference.
+            if not fast:
+                raise
+            systems, passes = propensity_fits(False)
+    fitted = [
+        RowFunctional(c, method, trim, systems[ps] if systems else None, [models[k] for k in ix], K)
+        for c, ps, ix in zip(comparisons, ps_indices, model_indices, strict=True)
+    ]
+    if kernels is not None:
+        params = kernels.parameters(fitted, method, K)
+        acc = kernels.empty_moments(K)
+        for raw in batches():
+            kernels.moments(raw, method, trim, params, acc)
+        for j, f in enumerate(fitted):
+            f.load(acc, j)
+    else:
+        for raw in batches():
+            for f in fitted:
+                f.update(raw)
+    for f in fitted:
+        f.finish()
+    M = len(cells)
+    cohorts = np.array([g for g, _ in cells])
+    counts = np.array([store.gw[g] for g, _ in cells], float)
+    # Unadjusted sums of [scores, share] outer products: api.py owns the
+    # explicit CRV0/CRV1 factor for all covariance blocks.
+    accumulator = CohortScoreCovariance(cohorts, counts / store.n_units)
+    for raw in batches():
+        g, y, _, w = raw
+        anchor = w.sum(axis=1) / np.isfinite(y).sum(axis=1) if w.ndim == 2 else w
+        if kernels is not None:
+            scores = kernels.scores(raw, method, trim, params, fitted)
+        else:
+            scores = np.zeros((len(g), M))
+            for j, f in enumerate(fitted):
+                unit, phi = f.scores(raw)
+                np.add.at(scores[:, j], unit, phi)
+        accumulator.add(g, scores, anchor)
+        if multiplier is not None:
+            share = anchor[:, None] * ((g[:, None] == cohorts) - counts / store.n_units)
+            multiplier.update(scores, share)
+    V, cross, share_cov = accumulator.blocks()
+    if balance_collector is not None:
+        balance_collector.unbalanced(comparisons, fitted, batches)
+    metadata = dict(
+        est_method=method,
+        covariate_names=names,
+        covariate_convention="observed-date covariates; repeated-cross-section functionals",
+        nuisance_systems=len(linear_keys),
+        propensity_models=len(systems),
+        nuisance_source_passes=passes + 2 + int(method != "ipw"),
+        trim_level=trim,
+        comparisons=[
+            dict(
+                cohort=c.cohort,
+                time=int(store.times[c.target]),
+                base=int(store.times[c.base]),
+                trimmed_controls=f.trimmed,
+                capped_probabilities=f.capped,
+                n_control_pre=int(f.stratum_counts[0]),
+                n_control_post=int(f.stratum_counts[1]),
+                n_treated_pre=int(f.stratum_counts[2]),
+                n_treated_post=int(f.stratum_counts[3]),
+                propensity_range=[f.probability_min, f.probability_max]
+                if method != "reg"
+                else None,
+                effective_controls=(f.means.weight[[3, 2]] ** 2 / f.control_weight_squares).tolist()
+                if method != "reg"
+                else (f.sample_weight[:2] ** 2 / f.sample_weight_squares[:2]).tolist(),
+                effective_treated=(f.sample_weight[2:] ** 2 / f.sample_weight_squares[2:]).tolist(),
+                logit_iterations=f.system.iterations if method != "reg" else None,
+            )
+            for c, f in zip(comparisons, fitted, strict=True)
+        ],
+        warnings=[],
+    )
+    if any(f.trimmed for f in fitted):
+        metadata["warnings"].append(
+            "Propensity trimming removed control contributions; inspect comparisons."
+        )
+    if any(f.capped for f in fitted):
+        metadata["warnings"].append(
+            "Some propensity probabilities were capped at 1-1e-6 under the R DRDID convention."
+        )
+    return (
+        np.array([f.estimate for f in fitted]),
+        counts,
+        V,
+        cross,
+        share_cov,
+        metadata,
+    )

diff_stream/effective_sample.py

@@ -0,0 +1,73 @@
+"""CS comparison support, warning policy and explicit coefficient suppression."""
+
+import warnings
+from numbers import Real
+
+import numpy as np
+import polars as pl
+
+
+class LowEffectiveSampleWarning(UserWarning):
+    """A CS comparison relies on a low effective treated/control sample."""
+
+
+def threshold(value, name, *, optional=True):
+    if value is None and optional:
+        return value
+    if (
+        isinstance(value, (bool, np.bool_))
+        or not isinstance(value, Real)
+        or not np.isfinite(value)
+        or value < 0
+    ):
+        raise ValueError(
+            f"{name} must be a finite nonnegative number" + (" or None" if optional else "")
+        )
+    return float(value)
+
+
+def finish_result(result, warn_threshold, suppress_threshold):
+    if result.estimator != "cs":
+        return result
+    rows = []
+    cell_index = {cell: j for j, cell in enumerate(result.cells)}
+    for comparison in result.diagnostics["comparisons"]:
+        indices = comparison.get("cells")
+        if indices is None:
+            indices = [cell_index[(comparison["cohort"], comparison["time"])]]
+        for index in indices:
+            for group, key in (("treated", "effective_treated"), ("control", "effective_controls")):
+                values = np.atleast_1d(comparison[key])
+                labels = [group] if len(values) == 1 else [group + "_pre", group + "_post"]
+                for label, value in zip(labels, values, strict=True):
+                    if not np.isfinite(value) or value <= 0:
+                        raise ValueError("CS comparison has undefined effective sample size")
+                    rows.append(
+                        dict(
+                            coefficient=result.coefnames[index],
+                            cell=result.coefnames[index],
+                            cohort=result.cells[index][0],
+                            time=result.cells[index][1],
+                            stratum=label,
+                            effective_n=float(value),
+                        )
+                    )
+    result._ess_table = pl.DataFrame(rows)
+    result.diagnostics["ess_warn_threshold"] = warn_threshold
+    result.diagnostics["ess_suppress_threshold"] = suppress_threshold
+    result.diagnostics["ess_definition"] = (
+        "Kish weight ESS within each required treated/control comparison stratum; not regression-information or influence-score ESS"
+    )
+    minima = result._minimum_ess()
+    bad = minima < (warn_threshold or 0)
+    if bad.any():
+        message = (
+            f"Low effective sample size: {int(bad.sum())} CS coefficients have a required "
+            f"treated/control stratum below {warn_threshold:g} (minimum {minima.min():.3g}). "
+            "Inspect effective_sample_size(); suppress(min_effective_sample=...) masks coefficients without changing aggregation weights."
+        )
+        result.diagnostics["warnings"] = [*result.diagnostics.get("warnings", []), message]
+        warnings.warn(message, LowEffectiveSampleWarning, stacklevel=3)
+    if suppress_threshold is not None:
+        return result.suppress(min_effective_sample=suppress_threshold)
+    return result