Use the real-input FFT for registration phase correlation

paper/figures.py

@@ -778,7 +778,8 @@ def registration_fig(frand, freg, refImg, cc_ex, yoff, xoff, yblock, xblock, nbl
     ax.set_position([pos[0]+0.02, pos[1], pos[2], pos[3]])
     from suite2p.registration.utils import ref_smooth_fft
     import torch
-    ref_img_w = torch.real(torch.fft.ifft2(ref_smooth_fft(torch.from_numpy(ref_img), smooth_sigma=0.85))).numpy()
+    ref_img_w = torch.fft.irfft2(ref_smooth_fft(torch.from_numpy(ref_img), smooth_sigma=0.85),
+                                 s=ref_img.shape[-2:]).numpy()
     ref_img_w = ref_img_w[::-1][:,::-1]
     # ref_img_w = np.fft.fft2(ref_img)
     # ref_img_w /= np.abs(ref_img_w)

suite2p/registration/nonrigid.py

@@ -156,14 +156,16 @@ def compute_masks_ref_smooth_fft(refImg0, maskSlope, smooth_sigma,
         computed as block_mean * (1 - maskMul_block) so that masked regions are
         filled with the local block mean scaled by the complement of the taper.
     cfRefImg_block : torch.Tensor (complex64)
-        Complex32 tensor of shape (nb, Ly, Lx). Frequency-domain (FFT) representation
-        of the Gaussian-smoothed reference blocks (output of ref_smooth_fft). These
+        Complex64 tensor of shape (nb, Ly, Lx // 2 + 1). Frequency-domain (FFT)
+        representation of the Gaussian-smoothed reference blocks (output of
+        ref_smooth_fft), holding only the non-redundant half of the spectrum. These
         are intended for use in phase-correlation registration.
 
     """
     nb, Ly, Lx = len(yblock), yblock[0][1] - yblock[0][0], xblock[0][1] - xblock[0][0]
     dims = (nb, Ly, Lx)
-    cfRef_dims = dims
+    # cfRefImg holds only the non-redundant half of the spectrum (real-input FFT)
+    cfRef_dims = (nb, Ly, Lx // 2 + 1)
     cfRefImg1 = torch.zeros(cfRef_dims, dtype=torch.complex64)
 
     maskMul = spatial_taper(maskSlope, *refImg0.shape)
@@ -295,7 +297,8 @@ def phasecorr(data, blocks, maskMul, maskOffset, cfRefImg, snr_thresh,
     device = data.device
 
     nimg = data.shape[0]
-    ly, lx = cfRefImg.shape[-2:]
+    # from maskMul, not cfRefImg: the latter holds a half-width spectrum
+    ly, lx = maskMul.shape[-2:]
 
     # maximum registration shift allowed
     lcorr = int(
@@ -308,17 +311,16 @@ def phasecorr(data, blocks, maskMul, maskOffset, cfRefImg, snr_thresh,
     for n in range(nb):
         yind, xind = yblock[n], xblock[n]
         Y[:, n] = data[:, yind[0]:yind[-1], xind[0]:xind[-1]]
-    Y = (Y.float() * maskMul + maskOffset).type(torch.complex64)
+    Y = Y.float() * maskMul + maskOffset
     batch = min(64, Y.shape[1])  #16
     for n in np.arange(0, nb, batch):
         nend = min(Y.shape[1], n + batch)
         Y[:, n:nend] = convolve(mov=Y[:, n:nend], img=cfRefImg[n:nend])
 
     # calculate ccsm
     lhalf = lcorr + lpad
-    cc0 = torch.cat((torch.cat((Y[..., -lhalf:, -lhalf:], Y[..., -lhalf:, :lhalf + 1]), axis=-1),   
+    cc0 = torch.cat((torch.cat((Y[..., -lhalf:, -lhalf:], Y[..., -lhalf:, :lhalf + 1]), axis=-1),
                     torch.cat((Y[..., :lhalf + 1, -lhalf:], Y[..., :lhalf + 1, :lhalf + 1]), axis=-1)), axis=-2)
-    cc0 = torch.real(cc0)
     cc0 = cc0.permute(1, 0, 2, 3)
     cc0 = cc0.reshape(cc0.shape[0], -1)
     cc0 = cc0.cpu().numpy()

suite2p/registration/rigid.py

@@ -51,11 +51,11 @@ def phasecorr(frames, cfRefImg, maskMul, maskOffset, maxregshift, smooth_sigma_t
     ----------
     frames : torch.Tensor
         Input image sequence, expected shape (N, Ly, Lx) where N is the number of frames.
-        The tensor may be on CPU or CUDA; it is converted to float and then to complex for the
-        Fourier-domain operations performed by the helper `convolve`.
+        The tensor may be on CPU or CUDA; it is converted to float for the Fourier-domain
+        operations performed by the helper `convolve`.
     cfRefImg : torch.Tensor
-        Complex-valued reference of shape (Ly, Lx) in the Fourier domain used to compute 
-        cross-correlation with each frame
+        Complex-valued half-spectrum reference of shape (Ly, Lx // 2 + 1) in the Fourier
+        domain used to compute cross-correlation with each frame
     maskMul : torch.Tensor
         Multiplicative mask applied to `frames` before correlation. Broadcasted over frames.
     maskOffset : torch.Tensor
@@ -87,15 +87,14 @@ def phasecorr(frames, cfRefImg, maskMul, maskOffset, maxregshift, smooth_sigma_t
     """
 
     device = frames.device
-    data = (frames.float() * maskMul + maskOffset).type(torch.complex64)
+    data = frames.float() * maskMul + maskOffset
     min_dim = min(data.shape[1], data.shape[2])  # maximum registration shift allowed
     lcorr = int(np.minimum(np.round(maxregshift * min_dim), min_dim // 2))
 
     data = convolve(data, cfRefImg)
-    cc = torch.cat((torch.cat((data[:, -lcorr:, -lcorr:], data[:, -lcorr:, :lcorr + 1]), axis=2),   
+    cc = torch.cat((torch.cat((data[:, -lcorr:, -lcorr:], data[:, -lcorr:, :lcorr + 1]), axis=2),
                     torch.cat((data[:, :lcorr + 1, -lcorr:], data[:, :lcorr + 1, :lcorr + 1]), axis=2)), axis=1)
-    cc = torch.real(cc)
-
+
     cc = temporal_smooth(cc, smooth_sigma_time) if smooth_sigma_time > 0 else cc
 
     imax = torch.stack([torch.argmax(cc[t]) for t in range(data.shape[0])], dim=0)

suite2p/registration/utils.py

@@ -9,37 +9,42 @@
 
 try:
     # pytorch > 1.7
-    from torch.fft import fft, fft2, ifft, ifft2, fftshift, ifftshift
+    from torch.fft import (fft, fft2, ifft, ifft2, rfft2, irfft2, fftshift,
+                           ifftshift)
 except:
     # pytorch <= 1.7
     raise ImportError("pytorch version > 1.7 required")
 
-eps = torch.complex(torch.tensor(1e-5), torch.tensor(0.0))
 
 def convolve(mov: np.ndarray, img: np.ndarray) -> np.ndarray:
     """
     Convolve a 3D frame sequence by a 2D image in the Fourier domain using phase-correlation.
 
     Applies FFT to each frame, normalizes by magnitude, multiplies by `img`, and returns the
-    inverse FFT (real part).
+    inverse FFT.
+
+    `mov` is real-valued, so the real-input FFT is used throughout: the spectrum is
+    Hermitian and the redundant half carries no information. `img` must therefore be a
+    half-spectrum kernel (see `ref_smooth_fft`).
 
     Parameters
     ----------
     mov : torch.Tensor
-        Input frames of shape (nImg, Ly, Lx).
+        Real-valued input frames of shape (..., Ly, Lx).
     img : torch.Tensor
-        2D complex-valued convolution kernel of shape (Ly, Lx), typically a conjugate FFT
-        of a reference image.
+        Complex-valued half-spectrum convolution kernel of shape (..., Ly, Lx // 2 + 1),
+        typically the conjugate FFT of a reference image.
 
     Returns
     -------
     convolved_data : torch.Tensor
-        Real-valued convolution result of shape (nImg, Ly, Lx).
+        Real-valued convolution result of shape (..., Ly, Lx).
     """
-    mov = fft2(mov)
-    mov /= (eps + torch.abs(mov))
+    Ly, Lx = mov.shape[-2], mov.shape[-1]
+    mov = rfft2(mov)
+    mov /= (1e-5 + torch.abs(mov))
     mov *= img
-    mov = torch.real(ifft2(mov))
+    mov = irfft2(mov, s=(Ly, Lx))
     return mov
 
 def spatial_taper(sig, Ly, Lx):
@@ -62,11 +67,11 @@ def spatial_taper(sig, Ly, Lx):
     Returns
     -------
     maskMul : torch.Tensor
-        Floating-point multiplicative mask of shape (Ly, Lx), with values near 1.0
+        Float32 multiplicative mask of shape (Ly, Lx), with values near 1.0
         in the center and smoothly decaying to 0.0 at the edges.
     """
-    y = torch.arange(0, Ly, dtype=torch.double)
-    x = torch.arange(0, Lx, dtype=torch.double)
+    y = torch.arange(0, Ly, dtype=torch.float32)
+    x = torch.arange(0, Lx, dtype=torch.float32)
     x = (x - x.mean()).abs()
     y = (y - y.mean()).abs()
     mY = ((Ly - 1) / 2) - 2 * sig
@@ -186,13 +191,15 @@ def ref_smooth_fft(refImg: np.ndarray, smooth_sigma=None) -> np.ndarray:
     Returns
     -------
     cfRefImg : torch.Tensor
-        Complex64 tensor of shape (Ly, Lx) containing the smoothed, whitened
-        complex-conjugate FFT of the reference image.
+        Complex64 tensor of shape (Ly, Lx // 2 + 1) containing the smoothed, whitened
+        complex-conjugate FFT of the reference image. Only the non-redundant half of
+        the spectrum is returned, matching the real-input FFT used by `convolve`.
     """
-    cfRefImg = complex_fft2(img=refImg)
+    Ly, Lx = refImg.shape[-2], refImg.shape[-1]
+    cfRefImg = torch.conj(rfft2(refImg.float()))
     cfRefImg /= (1e-5 + torch.abs(cfRefImg))
     if smooth_sigma is not None:
-        cfRefImg *= gaussian_fft(smooth_sigma, cfRefImg.shape[0], cfRefImg.shape[1])
+        cfRefImg *= gaussian_fft(smooth_sigma, Ly, Lx)[:, :Lx // 2 + 1]
     return cfRefImg.type(torch.complex64)
 
 def kernelD(xs: np.ndarray, ys: np.ndarray, sigL: float = 0.85) -> np.ndarray:

tests/test_registration.py

@@ -74,4 +74,67 @@ def test_transform_data_mps_cpu_consistency():
     max_diff = np.abs(cpu_np - mps_np).max()
 
     assert correlation > 0.99, f"Correlation: {correlation}"
-    assert max_diff < 2, f"Max diff: {max_diff}"
+    assert max_diff < 2, f"Max diff: {max_diff}"
+
+def test_convolve_matches_full_spectrum_reference():
+    """convolve uses a real-input FFT; check it against a full complex-FFT reference.
+
+    The phase-correlation spectrum of two real images is Hermitian, so the half
+    spectrum carries the same information. This pins that equivalence.
+    """
+    from suite2p.registration.utils import convolve, ref_smooth_fft
+
+    np.random.seed(0)
+    Ly, Lx = 64, 64
+    mov = torch.from_numpy(np.random.rand(4, Ly, Lx).astype(np.float32))
+    ref = torch.from_numpy(np.random.rand(Ly, Lx).astype(np.float32))
+
+    got = convolve(mov.clone(), ref_smooth_fft(ref, smooth_sigma=1.15))
+
+    # reference: full complex spectrum, as suite2p computed it before
+    cf_full = torch.conj(torch.fft.fft2(ref))
+    cf_full /= (1e-5 + torch.abs(cf_full))
+    from suite2p.registration.utils import gaussian_fft
+    cf_full *= gaussian_fft(1.15, Ly, Lx)
+    m = torch.fft.fft2(mov.clone().type(torch.complex64))
+    m /= (1e-5 + torch.abs(m))
+    m *= cf_full.type(torch.complex64)
+    expected = torch.real(torch.fft.ifft2(m))
+
+    assert got.shape == expected.shape
+    assert got.dtype == torch.float32
+    assert torch.allclose(got, expected, atol=1e-5)
+
+
+def test_ref_smooth_fft_returns_half_spectrum():
+    from suite2p.registration.utils import ref_smooth_fft
+
+    ref = torch.from_numpy(np.random.rand(64, 48).astype(np.float32))
+    cf = ref_smooth_fft(ref, smooth_sigma=1.15)
+    assert cf.shape == (64, 48 // 2 + 1)
+    assert cf.dtype == torch.complex64
+
+
+def test_rigid_phasecorr_recovers_known_shifts():
+    from suite2p.registration import rigid
+
+    np.random.seed(0)
+    Ly = Lx = 128
+    ref = np.random.rand(Ly, Lx).astype(np.float32) * 500
+    shifts = [(0, 0), (3, -2), (-5, 4)]
+    frames = np.stack([np.roll(ref, s, (0, 1)) for s in shifts])
+
+    maskMul, maskOffset, cfRefImg = rigid.compute_masks_ref_smooth_fft(
+        torch.from_numpy(ref), maskSlope=3.45, smooth_sigma=1.15)
+    ymax, xmax, cmax, _ = rigid.phasecorr(
+        torch.from_numpy(frames), cfRefImg, maskMul, maskOffset,
+        maxregshift=0.1, smooth_sigma_time=0)
+
+    assert [(int(y), int(x)) for y, x in zip(ymax, xmax)] == shifts
+
+
+def test_spatial_taper_is_float32():
+    """maskMul must stay float32: a float64 mask would upcast the FFT to complex128."""
+    from suite2p.registration.utils import spatial_taper
+
+    assert spatial_taper(3.45, 64, 64).dtype == torch.float32