This document outlines the strategy for improving diff-diff's performance on large datasets, particularly for BasicDiD/TWFE and CallawaySantAnna estimators.

Earlier sections of this document (v1.4.0, v2.0.3) measured isolated fit() calls on synthetic panels for R-parity. This section measures end-to-end practitioner chains - Bacon decomposition, fit, event-study pre-trend inspection, HonestDiD sensitivity grids, cross-estimator robustness refits, and reporting - at data shapes anchored to applied-econ papers and industry writeups. The six scenarios are defined in docs/performance-scenarios.md; scripts live in benchmarks/speed_review/bench_*.py; raw results in benchmarks/speed_review/baselines/*.json and flame profiles in benchmarks/speed_review/baselines/profiles/.

Environment: macOS darwin 25.3 on Apple Silicon M4, Python 3.9, numpy 2.x, diff_diff 3.1.3. Each multi-scale scenario runs at three data scales under both DIFF_DIFF_BACKEND=python and DIFF_DIFF_BACKEND=rust, with one intentional exception: the SDiD few-markets scenario at its large scale runs Rust only, because the pure-numpy jackknife at n=500 would exceed four minutes per run without changing the already-clear Python-vs-Rust conclusion established at small and medium. The numerical tables below are auto-generated from the committed JSON baselines by benchmarks/speed_review/gen_findings_tables.py; narrative prose is hand-written and must be re-read when numbers shift.

Four of the six scenarios run at three scales (small / medium / large). The small scale matches tutorial data shapes; medium reflects typical practitioner workloads; large stretches toward the upper end of what an analyst might bring (1M-row BRFSS microdata, 1,500-unit county-level staggered panel, 1,000-unit multi-region brand survey, 500-unit zip-level geo-experiment). Dose-response and reversible-dCDH run at a single mid-range scale. Data-shape details are in docs/performance-scenarios.md.

Three findings are load-bearing for the optimization priority list:

1. BRFSS aggregate_survey is now practitioner-fast at every measured scale. Prior to the precompute-scaffolding fix (see "Optimization landed" below), the full chain at 1M rows took ~24 seconds and was essentially all inside _compute_stratified_psu_meat. After the fix, the chain is sub-2s at every measured scale; aggregate_survey continues to dominate its own (now-cheap) chain share, but in absolute time the entire workflow is well under a practitioner- perceptible threshold at realistic pooled-multi-year BRFSS volume. The path is entirely Python, so Python and Rust backends track each other within noise.
2. Staggered CS chain stays cheap across scales. A 10x unit increase (150 -> 1,500) is a small-single-digit multiplier on total time. ImputationDiD and SunAbraham together consistently account for ~70-80% of the chain; either can be the single top phase at a given (scale, backend) cell, which is a per-cell ranking detail not a stable pattern to optimize against.
3. SDiD Rust gap is stable across scales, not emergent. Python SDiD has a fixed per-jackknife-refit overhead that dominates even at small n. Rust stays sub-second through 500 units.

Two findings hold across scales:

4. Brand-awareness survey total scales roughly linearly in n_units, but the JK1 replicate path inside it scales closer to n_units x n_replicates - faster growth than the chain total, so it increasingly dominates at large n.
5. Rust backend gives large uplift only for SDiD (order-of-magnitude and up). Elsewhere the gap is modest across all measured (scenario, scale) cells - see the scale-sweep table for exact ratios. The primary bottlenecks live in Python code the Rust backend does not touch (aggregate_survey, JK1 replicate fit), and paths that Rust does touch (CS bootstrap, ImputationDiD, Survey TSL) are already well-vectorized in Python.

Per-scenario phase narrative (cross-check against the table above after any rerun):

• Staggered campaign. ImputationDiD robustness and SunAbraham consistently account for ~70-80% of the chain at every scale. They sit in a narrow phase-share band (each typically ~25-50%) and which one leads varies by (scale, backend) and can flip across reruns at medium scale where the two are close; see the table for the exact ordering per cell. CS fit with n_bootstrap=999 (both with and without covariates) is well-vectorized and sits well below both in the ranking. Either phase is a legitimate optimization target; the aggregate share is what drives the "next hotspot" priority.
• Brand awareness survey. At small scale HonestDiD dominates. From medium onwards JK1 is the single largest phase under both backends; see the table for the exact share per cell. Python and Rust totals stay close across the sweep (within ~1.1x at any measured scale, see scale-sweep table); the JK1 replicate-fit loop is not Rust-accelerated, so the backends neither help nor hurt each other meaningfully on this chain.
• BRFSS. aggregate_survey remains the single largest chain share under both backends at every scale, but the absolute chain total is sub-2s at 1M rows after the precompute-scaffolding fix. Downstream phases (CS fit, SunAbraham, HonestDiD) are a fraction of a second combined - see the scale-sweep table for the current totals.
• SDiD few markets. sensitivity_to_zeta_omega and in_time_placebo are the two largest phases under Python at every scale and under Rust at medium/large (together ~70% of the chain). At Rust small the absolute cost collapses so far that per-phase fixed overhead dominates and 2_sdid_bootstrap_variance_200 slightly edges the other two. The difference across backends is absolute: under Python these phases drive a multi-second chain, under Rust they stay in the top ranks but of a sub-second total runtime. That is the Python-vs-Rust story for this scenario.
• Reversible dCDH. Main fit and heterogeneity refit are the two largest phases by design - together effectively the whole chain, with the remainder on HonestDiD at <2%. The two phases sit within a few percentage points of each other at this shape and the leader can flip across reruns under either backend. Both fits run under the same SurveyDesign and rebuild shared TSL scaffolding - that is the optimization opportunity, independent of which side is slightly larger on a given measurement.
• Pricing dose-response. Four spline fits account for essentially all runtime; linear scaling in variant count.

End-to-end peak RSS and per-scenario growth are captured in each JSON baseline under the memory field, recorded via a psutil background sampler at 10 ms. A standalone tracemalloc-based allocator attribution pass for the BRFSS-1M scenario lives at benchmarks/speed_review/mem_profile_brfss.py; its scrubbed output is in benchmarks/speed_review/baselines/mem_profile_brfss_large_<backend>.txt.

The ~115-130 MB floor is the Python + diff-diff + numpy import footprint; the "growth" columns are the practitioner-meaningful numbers.

1. BRFSS aggregate_survey was compute-bound, not memory-bound - and the compute side is now addressed. Working-memory growth stayed in the low tens of MB across the 20x data-growth sweep (50K -> 1M rows); the pre-fix tracemalloc pass confirmed net retained allocation under 1 MB and identified tracemalloc's own linecache overhead as the top allocation site (smoking gun that nothing else was allocating meaningfully). The precompute-scaffolding fix in this PR is a pure CPU win - no change to the function's memory profile, which was already Lambda-friendly.
2. Staggered CS chain is memory-heavier than wall-clock suggested. At 1,500 units the chain's peak RSS sits in the high-400s to high-500s MB depending on backend. Fine for workstations, tight for 512 MB Lambda tier. Bootstrap-999 in CS and ImputationDiD's saturated regression are plausible drivers. Rust uses slightly more memory here (likely FFI-held temporary array copies); not worth optimizing.
3. JK1 replicate path is allocation-heavy at large replicate count. At 1,000 units × 160 replicates the chain's growth during run sits in the mid-100s of MB (see memory table). Each replicate refit plus the n × n_replicates weight matrix drives this. A Rust port would save memory even though time is within noise today - the dual benefit strengthens the case for the port if replicate counts grow.
4. SDiD Rust path is essentially memory-free (growth at or below a single MB across scales). Rust does the work in native memory without round-tripping through the Python allocator. Confirms the existing Rust port is well-behaved on both axes.
5. No scenario hits OOM territory at measured scales. Peak RSS across the whole sweep stays under 600 MB. 1 GB is a comfortable ceiling for every scenario measured.

#1 shipped in this PR. diff_diff/survey.py now precomputes a per-design _PsuScaffolding (strata codes, global PSU codes, per- stratum counts and FPC ratios, singleton mask, lonely-PSU-aware variance-computable flag). aggregate_survey builds it once per call and threads it through _cell_mean_variance so each per-cell variance reduction uses two vectorized np.bincount passes instead of a per-stratum pandas groupby loop. Numerics are preserved to sub-ULP tolerance; equivalence tests across seven design cases (TestAggregateSurveyScaffolding) enforce assert_allclose(atol=1e-14, rtol=1e-14) between fast and legacy paths.

Replicate-weight designs (JK1 etc.) continue to use the legacy compute_replicate_if_variance code path and are unaffected.

Bottom line: no practitioner-perceptible bottleneck remains in the six measured workflows; four optional items stand by. Items #2-5 above should be prioritized by concrete deployment-environment signal (Lambda OOMs, practitioner reports of slowness at specific shapes), not proactively.

These are developer-ergonomics / API-consistency smells surfaced during scenario development. None are silent-failures and none belong in this PR or in the silent-failures audit; logging here for awareness.

1. aggregate / level parameter naming is inconsistent. CS accepts aggregate="event_study"; ContinuousDiD requires aggregate="eventstudy" on fit() but level="event_study" on to_dataframe(). Two different spellings within one estimator plus a third cross-estimator spelling. Surfaced when the P1 exit-propagation fix stopped silently swallowing the resulting ValueError in the dose-response benchmark. Route: API-consistency cleanup, minor.
2. generate_survey_did_data(panel=True) treated column. Row-level active-treatment indicator that is zero in pre-periods, which makes it quietly incompatible with check_parallel_trends (expects unit-level treatment group membership) and pre-period placebo tests. Tutorial 17 does not hit this because it uses a 2x2 design where post discriminates the comparison. Suggest adding a treat_unit column alongside treated for generator output clarity. Route: DGP cleanup, minor.
3. SurveyDesign.replicate_method case sensitivity. "jk1" raises ValueError("must be one of {'Fay', 'SDR', 'BRR', 'JKn', 'JK1'}"); "JK1" works. Either normalize the input or mention the expected casing in the error message. Route: API-ergonomics, minor.

• OOM behaviour at the edge: the sweep captures peak RSS up to ~600 MB (staggered CS large under Rust). Behaviour under a hard memory ceiling (512 MB Lambda, 1 GB container) is not exercised; if deployment signal emerges that practitioners hit those ceilings, a ceiling-test pass should be added.
• Pure-Rust profiles: scenarios run the Rust backend as a black box. Optimizing inside rust/ is a separate concern owned by the crate maintainers and is not in scope here.
• Real-data shapes: the scenarios use synthetic DGPs. The BRFSS scenario uses a BRFSS-shaped synthetic panel, not actual BRFSS microdata. If a real-data calibration becomes relevant, CDC BRFSS annual files are public.

pip install pyinstrument                  # one-time, dev-only
python benchmarks/speed_review/run_all.py # both backends, all scenarios

# Single scenario, single backend:
DIFF_DIFF_BACKEND=rust python benchmarks/speed_review/bench_campaign_staggered.py

Raw JSON is written under benchmarks/speed_review/baselines/ for scenario-level diffing as the library evolves; flame HTMLs are written alongside under baselines/profiles/ (gitignored; regenerated on each run).

v2.0.3 includes Rust backend optimizations that further improve SyntheticDiD performance:

20K Scale Results (new in v2.0.3 benchmarks):

1. Cholesky factorization for symmetric positive-definite matrix inversion (~2x faster for well-conditioned matrices)
2. Reduced bootstrap allocations - Direct Array2 allocation eliminates Vec<Vec> intermediate
3. Vectorized variance computation - HC1 meat uses BLAS-accelerated matrix operations
4. Webb lookup table - Faster Webb distribution weight generation
5. Rayon chunk size tuning - Reduced parallel scheduling overhead

Phase 1 is complete. Pure Python optimizations exceeded all targets:

1. Unified linalg.py backend (diff_diff/linalg.py)

   • solve_ols() - scipy lstsq with gelsy LAPACK driver
   • compute_robust_vcov() - Vectorized cluster-robust SE via pandas groupby
   • Single optimization point for all estimators

2. CallawaySantAnna optimizations (staggered.py)

   • _precompute_structures() - Pre-computed wide-format outcome matrix, cohort masks
   • _compute_att_gt_fast() - Vectorized ATT(g,t) using numpy (23x faster)
   • _generate_bootstrap_weights_batch() - Batch weight generation
   • Vectorized bootstrap using matrix operations (26x faster)

3. TWFE optimization (twfe.py)

   • Cached groupby indexes for within-transformation

4. All estimators migrated to unified backend

   • estimators.py, twfe.py, staggered.py, triple_diff.py, synthetic_did.py, sun_abraham.py, utils.py

Benchmark comparisons showed that while diff-diff was competitive or faster than R for small datasets, performance degraded significantly at scale:

Note: SyntheticDiD is already 37-1600x faster than R's synthdid package.

Current implementation uses np.linalg.lstsq with default settings:

• General-purpose LAPACK driver (gelsd) rather than faster alternatives
• Preceded by expensive matrix_rank() check (O(min(n,k)^3))
• NumPy may not link to optimized BLAS

Loop-based implementation:

for cluster in unique_clusters:
    mask = cluster_ids == cluster  # O(n) per cluster
    ...

• O(n * n_clusters) complexity
• Creates boolean mask array for each cluster
• No vectorization or parallelization

Multiple groupby operations:

for var in variables:
    unit_means = data.groupby(unit)[var].transform("mean")
    time_means = data.groupby(time)[var].transform("mean")
    ...

• Multiple passes over data per variable
• No caching of groupby indexes
• Not using alternating projections algorithm

for g in treatment_groups:
    for t in valid_periods:
        att_gt = self._compute_att_gt(...)

• Repeated DataFrame indexing (.set_index(), .loc[], .isin()) for each (g,t)
• No pre-computation of outcome changes
• Influence function dictionaries created per (g,t)

Quick wins that improve performance without adding dependencies.

Replace loop with vectorized groupby:

scores = X * residuals[:, np.newaxis]
cluster_scores = pd.DataFrame(scores).groupby(cluster_ids).sum()
meat = cluster_scores.values.T @ cluster_scores.values

Expected speedup: 5-10x for SE computation

from scipy.linalg import lstsq
coefficients = lstsq(X, y, lapack_driver='gelsy',
                     overwrite_a=True, overwrite_b=True,
                     check_finite=False)[0]

Expected speedup: 1.2-1.5x for OLS

Create groupby objects once and reuse:

unit_grouper = data.groupby(unit, sort=False)
time_grouper = data.groupby(time, sort=False)

Expected speedup: 1.5-2x for demeaning

Pivot to wide format once, pre-compute all period changes:

outcome_wide = data.pivot(index=unit, columns=time, values=outcome)
changes = {(t0, t1): outcome_wide[t1] - outcome_wide[t0] for ...}

Expected speedup: 3-5x for CallawaySantAnna

Implement performance-critical components in a compiled language for maximum speed.

We have two viable options for a compiled backend. Both can achieve near-identical performance; the choice depends on team expertise and maintenance considerations.

Pros:

• Memory safety by design - No segfaults, buffer overflows, or data races; compiler catches these at build time
• Modern tooling - Cargo package manager + maturin makes wheel building straightforward
• Zero-copy NumPy interop - rust-numpy crate provides direct array access without copying
• Easy parallelism - rayon crate makes parallel iteration trivial (.par_iter())
• Growing ecosystem - Used by polars, pyfixest, cryptography, orjson, ruff
• Low per-call overhead - Research shows PyO3 has ~0.14ms overhead vs NumPy's ~3.5ms for simple operations
• Single toolchain - cargo build works the same on all platforms

Cons:

• Learning curve - Rust's ownership model takes time to learn
• Smaller scientific ecosystem - Fewer numerical libraries than C++ (though ndarray and faer are mature)
• Slower compilation - Rust compiles slower than C++
• Newer language - Less institutional knowledge, fewer Stack Overflow answers

Key dependencies: pyo3, rust-numpy, ndarray, faer (linear algebra), rayon (parallelism)

Pros:

• Mature ecosystem - Eigen, Armadillo, Intel MKL, OpenBLAS all native C++
• Familiar to more developers - Larger pool of contributors
• Proven in scientific Python - NumPy, SciPy, scikit-learn, pandas all use C/C++ extensions
• Excellent Eigen integration - pybind11 has built-in support for Eigen matrices
• Faster compilation - C++ compiles faster than Rust
• More optimization resources - Decades of C++ performance tuning knowledge

Cons:

• Memory safety risks - Segfaults, buffer overflows, use-after-free possible; harder to debug
• Manual memory management - Must carefully manage lifetimes, especially with Python GC interaction
• Complex build systems - CMake configuration, compiler flags, platform-specific issues
• Copy overhead by default - pybind11 copies arrays unless carefully configured with py::array_t
• Manual GIL management - Easy to deadlock or corrupt state if GIL not handled correctly
• Platform differences - MSVC vs GCC vs Clang have different behaviors and flags

Key dependencies: pybind11, Eigen (linear algebra), OpenMP or TBB (parallelism)

Rust with PyO3 is the recommended approach because:

1. pyfixest validates this for our exact domain - They use Rust/PyO3 for fixed effects econometrics
2. Memory safety prevents production bugs - No risk of segfaults in user code
3. maturin simplifies distribution - Single command builds wheels for all platforms
4. rayon makes parallelization trivial - Critical for bootstrap and cluster SE

However, C++ is a viable alternative if:

• Team has stronger C++ expertise
• Need to integrate with existing C++ econometrics code
• Want to leverage Eigen's mature linear algebra

try:
    from diff_diff._rust_backend import solve_ols_clustered
    _HAS_RUST = True
except ImportError:
    _HAS_RUST = False

def _fit_ols(self, X, y, cluster_ids=None):
    if _HAS_RUST and self.backend == 'rust':
        return solve_ols_clustered(X, y, cluster_ids)
    else:
        # Existing NumPy implementation
        ...

diff_diff/
├── estimators.py          # Python API (unchanged)
├── _rust_backend/         # Compiled Rust module
│   └── ...
└── _fallback.py           # Pure Python fallback

src/                       # Rust source (Cargo workspace)
├── Cargo.toml
├── lib.rs
├── ols.rs                 # OLS with cluster SE
├── demeaning.rs           # Alternating projections
├── bootstrap.rs           # Parallel bootstrap
└── staggered.rs           # ATT(g,t) computation

pyproject.toml             # maturin build config

diff_diff/
├── estimators.py          # Python API (unchanged)
├── _cpp_backend/          # Compiled C++ module
│   └── ...
└── _fallback.py           # Pure Python fallback

cpp/                       # C++ source
├── CMakeLists.txt
├── src/
│   ├── module.cpp         # pybind11 bindings
│   ├── ols.cpp            # OLS with cluster SE
│   ├── ols.hpp
│   ├── demeaning.cpp      # Within transformation
│   ├── demeaning.hpp
│   ├── bootstrap.cpp      # Parallel bootstrap
│   └── bootstrap.hpp
└── extern/
    └── eigen/             # Eigen submodule (or system install)

pyproject.toml             # scikit-build-core config

# .github/workflows/wheels.yml
- uses: PyO3/maturin-action@v1
  with:
    command: build
    args: --release --out dist

• Simple single-action CI configuration
• Use abi3 stable ABI for Python version-independent wheels
• Cross-compilation via --target flag

# .github/workflows/wheels.yml
- uses: pypa/cibuildwheel@v2
  env:
    CIBW_BUILD: "cp39-* cp310-* cp311-* cp312-*"

• More configuration required for CMake integration
• Need to handle OpenMP linking per-platform
• Consider vcpkg or conan for dependency management

Both approaches build wheels for:

• Linux (manylinux2014, x86_64 and aarch64)
• macOS (x86_64 and ARM64)
• Windows (x86_64)

Exceeded all targets:

• BasicDiD @ 10K: 0.835s → 0.011s (76x improvement, 4.2x faster than R)
• CallawaySantAnna @ 10K: 2.2s → 0.109s (20x improvement, 7.2x faster than R)
• Bootstrap inference: 26x faster via vectorization

No longer required for R parity. May be pursued for:

• Further optimization at extreme scales (100K+ units)
• Parallel bootstrap across CPU cores
• Memory efficiency for very large datasets

• PyO3 User Guide - Rust bindings for Python
• rust-numpy - Zero-copy NumPy interop
• maturin - Build and publish Rust Python packages
• faer - Pure Rust linear algebra (competitive with MKL)
• Polars - Example of Rust/Python hybrid architecture
• pyfixest - Rust backend for fixed effects econometrics

• pybind11 documentation - C++ bindings for Python
• pybind11 Eigen integration - Zero-copy with Eigen
• Eigen - C++ linear algebra library
• scikit-build-core - CMake integration for Python packages
• cibuildwheel - Build wheels for all platforms

• fixest demeaning algorithm - Reference implementation
• PyO3 vs C performance comparison - Academic benchmark
• Making Python 100x faster with Rust - Practical tutorial