[PERF] Vectorize the likelihood estimators across parallel datasets

[jdkent]

Summary

The two likelihood estimators looped over the second dimension of y in Python and ran scipy.optimize.minimize once per dataset, putting them 200–3000× behind the closed-form estimators. Measured at 30 observations and 12 groups, seconds per 1000 datasets: ~9 for VarianceBasedLikelihoodEstimator and ~135–205 for SampleSizeBasedLikelihoodEstimator, against ~0.03 for WLS, DL and Hedges. Extrapolated to a 200k-voxel image that is roughly half an hour and eleven hours respectively.

Both are now profiled down to a single bounded parameter and searched for every dataset at once, and the inversion that dominated every estimator in the library has been replaced.

What changed

Profiling the likelihoods down to one parameter.

• The coefficients have a closed form at a fixed variance component, so substituting it leaves a function of tau^2 alone. Exactly equivalent to minimizing the joint likelihood.
• The sample-size-based likelihood is invariant to the scale of (sigma^2, tau^2) up to a closed-form factor (Q/K for ML, Q/(K-P) for REML), so only their ratio, which lies in [0, 1], has to be searched for. Three parameters become one.
• tau^2 is searched through u / (1 - u), a monotone map from [0, 1) onto [0, inf), so a bounded search cannot truncate the parameter space. The DerSimonian-Laird moment estimate sets the scale of that map, which is the vectorized counterpart of warm-starting a per-dataset optimizer.
• stats.bounded_scalar_min scans coarsely to bracket each dataset's minimum, then refines every bracket in step by successive parabolic interpolation with Brent's safeguards. The whole fit costs a few dozen vectorized evaluations regardless of dataset count.

Not taking an SVD per objective evaluation. Profiling showed one call dominating every estimator: np.linalg.pinv, which decomposes each of the D copies of X'WX by SVD — 48–58% of a likelihood fit and 65% of a DerSimonian-Laird fit. stats._invert_stack uses a reciprocal when there is one predictor and an LU inverse otherwise, keeping pinv as the fallback where it is the only defined answer. On a stack of 20000 matrices: 74.7 ms for pinv, 7.1 ms for inv, 0.03 ms for 1 / x.

Not recomputing what the fit already has. The weights are formed once and handed to weighted_least_squares instead of being rebuilt inside it, and REML reads log|X'WX| off the returned covariance as -log|(X'WX)^-1| rather than forming X'WX a second time.

Timings

Seconds per 1000 datasets, 30 observations, 12 groups, D = 2000:

estimator	before	after	speedup
VarianceBased ML	1.707	0.0344	50x
VarianceBased REML	2.728	0.0404	68x
SampleSizeBased ML	19.429	0.0634	306x
SampleSizeBased REML	22.413	0.0770	291x
WLS	0.0017	0.0007	2.4x
Hedges	0.0037	0.0006	6.2x
DerSimonian-Laird	0.0036	0.0019	1.9x

A 200k-voxel image now takes 8–15 s for the likelihood estimators.

Correctness

• The metafor ground-truth values the tests pin (tau2 7.7649 for ML, 10.9499 for REML) still hold at atol=1e-4.
• Scored with the original joint objective, the new solution is never worse than what the per-dataset optimizer reached — largest difference 2e-13 across 200 datasets — and on 6 of 200 sample-size-based datasets it lands a substantially lower negative log-likelihood, up to 1.59. That is where its sigma2 and tau2 estimates change: the old three-parameter L-BFGS-B run was under-converging.
• The refinement was checked against a 20001-point grid on 400 random non-quadratic unimodal objectives, where it never returned a higher value.
• New tests cover the minimizer (per-dataset optima, optima exactly on a bound, a minimum off the scan grid, a constant objective, a minimum with almost no curvature, degenerate datasets, bound validation), the pseudo-inverse fallback on a collinear design, per-column equivalence of a 40-dataset fit against 40 single-dataset fits, that both estimators land on a minimum of their profile likelihood under ML and REML, and that the reparametrized sample-size fit is not beaten by a grid search over sigma^2 and tau^2.

Full suite: 331 passed, 1 skipped. flake8 clean.

API notes

Three deliberate changes worth flagging:

1. **kwargs on the two estimators now goes to stats.bounded_scalar_min (xtol, ftol, maxiter) rather than to scipy.optimize.minimize. Code passing SciPy minimizer options will raise TypeError at fit time.
2. The warning about looping over more than ten parallel datasets is gone along with the loop, and its test has been replaced by one asserting that a 40-dataset fit is silent and column-wise identical to fitting each alone.
3. wrapt is dropped from install_requires; it existed only for the _loopable decorator. An ill-conditioned but nonsingular design now gets the ordinary inverse rather than one with its smallest singular values truncated, so a nearly collinear design reports a large covariance instead of a quietly regularized one.

The benchmark suite's likelihood entries move to the shared dataset size, since they no longer need a smaller second dimension, and N_DATASETS_LOOPED is removed.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 99.24242% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 92.35%. Comparing base (e2df937) to head (a16c648).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
pymare/stats.py	98.55%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #137      +/-   ##
==========================================
+ Coverage   92.18%   92.35%   +0.17%     
==========================================
  Files          13       13              
  Lines        1817     1845      +28     
==========================================
+ Hits         1675     1704      +29     
+ Misses        142      141       -1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Vectorizes likelihood estimators across parallel datasets and accelerates matrix inversion.

Changes:

• Adds a vectorized bounded scalar minimizer.
• Profiles likelihood estimation to one parameter.
• Updates tests, benchmarks, documentation, and dependencies.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
setup.cfg	Removes the unused wrapt dependency.
pymare/stats.py	Adds vectorized minimization and faster inversion.
pymare/estimators/estimators.py	Vectorizes likelihood estimators.
pymare/tests/test_stats.py	Tests minimization and inversion behavior.
pymare/tests/test_estimators.py	Tests parallel likelihood fitting.
docs/api.rst	Exposes the new minimizer API.
benchmarks/common.py	Removes loop-specific dataset sizing.
benchmarks/bench_estimators.py	Benchmarks vectorized likelihood fits.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jdkent]

LGTM, will stress test with real data.