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

diff_stream/__init__.py

@@ -0,0 +1,52 @@
+"""Disk-backed staggered difference-in-differences and TWFE diagnostics."""
+
+import stream_limits as _limits
+
+_limits.apply_scheduler_defaults()
+# The compiled passes make many short parallel calls, which run faster when idle
+# OpenMP threads sleep than when they spin. A policy the user has set is kept.
+_limits.set_openmp_wait_policy()
+
+from hdfe_stream.core.profile import Profile  # noqa: E402
+from hdfe_stream.core.reporting import etable  # noqa: E402
+from stream_limits import limit_threads  # noqa: E402
+
+from .absorbed import UnimputedObservationsWarning  # noqa: E402
+from .api import did_bjs, did_callaway_santanna, did_sun_abraham  # noqa: E402
+from .bacon import BaconResult, bacon_decomp  # noqa: E402
+from .bjs_placebo import BJSPlaceboResult  # noqa: E402
+from .bootstrap import BootstrapFailureWarning, RefitBootstrapResult, refit_bootstrap  # noqa: E402
+from .effective_sample import LowEffectiveSampleWarning  # noqa: E402
+from .etwfe import did_extended_twfe  # noqa: E402
+from .gardner import TwoStageResult, did_two_stage  # noqa: E402
+from .reference import ReferenceWindowWarning  # noqa: E402
+from .results import DIDResult  # noqa: E402
+from .stacked import StackedResult, did_stacked  # noqa: E402
+from .twfe_weights import TWFEWeightResult, twfe_weights  # noqa: E402
+
+__version__ = "0.1.0.dev0"
+__all__ = [
+    "did_callaway_santanna",
+    "did_sun_abraham",
+    "did_bjs",
+    "did_two_stage",
+    "did_extended_twfe",
+    "did_stacked",
+    "bacon_decomp",
+    "twfe_weights",
+    "refit_bootstrap",
+    "DIDResult",
+    "TwoStageResult",
+    "StackedResult",
+    "BJSPlaceboResult",
+    "BaconResult",
+    "TWFEWeightResult",
+    "RefitBootstrapResult",
+    "BootstrapFailureWarning",
+    "LowEffectiveSampleWarning",
+    "ReferenceWindowWarning",
+    "UnimputedObservationsWarning",
+    "Profile",
+    "etable",
+    "limit_threads",
+]

diff_stream/_bootstrap_worker.py

@@ -0,0 +1,48 @@
+"""Private fresh-process refit entry point; configure limits before numerics."""
+
+import os
+import pickle
+import sys
+import traceback
+from pathlib import Path
+
+
+def main():
+    job, output, threads, memory = sys.argv[1:]
+    import stream_limits
+
+    spill = Path(job).parent / "spill"
+    spill.mkdir()
+    stream_limits.set_limits(
+        threads=int(threads), memory_gb=float(memory), spill_dir=spill, override=True
+    )
+    # Import this same checkout/wheel, including when invoked from a notebook
+    # or from a working directory outside the package repository.
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    try:
+        # Let the package apply scheduler defaults before importing Polars.
+        from diff_stream.bootstrap import _refit_attempt  # noqa: I001
+
+        import polars as pl
+
+        with open(job, "rb") as f:
+            task = pickle.load(f)  # Private, locally written task; never external input.
+        result = _refit_attempt(task)
+        payload = dict(
+            result=result,
+            resources=dict(pid=os.getpid(), polars_threads=pl.thread_pool_size()),
+        )
+    except BaseException as error:
+        # Preserve resource/programming exceptions instead of treating them
+        # as statistical failures or returning a partial bootstrap result.
+        try:
+            pickle.dumps(error)
+        except Exception:
+            error = RuntimeError(f"{type(error).__name__}: {error}")
+        payload = dict(error=error, traceback=traceback.format_exc())
+    with open(output, "wb") as f:
+        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+if __name__ == "__main__":
+    main()

diff_stream/_core/__init__.py

