Claude/add rust backend

.github/workflows/publish.yml

@@ -1,47 +1,124 @@
-name: Publish to PyPI
+name: Build and Publish
 
 on:
   release:
     types: [published]
 
 jobs:
-  build:
+  # Build wheels on Linux
+  build-linux:
+    name: Build Linux wheels
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        target: [x86_64, aarch64]
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
         with:
-          python-version: "3.11"
+          target: ${{ matrix.target }}
+          args: --release --out dist
+          manylinux: auto
 
-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install build
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-linux-${{ matrix.target }}
+          path: dist/*.whl
+
+  # Build wheels on macOS (x86_64)
+  build-macos-x86:
+    name: Build macOS x86_64 wheels
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: x86_64-apple-darwin
+          args: --release --out dist
+
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-macos-x86_64
+          path: dist/*.whl
+
+  # Build wheels on macOS (ARM64)
+  build-macos-arm:
+    name: Build macOS ARM64 wheels
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: aarch64-apple-darwin
+          args: --release --out dist
 
-      - name: Build package
-        run: python -m build
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-macos-arm64
+          path: dist/*.whl
+
+  # Build wheels on Windows
+  build-windows:
+    name: Build Windows wheels
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: x64
+          args: --release --out dist
+
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-windows
+          path: dist/*.whl
+
+  # Build source distribution
+  build-sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build sdist
+        uses: PyO3/maturin-action@v1
+        with:
+          command: sdist
+          args: --out dist
 
-      - name: Upload build artifacts
+      - name: Upload sdist
         uses: actions/upload-artifact@v4
         with:
-          name: dist
-          path: dist/
+          name: sdist
+          path: dist/*.tar.gz
 
+  # Publish all artifacts to PyPI
   publish:
-    needs: build
+    name: Publish to PyPI
+    needs: [build-linux, build-macos-x86, build-macos-arm, build-windows, build-sdist]
     runs-on: ubuntu-latest
     environment: pypi
     permissions:
       id-token: write  # Required for trusted publishing
 
     steps:
-      - name: Download build artifacts
+      - name: Download all artifacts
         uses: actions/download-artifact@v4
         with:
-          name: dist
-          path: dist/
+          path: dist
+          merge-multiple: true
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1

.github/workflows/rust-test.yml

@@ -0,0 +1,128 @@
+name: Rust Backend Tests
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - 'rust/**'
+      - 'diff_diff/**'
+      - 'tests/**'
+      - 'pyproject.toml'
+      - '.github/workflows/rust-test.yml'
+  pull_request:
+    branches: [main]
+    paths:
+      - 'rust/**'
+      - 'diff_diff/**'
+      - 'tests/**'
+      - 'pyproject.toml'
+      - '.github/workflows/rust-test.yml'
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  # Run Rust unit tests
+  rust-tests:
+    name: Rust Unit Tests
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Install OpenBLAS
+        run: sudo apt-get update && sudo apt-get install -y libopenblas-dev
+
+      - name: Run Rust tests
+        working-directory: rust
+        run: cargo test --verbose
+
+  # Build and test with Python on multiple platforms
+  python-tests:
+    name: Python Tests (${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        # Windows excluded due to Intel MKL build complexity
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install OpenBLAS (Ubuntu)
+        if: matrix.os == 'ubuntu-latest'
+        run: sudo apt-get update && sudo apt-get install -y libopenblas-dev
+
+      - name: Install OpenBLAS (macOS)
+        if: matrix.os == 'macos-latest'
+        run: brew install openblas
+
+      - name: Set OpenBLAS paths (macOS)
+        if: matrix.os == 'macos-latest'
+        run: |
+          echo "OPENBLAS_DIR=$(brew --prefix openblas)" >> $GITHUB_ENV
+          echo "PKG_CONFIG_PATH=$(brew --prefix openblas)/lib/pkgconfig" >> $GITHUB_ENV
+
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Install test dependencies
+        run: pip install pytest numpy pandas scipy
+
+      - name: Build and install with maturin
+        run: |
+          pip install maturin
+          maturin build --release -o dist
+          echo "=== Built wheels ==="
+          ls -la dist/
+          # --no-index ensures we install from local wheel, not PyPI
+          pip install --no-index --find-links=dist diff-diff
+
+      - name: Verify Rust backend is available
+        # Run from /tmp to avoid source directory shadowing installed package
+        working-directory: /tmp
+        run: |
+          python -c "import diff_diff; print('Location:', diff_diff.__file__)"
+          python -c "from diff_diff import HAS_RUST_BACKEND; print('HAS_RUST_BACKEND:', HAS_RUST_BACKEND); assert HAS_RUST_BACKEND, 'Rust backend not available'"
+
+      - name: Copy tests to isolated location
+        run: cp -r tests /tmp/tests
+
+      - name: Run Rust backend tests
+        working-directory: /tmp
+        run: pytest tests/test_rust_backend.py -v
+
+      - name: Run tests with Rust backend
+        working-directory: /tmp
+        run: DIFF_DIFF_BACKEND=rust pytest tests/ -x -q
+
+  # Test pure Python fallback (without Rust extension)
+  python-fallback:
+    name: Pure Python Fallback
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: pip install numpy pandas scipy pytest
+
+      - name: Verify pure Python mode
+        run: |
+          # Use PYTHONPATH to import directly (skips maturin build)
+          PYTHONPATH=. python -c "from diff_diff import HAS_RUST_BACKEND; print(f'HAS_RUST_BACKEND: {HAS_RUST_BACKEND}'); assert not HAS_RUST_BACKEND"
+
+      - name: Run tests in pure Python mode
+        run: PYTHONPATH=. DIFF_DIFF_BACKEND=python pytest tests/ -x -q --ignore=tests/test_rust_backend.py

.gitignore

@@ -56,3 +56,13 @@ Thumbs.db
 # Benchmarks - generated data and results (can be regenerated)
 benchmarks/data/synthetic/*.csv
 benchmarks/results/
+
+# Rust build artifacts
+rust/target/
+Cargo.lock
+*.so
+*.pyd
+*.dylib
+
+# Maturin build artifacts
+target/

CLAUDE.md

@@ -31,6 +31,28 @@ ruff check diff_diff tests
 mypy diff_diff
 ```
 
+### Rust Backend Commands
+
+```bash
+# Build Rust backend for development (requires Rust toolchain)
+maturin develop
+
+# Build with release optimizations
+maturin develop --release
+
+# Run Rust unit tests
+cd rust && cargo test
+
+# Force pure Python mode (disable Rust backend)
+DIFF_DIFF_BACKEND=python pytest
+
+# Force Rust mode (fail if Rust not available)
+DIFF_DIFF_BACKEND=rust pytest
+
+# Run Rust backend equivalence tests
+pytest tests/test_rust_backend.py -v
+```
+
 ## Architecture
 
 ### Module Structure
@@ -81,6 +103,20 @@ mypy diff_diff
   - Single optimization point for all estimators (reduces code duplication)
   - Cluster-robust SEs use pandas groupby instead of O(n × clusters) loop
 
+- **`diff_diff/_backend.py`** - Backend detection and configuration (v2.0.0):
+  - Detects optional Rust backend availability
+  - Handles `DIFF_DIFF_BACKEND` environment variable ('auto', 'python', 'rust')
+  - Exports `HAS_RUST_BACKEND` flag and Rust function references
+  - Other modules import from here to avoid circular imports with `__init__.py`
+
+- **`rust/`** - Optional Rust backend for accelerated computation (v2.0.0):
+  - **`rust/src/lib.rs`** - PyO3 module definition, exports Python bindings
+  - **`rust/src/bootstrap.rs`** - Parallel bootstrap weight generation (Rademacher, Mammen, Webb)
+  - **`rust/src/linalg.rs`** - OLS solver and cluster-robust variance estimation
+  - **`rust/src/weights.rs`** - Synthetic control weights and simplex projection
+  - Uses ndarray-linalg with OpenBLAS (Linux/macOS) or Intel MKL (Windows)
+  - Provides 4-8x speedup for SyntheticDiD, minimal benefit for other estimators
+
 - **`diff_diff/results.py`** - Dataclass containers for estimation results:
   - `DiDResults`, `MultiPeriodDiDResults`, `SyntheticDiDResults`, `PeriodEffect`
   - Each provides `summary()`, `to_dict()`, `to_dataframe()` methods

TODO.md

@@ -102,9 +102,19 @@ From code review (PR #32):
 
 ---
 
+## Rust Backend Optimizations
+
+Deferred from PR #58 code review (can be done post-merge):
+
+- [ ] **Matrix inversion efficiency** (`rust/src/linalg.rs:180-194`): Use Cholesky factorization for symmetric positive-definite matrices instead of column-by-column solve
+- [ ] **Reduce bootstrap allocations** (`rust/src/bootstrap.rs`): Currently uses `Vec<Vec<f64>>` → flatten → `Array2` which allocates twice. Should allocate directly into ndarray.
+- [ ] **Consider static BLAS linking** (`rust/Cargo.toml`): Currently requires system BLAS libraries. Consider `openblas-static` or `intel-mkl-static` features for easier distribution.
+
+---
+
 ## Performance Optimizations
 
-No major performance issues identified. Potential future optimizations:
+Potential future optimizations:
 
 - [ ] JIT compilation for bootstrap loops (numba)
 - [ ] Parallel bootstrap iterations