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

diff_stream/balance.py

@@ -0,0 +1,149 @@
+"""CS-specific comparison selection around shared weighted balance kernels.
+
+Optional one-comparison-at-a-time staging avoids retaining all unit covariates
+or propensity scores on a result. Regression adjustment does not reweight controls.
+"""
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import polars as pl
+import pyarrow.parquet as pq
+from hdfe_stream.core.balance import effective_sizes, weighted_balance
+
+from ._core.contrasts import att_parts
+
+
+class CSBalanceCollector:
+    def __init__(self, store, ecdf):
+        self.store, self.ecdf = store, ecdf
+        names = store.design.names if store.design is not None else []
+        self.keep = [j for j, name in enumerate(names) if name != "Intercept"]
+        if not self.keep:
+            raise ValueError("covariate_balance=True requires covariates in the formula")
+        self.names = [names[j] for j in self.keep]
+        self.columns = [f"__x{j}" for j in range(len(self.keep))]
+        self.tables = []
+        self.source_passes = 0
+
+    def compute(self, rows, adjusted):
+        """Shared kernels consume a lazy one-comparison temporary Parquet file."""
+        with TemporaryDirectory(prefix="cs-balance-", dir=self.store.root) as run:
+            path = Path(run) / "rows.parquet"
+            writer = None
+            try:
+                for X, d, w, after in rows():
+                    if not len(d):
+                        continue
+                    self.store._guard(32e6 + X.nbytes * 4)
+                    frame = pl.DataFrame(
+                        {
+                            **{
+                                name: X[:, k]
+                                for name, k in zip(self.columns, self.keep, strict=True)
+                            },
+                            "__d": d,
+                            "__before": w,
+                            "__after": after,
+                        }
+                    )
+                    table = frame.to_arrow()
+                    if writer is None:
+                        writer = pq.ParquetWriter(path, table.schema)
+                    writer.write_table(table, row_group_size=self.store.read_rows)
+            finally:
+                if writer is not None:
+                    writer.close()
+            self.source_passes += 1
+            if writer is None:
+                raise ValueError("CS balance comparison has no observations")
+            lf = pl.scan_parquet(path, low_memory=True)
+            before = weighted_balance(lf, "__d", self.columns, "__before", ecdf=self.ecdf)
+            scales = {
+                row["term"]: float(np.sqrt((row["var1"] + row["var0"]) / 2))
+                for row in before.iter_rows(named=True)
+            }
+            tables = []
+            for stage, table, weight in [
+                ("before", before, "__before"),
+                *([("after", None, "__after")] if adjusted else []),
+            ]:
+                if table is None:
+                    table = weighted_balance(
+                        lf, "__d", self.columns, weight, scale=scales, ecdf=self.ecdf
+                    )
+                sizes = effective_sizes(lf, "__d", weight)
+                tables.append(
+                    table.with_columns(
+                        pl.col("term").replace_strict(
+                            dict(zip(self.columns, self.names, strict=True))
+                        ),
+                        pl.lit(stage).alias("stage"),
+                        pl.lit(sizes.n1, dtype=pl.Int64).alias("n_treated"),
+                        pl.lit(sizes.n0, dtype=pl.Int64).alias("n_controls"),
+                        pl.lit(sizes.ess1).alias("ess_treated"),
+                        pl.lit(sizes.ess0).alias("ess_controls"),
+                    )
+                )
+            return pl.concat(tables)
+
+    def balanced(self, comparisons, fitted, cells, batches):
+        for c, f in zip(comparisons, fitted, strict=True):
+
+            def rows(c=c, f=f):
+                for raw in batches():
+                    mask, X, d, _ = c.select(raw)
+                    sw = c.weights(raw, mask)
+                    wc = (
+                        att_parts(X, d, f.nuisance.beta, f.control_max)[1]
+                        if f.method != "reg"
+                        else np.ones(len(d))
+                    )
+                    yield X, d, sw, sw * np.where(d, 1, wc)
+
+            table = self.compute(rows, f.method != "reg")
+            for k, j in enumerate(c.cells):
+                self.tables.append(
+                    table.with_columns(
+                        pl.lit(c.cohort, dtype=pl.Int64).alias("cohort"),
+                        pl.lit(cells[j][1], dtype=pl.Int64).alias("time"),
+                        pl.lit(int(self.store.times[c.bases[k]]), dtype=pl.Int64).alias("base"),
+                        pl.lit("comparison").alias("stratum"),
+                        pl.lit(int(self.store.times[c.xperiod]), dtype=pl.Int64).alias(
+                            "covariate_time"
+                        ),
+                    )
+                )
+
+    def unbalanced(self, comparisons, fitted, batches):
+        for c, f in zip(comparisons, fitted, strict=True):
+            for use_post, stratum in [(False, "pre"), (True, "post")]:
+
+                def rows(c=c, f=f, use_post=use_post):
+                    for raw in batches():
+                        unit, X, d, _, post = c.select(raw)
+                        sw = c.weights(raw, unit, post)
+                        wc = (
+                            att_parts(X, d, f.system.beta, f.trim)[1]
+                            if f.method != "reg"
+                            else np.ones(len(d))
+                        )
+                        mask = post == use_post
+                        yield X[mask], d[mask], sw[mask], (sw * np.where(d, 1, wc))[mask]
+
+                table = self.compute(rows, f.method != "reg")
+                self.tables.append(
+                    table.with_columns(
+                        pl.lit(c.cohort, dtype=pl.Int64).alias("cohort"),
+                        pl.lit(int(self.store.times[c.target]), dtype=pl.Int64).alias("time"),
+                        pl.lit(int(self.store.times[c.base]), dtype=pl.Int64).alias("base"),
+                        pl.lit(stratum).alias("stratum"),
+                        pl.lit(int(self.store.times[c.target if use_post else c.base])).alias(
+                            "covariate_time"
+                        ),
+                    )
+                )
+
+    def table(self):
+        return pl.concat(self.tables).sort("cohort", "time", "stratum", "stage", "term")

