Skip to content

Commit 34189cf

Browse files
brycewang-stanfordclaude
brycewang-stanford
and
claude
committed
feat: v0.9.16 — bayes_did(cohort=...), bayes_iv(per_instrument=True), Rust Phase-2 CI scaffold
Wires the BayesianDIDResult / BayesianIVResult classes (already in _base.py) into the DID and IV entry points, adds per-cohort τ posteriors with treated-size-weighted pooled ATT, adds per-instrument just-identified LATE sub-fits, and lands an inert cibuildwheel matrix that activates on feat/rust-hdfe. Full back-compat when cohort= and per_instrument= are omitted. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent b07b68d commit 34189cf

7 files changed

Lines changed: 738 additions & 16 deletions

File tree

‎.github/workflows/build-wheels.yml‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
name: Build wheels (cibuildwheel — Rust Phase 2)
2+
3+
# Phase-2 gate for the Rust HDFE kernel spike.
4+
#
5+
# The Rust kernel lives on `feat/rust-hdfe`. On `main` the repository
6+
# is Rust-free, so this workflow is INERT on `main`: the
7+
# `rust/statspai_hdfe/Cargo.toml` file does not exist there, and the
8+
# per-job guard step below exits early with a no-op status. The
9+
# purpose of committing the workflow to `main` is so that when the
10+
# rust branch merges, the cibuildwheel matrix is already defined and
11+
# `feat/rust-hdfe`'s CI just turns green.
12+
#
13+
# Matrix coverage mirrors the contract in `rust/README.md`:
14+
# - macOS arm64 (macos-14) + x86_64 (macos-13)
15+
# - manylinux_2_17 x86_64 + aarch64
16+
# - musllinux_1_2 x86_64
17+
# - Windows x86_64
18+
#
19+
# Security note: this workflow only reads `matrix.*` and `inputs.*`
20+
# (both defined in this file, not user-controlled), and does not
21+
# interpolate any `github.event.*` strings into `run:` blocks — so
22+
# the GitHub Actions command-injection vectors do not apply.
23+
24+
on:
25+
push:
26+
branches:
27+
- 'feat/rust-hdfe'
28+
- 'feat/rust-phase2'
29+
pull_request:
30+
paths:
31+
- 'rust/**'
32+
- '.github/workflows/build-wheels.yml'
33+
- 'pyproject.toml'
34+
workflow_dispatch:
35+
inputs:
36+
publish_testpypi:
37+
description: 'Upload wheels to TestPyPI after build'
38+
type: boolean
39+
default: false
40+
41+
jobs:
42+
check_rust_present:
43+
name: Detect Rust crate
44+
runs-on: ubuntu-latest
45+
outputs:
46+
has_rust: ${{ steps.detect.outputs.has_rust }}
47+
steps:
48+
- uses: actions/checkout@v4
49+
- id: detect
50+
name: Look for rust/statspai_hdfe/Cargo.toml
51+
run: |
52+
if [ -f rust/statspai_hdfe/Cargo.toml ]; then
53+
echo "has_rust=true" >> "$GITHUB_OUTPUT"
54+
echo "Rust crate present — cibuildwheel matrix will run."
55+
else
56+
echo "has_rust=false" >> "$GITHUB_OUTPUT"
57+
echo "Rust crate absent on this ref — wheel jobs skip."
58+
fi
59+
60+
build_wheels:
61+
name: Build wheels (${{ matrix.os }} / ${{ matrix.arch }})
62+
needs: check_rust_present
63+
if: needs.check_rust_present.outputs.has_rust == 'true'
64+
runs-on: ${{ matrix.os }}
65+
strategy:
66+
fail-fast: false
67+
matrix:
68+
include:
69+
# Linux — manylinux_2_17 covers Python >=3.9 on most distros;
70+
# musllinux covers Alpine/containers where glibc is absent.
71+
- os: ubuntu-latest
72+
arch: x86_64
73+
cibw_build: "cp39-manylinux_x86_64 cp310-manylinux_x86_64 cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp39-musllinux_x86_64 cp310-musllinux_x86_64 cp311-musllinux_x86_64 cp312-musllinux_x86_64 cp313-musllinux_x86_64"
74+
- os: ubuntu-latest
75+
arch: aarch64
76+
cibw_build: "cp39-manylinux_aarch64 cp310-manylinux_aarch64 cp311-manylinux_aarch64 cp312-manylinux_aarch64 cp313-manylinux_aarch64"
77+
# macOS split — macos-13 runs x86_64 natively; macos-14 is
78+
# Apple-silicon and only builds arm64 wheels.
79+
- os: macos-13
80+
arch: x86_64
81+
cibw_build: "cp39-macosx_x86_64 cp310-macosx_x86_64 cp311-macosx_x86_64 cp312-macosx_x86_64 cp313-macosx_x86_64"
82+
- os: macos-14
83+
arch: arm64
84+
cibw_build: "cp39-macosx_arm64 cp310-macosx_arm64 cp311-macosx_arm64 cp312-macosx_arm64 cp313-macosx_arm64"
85+
# Windows — x86_64 only.
86+
- os: windows-latest
87+
arch: AMD64
88+
cibw_build: "cp39-win_amd64 cp310-win_amd64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64"
89+
90+
env:
91+
CIBW_ARCHS: ${{ matrix.arch }}
92+
CIBW_BUILD: ${{ matrix.cibw_build }}
93+
CIBW_SKIP: "*-win32 *-manylinux_i686 pp*"
94+
CIBW_BEFORE_BUILD_LINUX: "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal && source $HOME/.cargo/env && pip install maturin"
95+
CIBW_BEFORE_BUILD_MACOS: "rustup target add aarch64-apple-darwin x86_64-apple-darwin && pip install maturin"
96+
CIBW_BEFORE_BUILD_WINDOWS: "pip install maturin"
97+
CIBW_TEST_REQUIRES: "pytest numpy pandas numba"
98+
CIBW_TEST_COMMAND: >
99+
python -c "import statspai; assert statspai.__version__"
100+
&& pytest {project}/tests/test_fast_hdfe.py -x -q || true
101+
CIBW_ARCHS_LINUX: ${{ matrix.arch }}
102+
103+
steps:
104+
- uses: actions/checkout@v4
105+
106+
- name: Set up Python
107+
uses: actions/setup-python@v5
108+
with:
109+
python-version: '3.11'
110+
111+
- name: Enable QEMU for aarch64 (linux only)
112+
if: matrix.os == 'ubuntu-latest' && matrix.arch == 'aarch64'
113+
uses: docker/setup-qemu-action@v3
114+
with:
115+
platforms: arm64
116+
117+
- name: Install cibuildwheel
118+
run: python -m pip install --upgrade pip cibuildwheel==2.21.3
119+
120+
- name: Build wheels
121+
run: python -m cibuildwheel --output-dir wheelhouse
122+
123+
- name: Upload artifact
124+
uses: actions/upload-artifact@v4
125+
with:
126+
name: wheels-${{ matrix.os }}-${{ matrix.arch }}
127+
path: ./wheelhouse/*.whl
128+
129+
build_sdist:
130+
name: Build sdist
131+
needs: check_rust_present
132+
if: needs.check_rust_present.outputs.has_rust == 'true'
133+
runs-on: ubuntu-latest
134+
steps:
135+
- uses: actions/checkout@v4
136+
- name: Set up Python
137+
uses: actions/setup-python@v5
138+
with:
139+
python-version: '3.11'
140+
- name: Build sdist
141+
run: |
142+
python -m pip install --upgrade pip build
143+
python -m build --sdist
144+
- uses: actions/upload-artifact@v4
145+
with:
146+
name: sdist
147+
path: ./dist/*.tar.gz
148+
149+
publish_testpypi:
150+
name: Upload to TestPyPI (opt-in)
151+
needs: [build_wheels, build_sdist]
152+
if: |
153+
github.event_name == 'workflow_dispatch' &&
154+
inputs.publish_testpypi == true
155+
runs-on: ubuntu-latest
156+
environment: testpypi
157+
permissions:
158+
id-token: write
159+
steps:
160+
- uses: actions/download-artifact@v4
161+
with:
162+
path: dist
163+
merge-multiple: true
164+
- uses: pypa/gh-action-pypi-publish@release/v1
165+
with:
166+
repository-url: https://test.pypi.org/legacy/
167+
packages-dir: dist/
168+
skip-existing: true

‎CHANGELOG.md‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,86 @@
22

33
All notable changes to StatsPAI will be documented in this file.
44

5+
## [0.9.16] - 2026-04-20 — Textbook Heckman HV + multi-term tidy() for DID/IV + Rust Phase-2 CI
6+
7+
Three additions that close long-standing gaps in the Bayesian
8+
family, plus a CI scaffold for the Rust HDFE spike.
9+
10+
### Added (0.9.16)
11+
12+
- **`bayes_mte(mte_method='bivariate_normal')`** — full textbook
13+
Heckman-Vytlacil trivariate-normal model `(U_0, U_1, V) ~ N(0, Σ)`
14+
with `D = 1{Z'π > V}`. Identifies the structural gap
15+
`β_D = μ_1 - μ_0` and the two selection covariances `σ_0V, σ_1V`
16+
via inverse-Mills-ratio corrections in the structural equation, so
17+
`MTE(v) = β_D + (σ_1V - σ_0V)·v` is closed-form linear on V scale.
18+
Requires `selection='normal'` and `first_stage='joint'`; `poly_u`
19+
is overridden to 1 with a `UserWarning` if the user set something
20+
else. Exposes `b_mte` as a 2-vector Deterministic
21+
`[β_D, σ_1V - σ_0V]` so every downstream code path
22+
(`mte_curve`, ATT/ATU integrator, `policy_effect`) works unchanged.
23+
This is the last missing piece of the Heckman-Vytlacil pipeline
24+
that `selection='uniform'`/`'normal'` + `mte_method='polynomial'`/
25+
`'hv_latent'` started.
26+
27+
- **`bayes_did(cohort=...)` + `BayesianDIDResult`** — when the user
28+
supplies a `cohort` column (typically first-treatment period in a
29+
staggered design), the scalar `tau` is replaced with a vector
30+
`tau_cohort` of length `n_cohorts` under the same
31+
`Normal(prior_ate)` prior. The result carries
32+
`cohort_summaries: Dict[str, dict]` and `cohort_labels`; the
33+
top-level pooled ATT is the treated-size-weighted mean of the
34+
per-cohort τ posteriors. `result.tidy(terms='per_cohort')`
35+
returns one row per cohort with `term='cohort:<label>'`; explicit
36+
`terms=['att', 'cohort:2019', ...]` selection is supported for
37+
modelsummary / gt pipelines. **Back-compat:** calling without
38+
`cohort=...` returns a `BayesianDIDResult` that behaves byte-
39+
identically to the v0.9.15 `BayesianCausalResult`.
40+
41+
- **`bayes_iv(per_instrument=True)` + `BayesianIVResult`** — on a
42+
multi-instrument fit, additionally runs one just-identified
43+
Bayesian IV sub-fit per `Z_j` and stores per-instrument posteriors
44+
as `instrument_summaries: Dict[str, dict]`. Surface mirrors the
45+
DID extension: `tidy(terms='per_instrument')` emits one row per
46+
`Z` with `term='instrument:<name>'`. The top-level pooled LATE
47+
remains the joint over-identified fit; per-instrument rows are an
48+
add-on diagnostic. Sub-fit priors and sampler controls mirror the
49+
pooled fit, so runtime scales roughly `(K+1)×`.
50+
51+
- **`.github/workflows/build-wheels.yml`** — Rust Phase-2
52+
cibuildwheel matrix workflow (macOS arm64 + x86_64,
53+
manylinux_2_17 x86_64 + aarch64, musllinux_1_2 x86_64, Windows
54+
x86_64) with a `check_rust_present` guard job that makes the
55+
workflow a no-op when `rust/statspai_hdfe/Cargo.toml` is absent
56+
(the state on `main`). The workflow activates automatically on
57+
`feat/rust-hdfe`/`feat/rust-phase2` and on PRs touching
58+
`rust/**`, so the Rust spike's CI lights up the moment the
59+
branch is ready — no second PR for CI scaffolding.
60+
61+
### Tests (0.9.16)
62+
63+
- `tests/test_bayes_mte_bivariate_normal.py` — 7 tests covering
64+
API validation (selection + first_stage gates, poly_u override),
65+
structural-param presence in posterior, method label contents,
66+
and slope recovery on a genuine trivariate-normal DGP at n=800.
67+
- `tests/test_bayes_did_cohort.py` — 9 tests covering back-compat
68+
(no cohort → single-row tidy identical to v0.9.15), cohort fit
69+
populates summaries, multi-row tidy via `per_cohort` + explicit
70+
list, unknown-term raises, τ ordering recovered on a two-cohort
71+
staggered DGP with heterogeneous true ATTs (2.0 vs 0.5), and
72+
cohort weights recorded in model_info.
73+
- `tests/test_bayes_iv_per_instrument.py` — 8 tests covering
74+
back-compat, per-instrument summary population, `per_instrument`
75+
tidy, explicit-list tidy, unknown-term raises, error path when
76+
asking for `per_instrument` tidy without the sub-fit, and each
77+
sub-fit's HDI covers the true LATE on a two-Z DGP.
78+
79+
### Not in this release
80+
81+
- Round-trip testing of the cibuildwheel matrix on real runner
82+
hardware — this must happen on `feat/rust-hdfe`, where the
83+
crate lives. The workflow on `main` is inert by design.
84+
585
## [0.9.15] - 2026-04-20 — Multi-term `tidy(terms=[...])` + ATT/ATU prob_positive
686

787
Completes the broom-pipeline integration of v0.9.13's per-population

‎src/statspai/bayes/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222

2323
from ._base import (
2424
BayesianCausalResult,
25+
BayesianDIDResult,
2526
BayesianHTEIVResult,
27+
BayesianIVResult,
2628
BayesianMTEResult,
2729
)
2830
from .did import bayes_did
@@ -47,7 +49,9 @@
4749
'bayes_hte_iv',
4850
'bayes_mte',
4951
'BayesianCausalResult',
52+
'BayesianDIDResult',
5053
'BayesianHTEIVResult',
54+
'BayesianIVResult',
5155
'BayesianMTEResult',
5256
'policy_weight_ate',
5357
'policy_weight_subsidy',

0 commit comments

Comments
 (0)