pyrollmatch 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,102 @@
1
+ name: Build & Publish
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+ workflow_dispatch:
7
+
8
+ jobs:
9
+ build:
10
+ runs-on: ubuntu-latest
11
+ steps:
12
+ - uses: actions/checkout@v4
13
+
14
+ - name: Set up Python
15
+ uses: actions/setup-python@v5
16
+ with:
17
+ python-version: "3.12"
18
+
19
+ - name: Install build tools
20
+ run: pip install build twine
21
+
22
+ - name: Build wheel and sdist
23
+ run: python -m build
24
+
25
+ - name: Check distributions
26
+ run: twine check dist/*
27
+
28
+ - name: Upload artifacts
29
+ uses: actions/upload-artifact@v4
30
+ with:
31
+ name: dist
32
+ path: dist/
33
+
34
+ # pyrollmatch is pure Python — one wheel works on all platforms.
35
+ verify:
36
+ needs: build
37
+ strategy:
38
+ matrix:
39
+ os: [ubuntu-latest, macos-latest, windows-latest]
40
+ python-version: ["3.10", "3.12"]
41
+ runs-on: ${{ matrix.os }}
42
+ steps:
43
+ - uses: actions/setup-python@v5
44
+ with:
45
+ python-version: ${{ matrix.python-version }}
46
+
47
+ - name: Download artifacts
48
+ uses: actions/download-artifact@v4
49
+ with:
50
+ name: dist
51
+ path: dist/
52
+
53
+ - name: Install wheel
54
+ shell: bash
55
+ run: pip install dist/*.whl
56
+
57
+ - name: Smoke test
58
+ shell: bash
59
+ run: |
60
+ python -c "
61
+ import pyrollmatch
62
+ print(f'pyrollmatch {pyrollmatch.__version__} loaded successfully')
63
+ from pyrollmatch.score import SUPPORTED_MODELS
64
+ print(f'Models: {SUPPORTED_MODELS}')
65
+ "
66
+
67
+ upload-release:
68
+ needs: [build, verify]
69
+ runs-on: ubuntu-latest
70
+ if: github.event_name == 'release'
71
+ permissions:
72
+ contents: write
73
+ steps:
74
+ - name: Download artifacts
75
+ uses: actions/download-artifact@v4
76
+ with:
77
+ name: dist
78
+ path: dist/
79
+
80
+ - name: Upload to GitHub Release
81
+ env:
82
+ GH_TOKEN: ${{ github.token }}
83
+ run: gh release upload "${{ github.event.release.tag_name }}" dist/* --clobber --repo "${{ github.repository }}"
84
+
85
+ publish-pypi:
86
+ needs: [build, verify]
87
+ runs-on: ubuntu-latest
88
+ if: github.event_name == 'release'
89
+ permissions:
90
+ id-token: write
91
+ environment:
92
+ name: pypi
93
+ url: https://pypi.org/p/pyrollmatch
94
+ steps:
95
+ - name: Download artifacts
96
+ uses: actions/download-artifact@v4
97
+ with:
98
+ name: dist
99
+ path: dist/
100
+
101
+ - name: Publish to PyPI
102
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,50 @@
1
+ name: Tests
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ test:
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ matrix:
14
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
15
+
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+
19
+ - name: Set up Python ${{ matrix.python-version }}
20
+ uses: actions/setup-python@v5
21
+ with:
22
+ python-version: ${{ matrix.python-version }}
23
+
24
+ - name: Install dependencies
25
+ run: |
26
+ python -m pip install --upgrade pip
27
+ pip install -e ".[dev]"
28
+
29
+ - name: Run tests
30
+ run: |
31
+ python -m pytest tests/test_smoke.py tests/test_robust.py -v
32
+
33
+ test-stress:
34
+ runs-on: ubuntu-latest
35
+ steps:
36
+ - uses: actions/checkout@v4
37
+
38
+ - name: Set up Python
39
+ uses: actions/setup-python@v5
40
+ with:
41
+ python-version: "3.12"
42
+
43
+ - name: Install dependencies
44
+ run: |
45
+ python -m pip install --upgrade pip
46
+ pip install -e ".[dev]"
47
+
48
+ - name: Run stress tests
49
+ run: |
50
+ python -m pytest tests/test_stress.py -v
@@ -0,0 +1,7 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .venv/
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+ .pytest_cache/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Alan Huang
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,278 @@
1
+ Metadata-Version: 2.4
2
+ Name: pyrollmatch
3
+ Version: 0.1.0
4
+ Summary: Fast rolling entry matching for staggered adoption studies using polars + numpy
5
+ Author: Alan Huang
6
+ License-Expression: MIT
7
+ License-File: LICENSE
8
+ Keywords: causal-inference,difference-in-differences,matching,propensity-score,rolling-entry,staggered-adoption,treatment-effects
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.10
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
17
+ Requires-Python: >=3.10
18
+ Requires-Dist: numpy>=1.24
19
+ Requires-Dist: polars>=1.0
20
+ Requires-Dist: scikit-learn>=1.3
21
+ Requires-Dist: scipy>=1.10
22
+ Provides-Extra: dev
23
+ Requires-Dist: pytest-xdist; extra == 'dev'
24
+ Requires-Dist: pytest>=7.0; extra == 'dev'
25
+ Provides-Extra: r-comparison
26
+ Requires-Dist: rpy2>=3.5; extra == 'r-comparison'
27
+ Description-Content-Type: text/markdown
28
+
29
+ # pyrollmatch
30
+
31
+ > **Alpha** (v0.1.x) — This package is in early development. Results have been validated against R rollmatch on synthetic and real data, but edge cases may remain. APIs are not stable and may change without notice. **Please verify critical results independently.**
32
+
33
+ Fast rolling entry matching for staggered adoption studies in Python.
34
+
35
+ High-performance reimplementation of the R [rollmatch](https://github.com/RTIInternational/rollmatch) package using **polars** + **numpy**, with 8 propensity score models and comprehensive post-matching diagnostics.
36
+
37
+ | | R rollmatch | pyrollmatch |
38
+ |---|---|---|
39
+ | Max scale | ~5K treated (crashes) | **90K+ treated** |
40
+ | 10K × 30K | OOM (5.8B row join) | **1.7 seconds** |
41
+ | Scoring models | Logistic only | **8 models** |
42
+ | Diagnostics | SMD only | SMD + t-test + VR + KS + TOST |
43
+ | Data library | dplyr (deprecated APIs) | polars |
44
+
45
+ ## Installation
46
+
47
+ ```bash
48
+ pip install pyrollmatch
49
+ ```
50
+
51
+ **From source:**
52
+ ```bash
53
+ git clone https://github.com/AlanHuang99/pyrollmatch.git
54
+ cd pyrollmatch
55
+ pip install -e ".[dev]"
56
+ ```
57
+
58
+ ## Quick Start
59
+
60
+ ```python
61
+ import polars as pl
62
+ from pyrollmatch import rollmatch, alpha_sweep
63
+
64
+ # Panel data: unit × time with treatment indicator
65
+ data = pl.read_parquet("panel_data.parquet")
66
+
67
+ # Rolling entry matching
68
+ result = rollmatch(
69
+ data,
70
+ treat="treat", # binary treatment group indicator (1=treated, 0=control)
71
+ tm="time_period", # time period (integer)
72
+ entry="entry_time", # treatment onset period for treated units;
73
+ # for controls, set to any value > max(time_period)
74
+ id="unit_id", # unit identifier
75
+ covariates=["x1", "x2"], # matching covariates
76
+ lookback=1, # periods to look back for baseline
77
+ alpha=0.1, # caliper = alpha × pooled_SD
78
+ num_matches=3, # controls per treated
79
+ model_type="logistic", # propensity score model
80
+ )
81
+
82
+ # Results
83
+ result.balance # SMD table (polars DataFrame)
84
+ result.weights # unit_id → weight (polars DataFrame)
85
+ result.matched_data # treat_id, control_id, difference
86
+ ```
87
+
88
+ ### Alpha Sweep
89
+
90
+ Find the optimal caliper automatically:
91
+
92
+ ```python
93
+ summary, best = alpha_sweep(
94
+ data, treat="treat", tm="time", entry="entry_time", id="unit_id",
95
+ covariates=["x1", "x2", "x3"],
96
+ alphas=[0.01, 0.02, 0.05, 0.1, 0.15, 0.2],
97
+ )
98
+ # summary: alpha, match_rate, max|SMD|, all_pass
99
+ # best: RollmatchResult from the best alpha
100
+ ```
101
+
102
+ ### Post-Matching Diagnostics
103
+
104
+ ```python
105
+ from pyrollmatch import balance_test, equivalence_test
106
+ from pyrollmatch import reduce_data, score_data
107
+
108
+ reduced = reduce_data(data, "treat", "time", "entry_time", "unit_id")
109
+ scored = score_data(reduced, ["x1", "x2", "x3"], "treat")
110
+
111
+ # SMD + t-test + variance ratio + KS test
112
+ diag = balance_test(scored, result.matched_data,
113
+ "treat", "unit_id", "time", ["x1", "x2", "x3"])
114
+
115
+ # TOST equivalence test (Hartman & Hidalgo 2018)
116
+ equiv = equivalence_test(scored, result.matched_data,
117
+ "treat", "unit_id", "time", ["x1", "x2", "x3"])
118
+ ```
119
+
120
+ ## Propensity Score Models
121
+
122
+ 8 scoring methods, all accessible via the `model_type` parameter:
123
+
124
+ | Model | `model_type` | Description | Best for |
125
+ |---|---|---|---|
126
+ | **Logistic** | `"logistic"` | Standard logistic regression (default) | Most cases |
127
+ | **Probit** | `"probit"` | Probit model (Φ⁻¹ transform) | Robustness check |
128
+ | **GBM** | `"gbm"` | Gradient boosting (HistGradientBoosting) | Non-linear relationships |
129
+ | **Random Forest** | `"rf"` | Random forest classifier | High-dimensional, interactions |
130
+ | **Lasso** | `"lasso"` | L1-regularized logistic | Variable selection |
131
+ | **Ridge** | `"ridge"` | L2-regularized logistic | Multicollinearity |
132
+ | **ElasticNet** | `"elasticnet"` | L1+L2 regularized logistic | Combined regularization |
133
+ | **Mahalanobis** | `"mahalanobis"` | No propensity model — covariate distance | Direct covariate matching |
134
+
135
+ ```python
136
+ # Use gradient boosting for propensity scores
137
+ result = rollmatch(data, ..., model_type="gbm")
138
+
139
+ # Use Mahalanobis distance (no propensity model)
140
+ result = rollmatch(data, ..., model_type="mahalanobis")
141
+ ```
142
+
143
+ ## How It Works
144
+
145
+ ### Rolling Entry Matching
146
+
147
+ For staggered adoption, each treated unit enters at a different time. Rolling entry matching:
148
+
149
+ 1. **`reduce_data()`** — For each treated unit at entry time *t*, select covariates at *t − lookback*. Controls get one observation per treatment entry period.
150
+
151
+ 2. **`score_data()`** — Fit propensity model, compute logit-transformed scores.
152
+
153
+ 3. **`match_all_periods()`** — For each time period, match treated to closest controls within caliper. Uses **block-vectorized numpy broadcasting** (not full cross-join).
154
+
155
+ 4. **`compute_balance()`** — Compare covariate means/SDs before and after matching.
156
+
157
+ ### Key Optimization
158
+
159
+ R's bottleneck: `inner_join(treated, controls, by=time)` creates N_treated × N_controls rows.
160
+
161
+ pyrollmatch instead processes treated in **blocks** via numpy broadcasting:
162
+ ```
163
+ Block of 2000 treated × 64K controls = 128M distances = ~1 GB
164
+ ```
165
+ Never materializes the full cross-product as a DataFrame.
166
+
167
+ ## Data Format
168
+
169
+ Input must be a `polars.DataFrame` in **long panel format** (one row per unit per time period):
170
+
171
+ | Column | Type | Description |
172
+ |---|---|---|
173
+ | `id` | int or str | Unit identifier (e.g., individual, firm, repository) |
174
+ | `tm` | int | Time period, must be integer and monotonically increasing (e.g., 1, 2, ..., 20) |
175
+ | `treat` | int (0 or 1) | **Time-invariant treatment group indicator.** `1` = unit that eventually receives treatment, `0` = unit that never receives treatment. This is NOT a time-varying treatment status — it labels the *group*, not whether treatment is currently active. |
176
+ | `entry` | int | **Treatment onset period** for treated units (the time period when treatment begins). For control units, set this to any value **strictly greater** than the maximum time period in your data. For example, if your panel spans periods 1–20, use `entry=99` or `entry=999` for controls. The specific value does not matter as long as it exceeds `max(tm)`. |
177
+ | covariates | float | Matching variables (e.g., pre-computed rolling means of activity measures) |
178
+
179
+ **Example:**
180
+ ```
181
+ unit_id | time | treat | entry_time | x1 | x2
182
+ --------|------|-------|------------|------|-----
183
+ 1 | 1 | 1 | 5 | 2.3 | 1.1 <- treated, enters at period 5
184
+ 1 | 2 | 1 | 5 | 2.5 | 1.0
185
+ ...
186
+ 1 | 5 | 1 | 5 | 4.1 | 1.8 <- treatment starts here
187
+ ...
188
+ 101 | 1 | 0 | 99 | 1.8 | 0.9 <- control, entry=99 (sentinel)
189
+ 101 | 2 | 0 | 99 | 1.9 | 1.0
190
+ ```
191
+
192
+ ## API Reference
193
+
194
+ ### Core Functions
195
+
196
+ | Function | Description |
197
+ |---|---|
198
+ | `rollmatch()` | Full rolling entry matching pipeline |
199
+ | `alpha_sweep()` | Try multiple calipers, select best |
200
+ | `reduce_data()` | Construct quasi-panel for matching |
201
+ | `score_data()` | Compute propensity scores (8 models) |
202
+
203
+ ### Diagnostics
204
+
205
+ | Function | Description |
206
+ |---|---|
207
+ | `balance_test()` | SMD + t-test + variance ratio + KS test |
208
+ | `equivalence_test()` | TOST equivalence test (Hartman & Hidalgo 2018) |
209
+ | `compute_balance()` | Covariate balance table |
210
+ | `smd_table()` | Print formatted SMD table |
211
+
212
+ ### RollmatchResult
213
+
214
+ | Attribute | Type | Description |
215
+ |---|---|---|
216
+ | `matched_data` | `pl.DataFrame` | `treat_id`, `control_id`, `difference` |
217
+ | `balance` | `pl.DataFrame` | SMD for each covariate |
218
+ | `weights` | `pl.DataFrame` | `id`, `weight` |
219
+ | `n_treated_matched` | `int` | Number matched |
220
+ | `n_treated_total` | `int` | Total treated |
221
+ | `alpha` | `float` | Caliper used |
222
+
223
+ ## Validation Against R rollmatch
224
+
225
+ Tested on identical synthetic data across 15 configurations:
226
+
227
+ | Metric | Result |
228
+ |---|---|
229
+ | Propensity score correlation | **≥ 0.9999** across all configs |
230
+ | Match count agreement | Within ±3% |
231
+ | Balance quality (max\|SMD\|) | Both achieve < 0.05 |
232
+
233
+ Pair-level overlap varies by caliper tightness (3%–71%) due to different GLM solvers (sklearn vs R glm). This is expected — **balance quality is what matters, not pair identity**.
234
+
235
+ Validated on real data: [GithubRepo project](https://github.com/AlanHuang99/) (8K treated × 7.5K controls, 10 covariates, all diagnostics pass).
236
+
237
+ ## Performance
238
+
239
+ | Scale | Runtime | Match rate |
240
+ |---|---:|---:|
241
+ | 500 × 2,000 | 0.03s | 100% |
242
+ | 1,000 × 3,000 | 0.05s | 100% |
243
+ | 5,000 × 15,000 | 0.10s | 100% |
244
+ | 10,000 × 30,000 | 1.7s | 99.99% |
245
+
246
+ R rollmatch crashes at 10K treated due to dplyr's 2.1B row limit.
247
+
248
+ ## Design Principles
249
+
250
+ - **Polars-first**: Native polars DataFrames throughout. No pandas dependency.
251
+ - **Reproducible**: Deterministic with fixed `random_state`. Same input → same output.
252
+ - **Scalable**: Block-vectorized matching, O(block × N_controls) memory per iteration.
253
+ - **Validated**: 44 tests, R-validated across 15 configurations, real-data tested.
254
+ - **Modular**: Each step (reduce → score → match → balance) is independently usable.
255
+
256
+ ## References
257
+
258
+ - Rickles, J. H., & Seltzer, M. (2014). "A Two-Stage Matching Strategy for Estimating the Effects of Staggered Adoption of Interventions."
259
+ - RTI International. `rollmatch` R package. https://github.com/RTIInternational/rollmatch
260
+ - Hartman, E., & Hidalgo, F. D. (2018). "An Equivalence Approach to Balance and Placebo Tests." *American Journal of Political Science*.
261
+
262
+ ## Status & Contributing
263
+
264
+ This package is under **active development**. We welcome:
265
+
266
+ - **Bug reports** — if you encounter unexpected behavior, please open an issue with a minimal reproducible example
267
+ - **Feature requests** — suggestions for new matching methods, diagnostics, or API improvements
268
+ - **Contributions** — pull requests are welcome; please open an issue first to discuss
269
+
270
+ Please be cautious when using this package for published research. While we have validated against the R rollmatch package across multiple configurations (see [Validation](#validation-against-r-rollmatch)), this is alpha software. We strongly recommend:
271
+
272
+ 1. **Cross-checking** results against an established implementation (R rollmatch, R MatchIt) for critical analyses
273
+ 2. **Inspecting balance diagnostics** carefully before proceeding to estimation
274
+ 3. **Reporting any discrepancies** you find via [GitHub Issues](https://github.com/AlanHuang99/pyrollmatch/issues)
275
+
276
+ ## License
277
+
278
+ MIT