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

diff_stream/twostage_absorbed.py

@@ -0,0 +1,190 @@
+"""Gardner's two-stage estimator with the fixed effects absorbed iteratively.
+
+The first stage is BJS's: the untreated rows' model (unit, time and additional
+fixed effects, covariates) fitted by hdfe-stream, whose predictions give every
+row its first-stage residual. The second stage regresses those residuals on the
+treatment indicators, one per event time or a single one. The indicators do not
+overlap and there is no intercept, so a coefficient is the weighted mean of its
+rows' residuals.
+
+The covariance is Gardner's: each unit's second-stage score, net of the effect of
+the estimated first stage on it. A row's weight in the first stage's effect on
+coefficient m is -w_it z_it'u_m, with u_m = S^- Z'v_m the solution of the
+untreated model's normal equations for the rows of m
+(``HDFEResult.solve_normal_equations``), the computation that BJS uses per
+cohort-period cell. Here there is one column per coefficient, so the solve is
+small however many fixed-effect levels there are, and a single pass over the
+residuals gives each unit's score.
+"""
+
+import warnings
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import polars as pl
+from hdfe_stream.core.influence import ClusterScoreAccumulator
+
+from ._core.partition import whole_group_batches
+from .absorbed import (
+    _DROPPED,
+    absorbed_terms,
+    first_columns,
+    hdfe_options,
+    info,
+    solve_block,
+    staged,
+    untreated,
+)
+
+
+def _indicator(specification, horizons, anticipation):
+    """The row's second-stage coefficient (an index), or null for none."""
+    event = pl.col("t") - pl.col("g")
+    if specification == "static":
+        return (
+            pl.when((pl.col("g") != 0) & (event >= -anticipation))
+            .then(pl.lit(0, dtype=pl.Int64))
+            .otherwise(None)
+        )
+    return (
+        pl.when(pl.col("g") != 0)
+        .then(
+            event.replace_strict(
+                horizons, list(range(len(horizons))), default=None, return_dtype=pl.Int64
+            )
+        )
+        .otherwise(None)
+    )
+
+
+def fit_two_stage_absorbed(store, terms, specification, horizons, anticipation, *, threads=None):
+    """(beta, covariance sums, diagnostics) of the two-stage fit, CRV0 by unit.
+
+    ``horizons`` are the event times of a dynamic fit (empty for a static one).
+    """
+    from hdfe_stream import feols_stream
+
+    frame, covariates, xs, absorbed = staged(store, terms)
+    H = len(horizons) if specification == "dynamic" else 1
+    fml = f"y ~ {' + '.join(xs) or '1'} | " + " + ".join(["u", "t", *absorbed])
+    options = hdfe_options(store, threads, len(xs), vcov="iid", keep_intermediates=True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", _DROPPED)
+        first = feols_stream(
+            fml, frame.filter(untreated(anticipation)), workdir=store.root, **options
+        )
+    try:
+        with TemporaryDirectory(prefix="diff-stream-gardner-", dir=store.root) as workspace:
+            beta, V = _second_stage(
+                store, first, terms, frame, workspace, specification, horizons, anticipation, H
+            )
+        diagnostics = info(covariates, xs, list(first.coefnames), terms)
+    finally:
+        first.cleanup()
+    diagnostics["first_stage_columns"] = diagnostics["regression_columns"]
+    return beta, V, diagnostics
+
+
+def _second_stage(store, first, terms, frame, workspace, specification, horizons, anticipation, H):
+    weight = pl.col("w") if store.weighted else pl.lit(1.0)
+    _, _, xs, absorbed = staged(store, terms)
+
+    # Every row's first-stage residual: the outcome less the imputed untreated one.
+    residuals = f"{workspace}/residuals.parquet"
+    (
+        first.predict(frame, name="prediction")
+        .with_columns(pl.col("g").cast(pl.Int64), pl.col("t").cast(pl.Int64))
+        .with_columns(
+            (pl.col("y") - pl.col("prediction")).alias("r"),
+            untreated(anticipation).alias("untreated"),
+            _indicator(specification, horizons, anticipation).alias("m"),
+        )
+        .drop("prediction")
+        .sink_parquet(residuals, row_group_size=store.read_rows, engine="streaming")
+    )
+    rows = pl.scan_parquet(residuals).filter(pl.col("m").is_not_null())
+    unidentified = rows.select(pl.col("r").null_count()).collect().item()
+    if unidentified:
+        raise ValueError(
+            f"Gardner's second stage is not identified for {unidentified:,} observations: "
+            "a fixed-effect level of theirs has no untreated observation"
+        )
+    sums = (
+        rows.group_by("m")
+        .agg(
+            (weight * pl.col("r")).sum().alias("effect"),
+            (pl.col("w").sum() if store.weighted else pl.len().cast(pl.Float64)).alias("mass"),
+        )
+        .collect()
+        .sort("m")
+    )
+    if sums.height != H or not (sums["mass"].to_numpy() > 0).all():
+        raise ValueError("Gardner second-stage design is unidentified")
+    mass = sums["mass"].to_numpy()
+    beta = sums["effect"].to_numpy() / mass
+
+    # The untreated model's normal equations for each coefficient's rows.
+    names = [f"v{j}" for j in range(H)]
+    target = rows.select(
+        "u",
+        "t",
+        *[c for c in first_columns(first) if c not in ("u", "t")],
+        *[
+            pl.when(pl.col("m") == j).then(weight / mass[j]).otherwise(0.0).alias(n)
+            for j, n in enumerate(names)
+        ],
+    )
+    solution = first.solve_normal_equations(target, names, block=solve_block(store, H))
+    try:
+        V = _covariance(store, solution, residuals, names, xs, absorbed, beta, mass)
+    finally:
+        solution.cleanup()
+    return beta, V
+
+
+def _covariance(store, solution, residuals, names, xs, absorbed, beta, mass):
+    """Sum over units of the outer products of their corrected second-stage scores."""
+    H = len(names)
+    tables = []
+    for term in absorbed_terms(absorbed):
+        table = solution.fixef(term).collect()
+        keys = [c for c in table.columns if c not in names]
+        tables.append((keys, table.with_row_index("__level"), table.select(names).to_numpy()))
+    kept = solution.coef["coefficient"].to_list()
+    coef = solution.coef.select(names).to_numpy()
+    accumulator = ClusterScoreAccumulator()
+    store._guard(32e6 + 8 * H * 4 * store.batch_units * store.periods)
+    for d in whole_group_batches(residuals, "u", store.batch_units, store.read_rows):
+        w = d["w"].to_numpy() if store.weighted else np.ones(d.height)
+        r = d["r"].to_numpy()
+        rowscore = np.zeros((d.height, H))
+        # the first stage's part: a clean row's weight in each coefficient is
+        # -w z'u, so its share of the score is that times its residual
+        clean = np.flatnonzero(d["untreated"].to_numpy())
+        if clean.size:
+            c = d[clean]
+            q = np.zeros((clean.size, H))
+            for keys, table, values in tables:
+                level = (
+                    c.select(keys)
+                    .with_columns(pl.col(k).cast(table.schema[k]) for k in keys)
+                    .join(
+                        table.select(*keys, "__level"),
+                        on=keys,
+                        how="left",
+                        maintain_order="left",
+                    )["__level"]
+                    .to_numpy()
+                )
+                q += values[level]
+            if kept:
+                q += c.select(kept).to_numpy() @ coef
+            rowscore[clean] = -(w[clean] * r[clean])[:, None] * q
+        # the second stage's part: the rows of coefficient m, net of the mean
+        m = d["m"].fill_null(-1).to_numpy()
+        hit = np.flatnonzero(m >= 0)
+        rowscore[hit, m[hit]] += w[hit] * (r[hit] - beta[m[hit]]) / mass[m[hit]]
+        u = d["u"].to_numpy()
+        starts = np.r_[0, np.flatnonzero(u[1:] != u[:-1]) + 1]
+        accumulator.add_clusters(np.add.reduceat(rowscore, starts, axis=0))
+    return accumulator.sums

diff_stream/unbalanced.py

@@ -0,0 +1,297 @@
+"""Observed-row Sun–Abraham and BJS systems for unbalanced panels."""
+
+import shutil
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+from ._core import covariate_kernels
+from ._core.nuisance import solve_checked
+
+
+def fit_sunab(store, cells, anticipation=0):
+    """Exact saturated SA regression via independent cohort calendar effects.
+
+    Unit intercepts absorb each cohort's calendar-effect normalization. Since
+    every other observed cohort/date has an interaction, the regression splits
+    into one unit/time FE fit per cohort. Subtract the never-treated calendar
+    change from each treated cohort's change relative to g-1. This identity
+    holds with arbitrary observed-row patterns; no balanced change shortcut or
+    all-row interaction design is used.
+    """
+    T, M = store.periods, len(cells)
+    systems = {}
+    for g, counts in store.cell_counts.items():
+        reference = g - anticipation - 1 - store.times[0] if g else 0
+        if counts[reference] == 0 or (g == 0 and np.any(counts == 0)):
+            raise ValueError(
+                "Sun–Abraham interactions are unidentified: missing cohort reference or never-treated date"
+            )
+        free = np.flatnonzero((counts > 0) & (np.arange(T) != reference))
+        Z = np.eye(T)[:, free]
+        systems[g] = [Z, np.zeros((len(free), len(free))), np.zeros(len(free))]
+    for g, y, w in store.batches(include_weights=True):
+        for cg in np.unique(g):
+            Z, Q, q = systems[cg]
+            values = y[g == cg]
+            sw = w[g == cg]
+            obs = np.isfinite(values)
+            n = obs.sum(axis=1)
+            zbar = obs @ Z / n[:, None]
+            raw = np.where(obs, values, 0)
+            centered = np.where(obs, raw - raw.sum(axis=1)[:, None] / n[:, None], 0)
+            Q += Z.T @ ((sw @ obs)[:, None] * Z) - zbar.T @ ((sw * n)[:, None] * zbar)
+            q += Z.T @ (sw @ centered)
+    coefficients, derivatives = {}, {}
+    for g, (Z, Q, q) in systems.items():
+        context = f"Sun–Abraham interactions are unidentified in cohort {g} (disconnected observed design)"
+        b = solve_checked(Q, q, context) if len(q) else np.zeros(0)
+        coefficients[g] = Z @ b
+        ix = [j for j, (cg, t) in enumerate(cells) if cg == g] if g else list(range(M))
+        D = np.zeros((len(ix), Z.shape[1]))
+        for k, j in enumerate(ix):
+            cg, t = cells[j]
+            ti, bi = t - store.times[0], cg - anticipation - 1 - store.times[0]
+            D[k] = (Z[ti] - Z[bi]) * (1 if g else -1)
+        H = solve_checked(Q, D.T, context) if len(q) else np.zeros((0, len(ix)))
+        derivatives[g] = (ix, H)
+    beta = np.array(
+        [
+            coefficients[g][t - store.times[0]]
+            - coefficients[g][g - anticipation - 1 - store.times[0]]
+            - coefficients[0][t - store.times[0]]
+            + coefficients[0][g - anticipation - 1 - store.times[0]]
+            for g, t in cells
+        ]
+    )
+    counts = np.array([store.cell_weights[g][t - store.times[0]] for g, t in cells], float)
+    V = np.zeros((M, M))
+    for g, y, w in store.batches(include_weights=True):
+        scores = np.zeros((len(g), M))
+        for cg in np.unique(g):
+            mask = g == cg
+            obs = np.isfinite(y[mask])
+            raw = np.where(obs, y[mask] - coefficients[cg], 0)
+            residual = np.where(obs, raw - raw.sum(axis=1)[:, None] / obs.sum(axis=1)[:, None], 0)
+            ix, H = derivatives[cg]
+            scores[np.ix_(np.flatnonzero(mask), ix)] = w[mask, None] * (
+                residual @ systems[cg][0] @ H
+            )
+        V += scores.T @ scores
+    return beta, counts, V
+
+
+def residuals(g, y, times, Z, theta, anticipation=0):
+    observed = np.isfinite(y)
+    untreated = observed & ((g[:, None] == 0) | (times[None, :] < g[:, None] - anticipation))
+    n0 = untreated.sum(axis=1)
+    if np.any(n0 == 0):
+        raise ValueError(
+            "BJS requires at least one untreated observation per unit; exclude units without one explicitly"
+        )
+    zbar = untreated @ Z / n0[:, None]
+    raw = np.where(observed, y - Z @ theta, 0)
+    alpha = (raw * untreated).sum(axis=1) / n0
+    r = np.where(observed, raw - alpha[:, None], 0)
+    return observed, untreated, zbar, r
+
+
+def fit_bjs(store, cells, avgeffectsby, leaveout, workdir, anticipation=0):
+    T, M = store.periods, len(cells)
+    Z = np.eye(T)[:, 1:]
+    Q, q = np.zeros((T - 1, T - 1)), np.zeros(T - 1)
+    for g, y, w in store.batches(include_weights=True):
+        _, use, zm, centered = residuals(g, y, store.times, Z, np.zeros(T - 1), anticipation)
+        n0 = use.sum(axis=1)
+        Q += Z.T @ ((w @ use)[:, None] * Z) - zm.T @ ((w * n0)[:, None] * zm)
+        q += Z.T @ (w @ (centered * use))
+    theta = solve_checked(Q, q, "BJS untreated time effects (disconnected observed design)")
+    counts = np.array([store.cell_weights[g][t - store.times[0]] for g, t in cells], float)
+    beta, D = np.zeros(M), np.zeros((M, T - 1))
+    for g, y, w in store.batches(include_weights=True):
+        observed, _, zm, r = residuals(g, y, store.times, Z, theta, anticipation)
+        for j, (cg, t) in enumerate(cells):
+            ti = t - store.times[0]
+            mask = (g == cg) & observed[:, ti]
+            beta[j] += w[mask] @ r[mask, ti] / counts[j]
+            D[j] += w[mask] @ (Z[ti] - zm[mask]) / counts[j]
+    H = solve_checked(Q, D.T, "BJS untreated time effects").T
+    context = ObservedBJSInference(
+        store, cells, counts, Z, theta, H, avgeffectsby, leaveout, workdir, anticipation
+    )
+    try:
+        V = context.covariance(np.eye(M), allow_unidentified=avgeffectsby != "cohort_time")
+    except BaseException:
+        context.cleanup()
+        raise
+    return beta, counts, V, context
+
+
+class ObservedBJSInference:
+    """Disk-backed, bounded cluster inference including missingness patterns.
+
+    The result owns these sorted observed-row partitions. Each aggregation uses
+    two streaming passes; no unit-by-effect score matrix is saved or retained.
+    """
+
+    def __init__(
+        self,
+        store,
+        cells,
+        counts,
+        Z,
+        theta,
+        H,
+        avgeffectsby,
+        leaveout,
+        workdir,
+        anticipation=0,
+        *,
+        model=None,
+    ):
+        self.workspace = TemporaryDirectory(prefix="diff-stream-bjs-", dir=workdir)
+        self.store, self.cells, self.counts = store, cells, counts
+        self.Z, self.theta, self.H = Z, theta, H
+        self.model = model
+        self.avgeffectsby, self.leaveout = avgeffectsby, leaveout
+        self.anticipation = anticipation
+        self.closed = False
+        paths = []
+        for path in store.paths:
+            paths.append(shutil.move(str(path), self.workspace.name + "/" + path.name))
+        store.paths = paths
+
+    def cleanup(self):
+        self.workspace.cleanup()
+        self.closed = True
+
+    def covariance(self, contrasts, *, allow_unidentified=False):
+        if self.closed:
+            raise ValueError("BJS aggregation data were removed by cleanup(); refit to aggregate")
+        A = np.atleast_2d(contrasts)
+        P, M = A.shape
+        self.store._guard(
+            8
+            * (
+                8 * P * P
+                + 8 * min(self.store.batch_units, self.store.n_units) * (P + self.store.periods)
+            )
+            + 32e6
+            + (
+                8
+                * (
+                    8
+                    * min(self.store.batch_units, self.store.n_units)
+                    * self.store.periods
+                    * (len(self.model.keep) + 1)
+                    + 8 * P * len(self.model.keep)
+                )
+                if self.model is not None
+                else 0
+            )
+        )
+        labels, groups = {}, {}
+        for j, (g, t) in enumerate(self.cells):
+            label = {"cohort_time": (g, t), "event_time": t - g, "cohort": g, "all": 0}[
+                self.avgeffectsby
+            ]
+            r = labels.setdefault(label, len(labels))
+            groups.setdefault((g, r), []).append(j)
+        R = len(labels)
+        denom, numer = np.zeros((P, R)), np.zeros((P, R))
+
+        def batches():
+            if self.model is not None:
+                yield from self.model.batches()
+            else:
+                for g, y, sw in self.store.batches(include_weights=True):
+                    yield g, y, None, sw
+
+        def parts(g, y, X, sw):
+            if self.model is not None:
+                observed, untreated, _, resid, first_score = self.model.residuals(
+                    g, y, X, sw, self.theta, self.anticipation
+                )
+            else:
+                observed, untreated, _, resid = residuals(
+                    g, y, self.store.times, self.Z, self.theta, self.anticipation
+                )
+                first_score = sw[:, None] * ((resid * untreated) @ self.Z)
+                sw = np.broadcast_to(sw[:, None], y.shape)
+            yield first_score
+            for (cg, r), js in groups.items():
+                mask = g == cg
+                if not mask.any():
+                    continue
+                ti = np.array([self.cells[j][1] - self.store.times[0] for j in js])
+                w = A[:, js] / self.counts[js]
+                totals = (sw[mask][:, ti] * observed[mask][:, ti]) @ w.T
+                effects = (sw[mask][:, ti] * resid[mask][:, ti]) @ w.T
+                yield mask, r, totals, effects
+
+        compiled = self.model is not None and covariate_kernels.ENABLED
+        if compiled:
+            kernels = covariate_kernels.BJSKernels(
+                self.store, self.cells, self.model.keep, self.theta, self.anticipation
+            )
+            cell_labels = [
+                labels[
+                    {"cohort_time": (g, t), "event_time": t - g, "cohort": g, "all": 0}[
+                        self.avgeffectsby
+                    ]
+                ]
+                for g, t in self.cells
+            ]
+            grouping = kernels.grouping(cell_labels, self.counts, A)
+            denom, numer = kernels.totals(self.model.raw_batches, grouping, R)
+        else:
+            for g, y, X, sw in batches():
+                iterator = parts(g, y, X, sw)
+                next(iterator)
+                for _mask, r, totals, effects in iterator:
+                    denom[:, r] += (totals**2).sum(axis=0)
+                    numer[:, r] += (totals * effects).sum(axis=0)
+        averages = np.divide(numer, denom, out=np.zeros_like(numer), where=denom > 0)
+        V, unidentified = np.zeros((P, P)), np.zeros(P, bool)
+        derivative = A @ self.H
+        if compiled:
+            for raw in self.model.raw_batches():
+                scores, invalid = kernels.scores(
+                    raw, grouping, derivative, averages, denom, self.leaveout
+                )
+                if self.leaveout:
+                    unidentified |= invalid
+                    if invalid.any() and not allow_unidentified:
+                        raise ValueError(
+                            "BJS leave_out requires more than one contributing unit in each residual group; use a coarser residual_averaging"
+                        )
+                V += scores.T @ scores
+            V[unidentified, :] = np.nan
+            V[:, unidentified] = np.nan
+            return (V + V.T) / 2
+        for g, y, X, sw in batches():
+            iterator = parts(g, y, X, sw)
+            scores = -next(iterator) @ derivative.T
+            for mask, r, totals, effects in iterator:
+                contribution = effects - totals * averages[:, r]
+                if self.leaveout:
+                    lev = np.divide(
+                        totals**2, denom[:, r], out=np.zeros_like(totals), where=denom[:, r] > 0
+                    )
+                    invalid = (np.abs(totals) > 0) & (lev >= 1 - 1e-12)
+                    unidentified |= invalid.any(axis=0)
+                    if invalid.any() and not allow_unidentified:
+                        raise ValueError(
+                            "BJS leave_out requires more than one contributing unit in each residual group; use a coarser residual_averaging"
+                        )
+                    contribution = np.divide(
+                        contribution,
+                        1 - lev,
+                        out=np.zeros_like(contribution),
+                        where=lev < 1 - 1e-12,
+                    )
+                scores[mask] += contribution
+            V += scores.T @ scores
+        V[unidentified, :] = np.nan
+        V[:, unidentified] = np.nan
+        return (V + V.T) / 2

diff_stream-0.1.0.dev0.dist-info/METADATA

@@ -0,0 +1,192 @@
+Metadata-Version: 2.4
+Name: diff-stream
+Version: 0.1.0.dev0
+Summary: Staggered difference-in-differences on data larger than memory: Callaway-Sant'Anna, Sun-Abraham, Borusyak-Jaravel-Spiess imputation, two-stage DiD and extended TWFE, with TWFE diagnostics
+Author: Lee Tucker
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/leetucker/diff-stream
+Project-URL: Source, https://github.com/leetucker/diff-stream
+Project-URL: Issues, https://github.com/leetucker/diff-stream/issues
+Project-URL: Documentation, https://leetucker.github.io/diff-stream/
+Project-URL: Changelog, https://github.com/leetucker/diff-stream/blob/main/CHANGELOG.md
+Keywords: econometrics,difference-in-differences,staggered adoption,event study,Callaway-Sant'Anna,Sun-Abraham,imputation,two-way fixed effects,Goodman-Bacon,treatment effects,out-of-core
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: hdfe-stream[formula]>=0.10.0
+Requires-Dist: stream-limits>=0.3.0
+Requires-Dist: polars!=1.41.*,>=1.37
+Requires-Dist: numpy>=2.0
+Requires-Dist: scipy>=1.13
+Requires-Dist: numba>=0.60
+Requires-Dist: pyarrow>=16.0
+Provides-Extra: benchmark
+Requires-Dist: stream-limits[launcher]>=0.3.0; extra == "benchmark"
+Requires-Dist: matplotlib>=3.7; extra == "benchmark"
+Requires-Dist: psutil>=5.9; extra == "benchmark"
+Requires-Dist: pandas>=2.0; extra == "benchmark"
+Provides-Extra: docs
+Requires-Dist: sphinx>=8.0; extra == "docs"
+Requires-Dist: myst-parser>=4.0; extra == "docs"
+Requires-Dist: furo>=2024.8.6; extra == "docs"
+Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
+Requires-Dist: sphinx-design>=0.6; extra == "docs"
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == "test"
+Requires-Dist: stream-limits[testing]>=0.3.0; extra == "test"
+Dynamic: license-file
+
+# diff-stream
+
+Difference-in-differences with staggered treatment timing, on panels larger than
+memory. Includes Callaway and Sant'Anna, Sun and Abraham, Borusyak, Jaravel and
+Spiess imputation, Gardner's two-stage DiD, Wooldridge's extended two-way fixed
+effects, and stacked DiD. Two diagnostics show what a conventional two-way fixed
+effects estimate averages. Results are tested against the authors' R and Stata
+packages.
+
+It reads Parquet in a stream, sorts the panel once (in memory when it fits,
+in whole-unit pieces on disk when it does not), and computes standard errors
+clustered by unit without holding all units' scores.
+
+```python
+from diff_stream import did_callaway_santanna
+
+fit = did_callaway_santanna(
+    "earnings ~ age + C(region)",
+    "panel/*.parquet",
+    panel=("person", "year"),
+    cohort="first_treated_year",
+    control_group="not_yet_treated",
+)
+fit.aggregate().summary()           # event study
+fit.aggregate("overall").summary()  # overall ATT
+```
+
+**Status: in development** (version 0.1.0.dev0, not yet on PyPI). The
+estimators, inference and documentation are complete and tested; the
+[benchmarks](docs/benchmarks/index.md) compare them with other implementations.
+
+## What it does
+
+**Estimators** (each cites its paper in its docstring)
+- `did_callaway_santanna`: group-time effects, with outcome regression, IPW
+  or doubly robust covariate adjustment; never-treated or not-yet-treated
+  controls; varying or universal base periods; panels or repeated cross
+  sections.
+- `did_sun_abraham`: the interaction-weighted event study, with covariates
+  and additional fixed effects.
+- `did_bjs`: imputation, with the authors' standard errors (residual
+  averaging groups, leave-out) and their pre-trend test.
+- `did_two_stage`: Gardner's two-stage DiD, event study or overall ATT.
+- `did_extended_twfe`: Wooldridge's extended TWFE (linear), with
+  covariate-interacted effects.
+- `did_stacked`: stacked DiD with corrective weights, without building the
+  stack.
+
+**Data**
+- Unbalanced panels, sampling weights (also time-varying), anticipation, and
+  panels where every unit is eventually treated (the last-treated cohort as
+  comparison).
+
+**Results and inference**
+- Aggregation by event time, cohort and calendar period, or overall, with
+  balanced event windows and the cells and cohorts behind every estimate.
+- Standard errors clustered by unit, including the uncertainty in the cohort
+  shares used to aggregate; pre-trend tests.
+- Callaway-Sant'Anna's multiplier bootstrap for simultaneous bands, and a
+  bootstrap that refits any estimator on unit resamples, in parallel.
+
+**Diagnostics**
+- `bacon_decomp`: the Goodman-Bacon decomposition (with covariates and
+  weights).
+- `twfe_weights`: de Chaisemartin and D'Haultfoeuille's weights and
+  sensitivity measures.
+
+**Resources**
+- `n_threads=` and `max_memory_gb=` respect a job's limits (stream-limits);
+  `profile=True` reports the memory and disk a fit needs, without fitting.
+- Results with hdfe-stream's reporting: `summary()`, `tidy()`, `confint()`,
+  `wald_test()`, `etable`.
+
+## Why it is fast
+
+- **Sorted once.** The panel is sorted by unit in one piece when it fits in
+  memory, in whole-unit pieces on disk otherwise; later passes read whole units
+  in large batches.
+- **Sufficient statistics.** Without covariates, the estimators reduce to
+  cohort-by-period means plus unit scores computed in parallel.
+- **Covariance by cohort.** A unit's scores touch only its cohort's cells, so
+  the joint covariance of hundreds of effects is accumulated cohort by cohort.
+- **Shared passes.** All comparisons' propensity and outcome models are
+  fitted together.
+
+On a simulated panel of one million units over eight periods (8 million
+rows), each estimator fits in one to two seconds on a 16-vCPU cloud instance,
+and ten million units (80 million rows) in under half a minute, with a peak of
+about 4 GB. With `max_memory_gb=2`, a million units fit within about 0.8 GB.
+The benchmarks give the comparisons with R and Python implementations.
+
+## Validation against other implementations
+
+| Implementation | Compared on |
+|---|---|
+| R did, DRDID | Callaway-Sant'Anna: effects, full covariance, aggregations, multiplier bootstrap; the `mpdta` example |
+| R fixest (`sunab`) | Sun-Abraham: coefficients, covariance, aggregation; the Castle Doctrine panel |
+| Stata `did_imputation` (authors), R didimputation | BJS: estimates and full covariance, every averaging group, leave-out, pre-trend leads |
+| R did2s (authors) | two-stage: coefficients and covariance |
+| R etwfe, emfx (authors) | extended TWFE: effects and the covariance of every aggregation |
+| fixest on the stacked data | stacked DiD: coefficients and covariance |
+| R bacondecomp | Goodman-Bacon: every comparison; the `math_reform` example |
+| R TwoWayFEWeights (authors) | TWFE weights (`feTR`) and sensitivity measures |
+
+Agreement is within 2e-7 absolute, usually 1e-10. The frozen references and
+the R and Stata scripts that make them are in `tests/data/` and
+`benchmarks/external/`; `docs/implementations.md` lists the conventions on
+which these packages differ.
+
+## Install
+
+diff-stream is not on PyPI yet. Install it from the repository:
+
+```bash
+pip install git+https://github.com/leetucker/diff-stream
+```
+
+Python 3.10 or later. It depends on hdfe-stream (with its formula extra),
+stream-limits, Polars, NumPy 2, numba, SciPy and pyarrow.
+
+## Documentation
+
+Build the documentation locally with:
+
+```bash
+pip install -e ".[docs]"
+sphinx-build -W --keep-going -b html docs docs/_build/html
+```
+
+It covers choosing an estimator, your data, each estimator, results and
+inference, the diagnostics, worked examples with their output, how it works,
+memory and threads, the conventions that differ from other packages, and the
+validation. `examples/` has runnable scripts, and `benchmarks/` the
+performance comparisons (see `benchmarks/README.md`).
+
+A sibling of [hdfe-stream](https://github.com/leetucker/hdfe-stream),
+[logit-stream](https://github.com/leetucker/logit-stream),
+[dpd-stream](https://github.com/leetucker/dpd-stream) and
+[match-stream](https://github.com/leetucker/match-stream).
+
+## License
+
+MIT.