71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
# P372-PS14-FFT-animate.py
|
|
"""
|
|
A wavepacket is made to move. Based on P372-PS14-FFT.py
|
|
and uses matplotlib.animation library
|
|
"""
|
|
|
|
import numpy as np
|
|
import pylab as pl
|
|
import matplotlib.animation as animation
|
|
import numpy.fft as fft
|
|
|
|
# Some basic numbers, first. These are all design decisions you can change!
|
|
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] # 25% of highest k value <-- PICK THIS!
|
|
alpha = 0.1 # smaller alpha means wider phi_k, narrower psi_x <-- AND THIS!
|
|
# Now define the k-spectrum amplitudes
|
|
phi_k = np.exp(-0.5 * alpha * (ks - k0) ** 2) # the spectral amplitudes (complex)
|
|
# And generate the wavefunction from phi(k) using Inverse Real FFT
|
|
psi_x = fft.irfft(phi_k) # inverse real FFT gets the real wavefunction
|
|
|
|
# Now we want the wavepacket to move! Let's define a velocity:
|
|
v = 1.0
|
|
|
|
fig, ax = pl.subplots() # Create the graphics canvas
|
|
pl.title(r"Wavefunction $\Psi(x,t)$ animated")
|
|
pl.xlabel("position x")
|
|
pl.ylabel(r"wavefunction $\Psi(x,t)$")
|
|
|
|
T = 1.0 # total time duration
|
|
Nsteps = 200 # number of steps of animation
|
|
times = np.linspace(0, T, Nsteps + 1) # the times at which to calculate the shifts
|
|
# Draw the first plot and get the lines object
|
|
lines = ax.plot(x[:numPoints // 2], psi_x[:numPoints // 2], lw=1)[0]
|
|
|
|
|
|
def update(frame):
|
|
# This calculates the new values for the plot, called by the animation
|
|
time_now = times[frame]
|
|
# Calculate the phase shifts
|
|
phase_shift = np.exp(-1j * ks * 2 * v * time_now)
|
|
alpha = time_now
|
|
phi_k = np.exp(-0.5 * alpha * (ks - k0) ** 2)
|
|
# And apply them before doing the F.T.
|
|
psi_x_shift = fft.irfft(phi_k * phase_shift)
|
|
lines.set_ydata(psi_x_shift[:numPoints // 2])
|
|
return lines
|
|
|
|
|
|
# Connect the update function to the animation
|
|
ani = animation.FuncAnimation(fig=fig, func=update, frames=Nsteps, interval=30)
|
|
|
|
# Show the animation and run it until the window is closed
|
|
pl.show()
|
|
|
|
print("Done.")
|