diff_stream/bjs.py

@@ -0,0 +1,101 @@
+"""BJS smart residual averaging and cluster leave-out adjustments.
+
+The implementation uses cohort-sized first/second moments. It implements the
+published score algebra independently; no author implementation is vendored.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class BJSInference:
+    groups: np.ndarray
+    sizes: np.ndarray
+    times: np.ndarray
+    cells: list
+    means: np.ndarray
+    Z: np.ndarray
+    theta: np.ndarray
+    H: np.ndarray
+    scatter: np.ndarray
+    avgeffectsby: str
+    leaveout: bool
+    anticipation: int = 0
+
+    def covariance(self, contrasts, *, allow_unidentified=False):
+        """Covariance for contrasts of the original cell estimates.
+
+        Averaging is contrast-specific. Coarser residual groups therefore need
+        fresh smart averages, rather than a linear transform of the cell vcov.
+        """
+        A = np.atleast_2d(contrasts)
+        P, M = A.shape
+        T, G = len(self.times), len(self.groups)
+        weights = np.zeros((G, P, T))
+        labels = {}
+        cell_groups = []
+        for j, (g, t) in enumerate(self.cells):
+            k = int(np.searchsorted(self.groups, g))
+            weights[k, :, t - self.times[0]] = A[:, j] / self.sizes[k]
+            label = {"cohort_time": (g, t), "event_time": t - g, "cohort": g, "all": 0}[
+                self.avgeffectsby
+            ]
+            cell_groups.append(labels.setdefault(label, len(labels)))
+        R = len(labels)
+        residual_means = np.zeros((G, T))
+        denom, numer = np.zeros((P, R)), np.zeros((P, R))
+
+        def group_totals(k):
+            totals, effects = np.zeros((P, R)), np.zeros((P, R))
+            for j, (g, t) in enumerate(self.cells):
+                if g == self.groups[k]:
+                    ti, r = t - self.times[0], cell_groups[j]
+                    totals[:, r] += weights[k, :, ti]
+                    effects[:, r] += weights[k, :, ti] * residual_means[k, ti]
+            return totals, effects
+
+        for k, g in enumerate(self.groups):
+            use = self.times < g - self.anticipation if g else np.ones(T, dtype=bool)
+            mu = self.means[k] - self.Z @ self.theta
+            residual_means[k] = mu - mu[use].mean()
+            totals, effects = group_totals(k)
+            denom += self.sizes[k] * totals**2
+            numer += self.sizes[k] * totals * effects
+        averages = np.divide(numer, denom, out=np.zeros_like(numer), where=denom > 0)
+        target_derivative = A @ self.H
+        V = np.zeros((P, P))
+        unidentified = np.zeros(P, dtype=bool)
+        for k, g in enumerate(self.groups):
+            use = self.times < g - self.anticipation if g else np.ones(T, dtype=bool)
+            totals, _ = group_totals(k)
+            leverage = np.divide(totals**2, denom, out=np.zeros_like(totals), where=denom > 0)
+            raw = weights[k].copy()
+            raw[:, use] = -target_derivative @ self.Z[use].T
+            center = np.zeros(P)
+            for j, (cg, t) in enumerate(self.cells):
+                if cg != g:
+                    continue
+                ti, r = t - self.times[0], cell_groups[j]
+                w = weights[k, :, ti]
+                lev = leverage[:, r]
+                if self.leaveout:
+                    invalid = (w != 0) & (lev >= 1 - 1e-12)
+                    if np.any(invalid) and not allow_unidentified:
+                        raise ValueError(
+                            "BJS leave_out requires more than one contributing unit in each residual group; use a coarser residual_averaging"
+                        )
+                    unidentified |= invalid
+                    raw[:, ti] = np.divide(w, 1 - lev, out=np.zeros_like(w), where=lev < 1 - 1e-12)
+                center += raw[:, ti] * averages[:, r]
+            # Scatter was formed after local elimination of unit intercepts.
+            # A PSD root avoids cancellation when residual variation is tiny.
+            values, vectors = np.linalg.eigh(self.scatter[k])
+            linear = raw @ (vectors * np.sqrt(np.maximum(values, 0)))
+            mean_score = raw @ residual_means[k] - center
+            V += linear @ linear.T + self.sizes[k] * np.outer(mean_score, mean_score)
+        V = (V + V.T) / 2
+        V[unidentified, :] = np.nan
+        V[:, unidentified] = np.nan
+        return V

diff_stream/bjs_placebo.py

@@ -0,0 +1,167 @@
+"""Separate clean-sample BJS lead regression with bounded unit elimination.
+
+The authors' pre-trend exercise augments the untreated unit/time regression
+with lead indicators. It does not average residuals from the post-effect fit.
+No author code is copied; the regression is implemented via sufficient sums.
+"""
+
+from copy import copy
+from dataclasses import dataclass, field
+
+import numpy as np
+import polars as pl
+from hdfe_stream.core.results_base import RegressionResult
+
+from ._core.nuisance import solve_checked
+from ._present import label_frame, labels_of
+
+
+@dataclass(kw_only=True)
+class BJSPlaceboResult(RegressionResult):
+    """Pooled clean lead coefficients from a separate untreated-sample WLS fit.
+
+    Standard coefficient/interval/Wald methods are inherited from the family.
+    There is no estimated joint covariance with the main treatment effects.
+    """
+
+    n_units: int
+    _support: pl.DataFrame = field(repr=False)
+    _internal: dict | None = field(default=None, repr=False)
+    _periods: object = field(default=None, repr=False)
+    estimator = "bjs_placebo"
+    cells = ()
+    _normal_inference = True
+
+    def support(self):
+        """Observed cohort/date support and original-scale weight totals."""
+        return label_frame(self._support, labels_of(self))
+
+    def pretrend_test(self, *, level="auto", horizons=None):
+        """Joint clean-lead Wald test; horizon selection does not refit the model."""
+        from .pretrend import pretrend_test
+
+        return pretrend_test(self, level=level, horizons=horizons)
+
+    def with_vcov(self, vcov):
+        if vcov not in ("CRV0", "CRV1"):
+            raise ValueError("vcov must be 'CRV0' or 'CRV1'")
+        factor = self.n_units / (self.n_units - 1)
+        ratio = (factor if vcov == "CRV1" else 1) / (factor if self.vcov_type == "CRV1" else 1)
+        out = copy(self)
+        out.vcov = self.vcov * ratio
+        out.vcov_type = vcov
+        out.all_vcovs = {vcov: (out.vcov, out.df_t)}
+        return out
+
+
+def fit_placebos(store, pretrends, anticipation, *, fml, depvar, panel, vcov):
+    """Eliminate unit intercepts locally, solve time/lead coefficients jointly."""
+    if store.design is not None or not store.weights_time_invariant:
+        from .linear import fit_placebos as fit_rich_placebos
+
+        return fit_rich_placebos(
+            store, pretrends, anticipation, fml=fml, depvar=depvar, panel=panel, vcov=vcov
+        )
+    T, L = store.periods, pretrends
+    horizons = np.arange(-anticipation - L, -anticipation, dtype=np.int64)
+    rows = []
+    raw_scale = store.weight_scale / store.weight_normalizer
+    support = np.zeros(L, dtype=int)
+    for g, counts in sorted(store.actual_cell_counts.items()):
+        if not g:
+            continue
+        for ti, t in enumerate(store.times):
+            j = int(t) - g + anticipation + L
+            if 0 <= j < L and counts[ti]:
+                support[j] += counts[ti]
+                rows.append(
+                    dict(
+                        aggregate=f"event::{horizons[j]}",
+                        cohort=int(g),
+                        time=int(t),
+                        event_time=int(horizons[j]),
+                        observed_units=int(counts[ti]),
+                        weight_total=float(store.actual_cell_weights[g][ti] * raw_scale),
+                    )
+                )
+    if np.any(support == 0):
+        raise ValueError(
+            f"BJS requested clean placebo horizons are not observed: {horizons[support == 0].tolist()}; reduce pretrend_leads"
+        )
+    Z = np.eye(T)[:, 1:]
+    P = T - 1 + L
+    Q, q = np.zeros((P, P)), np.zeros(P)
+    context = "BJS clean placebo design is unidentified (missing earlier clean reference or disconnected/collinear design); reduce pretrend_leads"
+
+    def batches():
+        return store.batches(include_weights=True, actual_cohort=True)
+
+    def design(g):
+        lead = (store.times[:, None] - g == horizons) & (g != 0)
+        return np.column_stack([Z, lead])
+
+    def parts(g, values):
+        obs = np.isfinite(values)
+        use = obs & ((store.times < g - anticipation) if g else True)
+        n = use.sum(axis=1)
+        if np.any(n == 0):
+            raise ValueError(
+                "BJS requires at least one untreated observation per unit for clean placebo fitting"
+            )
+        return use, n, np.where(use, values, 0)
+
+    n_obs = 0
+    for g, y, w in batches():
+        for cg in np.unique(g):
+            mask = g == cg
+            use, n, raw = parts(cg, y[mask])
+            X, sw = design(cg), w[mask]
+            xm = use @ X / n[:, None]
+            centered = np.where(use, raw - raw.sum(axis=1)[:, None] / n[:, None], 0)
+            Q += X.T @ ((sw @ use)[:, None] * X) - xm.T @ ((sw * n)[:, None] * xm)
+            q += X.T @ (sw @ centered)
+            n_obs += int(n.sum())
+    coefficients = solve_checked(Q, q, context)
+    H = solve_checked(Q, np.eye(P)[:, -L:], context)
+    V = np.zeros((L, L))
+    for g, y, w in batches():
+        scores = np.zeros((len(g), L))
+        for cg in np.unique(g):
+            mask = g == cg
+            use, n, raw = parts(cg, y[mask])
+            X = design(cg)
+            residual = np.where(use, raw - X @ coefficients, 0)
+            residual = np.where(use, residual - residual.sum(axis=1)[:, None] / n[:, None], 0)
+            scores[mask] = w[mask, None] * (residual @ X @ H)
+        V += scores.T @ scores
+    factor = store.n_units / (store.n_units - 1) if vcov == "CRV1" else 1.0
+    V *= factor
+    return BJSPlaceboResult(
+        fml=fml,
+        depvar=depvar,
+        coefnames=[f"event::{h}" for h in horizons],
+        beta=coefficients[-L:],
+        vcov=(V + V.T) / 2,
+        vcov_type=vcov,
+        df_t=store.n_units - 1,
+        all_vcovs={vcov: (V, store.n_units - 1)},
+        n_obs=n_obs,
+        n_clusters={panel[0]: store.n_units},
+        n_units=store.n_units,
+        weights=None,
+        _support=pl.DataFrame(rows),
+        diagnostics=dict(
+            aggregation="event_time",
+            anticipation=anticipation,
+            reference_event_time=None,
+            fitted_horizons=horizons.tolist(),
+            pretrends=L,
+            base_period="all other clean observations; no fixed event-time reference",
+            inference="separate clean-sample WLS; unit-clustered analytic scores",
+            cohort_sizes={str(g): int(n) for g, n in store.actual_gc.items()},
+            reference_observations=n_obs - int(support.sum()),
+            normalized_future_cohorts="original timing retained for placebo indicators",
+            post_effect_covariance="not estimated",
+            warnings=[],
+        ),
+    )