Search code examples
pythonarraysnumpyfftimshow

Displaying a spatio-temporal spectrum with imshow: what order are the frequencies in after taking an FFT down columns?

---------- Summary ----------

If I have 2D data psi_hat_kxt and I take its FFT down columns and FFT-shift the result, np.fft.fftshift(np.fft.fft(psi_hat_kxt, axis=0)), where are the negative frequencies located? Shouldn't they be in the top half of the resulting array? If so, when I plot the result (squared element-wise to get real data) using imshow, is there an implicit up-down flip somewhere?

---------- Detailed ----------

I have a simulation that evolves a field in spatial Fourier space (k-space) and I want to show the frequency content for each k. That is to say, I want to plot the spatio-temporal spectrum, a.k.a. k-omega plot, by taking the FFT in the time direction, squaring, and plotting. I am doing all this in NumPy,

import numpy as np
from matplotlib import pyplot as plt
# other code that assigns variables like Nt, Nx, Deltat, dx etc.

My data [Edit: which is complex-valued] is arranged as rows containing the k-space data, and time evolves moving down each row. I read this in from an external binary file and then reshape:

fname = open('../output/ky0slices.bin','rb')
psi_hat_kxt_vec = np.fromfile(fname, dtype=np.complex_)
fname.close()
psi_hat_kxt  = np.reshape(psi_hat_kxt_vec , (Nt,Nx))

I then do the FFT down the columns, shift, and square to get a real number:

komega_spec = np.abs( np.fft.fftshift (np.fft.fft(psi_hat_kxt, axis=0)) )**2.0

Finally I plot using imshow:

om_ax = 2*np.pi * np.fft.fftshift( np.fft.fftfreq(Nt,d=Deltat) )
k_ax  = 2*np.pi * np.fft.fftshift( np.fft.fftfreq(Nx,d=dx) )
log_komega_spec = np.log(komega_spec)

extnt=[k_ax[0], k_ax[-1] , om_ax[0], om_ax[-1] ]

fig, ax = plt.subplots()
im = plt.imshow(log_komega_spec, extent=extnt , aspect='auto')

In the end I get an image that looks correct, correctly oriented k-omega plot

But I don't understand why it is actually correct.

Namely, from reading the docs I thought that after the fftshift, the negative temporal frequencies should have their Fourier coefficients in the top rows of the fft data, i.e. komega_spec[0,:] should contain a row with all the Fourier coefficients corresponding to frequency -Nt/2.

But from the shape of the plot it appears that this row corresponds to the positive frequency Nt/2-1. (It seems this way because the parabola is convex, as it should be for physical reasons, please ignore the omega-axis ticks as they are controlled by extnt.)

Is imshow maybe doing an implicit flipud?

In short: why does it seem that the top half of komega_spec contains the positive frequency data?

Solution

  • The answer is that we need to make the transformation omega -> - omega to convert between the physical omega = k^2 of the dispersion relation (in the sense that we expand the function in plane waves psi(x,t) = Sum psi_hat(k, omega) * exp[i(k*x-omega*t)]
    ) and the omega that is the argument of what the FFT spits out (in the sense that the FFT calculates F(omega) = Sum f(t) * exp[ +i(omega*t) ]

    So the physical dispersion relation should indeed show up in the "negative frequencies" as far as the FFT is concerned..