80 lines
3.3 KiB
Python
80 lines
3.3 KiB
Python
# P372-PS14-FFT
|
|
# Matthew Oros
|
|
"""
|
|
This script introduces the Real FFT and Inverse Real FFT capabilities.
|
|
"""
|
|
|
|
import numpy as np
|
|
import pylab as pl
|
|
import numpy.fft as fft
|
|
|
|
# Some constants, I left these all the same
|
|
numPoints = 1000 # number of points in x-waveform, how much resolution
|
|
numFreqs = numPoints / 2 + 1 # (numPoints/2 -1) complex & 2 real k-amplitudes
|
|
# This is still same numPoints of bits of information!
|
|
# The wavenumber k's are also called "spatial frequencies"
|
|
# We are going to look at the region between 0 and 2 in x-space,
|
|
# but suppose we don't want anything to repeat within twice this, so
|
|
windowLength = 4.0 # effective window length
|
|
dx = windowLength / numPoints # Step size in x-space
|
|
x = np.arange(0.0, windowLength, dx) # All the x-positions
|
|
|
|
# Now prepare the spectrum and k wavenumbers
|
|
dk = 2 * np.pi / windowLength # Step size in k-space
|
|
ks = np.arange(numFreqs) * dk # All the k's, highest k found at ks[-1]
|
|
|
|
# HERE IS WHERE THE REAL WORK GETS DONE, the above is all just setting up
|
|
# Choose the <k>
|
|
k0 = 0.1 * ks[-1] # I chose 0.1 to have lower frequency inside the wave packet
|
|
alpha = 20 # 0.1 # I chose 0.1 to make the wave packet wider
|
|
# Now define the k-spectrum amplitudes
|
|
# phi_k = np.exp(-0.5 * alpha * (ks - k0) ** 2) # the spectral amplitudes (complex)
|
|
# Upside-down parabola!
|
|
phi_k = np.piecewise(ks,
|
|
[(ks - k0 < -alpha), (-alpha <= ks - k0) & (ks - k0 <= alpha), (ks - k0 > alpha)],
|
|
[0, lambda k: (1 / alpha ** 2) * (k - k0) ** 2, 0])
|
|
|
|
# And generate the wavefunction from phi(k) using Inverse Real FFT
|
|
psi_x = fft.irfft(phi_k) # inverse real FFT gets the real wavefunction
|
|
|
|
# I will now normalize the wave function,
|
|
# which was not done in the original script
|
|
norm_factor = 1 / np.sqrt(np.sum(psi_x * psi_x))
|
|
psi_x *= norm_factor
|
|
|
|
pl.figure(1)
|
|
pl.plot(ks, phi_k) # phi_k could be complex, which would throw a warning here
|
|
# if so, could use: pl.plot(ks, np.abs(phi_k) )
|
|
pl.title(r"Spectrum $\phi(k)$")
|
|
pl.xlabel("wavenumber k")
|
|
pl.ylabel("spectral amplitude")
|
|
|
|
pl.figure(2)
|
|
pl.plot(x, psi_x)
|
|
pl.title(r"Wavefunction $\psi(x)$")
|
|
pl.xlabel("position x")
|
|
pl.ylabel(r"wavefunction $\psi(x)$")
|
|
pl.annotate("This piece also appears to left of x=0!", (3.8, -0.010),
|
|
xytext=(1.0, -0.03), arrowprops=dict(facecolor='black', shrink=0.05))
|
|
|
|
# Here's a fancier figure, without recalculating anything but noticing
|
|
# that the wavepacket right now is centered around x=0. Some of that is
|
|
# showing up on the right side of the window, because the whole signal is
|
|
# repeating with a period of 1 windowLength, so we will copy that and stitch
|
|
# it on the left side in a different color. (If you look closely, you will
|
|
# see that the last point of the "orange" needs to be connected to the first
|
|
# point of the "blue")
|
|
# In the NEXT script P372-PS14-FFT-shift.py, the wavepacket is shifted in x-space
|
|
pl.figure(3)
|
|
cutPoint = (3 * numPoints) // 4 # 3/4 of the way through the window
|
|
pl.plot(x[0:cutPoint], psi_x[0:cutPoint]) # as before, but only to cutPoint
|
|
pl.plot(x[cutPoint:] - windowLength, psi_x[cutPoint:]) # put last 1/4 on left side
|
|
pl.plot([-dx, 0], [psi_x[-1], psi_x[0]], 'k') # Final stitch
|
|
pl.title(r"Wavefunction $\psi(x)$")
|
|
pl.xlabel("position x")
|
|
pl.ylabel(r"wavefunction $\psi(x)$")
|
|
|
|
pl.show() # show the graph windows and loop here until they are closed
|
|
|
|
print("Done.")
|