Skip to content

[PERF] Vectorize the likelihood estimators across parallel datasets - #137

Open
jdkent wants to merge 3 commits into
neurostuff:masterfrom
jdkent:perf/vectorize-likelihood
Open

[PERF] Vectorize the likelihood estimators across parallel datasets#137
jdkent wants to merge 3 commits into
neurostuff:masterfrom
jdkent:perf/vectorize-likelihood

Conversation

@jdkent

@jdkent jdkent commented Aug 19, 2026

Copy link
Copy Markdown
Member

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

James Kent and others added 2 commits August 19, 2026 00:27
Both likelihood estimators looped over the second dimension of y in Python
and ran scipy.optimize.minimize once per dataset, which put them 200-3000x
behind the closed-form estimators: ~9 s per 1000 datasets for the
variance-based one and ~135-205 s for the sample-size-based one, against
~0.03 s for WLS, DL and Hedges.

Both likelihoods are now profiled down to a single bounded parameter and
searched for every dataset at once:

- The coefficients have a closed form at a fixed variance component, so
  substituting it leaves a function of tau^2 alone. That is 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, 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 the bounded search cannot truncate the parameter space. The
  D-L moment estimate sets the scale of that map, which is the vectorized
  counterpart of warm-starting a per-dataset optimizer.
- stats.bounded_scalar_min does a coarse scan to bracket each dataset's
  minimum and then refines every bracket in step by golden section, so the
  whole fit costs ~80 vectorized evaluations regardless of dataset count.

Measured at 30 observations, 12 groups, seconds per 1000 datasets:

                    before           after
                 ungrouped  resc.  ungrouped  resc.
  VBL   ML            8.93  11.75       0.54   0.54
  VBL   REML          8.39  10.38       0.57   0.51
  SSBL  ML          134.10 135.21       0.51   0.49
  SSBL  REML        204.92 131.54       0.56   0.53

Scored with the original joint objective, the new solution is never worse
(largest difference 1e-14) and on 6 of 200 datasets the sample-size-based
fit lands a substantially lower nll -- up to 1.59 -- than the per-dataset
optimizer reached, which is where its parameter estimates change.

The warning about looping over more than ten datasets goes away with the
loop, as does the wrapt dependency it was built on.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… them

Profiling the vectorized fits showed one call dominating every estimator in
the library: np.linalg.pinv, which decomposes each of the D copies of X'WX
by SVD. It was 48-58% of a likelihood fit and 65% of a DerSimonian-Laird
fit. Nothing else came close, and the Python-level overhead of the search is
7.5 ms per fit however large D is, so it was never dispatch-bound.

Three changes, in descending order of payoff:

- stats._invert_stack replaces the pseudo-inverse with a reciprocal when
  there is one predictor and an LU inverse otherwise, falling back to
  np.linalg.pinv where that is the only defined answer. Measured on a stack
  of 20000 matrices: 74.7 ms for pinv, 7.1 ms for inv, 0.03 ms for 1 / x.
  This one also speeds up every closed-form estimator.

- The objective no longer recomputes what the fit it just ran 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.

- bounded_scalar_min refines by successive parabolic interpolation with
  Brent's safeguards rather than by golden section, stops where the
  objective has gone flat to double precision, and leaves alone the
  datasets whose optimum the scan already placed at an end of the interval
  -- the very common tau^2 = 0. The variance-based fit went from 78
  objective evaluations to 49. Brent's step-size safeguard is load-bearing:
  without it the interpolation creeps by ever-smaller steps on a nearly flat
  minimum and stops converging.

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

                  before  vectorized   now    vs vectorized  vs before
  VBL   ML         1.707     0.147    0.0344      4.3x         50x
  VBL   REML       2.728     0.201    0.0404      5.0x         68x
  SSBL  ML        19.429     0.192    0.0634      3.0x        306x
  SSBL  REML      22.413     0.200    0.0770      2.6x        291x
  WLS              0.0017    0.0017   0.0007      2.4x          2.4x
  Hedges           0.0037    0.0031   0.0006      5.2x          6.2x
  DL               0.0036    0.0028   0.0019      1.5x          1.9x

Scored with the original joint objective the fits are unchanged: never
worse than the per-dataset optimizer anywhere (largest difference 2e-13)
and better on the same 6 of 200 sample-size-based datasets as before. The
refinement was also checked against a 20001-point grid on 400 random
non-quadratic objectives, where it never returned a higher value.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@codecov

codecov Bot commented Aug 19, 2026

Copy link
Copy Markdown

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 (4c976c6).

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 AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread pymare/estimators/estimators.py
Comment thread pymare/stats.py Outdated
Comment thread pymare/stats.py
All three reproduce, and each now has a test that fails without its fix.

The sample size identifiability guard was reducing over the whole K x D
array once the per-dataset loop went away, so it asked whether sample sizes
vary *anywhere* rather than within each column. Columns that are each
constant still differ from one another, so a set of datasets that are every
one of them unidentifiable passed the guard and came back with numbers: two
columns of constant n = 20 and n = 50 fitted to tau^2 = 1.91 and 1e-4. A
single constant column also slipped through beside a varying one. The old
loop raised for both. Reduced over observations now, and the message says
how many datasets are affected. Same for the near-equal warning.

bounded_scalar_min's flat test compared the two ends of the bracket with
each other rather than with the middle, which its own docstring says it
does. An objective steeper on one side of its minimum than the other can
hold both ends at equal height while the middle sits far below them; the
search read that as flat, stopped before its first refinement step and
returned the scan point. On a V with slopes 10 and 1, placed so the ends
come out equal, it returned 1/3 against a true minimum of 79/264 -- out by
3.4e-2, most of a scan cell, after 31 evaluations. Now 8.6e-9 after 81.
Costs nothing on real data: the evaluation counts and timings for all four
likelihood fits are unchanged.

The documented ``lower <= upper`` contract was not enforced. A reversed
interval descends, so the tolerance changes sign with it, the first
convergence test passes, and a scan point is returned as though it had been
refined -- 0.333 rather than 0.314 on a quadratic. It raises now. An
interval of one point is ordered, and still allowed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants