[hippke]

Summary

This PR vectorizes the accept/reject step in RedBlueMove.propose. The sampler's output is bit-identical, and it is 1.1–2.6× faster, most noticeably for large ensembles and vectorize=True.

RedBlueMove is the base class for StretchMove (the default), WalkMove, DEMove, DESnookerMove and KDEMove, so all of them get faster.

The change

It is a single block in src/emcee/moves/red_blue.py: 7 lines replaced by 4. No other library code changes, and no API, behavior or dependency changes.

-            # Loop over the walkers and update them accordingly.
-            for i, (j, f, nlp) in enumerate(
-                zip(all_inds[S1], factors, new_log_probs)
-            ):
-                lnpdiff = f + nlp - state.log_prob[j]
-                if lnpdiff > np.log(model.random.rand()):
-                    accepted[j] = True
+            # Accept or reject all walkers in this subset at once. Drawing
+            # ``rand(n)`` gives the same numbers as ``n`` calls to ``rand()``.
+            lnpdiff = factors + new_log_probs - state.log_prob[S1]
+            accepted[S1] = lnpdiff > np.log(model.random.rand(len(lnpdiff)))

Why the output is bit-identical

• Same random numbers. RandomState.rand(n) returns exactly the same numbers, in the same order, as n successive calls to rand(), and leaves the generator in the same state afterwards.
• Same order of walkers. state.log_prob[S1] lists walkers in the same ascending order as the old all_inds[S1] loop, and factors and new_log_probs are already in that order.
• Same arithmetic. Each element is still computed as (f + nlp) - log_prob, in the same order and at the same precision. -inf log-probabilities give the same nan → reject result as before.
• Same acceptance flags. Every walker belongs to exactly one subset and starts as False, so assigning accepted[S1] gives the same result as the old "set to True if accepted" loop.

Testing

• Byte-for-byte comparison against the unmodified code. 13 configurations, 300 steps each, comparing the stored chain, log-probabilities, blobs, acceptance counts and final RNG state:
  • Moves: StretchMove, WalkMove, DEMove, DESnookerMove, KDEMove, and a weighted mix of moves.
  • Settings: vectorize=True, an odd number of walkers, nsplits=3, randomize_split=False, a likelihood that returns -inf, blobs, and a longdouble backend.
  • Result: all 53 arrays are identical. The longdouble log-prob array matches in every value; its only byte differences are in the unused padding bytes of x86's 80-bit long double, which come from np.empty.
• New unit test test_vectorized_accept_matches_loop in tests/unit/test_stretch.py. It runs the new propose alongside a copy of the original loop for 50 steps and asserts identical acceptance flags, coordinates, log-probabilities and RNG state. It covers 32 and 33 walkers, nsplits 2 and 3, and -inf log-probabilities.
  • I checked that the test can fail: reversing the order of the random draws in the implementation makes all 4 new cases fail.
• Full test suite: 133 passed. One test was deselected: test_tqdm_modes, which needs ipywidgets and is unrelated to this change.
• black and isort (the versions pinned in .pre-commit-config.yaml) and pyflakes are clean on the changed files.

Performance

Iterations per second with a trivial Gaussian log-prob, so the numbers mostly reflect emcee's own overhead. Best of 3 runs, with the old and new versions run alternately. Python 3.13, NumPy 2.5, x86_64.

Walkers × dims	Mode	Before (it/s)	After (it/s)	Speedup
32 × 5	per-walker	3071	3275	1.07×
100 × 10	per-walker	1856	2270	1.22×
1024 × 20	per-walker	237	367	1.55×
32 × 5	vectorize=True	3250	4104	1.26×
100 × 10	vectorize=True	2295	3416	1.49×
1024 × 20	vectorize=True	401	880	2.19×
4096 × 20	vectorize=True	113	288	2.55×

The speedup grows with the number of walkers because the removed loop ran Python code once per walker per step. With an expensive likelihood the absolute saving per walker per step stays the same, but it becomes a smaller share of total runtime.

---

Generated by Claude Code