PyTorch's FFT not preserving Linearity

user_2407438 0
Viewed 100

I'm not sure if this is a bug or I'm misunderstanding something about the indexing/shifting. When I try to calculate the difference in the Fourier domain, something is going wrong in the calculation. For a simple case, consider a identity kernel K, such that for any signal X it follows that X - K*X = 0. Further, x - k o x = 0. However, when I compute the residual in the Fourier domain using pytorch, I don't get a numerically 0 solution.

import torch
from torch.fft import rfftn, irfftn
from PIL import Image
from torchvision.transforms import ToTensor
from torch.nn.functional import pad
from matplotlib import pyplot as plt

# Let X be an input image.
with Image.open("barbara.png") as im:
    X = ToTensor()(im)[0,...]

# Let K be a convolutional identity kernel.
K = torch.zeros((3,3))
K[1,1] = 1.0

# Full convolution shape.
cnvShape = (X.shape[0] + K.shape[0] - 1, X.shape[1] + K.shape[1] - 1)

# Apply DFT. Uncentered zero-padding to full shape.
Xhat = rfftn(X, s=cnvShape, dim=(0,1))
Khat = rfftn(K, s=cnvShape, dim=(0,1))

# Calculate convolution via Fourier domain.
XKhat = Xhat * Khat
XK = irfftn(XKhat, s=cnvShape)[1:-1,1:-1]
# Difference in spatial domain. Is numerically close to 0.
R = X - XK

# Difference in Fourier domain. This should be numerically close to 0.
Shat = Xhat - XKhat
S = irfftn(Shat)

plt.subplot(1,2,1)
plt.imshow(R)
plt.colorbar()
plt.subplot(1,2,2)
plt.imshow(S)
plt.colorbar()
plt.show()

# Center zeropad image FFT, produces numerically 0 solution.
That = rfftn(pad(X, (1,1)*2), s=cnvShape) - XKhat
T = irfftn(That)
plt.imshow(T)
plt.colorbar()
plt.show()

This results in Difference between the signal and the convolved signal. On the left is taking the difference in spatial domain. On the right is the Fourier domain. Using a centered zeropad around the signal (y) such that y - k * x approx 0.

So from the last plot, I can see there's some sort of shifting going on with operations in the Fourier domain, but it's not intuitive what. I'd also like to figure out how to implement consistent operations in the Fourier domain, as this is just a small sample of what I need to compute. I'm aware of fftshift, but that doesn't seem to be appropriate here, as the shift doesn't seem to be with negative and positive frequencies. So, what is going on in the indexing?

pytorch fft python
0 Answers
Related