Zoom in on np.fft2 result

Question

Is there a way to chose the x/y output axes range from np.fft2 ? I have a piece of code computing the diffraction pattern of an aperture. The aperture is defined in a 2k x 2k pixel array. The diffraction pattern is basically the inner part of the 2D FT of the aperture. The np.fft2 gives me an output array same size of the input but with some preset range of the x/y axes. Of course I can zoom in by using the image viewer, but I have already lost detail. What is the solution?

Thanks, Gert

import numpy             as np
import matplotlib.pyplot as plt

r= 500
s= 1000

y,x = np.ogrid[-s:s+1, -s:s+1]
mask = x*x + y*y <= r*r
aperture = np.ones((2*s+1, 2*s+1))
aperture[mask] = 0

plt.imshow(aperture)
plt.show()

ffta= np.fft.fft2(aperture)

plt.imshow(np.log(np.abs(np.fft.fftshift(ffta))**2))
plt.show()

Answer 1

Unfortunately, much of the speed and accuracy of the FFT come from the outputs being the same size as the input.

The conventional way to increase the apparent resolution in the output Fourier domain is by zero-padding the input: np.fft.fft2(aperture, [4 * (2*s+1), 4 * (2*s+1)]) tells the FFT to pad your input to be 4 * (2*s+1) pixels tall and wide, i.e., make the input four times larger (sixteen times the number of pixels).

Begin aside I say "apparent" resolution because the actual amount of data you have hasn't increased, but the Fourier transform will appear smoother because zero-padding in the input domain causes the Fourier transform to interpolate the output. In the example above, any feature that could be seen with one pixel will be shown with four pixels. Just to make this fully concrete, this example shows that every fourth pixel of the zero-padded FFT is numerically the same as every pixel of the original unpadded FFT:

# Generate your `ffta` as above, then
N = 2 * s + 1
Up = 4
fftup = np.fft.fft2(aperture, [Up * N, Up * N])

relerr = lambda dirt, gold: np.abs((dirt - gold) / gold)
print(np.max(relerr(fftup[::Up, ::Up] , ffta))) # ~6e-12.

(That relerr is just a simple relative error, which you want to be close to machine precision, around 2e-16. The largest error between every 4th sample of the zero-padded FFT and the unpadded FFT is 6e-12 which is quite close to machine precision, meaning these two arrays are nearly numerically equivalent.) End aside

Zero-padding is the most straightforward way around your problem. But it does cost you a lot of memory. And it is frustrating because you might only care about a tiny, tiny part of the transform. There's an algorithm called the chirp z-transform (CZT, or colloquially the "zoom FFT") which can do this. If your input is N (for you 2*s+1) and you want just M samples of the FFT's output evaluated anywhere, it will compute three Fourier transforms of size N + M - 1 to obtain the desired M samples of the output. This would solve your problem too, since you can ask for M samples in the region of interest, and it wouldn't require prohibitively-much memory, though it would need at least 3x more CPU time. The downside is that a solid implementation of CZT isn't in Numpy/Scipy yet: see the scipy issue and the code it references. Matlab's CZT seems reliable, if that's an option; Octave-forge has one too and the Octave people usually try hard to match/exceed Matlab.

But if you have the memory, zero-padding the input is the way to go.

Answer 2

Python's scipy module has a czt function. You can simply use czt as fast dft in your case. However it is 1d only. To perform 2d dft, I had to implement my own method. It's basically dft along x, and y sequentially. Here is a runnable code sample.

import numpy as np
import scipy
import matplotlib.pyplot as plt


def get_blob_pattern(size=100, sigma=30):
    # get an image of spherical aperture with smooth edge
    radius = 0.5 * size
    x_grid, y_grid = np.meshgrid(np.linspace(-radius, radius, size), np.linspace(-radius, radius, size))
    dist = (x_grid**2 + y_grid**2) ** 0.5
    img = 1 / (1 + np.exp(np.clip(3 * (dist - sigma), -20, 20)))
    return img


def get_dft2d(data, m, kmax):
    # perform 2d dft on image (using scipy 1d czt), get a 2d spectrum with frequency range of (-kmax, +kmax) on both axis
    # m: spectrum bin number
    w, h = data.shape
    aw = 2j * kmax / ((m - 1) * w)
    f1 = scipy.signal.czt(data, m=m, w=np.exp(aw), a=np.exp(0.5 * (m - 1) * aw), axis=0)
    data1 = f1 * np.exp(-0.5j * kmax * np.linspace(-1, 1, m))[:, np.newaxis]
    ah = 2j * kmax / ((m - 1) * h)
    f2 = scipy.signal.czt(data1, m=m, w=np.exp(ah), a=np.exp(0.5 * (m - 1) * ah), axis=1)
    data2 = f2 * np.exp(-0.5j * kmax * np.linspace(-1, 1, m))[np.newaxis, :]
    return data2


img = get_blob_pattern(size=100, sigma=20)
plt.imshow(img, vmin=0, vmax=1)
plt.show()

img_f = get_dft2d(img, 500, 100)
plt.imshow(np.abs(img_f) ** 2, vmin=0, vmax=10000)
plt.show()

the original round aperture and diffraction pattern as an airy disk.

I encountered this exact problem yesterday and had a hard time finding the solution, eventually having to do it myself. Sad that no answer in 8 years. Hope it would help some physics student like me in the future.