match-stream 0.1.0__tar.gz

match_stream-0.1.0/CHANGELOG.md

@@ -0,0 +1,72 @@
+# Changelog
+
+## 0.1.0 (2026-11-06)
+
+First release: matching for data too large for memory, with treatment effects
+and the standard errors that go with them.
+
+**Three functions, one for each kind of matching.** Each takes only its own
+options and returns its own result, with `matches()`, `weights()`, `counts()`,
+`balance()` and `effect()`.
+
+- **`match_nn`: nearest-neighbor matching**, on a propensity score or on the
+  covariates themselves.
+  - The score is fitted for you (logit or probit, with high-dimensional fixed
+    effects absorbed, out of core by hdfe-stream), or it is a column you
+    supply. Or match on Mahalanobis, inverse-variance or Euclidean distance.
+  - One match or k per treated unit, with or without replacement, ties kept
+    or broken, and in any of five orders without replacement.
+  - Calipers on the score and on covariates, exact strata, common-support
+    trimming, and Mahalanobis distance within a score caliper.
+  - The focal units are the treated (the ATT) or the controls (the ATC);
+    `effect()` can also match both ways for the ATE.
+- **`match_kernel`: kernel and ridge matching** on the score. Five kernels,
+  ridge weights, a default bandwidth from the 1:1 nearest-neighbor distances
+  (`PairDistance`) and cross-validated bandwidths (`CrossValidation`), without
+  forming pairs.
+- **`match_cem`: coarsened exact matching**, with cells, weights and balance.
+
+**Treatment effects with the standard errors the literature uses.**
+`effect()` gives the ATT, ATC and ATE, for one outcome or several in one pass.
+
+- Abadie-Imbens standard errors for matching with replacement, with the bias
+  adjustment, with heteroskedastic variances (`n_variance_matches=`), and with
+  the correction for an estimated propensity score (Abadie and Imbens 2016).
+- Influence-function standard errors, by unit or clustered on the pair, the
+  matched unit or your own columns, as Abadie and Spiess (2022).
+- Results follow hdfe-stream's conventions, so they print, tabulate and export
+  to pyfixest the same way.
+
+**Balance diagnostics.** A table of standardized differences, variance ratios
+and empirical-CDF distances before and after matching, with the sample sizes
+by arm.
+
+**Built for large data.**
+
+- Only a few narrow columns are held; the rest of the data stays on disk.
+  Large data are matched a batch of exact strata at a time
+  (`memory="by_stratum"`), with memory bounded by the largest stratum.
+- Matching with replacement and by covariates uses k-d trees, and searches run
+  in parallel. On a 16-core machine, 1:1 matching of 100 million units takes
+  about one to two minutes (see the benchmarks).
+- `profile=True` estimates the memory and disk a match would need, before
+  running it. Threads and memory follow your job's limits (PBS, Slurm,
+  containers).
+
+**Checked against other software.** On shared reference data the matches and
+weights are the same, unit for unit, as R's MatchIt; the effects and standard
+errors agree with R's Matching and Stata's teffects and kmatch; and the matches
+agree with psmatch2's. `docs/implementations.md` lists where packages differ
+and which option reproduces each.
+
+**Documentation and benchmarks.**
+
+- A guide to choosing a matching method, runnable examples, a page on how it
+  works, a function reference and the conventions compared with other
+  packages.
+- Benchmarks against MatchIt and R Matching, pymatchit, fastmatch, bare
+  nearest-neighbor searches and Polars, from 10,000 to 100 million units.
+- Memory tests (`pytest -m resource`) check that matching a few narrow columns
+  stays under a limit that holding the data would exceed.
+
+Requires hdfe-stream 0.9.0 or later.

match_stream-0.1.0/CITATION.cff

@@ -0,0 +1,19 @@
+cff-version: 1.2.0
+message: "If you use this software in your research, please cite it as below."
+type: software
+title: "match-stream"
+abstract: "Out-of-core matching for large data: nearest-neighbor matching with or without replacement, calipers, exact strata, Mahalanobis, kernel, ridge and coarsened exact matching, with Abadie-Imbens inference. Conventions follow R's MatchIt; results are tested against MatchIt and Matching."
+authors:
+  - family-names: Tucker
+    given-names: Lee
+    email: lee.c.tucker@census.gov
+version: 0.1.0
+date-released: 2026-11-06
+license: MIT
+repository-code: "https://github.com/leetucker/match-stream"
+url: "https://github.com/leetucker/match-stream"
+keywords:
+  - econometrics
+  - matching
+  - propensity score
+  - treatment effects

