Numpy fft2 frequency shift

Question

I don't understand how to make frequency shift in fft2 or higher dimensions. Could anyone explain it, please?

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

%matplotlib inline

numb_steps = 100
x = np.linspace(-10, 10, numb_steps)
y = np.linspace(-10, 10, numb_steps)

X,Y = np.meshgrid(x,y)

R = 1

Z = np.zeros((100, 100))
# circle
for i in range(numb_steps):
    for j in range(numb_steps):
        if x[i]**2+y[j]**2 > R:
            Z[i,j] = 0
        else:
            Z[i,j] = 1



fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_surface(X,Y,Z)

I only can get those plots

fig = plt.figure()
ax = fig.gca(projection='3d')

Z_fft = np.fft.fft2(Z)
FreqCompRows = np.fft.fftfreq(Z.shape[0],d=2)
FreqCompCols = np.fft.fftfreq(Z.shape[1],d=2)
FreqCompRows = np.fft.fftshift(FreqCompRows)
FreqCompCols = np.fft.fftshift(FreqCompCols)

S,D = np.meshgrid(FreqCompRows, FreqCompCols)

ax.plot_surface(S, D, np.abs(Z_fft))

plt.imshow(np.abs(Z_fft))

plt.xticks(range(len(x_ax)), np.round(x_ax))
plt.yticks(range(len(y_ax)), np.round(y_ax))

P.S. Yes, the ticks overlap is my problem as well

Answer 1

You can solve this using scipy.fftpack (sfft) instead of np.fft, because the sfft implementation can be directly used on 2-dimensionnal arrays so you don't have to do it in a convoluted way (ba dum tsss).

In your code header, add :

import scipy.fftpack as sfft

Then, to calculate the fft and shift the spectrum, use :

Z_fft = sfft.fft2(Z)
Z_shift = sfft.fftshift(Z_fft)

The obtained spectrum is then nicely arranged for image display :

plt.figure(4)
plt.imshow(np.abs(Z_shift))

Also, the way you are constructing the circle seems overly complicated, you can take advantage of python's syntax using boolean syntax :

X,Y = np.meshgrid(x,y)
Z = X**2+Y**2
R = 1
Z[Z > R] = 0;Z[Z > 0] = 1;

What I did here is using the meshgrid you created to calculate the circle equation result, in the form of Z. Then I test Z for all values above R, put them to zero, and the other values to 1.

My code is below :

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import scipy.fftpack as sfft # you have to import this package

plt.close('all')

numb_steps = 100
x = np.linspace(-10, 10, numb_steps)
y = np.linspace(-10, 10, numb_steps)

#constructing the disk in a simpler manner
X,Y = np.meshgrid(x,y)
Z = X**2+Y**2
R = 1
Z[Z > R] = 0;Z[Z > 0] = 1;

# Here, use sfft routines instead of np.fft
Z_fft = sfft.fft2(Z)
Z_shift = sfft.fftshift(Z_fft)
FreqCompRows = np.fft.fftfreq(Z.shape[0],d=2)
FreqCompCols = np.fft.fftfreq(Z.shape[1],d=2)
FreqCompRows = np.fft.fftshift(FreqCompRows)
FreqCompCols = np.fft.fftshift(FreqCompCols)

S,D = np.meshgrid(FreqCompRows, FreqCompCols)

# plotting your disk signal image
fig = plt.figure(1) 
ax = fig.gca(projection='3d')
ax.plot_surface(X,Y,Z)

# plotting "3D" fft non-shifted
fig = plt.figure(2)
ax = fig.gca(projection='3d')
ax.plot_surface(S, D, np.abs(Z_fft))

# plotting images of FFT's
plt.figure(3)
plt.imshow(np.abs(Z_fft))

plt.figure(4)
plt.imshow(np.abs(Z_shift))