Skip to content

Commit 9692b15

Browse files
igerberigerber
authored
Merge pull request #58 from igerber/claude/add-rust-backend
adds rust backend
2 parents 7ffccaa + 24db1ba commit 9692b15

23 files changed

Lines changed: 2447 additions & 249 deletions

‎.github/workflows/publish.yml‎

Lines changed: 95 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,124 @@
1-
name: Publish to PyPI
1+
name: Build and Publish
22

33
on:
44
release:
55
types: [published]
66

77
jobs:
8-
build:
8+
# Build wheels on Linux
9+
build-linux:
10+
name: Build Linux wheels
911
runs-on: ubuntu-latest
12+
strategy:
13+
matrix:
14+
target: [x86_64, aarch64]
1015
steps:
1116
- uses: actions/checkout@v4
1217

13-
- name: Set up Python
14-
uses: actions/setup-python@v5
18+
- name: Build wheels
19+
uses: PyO3/maturin-action@v1
1520
with:
16-
python-version: "3.11"
21+
target: ${{ matrix.target }}
22+
args: --release --out dist
23+
manylinux: auto
1724

18-
- name: Install build dependencies
19-
run: |
20-
python -m pip install --upgrade pip
21-
pip install build
25+
- name: Upload wheels
26+
uses: actions/upload-artifact@v4
27+
with:
28+
name: wheels-linux-${{ matrix.target }}
29+
path: dist/*.whl
30+
31+
# Build wheels on macOS (x86_64)
32+
build-macos-x86:
33+
name: Build macOS x86_64 wheels
34+
runs-on: macos-latest
35+
steps:
36+
- uses: actions/checkout@v4
37+
38+
- name: Build wheels
39+
uses: PyO3/maturin-action@v1
40+
with:
41+
target: x86_64-apple-darwin
42+
args: --release --out dist
43+
44+
- name: Upload wheels
45+
uses: actions/upload-artifact@v4
46+
with:
47+
name: wheels-macos-x86_64
48+
path: dist/*.whl
49+
50+
# Build wheels on macOS (ARM64)
51+
build-macos-arm:
52+
name: Build macOS ARM64 wheels
53+
runs-on: macos-latest
54+
steps:
55+
- uses: actions/checkout@v4
56+
57+
- name: Build wheels
58+
uses: PyO3/maturin-action@v1
59+
with:
60+
target: aarch64-apple-darwin
61+
args: --release --out dist
2262

23-
- name: Build package
24-
run: python -m build
63+
- name: Upload wheels
64+
uses: actions/upload-artifact@v4
65+
with:
66+
name: wheels-macos-arm64
67+
path: dist/*.whl
68+
69+
# Build wheels on Windows
70+
build-windows:
71+
name: Build Windows wheels
72+
runs-on: windows-latest
73+
steps:
74+
- uses: actions/checkout@v4
75+
76+
- name: Build wheels
77+
uses: PyO3/maturin-action@v1
78+
with:
79+
target: x64
80+
args: --release --out dist
81+
82+
- name: Upload wheels
83+
uses: actions/upload-artifact@v4
84+
with:
85+
name: wheels-windows
86+
path: dist/*.whl
87+
88+
# Build source distribution
89+
build-sdist:
90+
name: Build source distribution
91+
runs-on: ubuntu-latest
92+
steps:
93+
- uses: actions/checkout@v4
94+
95+
- name: Build sdist
96+
uses: PyO3/maturin-action@v1
97+
with:
98+
command: sdist
99+
args: --out dist
25100

26-
- name: Upload build artifacts
101+
- name: Upload sdist
27102
uses: actions/upload-artifact@v4
28103
with:
29-
name: dist
30-
path: dist/
104+
name: sdist
105+
path: dist/*.tar.gz
31106

107+
# Publish all artifacts to PyPI
32108
publish:
33-
needs: build
109+
name: Publish to PyPI
110+
needs: [build-linux, build-macos-x86, build-macos-arm, build-windows, build-sdist]
34111
runs-on: ubuntu-latest
35112
environment: pypi
36113
permissions:
37114
id-token: write # Required for trusted publishing
38115

39116
steps:
40-
- name: Download build artifacts
117+
- name: Download all artifacts
41118
uses: actions/download-artifact@v4
42119
with:
43-
name: dist
44-
path: dist/
120+
path: dist
121+
merge-multiple: true
45122

46123
- name: Publish to PyPI
47124
uses: pypa/gh-action-pypi-publish@release/v1

‎.github/workflows/rust-test.yml‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
name: Rust Backend Tests
2+
3+
on:
4+
push:
5+
branches: [main]
6+
paths:
7+
- 'rust/**'
8+
- 'diff_diff/**'
9+
- 'tests/**'
10+
- 'pyproject.toml'
11+
- '.github/workflows/rust-test.yml'
12+
pull_request:
13+
branches: [main]
14+
paths:
15+
- 'rust/**'
16+
- 'diff_diff/**'
17+
- 'tests/**'
18+
- 'pyproject.toml'
19+
- '.github/workflows/rust-test.yml'
20+
21+
env:
22+
CARGO_TERM_COLOR: always
23+
24+
jobs:
25+
# Run Rust unit tests
26+
rust-tests:
27+
name: Rust Unit Tests
28+
runs-on: ubuntu-latest
29+
steps:
30+
- uses: actions/checkout@v4
31+
32+
- name: Install Rust toolchain
33+
uses: dtolnay/rust-toolchain@stable
34+
35+
- name: Install OpenBLAS
36+
run: sudo apt-get update && sudo apt-get install -y libopenblas-dev
37+
38+
- name: Run Rust tests
39+
working-directory: rust
40+
run: cargo test --verbose
41+
42+
# Build and test with Python on multiple platforms
43+
python-tests:
44+
name: Python Tests (${{ matrix.os }})
45+
runs-on: ${{ matrix.os }}
46+
strategy:
47+
fail-fast: false
48+
matrix:
49+
os: [ubuntu-latest, macos-latest]
50+
# Windows excluded due to Intel MKL build complexity
51+
52+
steps:
53+
- uses: actions/checkout@v4
54+
55+
- name: Set up Python
56+
uses: actions/setup-python@v5
57+
with:
58+
python-version: '3.11'
59+
60+
- name: Install OpenBLAS (Ubuntu)
61+
if: matrix.os == 'ubuntu-latest'
62+
run: sudo apt-get update && sudo apt-get install -y libopenblas-dev
63+
64+
- name: Install OpenBLAS (macOS)
65+
if: matrix.os == 'macos-latest'
66+
run: brew install openblas
67+
68+
- name: Set OpenBLAS paths (macOS)
69+
if: matrix.os == 'macos-latest'
70+
run: |
71+
echo "OPENBLAS_DIR=$(brew --prefix openblas)" >> $GITHUB_ENV
72+
echo "PKG_CONFIG_PATH=$(brew --prefix openblas)/lib/pkgconfig" >> $GITHUB_ENV
73+
74+
- name: Install Rust toolchain
75+
uses: dtolnay/rust-toolchain@stable
76+
77+
- name: Install test dependencies
78+
run: pip install pytest numpy pandas scipy
79+
80+
- name: Build and install with maturin
81+
run: |
82+
pip install maturin
83+
maturin build --release -o dist
84+
echo "=== Built wheels ==="
85+
ls -la dist/
86+
# --no-index ensures we install from local wheel, not PyPI
87+
pip install --no-index --find-links=dist diff-diff
88+
89+
- name: Verify Rust backend is available
90+
# Run from /tmp to avoid source directory shadowing installed package
91+
working-directory: /tmp
92+
run: |
93+
python -c "import diff_diff; print('Location:', diff_diff.__file__)"
94+
python -c "from diff_diff import HAS_RUST_BACKEND; print('HAS_RUST_BACKEND:', HAS_RUST_BACKEND); assert HAS_RUST_BACKEND, 'Rust backend not available'"
95+
96+
- name: Copy tests to isolated location
97+
run: cp -r tests /tmp/tests
98+
99+
- name: Run Rust backend tests
100+
working-directory: /tmp
101+
run: pytest tests/test_rust_backend.py -v
102+
103+
- name: Run tests with Rust backend
104+
working-directory: /tmp
105+
run: DIFF_DIFF_BACKEND=rust pytest tests/ -x -q
106+
107+
# Test pure Python fallback (without Rust extension)
108+
python-fallback:
109+
name: Pure Python Fallback
110+
runs-on: ubuntu-latest
111+
steps:
112+
- uses: actions/checkout@v4
113+
114+
- name: Set up Python
115+
uses: actions/setup-python@v5
116+
with:
117+
python-version: '3.11'
118+
119+
- name: Install dependencies
120+
run: pip install numpy pandas scipy pytest
121+
122+
- name: Verify pure Python mode
123+
run: |
124+
# Use PYTHONPATH to import directly (skips maturin build)
125+
PYTHONPATH=. python -c "from diff_diff import HAS_RUST_BACKEND; print(f'HAS_RUST_BACKEND: {HAS_RUST_BACKEND}'); assert not HAS_RUST_BACKEND"
126+
127+
- name: Run tests in pure Python mode
128+
run: PYTHONPATH=. DIFF_DIFF_BACKEND=python pytest tests/ -x -q --ignore=tests/test_rust_backend.py

‎.gitignore‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,13 @@ Thumbs.db
5656
# Benchmarks - generated data and results (can be regenerated)
5757
benchmarks/data/synthetic/*.csv
5858
benchmarks/results/
59+
60+
# Rust build artifacts
61+
rust/target/
62+
Cargo.lock
63+
*.so
64+
*.pyd
65+
*.dylib
66+
67+
# Maturin build artifacts
68+
target/

‎CLAUDE.md‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,28 @@ ruff check diff_diff tests
3131
mypy diff_diff
3232
```
3333

34+
### Rust Backend Commands
35+
36+
```bash
37+
# Build Rust backend for development (requires Rust toolchain)
38+
maturin develop
39+
40+
# Build with release optimizations
41+
maturin develop --release
42+
43+
# Run Rust unit tests
44+
cd rust && cargo test
45+
46+
# Force pure Python mode (disable Rust backend)
47+
DIFF_DIFF_BACKEND=python pytest
48+
49+
# Force Rust mode (fail if Rust not available)
50+
DIFF_DIFF_BACKEND=rust pytest
51+
52+
# Run Rust backend equivalence tests
53+
pytest tests/test_rust_backend.py -v
54+
```
55+
3456
## Architecture
3557

3658
### Module Structure
@@ -81,6 +103,20 @@ mypy diff_diff
81103
- Single optimization point for all estimators (reduces code duplication)
82104
- Cluster-robust SEs use pandas groupby instead of O(n × clusters) loop
83105

106+
- **`diff_diff/_backend.py`** - Backend detection and configuration (v2.0.0):
107+
- Detects optional Rust backend availability
108+
- Handles `DIFF_DIFF_BACKEND` environment variable ('auto', 'python', 'rust')
109+
- Exports `HAS_RUST_BACKEND` flag and Rust function references
110+
- Other modules import from here to avoid circular imports with `__init__.py`
111+
112+
- **`rust/`** - Optional Rust backend for accelerated computation (v2.0.0):
113+
- **`rust/src/lib.rs`** - PyO3 module definition, exports Python bindings
114+
- **`rust/src/bootstrap.rs`** - Parallel bootstrap weight generation (Rademacher, Mammen, Webb)
115+
- **`rust/src/linalg.rs`** - OLS solver and cluster-robust variance estimation
116+
- **`rust/src/weights.rs`** - Synthetic control weights and simplex projection
117+
- Uses ndarray-linalg with OpenBLAS (Linux/macOS) or Intel MKL (Windows)
118+
- Provides 4-8x speedup for SyntheticDiD, minimal benefit for other estimators
119+
84120
- **`diff_diff/results.py`** - Dataclass containers for estimation results:
85121
- `DiDResults`, `MultiPeriodDiDResults`, `SyntheticDiDResults`, `PeriodEffect`
86122
- Each provides `summary()`, `to_dict()`, `to_dataframe()` methods

‎TODO.md‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,19 @@ From code review (PR #32):
102102

103103
---
104104

105+
## Rust Backend Optimizations
106+
107+
Deferred from PR #58 code review (can be done post-merge):
108+
109+
- [ ] **Matrix inversion efficiency** (`rust/src/linalg.rs:180-194`): Use Cholesky factorization for symmetric positive-definite matrices instead of column-by-column solve
110+
- [ ] **Reduce bootstrap allocations** (`rust/src/bootstrap.rs`): Currently uses `Vec<Vec<f64>>` → flatten → `Array2` which allocates twice. Should allocate directly into ndarray.
111+
- [ ] **Consider static BLAS linking** (`rust/Cargo.toml`): Currently requires system BLAS libraries. Consider `openblas-static` or `intel-mkl-static` features for easier distribution.
112+
113+
---
114+
105115
## Performance Optimizations
106116

107-
No major performance issues identified. Potential future optimizations:
117+
Potential future optimizations:
108118

109119
- [ ] JIT compilation for bootstrap loops (numba)
110120
- [ ] Parallel bootstrap iterations

0 commit comments

Comments
 (0)