如何用NumPy在傅里叶空间平移奇数尺寸图像？代码报错求助

我正在开发一个可在傅里叶空间平移图像的函数，要求同时支持偶数和奇数尺寸的图像。目前该代码对偶数尺寸图像运行正常，但对奇数尺寸图像失效。点源在y轴的位置符合预期，但x轴位置不符合。我推测这与使用rfft2生成的傅里叶图像形状（N, N/2+1）有关，但不知如何修复。

现有代码 #

def phase_shift_2d(fourier_image_data: np.ndarray, x0: float, y0: float):
    # Shift image to position x0, y0 in Fourier space.
    image_shape = fourier_image_data.shape

    if image_shape[0] != image_shape[1]:
        u_data = np.arange(0, image_shape[1])
    else:
        u_data = np.arange(-image_shape[1] // 2, image_shape[1] // 2)

    v_data = np.arange(-image_shape[0] // 2, image_shape[0] // 2)

    u, v = np.meshgrid(u_data, v_data)

    ku =  u / image_shape[1]
    kv =  v / image_shape[0]

    shifted_fimage = fourier_image_data * np.exp(-2. * 1j * np.pi * (ku * x0 + kv * y0))

    return shifted_fimage

import numpy as np
N=256

image = np.zeros((N,N), dtype=float)
image[N//2,N//2] = 1.0
image[N//2+35,N//2+35] = 0.5


fourier_image = np.fft.fftshift(np.fft.rfft2(np.fft.ifftshift(image, axes=(-2,-1))), axes=-2)

x0_shift = np.random.randint(low=-N//2, high=N//2, size=1, dtype=int)
y0_shift = np.random.randint(low=-N//2, high=N//2, size=1, dtype=int)

shifted_fourier_image = phase_shift_2d(fourier_image, x0_shift[0], y0_shift[0])

shifted_image = np.fft.fftshift(np.fft.irfft2(np.fft.ifftshift(shifted_fourier_image, axes=-2)), axes=(-2,-1))

final_y, final_x = np.unravel_index(np.argmax(shifted_image.real), shifted_image.real.shape)[0], np.unravel_index(np.argmax(shifted_image.real), shifted_image.real.shape)[1]

assert final_x == x0_shift + N//2
assert final_y == y0_shift + N//2

运行错误 #

AssertionError                            Traceback (most recent call last)
Cell In[78], line 1
----> 1 assert final_x == x0_shift + N//2
      2 assert final_y == y0_shift + N//2

AssertionError: 

内容的提问来源于stack exchange，提问作者Miguel Cárcamo