ZoomFFT for speeding up off-axis computation (and more)

Sébastien M Popoff¹  and Rodrigo Gutiérrez-Cuevas¹

When performing the computation for some tasks such as off-axis holography, we often have to compute the entire FFT of a signal or an image while being only interested in a very small part of the spectrum. The rest of the information is just discarded. While the FFT algorithm and its implementation in standard computing libraries are very efficient, we can still take advantage of a slower approach which only computes what we need. The ZoomFFT algorithm does just that! And it's already available in standard packages such as SciPy for Python or in MATLAB. Using off-axis holography as an example, I will show how to save time – or not – using ZoomFFT.

What is the issue?

In many situations, and in particular for off-axis holography (see tutorial here), we perform FFTs that compute the whole spectral range while we are only interested in a narrow range of the spatial frequencies. As an illustration, I show in Fig.1 a typical FFT on an interferogram obtained when doing off-axis holography (left) and the mask representing the frequencies we are interested in (right). The rest is discareded, so most of the computation is done for nothing.

Figure 1. Left, the absolute value (in log scale) of the FFT of an typical interferogram obtained when performing off-axis holography (simulations). Right, the frequency window we select, the rest is useless and discarded.

What is ZoomFFT?

The ZoomFFT implementation in the scipy.signal package is presented this way:

This is a specialization of the chirp z-transform (CZT) for a set of equally spaced frequencies around the unit circle, used to calculate a section of the FFT more efficiently than calculating the entire FFT and truncating.

For repeated use of the same computation, we can create a callable that will be later used to compute the FFT. The difference with the standard FFT is that you specify the range of frequencies you want to compute. In addition, you have to specify the size of the input—the data you will feed the algorithm—and the size of the output, which will define your resolution in the Fourier domain.

transform = ZoomFFT(len(x), [f1, f2], len(x))

We can then call it to perform the FFT using

fft_of_f = transform(f)

ZoomFFT for 2D arrays

Note: This section explains how to use ZoomFFT for 2D arrays similarly to the standard FFT in Python. I have already written a class for you; if you are not interested, you can skip to the next section.

Cool! However, ZoomFFT is not really designed for 2D images. Although this is not a big deal, there are some precautions to take.

First, it deals with only one dimension, so we need to perform the FFT once for each of the two dimensions:

nx, ny = img.shape[0]
fftx = ZoomFFT(nx, nx, fn = [fx_min, fx_max])
ffty = ZoomFFT(ny, ny, fn = [fy_min, fy_max])
FFT = f(img, axis = 0)
FFT = f(img, axis = 1)

If your image is square and the spatial frequency range is the same for both dimensions, you can use the same ZoomFFT object for both axes.

Another important difference with the DFT is that ZoomFFT does not know about negative positions. That means the position 0 is always the beginning of the dimension of the array you use ZoomFFT on. You cannot directly indicate that the center of the image is the origin of your coordinate system (with the standard `np.fft` package, this would be done using `fftshift`).

The consequence is that when you perform the FFT of an image with only a spot at the center, the result of ZoomFFT is not a constant-phase image, but exhibits a phase slope. When using FFT for optics, we usually consider that the image is centered around the optical axis and then we expect a spot in the center to give rise to a constant-phase plane wave in the Fourier plane.

No big problem, since the expected result is just multiplied by a phase slope which can be easily removed. One way would be to compute this phase slope and remove it. But it is easier, and more flexible, to simply calibrate it by computing the FFT of a spot at the position that we want to be the origin of our coordinate system, storing the phase slope, multiplying it by -1, and then removing it from any subsequent computation of the FFT. An advantage is that you can set the reference of your coordinate system (i.e. the position of your optical axis) to any position, not necessarily at the center of the image, to fit your experimental data.

We can create a class to do that (it is simplified here, but a more complete code is available here ).

from scipy.signal import ZoomFFT
import numpy as np

