Is it possible to do fftshift in 0.4.0

Hi guys, the new release 0.4.0 of pytorch is available. The new feature fft is equiped by torch.fft. As far as I know, there is not torch.fftshift to shift the zero frequency to the center of planes. Did I miss something? Any suggestion about how to do fftshift?
Thanks a lot.
Regards.

There is not torch.fftshift yet, and SsnL (issue3775) is working on it. (And you can reference fft_extend and fft_extend_issue9 to achieve it by yourself)

1 Like

A workaround could be to element wise multiply each location by (-1)^(x+y) before doing the FFT, and after the IFFT. The -1 or +1 multiplications should cancel each other out in the back propagation I believe.

import numpy as np
from scipy import misc

def shift(img):
    tmpImg = img.copy()
    n,m,p = tmpImg.shape
    x = np.arange(0, m, 1)
    y = np.arange(0, n, 1)
    xm, ym  = np.meshgrid(x,y)
    shifter = (-1)**(xm + ym)
    shifter = np.expand_dims(shifter,2)
    shifter = np.repeat(shifter,3,2)
    tmpImg = np.multiply(tmpImg,shifter)

    return tmpImg

img = misc.face().astype(float)

tmpImg = shift(img)
altShift = np.fft.fft2(tmpImg,axes=(0,1))
npShift = np.fft.fftshift(np.fft.fft2(img,axes=(0,1)),axes=(0,1))

print(np.allclose(altShift,npShift))

altIFFT_Shift = np.fft.ifft2(altShift,axes=(0,1))
npIFFT_Shift = np.fft.ifft2(np.fft.ifftshift(npShift,axes=(0,1)),axes=(0,1))
altIFFT_Shift = shift(altIFFT_Shift)

print(np.allclose(altIFFT_Shift,npIFFT_Shift))

I have written fftshift and ifftshift operations as part of my Complex Steerable Pyramid implementation: https://github.com/tomrunia/PyTorchSteerablePyramid

(see the steerable/math_utils.py file)

1 Like