diff-stream 0.1.0.dev0__tar.gz

diff_stream-0.1.0.dev0/CHANGELOG.md

@@ -0,0 +1,60 @@
+# Changelog
+
+## Unreleased (0.1.0)
+
+The first release.
+
+- **Estimators**, one function each: `did_callaway_santanna` (regression,
+  IPW and doubly robust covariate adjustment; never-treated or not-yet-treated
+  controls; varying or universal base periods; panels or repeated cross
+  sections), `did_sun_abraham`, `did_bjs`
+  (the authors' standard errors, residual averaging groups, leave-out, and the
+  pre-trend test of `pretrend_leads`), `did_two_stage` (event study or overall
+  ATT), `did_extended_twfe` (linear; conditional or unconditional inference)
+  and `did_stacked` (stacked DiD with Wing, Freedman and Hollingsworth's
+  corrective weights, without building the stack).
+  Sun-Abraham, BJS and the two-stage estimator take covariates and additional
+  fixed effects (`"y ~ x | region^year"`) and absorb fixed effects with any
+  number of levels (county-by-year) through hdfe-stream
+  (`fixed_effect_solver`). BJS drops treated observations whose untreated
+  outcome cannot be imputed, with an `UnimputedObservationsWarning` and the
+  affected cells in its diagnostics.
+- **Speed**: Callaway-Sant'Anna with covariates (balanced and unbalanced panels, and
+  repeated cross sections), the two-stage estimator, extended TWFE, and Sun-Abraham and BJS with
+  covariates run their passes over the data in compiled code, in parallel, as do the
+  validation and summing when the panel is first staged, and the next batch of units is read in
+  the background. Importing the package sets `OMP_WAIT_POLICY=passive` (unless set), which
+  speeds the compiled passes. Results are identical for any number of threads. The first fit after installation compiles for a few seconds, once.
+- **Benchmarks** against R's did, fixest, didimputation, did2s, etwfe and
+  bacondecomp and Python's diff-diff and pyfixest, from 10,000 to 10 million
+  units, with covariates and with fixed effects of millions of levels
+  (`docs/benchmarks/`).
+- **Data**: Polars frames or Parquet larger than memory; treatment timing as
+  a first-treated-period column (`cohort=`, 0 or null for never treated) or a
+  0/1 column (`treatment=`); evenly spaced periods (every second year, ...),
+  with event times counted in periods and the original labels shown; unbalanced panels
+  (`incomplete_units="error" | "drop" | "keep"`), sampling weights (optionally
+  time-varying), anticipation, and panels where every unit is eventually
+  treated (`last_treated_controls`).
+- **Results**: cohort-period effects (`cohort_time()`) and their aggregation
+  by event time, cohort, calendar period or overall (`aggregate()`), with
+  balanced windows (`balanced_horizon`, `balance_window`) and the cells and
+  cohorts behind each estimate (`contributions()`, `composition()`);
+  pre-trend tests; effective sample sizes and covariate balance for
+  Callaway-Sant'Anna.
+- **Inference**: standard errors clustered by unit (CRV0, CRV1) that include
+  the uncertainty in aggregation shares; Callaway-Sant'Anna's multiplier
+  bootstrap with simultaneous bands; `refit_bootstrap`, which refits any
+  estimator on unit resamples, serially or in parallel.
+- **Diagnostics**: `bacon_decomp` (with covariates, weights and
+  always-treated units) and `twfe_weights` (`feTR`, with switching treatment).
+- **Resources**: memory and thread limits from the job (stream-limits);
+  staging in memory when the panel fits and in whole-unit pieces on disk
+  otherwise, with batch sizes that adapt to the memory left; `profile=True`.
+- Validated against R did/DRDID, fixest, didimputation, did2s, etwfe/emfx,
+  bacondecomp and TwoWayFEWeights, and the authors' Stata `did_imputation`.
+
+Known limitation: on macOS, the propensity logit of IPW and doubly robust
+Callaway-Sant'Anna can fail to converge when two covariates are almost
+identical (a condition number near 1e10); it converges on Linux and Windows.
+The test for that design is marked as an expected failure on macOS.

diff_stream-0.1.0.dev0/CITATION.cff

@@ -0,0 +1,20 @@
+cff-version: 1.2.0
+message: "If you use this software in your research, please cite it as below."
+type: software
+title: "diff-stream"
+abstract: "Out-of-core difference-in-differences with staggered treatment timing: Callaway-Sant'Anna, Sun-Abraham, Borusyak-Jaravel-Spiess imputation, Gardner's two-stage estimator, Wooldridge's extended two-way fixed effects and stacked DiD, with the Goodman-Bacon decomposition and TWFE weights as diagnostics. Standard errors are clustered by unit and include the uncertainty in aggregation shares; fixed effects with many levels are absorbed with hdfe-stream. Results are tested against the authors' R and Stata packages."
+authors:
+  - family-names: Tucker
+    given-names: Lee
+    email: lee.c.tucker@census.gov
+license: MIT
+repository-code: "https://github.com/leetucker/diff-stream"
+url: "https://github.com/leetucker/diff-stream"
+keywords:
+  - econometrics
+  - difference-in-differences
+  - staggered adoption
+  - event study
+  - treatment effects
+  - two-way fixed effects
+  - out-of-core

diff_stream-0.1.0.dev0/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Lee Tucker
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diff_stream-0.1.0.dev0/MANIFEST.in

@@ -0,0 +1,2 @@
+include LICENSE.md README.md CHANGELOG.md CITATION.cff
+include diff_stream/py.typed

diff_stream-0.1.0.dev0/PKG-INFO

@@ -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.

diff_stream-0.1.0.dev0/README.md

@@ -0,0 +1,143 @@
+# 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.

diff_stream-0.1.0.dev0/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-0.1.0.dev0/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-0.1.0.dev0/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.
+"""