class ZoomFFT2D:
    def __init__(self, n, m, f_min, f_max):
        nx, ny = n
        mx, my = m
        self.n = n
        self.m = m
        self.f_min = f_min
        self.f_max = f_max
        self.f1 = ZoomFFT(nx, m=mx, fn=[f_min[0], f_max[0]])
        self.f2 = ZoomFFT(ny, m=my, fn=[f_min[1], f_max[1]])
        self.ref = None
        self._get_phase_ref()

    def _get_phase_ref(self):
        foc = np.zeros(self.n)

        foc[self.n[0] // 2, self.n[1] // 2] = 1

        ref = self(foc)
        self.ref = np.exp(-1j * np.angle(ref))

    def __call__(self, A):
        FFT = self.f1(A, axis=-2)
        FFT = self.f2(FFT, axis=-1)
        if self.ref is not None:
            return FFT * self.ref
        else:
            return FFT

When initialized, we create the ZoomFFT callables for the `x` and `y` direction FFT. Then we call it once for a single spot image and store the calibration phase. When called again by the user, the reference phase is used to remove the additional phase slope. Here is a simple usage example:

N = 100
fx_min, fx_max = -0.25,0.25
fy_min, fy_max = -0.15,0.35
img = np.random.randn(N,N)
fft = ZoomFFT2D((N,N), (N,N), (fx_min,fy_min), (fx_max,fy_max))
FT_img = fft(img)

Install ZoomFFT2D

You can install the simple module I wrote using pip with

pip install git+https://github.com/wavefrontshaping/zoomfft2d.git

then import it with

from zoomfft2d import ZoomFFT2D

Does it save time? Test with Off-axis holography

The answer to this question seems to be heavily reliant on your hardware, the size of the images, and the frequency range you want to recover. I wrote some simple pieces of code to simulate the off-axis procedure and compare standard FFT with ZoomFFT for different image resolutions. I will not go over the code, but the full Jupyter Notebook is available here. You will see that ZoomFFT is not always faster but can give you a significant boost under certain conditions. Moreover, it simplifies the procedure and thus the code.

The Off-axis procedure

I will not go into the details as you can look at this tutorial on the off-axis holography technique, but basically, the steps are as follows:

- Perform the FFT of the interferogram,
- Select the window in the FFT corresponding to the desired order,
- Recenter it to zero frequency,
- Do the inverse FFT.

Using np.fft.fft2

Using the standard `np.fft.fft2`, the function performing this operation looks like this

def do_offaxis_FFT(A, mask, f_center, fft_size):
Fh = np.fft.fftshift(np.fft.fft2(A, s = fft_size).astype(complex))
## Apply mask
Fh_masked = Fh*mask
offset = [-int(f*n) for f,n in zip(f_center,fft_size)]
# Recenter the FFT
h_recentered = np.roll(Fh_masked, offset[0], axis=0)
Fh_recentered = np.roll(Fh_recentered, offset[1], axis=1) 
## Inverse FFT
fft_recovered = np.fft.ifft2(np.fft.ifftshift(Fh_recentered))
return fft_recovered[:A.shape[0], :A.shape[1]]

A is your interferogram image, mask is a boolean image that is equal to 1 in the region you want to select and 0 elsewhere, f_center is the central position of the region to select in the Fourier plane, and fft_size is the resolution of the image in the Fourier plane.

Note that in the general case, the resolution in the FFT plane may not be good enough; that is why we have to zero-pad your image first. This is done by specifying a larger resolution in the Fourier plane using np.fft.fft2(A, s=fft_size). After the inverse FFT, you will end up with a larger image we have to crop it using fft_recovered[:A.shape[0], :A.shape[1]].

Using ZoomFFT

ftzoom = ZoomFFT(...)

ft_small = ZoomFFT2D(
    img.shape, 
    window_resolution,
    f_center = f_center,
    f_range = f_range,
    )

ift = ZoomFFT2D(
    window_resolution,
    img.shape,
    f_center = 0.,
    f_range = 1., 
    direction='backward'
    )

def do_offaxis_FFT_zoom(A):
    ft_img = ftzoom(A)
    recovered_complex_pattern = ift(ft_img)
    return recovered_complex_pattern

We first define the objects for the direct and inverse Fourier transforms. Once done, the code is quite simpler as we already select the desired window by choosing the range of frequencies in the Fourier domain, so there is no need for a mask nor to recenter the spectrum. We then perform the inverse Fourier transform and voilà.

Speed tests

To compare the two approaches on the same basis, we make sure the `fft_size` given to the do_offaxis_FFT() function is set so that the selected window is equal to the window_resolution used in do_offaxis_FFT_zoom().

We compute the average computation time for the two techniques using different parameters of 

• the frequency range f_range, corresponding to the fraction of the range frequencies selected compared to the total range computed by the FFT,

• the image resolution N, 

• and of the resolution in the Fourier domain window_size.

We first fix f_range to 0.075, i.e. in each direction we select 7.5% of the frequencies. In 2D, this represents selecting only ~0.5% of the whole spectrum.

We then fix f_range to 0.15, i.e. in each direction we select 15% of the frequencies. In 2D, this represents selecting ~2.5% of the whole spectrum.

First, we see that in almost all configurations, the ZoomFFT approach is faster. As expected, the gain in time increases when the fraction of selected frequencies (f_range) decreases. 

We also observe that for other paremeters fixed, the speed enhancement of ZoomFFT become smaller when N increases, which I guess is due to the optimized performance of the DFT algotithm for large images.

Interestingly, for other parameters fixed, the speed enhancement gets better when the resolution of the FFT windows gets bigger. This is explained by the fact that the standard FFT has to compute the whole spectral range, so for the full size of the FT image which is quite huge, while most of it is discarded.

While these results are not to be taken for a fact, as it seems to depend heavily on your machine, ZoomFFT does save singificant time in our experiments.

Full code