Use the real-input FFT for registration phase correlation

[bendichter]

How This Was Found

I used Claude Code to profile the registration path. It surfaced the convolve hotspot and drafted the change. I am stating that at the top rather than leaving you to infer it, because I know AI-generated pull requests are frequently a net cost to the people maintaining a project, and you are entitled to not want them. If you would rather not take AI-assisted contributions regardless, say so and I will close this, no argument. If the change is welcome in principle but you want it structured differently, for instance the float32 taper split out as its own commit, I am glad to do that.

Summary

The frames passed to convolve in registration/utils.py are real-valued, but the function ran a full complex fft2 and then discarded the imaginary half of the inverse transform with torch.real(). Because both the frames and the reference image are real, the phase-correlation spectrum is Hermitian, so the redundant half carries no information. Using rfft2 / irfft2 computes the same quantity with half the transform work and half the spectrum memory.

The reason to look here at all is that convolve dominates the stage. On a 500 frame, 512x512 synthetic movie it accounts for roughly two thirds of registration time (7.75s of 11.6s) across only ten calls.

What Changed

ref_smooth_fft now returns the non-redundant half of the spectrum, of shape (Ly, Lx // 2 + 1), so that the kernel matches the transform. The callers in rigid.phasecorr and nonrigid.phasecorr keep their data real instead of casting to complex64. nonrigid.phasecorr had been storing real convolution output in a complex64 tensor, which doubled the memory traffic of the slicing that follows it.

spatial_taper changes from float64 to float32. This is required because a float64 mask promotes the masked frames to float64 and the transform to complex128, which would make this change slower than the code it replaces. The float64 intermediate was discarded by the cast to complex64 on the following line in any case, and the MPS backend has always run this path in float32, so float32 masks are not new behavior.

nonrigid.phasecorr previously derived the block spatial size from cfRefImg.shape, which is now half width. It reads from maskMul instead, which retains the full dimensions.

One caller outside the package, paper/figures.py, unpacked ref_smooth_fft with ifft2 and is updated to use irfft2.

Reproducing

The script I used to verify this, along with its reference output, is here:

https://gist.github.com/bendichter/6d2ecc2e01e6c03dfe42c6e499c054fd

It runs register_frames from two checkouts on byte-identical input and compares every output: the registered movie, the mean image, and all offsets. To use it against an unmodified tree:

git worktree add /tmp/s2p_main main
python verify_rfft2_registration.py run /tmp/s2p_main old.npz
python verify_rfft2_registration.py run /path/to/this/branch new.npz
python verify_rfft2_registration.py compare old.npz new.npz

Performance

Measured on 300 frames of 512x512 on CPU, median of four runs. I do not have a CUDA device available, so these numbers are CPU only. I would expect the float64 removal in particular to matter more on CUDA, where fp64 throughput is a fraction of fp32 on consumer cards, but I have not measured that and am not claiming it.

	before	after
rigid	1.15s	0.54s	2.1x
nonrigid	3.67s	1.90s	1.9x

Run-to-run variance on this machine is a few percent, mostly in the nonrigid numbers.

Numerical Equivalence

Rigid output is bit-identical. The registered frames, the mean image, and the integer offsets all match exactly. Only the correlation peak values differ, at 1e-7, which is float32 rounding.

With nonrigid enabled the output is not bit-identical. One block shift out of 10800 moved by a single subpixel step, 0.1 px at the default subpixel=10, which affected 1 frame out of 300. The correlation values themselves agree to 1e-8. This is float32 rounding in the transform flipping an argmax tie between two nearly equal peaks, not a systematic change in the shifts. The knock-on effect is that the nonrigid mean image moves by at most 0.13 on a movie whose values run to roughly 12000. Reviewers who want exact reproduction of existing nonrigid outputs on a pinned dataset should know about it.

Tests

This numerical path had no direct test coverage, so I added four tests to tests/test_registration.py:

• test_convolve_matches_full_spectrum_reference checks convolve against a full complex-FFT reference implementation, pinning the Hermitian half-spectrum equivalence that this change rests on.
• test_ref_smooth_fft_returns_half_spectrum pins the returned shape and dtype.
• test_rigid_phasecorr_recovers_known_shifts checks exact recovery of known integer shifts.
• test_spatial_taper_is_float32 guards the complex128 regression described above, which is otherwise silent and would only show up as a slowdown.

All 7 tests in tests/test_registration.py pass. The two test_io.py NWB round-trip failures also fail on unmodified main and are unrelated to this change.

🤖 Generated with Claude Code