@@ -0,0 +1,5 @@
+"""Local bounded infrastructure, pending promotion where consumers need it.
+
+IPW, influence-function covariance and first-stage helpers are imported from
+hdfe_stream.core. Treatment/population conventions remain outside this package.
+"""

diff_stream/_core/contrasts.py

@@ -0,0 +1,199 @@
+"""Reusable ATT mean contrasts with regression/estimated-propensity corrections.
+
+Outcomes can be any columns, including changes supplied by a DiD consumer.
+Scores are on the estimate scale, so joint covariance is sum_i phi_i phi_i'.
+The array path supplements the shared expression-based IPW implementation.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+import polars as pl
+from hdfe_stream.core.ipw import ipw_weight
+from scipy.special import expit
+
+from .nuisance import solve_checked
+from .weights import weighted_system
+
+_ATT_WEIGHT = ipw_weight("d", pl.col("p"), "att")
+
+
+def att_parts(X, d, beta, control_max):
+    """R-compatible ATT convention: cap p at 1-1e-6, trim controls only.
+
+    The propensity and regression models still use all selected units.
+    Trim membership is held fixed for analytic inference.
+    """
+    p = np.minimum(expit(X @ beta), 1 - 1e-6)
+    raw = pl.DataFrame({"d": d, "p": p}).select(_ATT_WEIGHT.alias("w"))["w"].to_numpy()
+    control = np.where((~d) & (p < control_max), raw, 0.0)
+    return p, control
+
+
+@dataclass
+class AdjustedContrast:
+    nuisance: object
+    method: str
+    control_max: float
+    treated_total: float
+    control_total: float
+    treated_mean: np.ndarray
+    control_mean: np.ndarray
+    outcome_adjustment: np.ndarray
+    propensity_adjustment: np.ndarray
+    diagnostics: dict
+
+    @property
+    def estimate(self):
+        return self.treated_mean - self.control_mean
+
+    def scores(self, X, d, Y, weights=None):
+        if not len(d):
+            return np.zeros_like(Y)
+        if self.method == "reg":
+            p, wc = np.zeros(len(d)), np.zeros(len(d))
+        else:
+            p, wc = att_parts(X, d, self.nuisance.beta, self.control_max)
+        residual = Y if self.method == "ipw" else Y - X @ self.nuisance.regression
+        phi = d[:, None] * (residual - self.treated_mean) / self.treated_total
+        if self.method != "reg":
+            phi -= wc[:, None] * (residual - self.control_mean) / self.control_total
+            phi += ((d - p)[:, None] * X) @ self.propensity_adjustment
+        if self.method != "ipw":
+            phi += (~d)[:, None] * residual * (X @ self.outcome_adjustment)[:, None]
+        return phi if weights is None else phi * weights[:, None]
+
+
+def fit_contrasts(batches, transforms, systems, method, control_max, *, fast=None):
+    """One common moment pass after nuisance fitting, for any number of outcomes.
+
+    ``fast`` (``cs_kernels.BalancedCS``) does the pass in compiled code.
+    """
+    moments = []
+    for s in systems:
+        K, L = s.regression.shape
+        moments.append(
+            dict(
+                w0=0.0,
+                w02=0.0,
+                s0=np.zeros(L),
+                s1=np.zeros(L),
+                xt=np.zeros(K),
+                xc=np.zeros(K),
+                xcy=np.zeros((K, L)),
+                H=np.zeros((K, K)),
+                pmin=1.0,
+                pmax=0.0,
+                trimmed=0,
+                capped=0,
+            )
+        )
+    if fast is not None:
+        C = len(systems)
+        K = systems[0].regression.shape[0]
+        L = fast.Lmax
+        acc = dict(
+            s1=np.zeros((C, L)),
+            xt=np.zeros((C, K)),
+            w0=np.zeros(C),
+            w02=np.zeros(C),
+            s0=np.zeros((C, L)),
+            xc=np.zeros((C, K)),
+            xcy=np.zeros((C, K, L)),
+            hess=np.zeros((C, K, K)),
+            pmin=np.ones(C),
+            pmax=np.zeros(C),
+            trimmed=np.zeros(C, dtype=np.int64),
+            capped=np.zeros(C, dtype=np.int64),
+        )
+        beta = np.array([s.beta for s in systems])
+        reg = np.zeros((C, K, L))
+        for j, s in enumerate(systems):
+            reg[j, :, : s.regression.shape[1]] = s.regression
+        for raw in batches():
+            fast.moments(raw, method, beta, reg, control_max, K, acc)
+        for j, m in enumerate(moments):
+            n = systems[j].regression.shape[1]
+            m.update(
+                w0=float(acc["w0"][j]),
+                w02=float(acc["w02"][j]),
+                s1=acc["s1"][j, :n].copy(),
+                s0=acc["s0"][j, :n].copy(),
+                xt=acc["xt"][j].copy(),
+                xc=acc["xc"][j].copy(),
+                xcy=acc["xcy"][j, :, :n].copy(),
+                H=acc["hess"][j].copy(),
+                pmin=float(acc["pmin"][j]),
+                pmax=float(acc["pmax"][j]),
+                trimmed=int(acc["trimmed"][j]),
+                capped=int(acc["capped"][j]),
+            )
+    else:
+        for raw in batches():
+            for j, transform in enumerate(transforms):
+                X, d, Y, w = weighted_system(transform(raw))
+                if not len(d):
+                    continue
+                s, m = systems[j], moments[j]
+                residual = Y if method == "ipw" else Y - X @ s.regression
+                m["s1"] += w[d] @ residual[d]
+                m["xt"] += w[d] @ X[d]
+                if method != "reg":
+                    p, wc = att_parts(X, d, s.beta, control_max)
+                    wc = wc * w
+                    m["pmin"] = min(m["pmin"], float(p.min()))
+                    m["pmax"] = max(m["pmax"], float(p.max()))
+                    m["trimmed"] += int(np.sum((~d) & (p >= control_max)))
+                    m["capped"] += int(np.sum(p >= 1 - 1e-6))
+                    m["w0"] += wc.sum()
+                    m["w02"] += wc @ wc
+                    m["s0"] += wc @ residual
+                    m["xc"] += wc @ X
+                    m["xcy"] += X.T @ (wc[:, None] * residual)
+                    m["H"] += X.T @ ((w * p * (1 - p))[:, None] * X)
+    fitted = []
+    for s, m in zip(systems, moments, strict=True):
+        wt, wc = s.treated_weight, m["w0"]
+        if method != "reg" and wc <= 0:
+            raise ValueError("no controls remain after propensity trimming")
+        mt = m["s1"] / wt
+        mc = m["s0"] / wc if method != "reg" else np.zeros_like(mt)
+        jo = -m["xt"] / wt
+        if method == "dr":
+            jo += m["xc"] / wc
+        ao = (
+            solve_checked(s.control_gram, jo, "outcome regression")
+            if method != "ipw"
+            else np.zeros(len(jo))
+        )
+        jp = (
+            -(m["xcy"] - np.outer(m["xc"], mc)) / wc if method != "reg" else np.zeros_like(m["xcy"])
+        )
+        ap = solve_checked(m["H"], jp, "propensity information") if method != "reg" else jp
+        fitted.append(
+            AdjustedContrast(
+                s,
+                method,
+                control_max,
+                wt,
+                wc,
+                mt,
+                mc,
+                ao,
+                ap,
+                dict(
+                    n_treated=s.n_treated,
+                    n_controls=s.n_controls,
+                    effective_treated=s.treated_weight**2 / s.treated_weight_squares,
+                    trimmed_controls=m["trimmed"],
+                    capped_probabilities=m["capped"],
+                    effective_controls=wc * wc / m["w02"]
+                    if method != "reg"
+                    else s.control_weight**2 / s.control_weight_squares,
+                    propensity_min=m["pmin"] if method != "reg" else None,
+                    propensity_max=m["pmax"] if method != "reg" else None,
+                    iterations=s.iterations,
+                ),
+            )
+        )
+    return fitted