match_stream-0.1.0/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.

match_stream-0.1.0/MANIFEST.in

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

match_stream-0.1.0/PKG-INFO

@@ -0,0 +1,227 @@
+Metadata-Version: 2.4
+Name: match-stream
+Version: 0.1.0
+Summary: Out-of-core matching: nearest-neighbor (with or without replacement), calipers, exact strata, Mahalanobis, kernel and coarsened exact matching, with Abadie-Imbens inference
+Author: Lee Tucker
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/leetucker/match-stream
+Project-URL: Source, https://github.com/leetucker/match-stream
+Project-URL: Issues, https://github.com/leetucker/match-stream/issues
+Project-URL: Documentation, https://leetucker.github.io/match-stream/
+Project-URL: Changelog, https://github.com/leetucker/match-stream/blob/main/CHANGELOG.md
+Keywords: econometrics,matching,propensity score,nearest neighbor,Mahalanobis,coarsened exact matching,Abadie-Imbens,treatment effects,out-of-core
+Classifier: Development Status :: 3 - 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: Programming Language :: Python :: 3.15
+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.9.0
+Requires-Dist: stream-limits>=0.1.1
+Requires-Dist: polars!=1.41.*,>=1.37
+Requires-Dist: numpy>=1.24
+Requires-Dist: numba>=0.59
+Requires-Dist: pyarrow>=14.0
+Requires-Dist: scipy>=1.11
+Requires-Dist: threadpoolctl>=3.1
+Provides-Extra: benchmark
+Requires-Dist: matplotlib>=3.7; extra == "benchmark"
+Requires-Dist: pyarrow>=14; 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.2.1; extra == "test"
+Requires-Dist: pandas>=2.0; extra == "test"
+Dynamic: license-file
+
+# match-stream
+
+Matching on data larger than memory. It does nearest-neighbor matching, on
+a propensity score or on the covariates themselves, with or without
+replacement, calipers on the score and on covariates, exact strata, and
+kernel, ridge and coarsened exact matching. It also
+gives balance diagnostics, and treatment effects with Abadie–Imbens or
+clustered standard errors.
+
+It reads Parquet in a stream, holds only a few narrow columns while it
+matches, and fits the propensity score out of core, with absorbed fixed
+effects if you want them. Its results are tested against R's MatchIt and
+Matching and Stata's teffects and kmatch (see below).
+
+```python
+from match_stream import match_nn
+
+m = match_nn("treat ~ age + tenure + log_wage0", "extract/*.parquet",
+             exact=["state", "naics2"], replace=True, caliper=0.2, id="person_id")
+m.balance().summary()
+m.effect("earnings").summary()          # the ATT, Abadie-Imbens standard errors
+m.matches().sink_parquet("pairs.parquet")
+```
+
+**Status: version 0.1.0, the first release.** The designs, the inference and
+the documentation are tested against MatchIt, Matching, Stata's teffects, kmatch and
+psmatch2, and the benchmarks are in the documentation (see Performance).
+
+## What it does
+
+Three functions: `match_nn` (nearest neighbor, on a score or a multivariate
+distance), `match_kernel` (kernel and ridge) and `match_cem` (coarsened exact
+matching). Each returns the matches and weights; `effect()` and `balance()`
+are methods of the result.
+
+**Designs**
+- k:1 nearest neighbor on a propensity score, with or without replacement.
+  Without replacement it is greedy in any of five orders (`largest`,
+  `smallest`, `closest`, `random`, `data`).
+- Calipers on the score (in standard deviations) and on covariates, and
+  exact strata.
+- Mahalanobis, inverse-variance or Euclidean distance on covariates, alone
+  or within a score caliper (`mahalanobis_vars=`).
+- Kernel and ridge matching on the score: five kernels, a
+  bandwidth rule or cross-validation.
+- Coarsened exact matching.
+- Common support (`discard=`, `reestimate_score=`), ties, a minimum number of
+  matches per unit, and one-pass k:1 matching without replacement.
+
+**The score**
+- Logit or probit, fitted by hdfe-stream, with absorbed fixed effects
+  (`score_model="treat ~ x | firm"`), or a column you supply.
+
+**Effects and inference**
+- ATT, ATC and ATE.
+- Abadie–Imbens (2006) standard errors, with or without heteroskedastic
+  conditional variances (`n_variance_matches`).
+- The Abadie–Imbens (2011) bias adjustment.
+- The Abadie–Imbens (2016) correction for an estimated score, also with
+  absorbed fixed effects.
+- Influence-function standard errors with the matching weights fixed, by unit or clustered one or several ways.
+- Several outcomes in one pass (`effect(["y1", "y2"])`), with their joint
+  covariance matrix, for every kind of standard error.
+
+**Diagnostics**
+- A balance table (means, standardized differences, variance ratios, eCDF
+  gaps) and sample sizes.
+
+**Resources**
+- `profile=True` reports the memory, disk and candidate pairs a match would
+  need, without matching, and a resource request for a scheduler.
+- `n_threads=` and `max_memory_gb=` respect a job's limits (stream-limits).
+- Large inputs are matched a batch of exact strata at a time
+  (`memory="by_stratum"`), with the same matches.
+
+**Interface**
+- The matches, weights and subclasses as Polars LazyFrames, ready to write to
+  Parquet.
+- Results with hdfe-stream's reporting: `summary()`, `tidy()`, `confint()`,
+  `etable`, `summary_json()`.
+
+## Why it is fast
+
+- **Windows, not pairs.** Within exact strata and sorted by the score, each
+  unit's candidates are a window found by binary search. Nothing forms all
+  the pairs.
+- **Cheap skipping.** Matching without replacement skips taken controls
+  with path-compressed pointers.
+- **Trees for covariates.** Mahalanobis matching uses k-d trees, and
+  candidate lists for matching without replacement and for calipers.
+- **Window sums.** Kernel and ridge weights come from running sums, with no
+  pairs formed.
+- **Narrow columns.** Only tens of bytes a unit are held. The score model is
+  fitted out of core, and the numba kernels run strata in parallel.
+
+## Validation against other implementations
+
+| Implementation | Compared on | Agreement |
+|---|---|---|
+| MatchIt 4.8.2 | nearest neighbor (every order, k:1, calipers), Mahalanobis, `mahalanobis_vars`, common support, CEM, weights, balance | pairs identical, weights and balance to 1e-8 |
+| R Matching 4.10-15 | Abadie–Imbens SEs (`Var.calc`), the bias adjustment | estimates to 1e-10, SEs to 1e-8 |
+| Stata teffects | nnmatch, psmatch with the AI (2016) correction, ipw | to 1e-8 |
+| Stata kmatch 1.1.6 | nearest neighbor, kernel, ridge, bandwidths, influence-function SEs | matches identical, estimates to 1e-12 |
+| MatchIt + sandwich, fixest | clustered influence-function SEs | to 1e-13 |
+
+The R and Stata scripts that make these references are in `tests/reference/`.
+`docs/implementations.md` lists the conventions on which these packages
+differ, and which option reproduces each.
+
+## Performance
+
+The benchmark suites (`benchmarks/`) compare match-stream with MatchIt and
+Matching (R), pymatchit, fastmatch and Polars' `join_asof` and `group_by`,
+and with bare k-d tree and faiss searches as a floor. They cover
+eight designs and the standard errors, from 10,000 to 100 million units.
+On a 16-vCPU, 32 GB cloud machine, at 10 million units:
+
+- match-stream is 6 to 35 times faster than MatchIt where MatchIt finishes,
+  with the same matches (estimates agree to about 1e-14). At 100 million
+  units MatchIt needs 28 to 30 GB where match-stream needs 8 to 21 GB, and
+  in two designs it runs out of memory.
+- With no strata, 1:1 Mahalanobis matching takes 6.1 s, as long as a bare
+  SciPy k-d tree search of the same data (6.7 s). MatchIt exceeds 30 minutes.
+- pymatchit and R's Matching are slower still, and mostly do not finish
+  at 10 million units or more.
+- Polars' `group_by` (coarsened exact matching) and `join_asof` (1:1 on a
+  score) are faster, as they do less.
+- At 100 million units, 1:1 Mahalanobis matching without replacement and
+  with calipers needs `memory="by_stratum"` on this machine (187 s, 9.2 GB).
+
+The results, figures and caveats are in the
+[benchmarks](https://leetucker.github.io/match-stream/benchmarks/).
+
+Small runs, for correctness only, give these results:
+
+- MatchIt's estimates agree to about 1e-15 on every design.
+- At 100,000 units, R Matching's Abadie–Imbens standard error (with
+  `Var.calc = 1`) took 307 s; match-stream's, the same to 1e-15, took 0.07 s.
+
+## Install
+
+```bash
+pip install match-stream
+```
+
+Python 3.10 or later. It depends on hdfe-stream (with its formula extra),
+stream-limits, Polars, NumPy, numba, SciPy and pyarrow. On Python 3.15,
+pyfixest (which hdfe-stream's formula extra brings in) has no wheels yet: it
+builds from source with a Rust toolchain and
+`PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1` set; everything else installs from wheels.
+
+## Documentation
+
+Build the documentation locally with:
+
+```bash
+pip install -e ".[docs]"
+sphinx-build -W --keep-going -b html docs docs/_build/html
+```
+
+The guide covers:
+
+- how it works;
+- the conventions that differ from other matching packages;
+- the resources a match needs;
+- the API reference.
+
+`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) and
+[dpd-stream](https://github.com/leetucker/dpd-stream).
+
+## License
+
+MIT.

match_stream-0.1.0/README.md

@@ -0,0 +1,177 @@
+# match-stream
+
+Matching on data larger than memory. It does nearest-neighbor matching, on
+a propensity score or on the covariates themselves, with or without
+replacement, calipers on the score and on covariates, exact strata, and
+kernel, ridge and coarsened exact matching. It also
+gives balance diagnostics, and treatment effects with Abadie–Imbens or
+clustered standard errors.
+
+It reads Parquet in a stream, holds only a few narrow columns while it
+matches, and fits the propensity score out of core, with absorbed fixed
+effects if you want them. Its results are tested against R's MatchIt and
+Matching and Stata's teffects and kmatch (see below).
+
+```python
+from match_stream import match_nn
+
+m = match_nn("treat ~ age + tenure + log_wage0", "extract/*.parquet",
+             exact=["state", "naics2"], replace=True, caliper=0.2, id="person_id")
+m.balance().summary()
+m.effect("earnings").summary()          # the ATT, Abadie-Imbens standard errors
+m.matches().sink_parquet("pairs.parquet")
+```
+
+**Status: version 0.1.0, the first release.** The designs, the inference and
+the documentation are tested against MatchIt, Matching, Stata's teffects, kmatch and
+psmatch2, and the benchmarks are in the documentation (see Performance).
+
+## What it does
+
+Three functions: `match_nn` (nearest neighbor, on a score or a multivariate
+distance), `match_kernel` (kernel and ridge) and `match_cem` (coarsened exact
+matching). Each returns the matches and weights; `effect()` and `balance()`
+are methods of the result.
+
+**Designs**
+- k:1 nearest neighbor on a propensity score, with or without replacement.
+  Without replacement it is greedy in any of five orders (`largest`,
+  `smallest`, `closest`, `random`, `data`).
+- Calipers on the score (in standard deviations) and on covariates, and
+  exact strata.
+- Mahalanobis, inverse-variance or Euclidean distance on covariates, alone
+  or within a score caliper (`mahalanobis_vars=`).
+- Kernel and ridge matching on the score: five kernels, a
+  bandwidth rule or cross-validation.
+- Coarsened exact matching.
+- Common support (`discard=`, `reestimate_score=`), ties, a minimum number of
+  matches per unit, and one-pass k:1 matching without replacement.
+
+**The score**
+- Logit or probit, fitted by hdfe-stream, with absorbed fixed effects
+  (`score_model="treat ~ x | firm"`), or a column you supply.
+
+**Effects and inference**
+- ATT, ATC and ATE.
+- Abadie–Imbens (2006) standard errors, with or without heteroskedastic
+  conditional variances (`n_variance_matches`).
+- The Abadie–Imbens (2011) bias adjustment.
+- The Abadie–Imbens (2016) correction for an estimated score, also with
+  absorbed fixed effects.
+- Influence-function standard errors with the matching weights fixed, by unit or clustered one or several ways.
+- Several outcomes in one pass (`effect(["y1", "y2"])`), with their joint
+  covariance matrix, for every kind of standard error.
+
+**Diagnostics**
+- A balance table (means, standardized differences, variance ratios, eCDF
+  gaps) and sample sizes.
+
+**Resources**
+- `profile=True` reports the memory, disk and candidate pairs a match would
+  need, without matching, and a resource request for a scheduler.
+- `n_threads=` and `max_memory_gb=` respect a job's limits (stream-limits).
+- Large inputs are matched a batch of exact strata at a time
+  (`memory="by_stratum"`), with the same matches.
+
+**Interface**
+- The matches, weights and subclasses as Polars LazyFrames, ready to write to
+  Parquet.
+- Results with hdfe-stream's reporting: `summary()`, `tidy()`, `confint()`,
+  `etable`, `summary_json()`.
+
+## Why it is fast
+
+- **Windows, not pairs.** Within exact strata and sorted by the score, each
+  unit's candidates are a window found by binary search. Nothing forms all
+  the pairs.
+- **Cheap skipping.** Matching without replacement skips taken controls
+  with path-compressed pointers.
+- **Trees for covariates.** Mahalanobis matching uses k-d trees, and
+  candidate lists for matching without replacement and for calipers.
+- **Window sums.** Kernel and ridge weights come from running sums, with no
+  pairs formed.
+- **Narrow columns.** Only tens of bytes a unit are held. The score model is
+  fitted out of core, and the numba kernels run strata in parallel.
+
+## Validation against other implementations
+
+| Implementation | Compared on | Agreement |
+|---|---|---|
+| MatchIt 4.8.2 | nearest neighbor (every order, k:1, calipers), Mahalanobis, `mahalanobis_vars`, common support, CEM, weights, balance | pairs identical, weights and balance to 1e-8 |
+| R Matching 4.10-15 | Abadie–Imbens SEs (`Var.calc`), the bias adjustment | estimates to 1e-10, SEs to 1e-8 |
+| Stata teffects | nnmatch, psmatch with the AI (2016) correction, ipw | to 1e-8 |
+| Stata kmatch 1.1.6 | nearest neighbor, kernel, ridge, bandwidths, influence-function SEs | matches identical, estimates to 1e-12 |
+| MatchIt + sandwich, fixest | clustered influence-function SEs | to 1e-13 |
+
+The R and Stata scripts that make these references are in `tests/reference/`.
+`docs/implementations.md` lists the conventions on which these packages
+differ, and which option reproduces each.
+
+## Performance
+
+The benchmark suites (`benchmarks/`) compare match-stream with MatchIt and
+Matching (R), pymatchit, fastmatch and Polars' `join_asof` and `group_by`,
+and with bare k-d tree and faiss searches as a floor. They cover
+eight designs and the standard errors, from 10,000 to 100 million units.
+On a 16-vCPU, 32 GB cloud machine, at 10 million units:
+
+- match-stream is 6 to 35 times faster than MatchIt where MatchIt finishes,
+  with the same matches (estimates agree to about 1e-14). At 100 million
+  units MatchIt needs 28 to 30 GB where match-stream needs 8 to 21 GB, and
+  in two designs it runs out of memory.
+- With no strata, 1:1 Mahalanobis matching takes 6.1 s, as long as a bare
+  SciPy k-d tree search of the same data (6.7 s). MatchIt exceeds 30 minutes.
+- pymatchit and R's Matching are slower still, and mostly do not finish
+  at 10 million units or more.
+- Polars' `group_by` (coarsened exact matching) and `join_asof` (1:1 on a
+  score) are faster, as they do less.
+- At 100 million units, 1:1 Mahalanobis matching without replacement and
+  with calipers needs `memory="by_stratum"` on this machine (187 s, 9.2 GB).
+
+The results, figures and caveats are in the
+[benchmarks](https://leetucker.github.io/match-stream/benchmarks/).
+
+Small runs, for correctness only, give these results:
+
+- MatchIt's estimates agree to about 1e-15 on every design.
+- At 100,000 units, R Matching's Abadie–Imbens standard error (with
+  `Var.calc = 1`) took 307 s; match-stream's, the same to 1e-15, took 0.07 s.
+
+## Install
+
+```bash
+pip install match-stream
+```
+
+Python 3.10 or later. It depends on hdfe-stream (with its formula extra),
+stream-limits, Polars, NumPy, numba, SciPy and pyarrow. On Python 3.15,
+pyfixest (which hdfe-stream's formula extra brings in) has no wheels yet: it
+builds from source with a Rust toolchain and
+`PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1` set; everything else installs from wheels.
+
+## Documentation
+
+Build the documentation locally with:
+
+```bash
+pip install -e ".[docs]"
+sphinx-build -W --keep-going -b html docs docs/_build/html
+```
+
+The guide covers:
+
+- how it works;
+- the conventions that differ from other matching packages;
+- the resources a match needs;
+- the API reference.
+
+`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) and
+[dpd-stream](https://github.com/leetucker/dpd-stream).
+
+## License
+
+MIT.

match_stream-0.1.0/match_stream/__init__.py

@@ -0,0 +1,45 @@
+"""match-stream: matching on data larger than memory.
+
+Nearest-neighbor matching with or without replacement (on a score or a
+multivariate distance), calipers on the score and on covariates, exact
+strata, kernel and ridge matching, coarsened exact matching, balance
+diagnostics, and treatment effects with Abadie-Imbens or influence-function
+standard errors. Conventions follow R's MatchIt. See
+https://leetucker.github.io/match-stream/.
+
+Layout
+    api.py          match_nn(), match_kernel(), match_cem()
+    bandwidth.py    PairDistance, CrossValidation (kernel bandwidth rules)
+    _spec.py        argument parsing and checks -> Spec
+    _prep.py        pass 0: the score model (feglm_stream), the narrow columns
+    _match.py       nearest neighbor: sort, order, run the kernel, both memory paths
+    _kernel.py      kernel and ridge matching: bandwidths, weights
+    _cem.py         coarsened exact matching (streaming group-by)
+    kernels/        numba: nn.py (score), metric.py (Mahalanobis etc.),
+                    kernel.py (window moments), var.py (conditional variances)
+    matches.py      MatchSet, CEMSet, KernelSet: pairs, weights, counts
+    balance.py      BalanceTable: MatchIt's summary(m, un = TRUE)
+    effects.py      effect(): estimates, AI and influence-function variances; MatchResult
+    simulate.py     simulate_matching()
+    _core/          code bound for hdfe_stream.core (influence functions, IPW,
+                    the first-stage correction, the weighted balance table)
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from .api import match_cem, match_kernel, match_nn
+from .balance import BalanceTable
+from .bandwidth import CrossValidation, PairDistance
+from .effects import MatchResult
+from .matches import CEMSet, KernelSet, MatchSet
+from .simulate import simulate_matching
+
+try:
+    __version__ = _version("match-stream")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.1.0"
+
+__all__ = ["BalanceTable", "CEMSet", "CrossValidation", "KernelSet", "MatchResult", "MatchSet",
+           "PairDistance", "__version__", "match_cem", "match_kernel", "match_nn",
+           "simulate_matching"]

match_stream-0.1.0/match_stream/_cem.py

@@ -0,0 +1,107 @@
+"""Coarsened exact matching (Iacus, King and Porro), as MatchIt's method="cem".
+
+CEM is a group-by: each matching covariate is cut into bins, a cell is a
+combination of bins and exact-matching values, and a cell is kept when it has
+both treated and control units. Two streaming passes over the source (the
+covariates' ranges, then the cells' counts) and a third that writes the
+matched units' weights; nothing the size of the source is held.
+
+Cutpoints follow MatchIt for a number of bins: equal-width bins over the
+covariate's range among the units used (R's ``cut(x, breaks = k)``; a value on
+a boundary goes to the upper bin here and to the lower one in R, which matters
+only for values exactly on a boundary, of measure zero for continuous data).
+A list gives the cut points themselves (bins are left-closed).
+
+Weights (ATT): a treated unit in a matched cell weighs 1; a control weighs
+(treated in its cell / controls in its cell) x (matched controls / matched
+treated), MatchIt's weights. The ATC swaps the roles.
+"""
+
+from __future__ import annotations
+
+import polars as pl
+from hdfe_stream.core.utils import as_lazy
+
+from ._spec import Spec
+
+ROW = "__ms_row"
+CELL = "__ms_cell"
+
+
+def _bins(x: str, cut, lo: float, hi: float) -> pl.Expr:
+    if isinstance(cut, int):
+        if cut < 1:
+            raise ValueError(f"cutpoints for {x!r} must be a positive number of bins")
+        span = hi - lo
+        if span <= 0:
+            return pl.lit(0, pl.Int32)
+        return ((pl.col(x) - lo) / span * cut).floor().clip(0, cut - 1).cast(pl.Int32)
+    edges = sorted(float(e) for e in cut)
+    return pl.col(x).cut(edges, left_closed=True).to_physical().cast(pl.Int32)
+
+
+def cem(spec: Spec, data, cutpoints) -> tuple[pl.DataFrame, dict, pl.LazyFrame]:
+    """The matched units' weights (row, id, treated, weight, subclass = cell)
+    and the counts."""
+    covs = spec.covariates
+    if not covs and not spec.exact:
+        raise ValueError("CEM needs matching covariates in the formula or exact=")
+    if cutpoints is None:
+        cutpoints = {}
+    elif isinstance(cutpoints, int):
+        cutpoints = {c: cutpoints for c in covs}
+    unknown = set(cutpoints) - set(covs)
+    if unknown:
+        raise ValueError("cutpoints names columns that are not matching covariates: "
+                         f"{sorted(unknown)}")
+    missing_cut = [c for c in covs if c not in cutpoints]
+    if missing_cut:
+        # MatchIt's default is Sturges' rule through R's hist(); not reproduced
+        # exactly, so ask for the number of bins instead of guessing
+        raise ValueError(f"give cutpoints for every covariate (a number of bins or the cut "
+                         f"points): missing {missing_cut}")
+
+    lf = as_lazy(data).with_row_index(ROW).with_columns(pl.col(ROW).cast(pl.Int64))
+    keep = list(dict.fromkeys([ROW, spec.treatment, *spec.exact, *covs,
+                               *([spec.id] if spec.id is not None else [])]))
+    schema = lf.collect_schema()
+    missing = [c for c in keep if c not in schema]
+    if missing:
+        raise ValueError(f"columns not in the data: {missing}")
+    floats = [c for c in keep if schema[c].is_float()]
+    ok = pl.all_horizontal([pl.col(c).is_not_null() for c in keep]
+                           + [pl.col(c).is_not_nan() for c in floats])
+    proj = lf.select(keep)
+    n_source = proj.select(pl.len()).collect(engine="streaming").item()
+    proj = proj.filter(ok).with_columns(pl.col(spec.treatment).cast(pl.Boolean))
+    rng = proj.select([pl.col(c).min().alias(f"lo:{c}") for c in covs]
+                      + [pl.col(c).max().alias(f"hi:{c}") for c in covs]
+                      + [pl.len().alias("n")]).collect(engine="streaming").row(0, named=True)
+
+    bins = [_bins(c, cutpoints[c], rng[f"lo:{c}"], rng[f"hi:{c}"]).alias(f"__b:{c}") for c in covs]
+    keys = [*spec.exact, *[f"__b:{c}" for c in covs]]
+    coarse = proj.with_columns(bins)
+    t = pl.col(spec.treatment)
+    cells = (coarse.group_by(keys)
+             .agg(t.sum().cast(pl.Int64).alias("n_t"), (~t).sum().cast(pl.Int64).alias("n_c"))
+             .filter((pl.col("n_t") > 0) & (pl.col("n_c") > 0))
+             .collect(engine="streaming")
+             .sort(keys).with_row_index(CELL, offset=1)
+             .with_columns(pl.col(CELL).cast(pl.Int64)))
+    mt, mc = int(cells["n_t"].sum()), int(cells["n_c"].sum())
+    focal_t = spec.estimand != "atc"
+    focal_n, other_n = ("n_t", "n_c") if focal_t else ("n_c", "n_t")
+    m_focal, m_other = (mt, mc) if focal_t else (mc, mt)
+    is_focal = t if focal_t else ~t
+    w = (pl.when(is_focal).then(1.0)
+         .otherwise(pl.col(focal_n) / pl.col(other_n) * (m_other / max(m_focal, 1))))
+    out = (coarse.join(cells.lazy(), on=keys, how="inner")
+           .select(pl.col(ROW).alias("row"), *([spec.id] if spec.id is not None else []),
+                   t.alias("treated"), w.alias("weight"), pl.col(CELL).alias("subclass"))
+           .collect(engine="streaming").sort("row"))
+    n_t = int(proj.select(t.cast(pl.Int64).sum()).collect(engine="streaming").item())
+    counts = {"treated": n_t, "control": int(rng["n"]) - n_t, "matched_treated": mt,
+              "matched_control": mc, "unmatched_treated": n_t - mt,
+              "unmatched_control": int(rng["n"]) - n_t - mc, "cells": cells.height,
+              "dropped_missing": n_source - int(rng["n"])}
+    return out, counts, proj