cool....
Building on the last post that evaluates differential equations in python. This is a study in art, math, physics and data visualization informed by code and concept generated by @Mocz.
When I learned about this in college, I was just grinding through textbooks and problem sets just to know enough to get the answer correct on the exam, but I never really understood the nuts and bolts of what was going on here. This is my quest back into the world of the mythical Schrodinger's equation and how to simulate it on a computer.
The 2D time-dependent Schrodinger's equation is an ordinary differential equation. @Mocz is a professional astrophysicist and I am a curious part-time arm-chair physicist, so I'm going to do my best to put this into my own understanding, learn something about how to create physical systems and visualizations using python, and then see if I can recreate something similar from scratch based on my understanding.
A 2D Fluid is going to be a 2D matrix of values with real and complex elements that are evaluated by applying a function to a 2D matrix of x, y wave position coordinates measured integer fractions of radians. The time dimension (t) is taken care of by the program looping through frames. Later, I will describe some techniques for defining initial values and shapes in that fluid!
Our wavefunction is a generic wave varying sinusoidally, represented as vectors moving in the complex plane as a function of position and time.... $$ Ψ(x,t) = Ae^{i(kx - ωt)} $$ The wavefunction can be derived from two traveling waves using Euler's identity (exercise left to reader!): $$ e^{iθ} = cos(θ) + isin(θ) $$ By multiplying Ψ by its complex conjugate "psi squared", we get the probability density function, which gives the probability of observing a particle in a specific spacetime location.
This is where the computational magic happens. Computers are really good at approximating forward and inverse Fourier Transforms with the famous Cooley-Tukey Fast Fourier Transform (FFT) algorithm. We can use this and a few tricks to help us evaluate the schrodinger equation and evaluate it.
Since ψ is an exponential wave function, ∇ becomes ik, and ∇^2 becomes -k^2 in any differential equation we apply a Fourier transform to. That means differentiation becomes a simple multiplication function. After solving, we can simply do the inverse Fourier Transform to return.
So, we've got a wave, and we have a 2D fluid, and in a bit we will have a force potential but we need some way of relating them! An equation to tie it all together. Something that can be evaluated to create an awesome visual. The Schrodinger equation (2) is just a fancy looking version of a very familiar energy equation from physics (1):
$$ (1) \;\;Energy(total) = E(kinetic) + E(potential) $$
$$ (2)\;\;iℏ \frac{\partialψ}{\partial t} = -\frac{h^2}{2m}\frac{\partial^2ψ}{\partial x^2}+Vψ $$
Ignoring the potential term, and assuming h = m = 1 for sake of this simulation, we take the Fourier Transform of the kinetic term and get the following: $$ iℏ \frac{\partialψ}{\partial t} = -\frac{1}{2}\frac{\partial^2ψ}{\partial x^2} = -\frac{1}{2}*-k^2 ψ = \frac{1}{2}k^2ψ $$
Now, while the function is still transformed (meaning there is no longer differential term to deal with on the right side), we can increment time by multiplying psi by e to -idt weighted by k-squared which will just be added into the exponential term in psi. $$ ψ = e^{-i\frac{k^2}{2}Δt} ψ $$ The k-squared divided by two constant aligns with the original relation: $$ iℏ \frac{\partialψ}{\partial t} = \frac{1}{2}k^2ψ $$
Then apply the Inverse Fourier Transform to return to the differential form in real space. We just evaluated the kinetic term!
Let's choose a force to act on the fluid. This means if we have two 2D blobs of fluid in our 2D grid, they will be atttracted (or repelled) to each other governed by some force. Here, we can just start by using the gravitational potential, which just says the force between any two points is a function of the mass at the two points and the distance between the two points.
@Mocz provides the Poisson Equation for a gravitational self-potential. I don't fully understand all the ins and outs of how this equation is derived, but it is a relation we can use to obtain a potential V. It says that the laplacian (double gradient) of the potential is a probability density function times the gravitational constant. $$ ∇^2V=4πG(|ψ|^2-1) $$
Side Adventure: I needed to understand this piece... And the equation was making sense to me other than I couldn't fully understand why there's the "-1" term. This was driving me nuts. I searched around and eventually found this paper). It says in "periodic cases", which I take to mean time-dependent, one must offset the density function term |ψ|^2 by its mean, a renormalization that prevents divergent solutions. I guess that means that if the PDF fluctuates in total amount as a function of time, this normalization prevents that from biasing a calculation of the potential. So in this case, @Mocz is just setting the mean density to 1, as we are assuming our system's total amount of density isn't changing periodically over time - it's conserved.
So, putting 2 and 2 together that sounds like in the case of our simulation, we don't really need it. And... removing the -1 from the equation seems to have no effect on the video output... :)
With our new differential equation describing the potential, we again apply our Fourier Transform trick to transform the Poisson Equation into a form which we can evaluate with the computer: $$ ∇^2V=4πG(|ψ|^2-1)→k^2V=4πG(|ψ|^2-1) → V = 4πG(|ψ|^2-1)/k^2 $$ We then apply the Inverse Fourier Transform to return to real space.
We increment time for the potential term by multiplying by e to the dt, adding it in the exponential term of psi, taking advantage of convenient properties of exponents. $$ ψ = e^{-i{V}{Δt}} ψ $$
@Mocz recommends using this kick-drift-kick scheme, where we evaluate the potential energy in two dt/2 increments (kick) and the kinetic energy in a full dt increment (drift). Since he's the expert, I'm going to go with that. Intuitively, this looks like it makes the simulation more accurate than just incrementing once using a full dt increment since the energy terms and the potential depend on each other for updating and evaluation. Here's some pseudocode:
dt = 0.01
for timeframe in sim:
evaluate psi PE at t + dt/2 (kick)
evaluate psi KE at t + dt (drift)
update potential based on new psi (update)
evaluate psi PE at t + dt/2 (kick)
increment t by dt (increment)
For this, we will use something called meshgrids. Meshgrids allow us to easily draw shapes in matrices. To start we will define an X matrix with shape (x, y) that just is columns containing each value's x-coordinate. We will repeat the process for the Y matrix, but with rows contain each value's y-coordinate. What's amazing about this is we can then create a circle or radial meshgrid by just saying R(radial) = X^2 + Y^2.
We can use meshgrids to define the both the strength/shape of the force potential, which in this case is radial.
We can also use meshgrid to define initial values of the system defining a set of 2D gaussian blobs represented by rho ρ, a variable that is often used to describe a mass density distribution. These are essentially 2D bell curves.
Here's the formula for the 2D bell curve distributions that are centered around coordinates x0 and y0. It's just a circle equation that drops off exponentially as a function of sigma (standard deviation) from a central point. $$ ρ = \frac{1}{2πσ^3}e^{((x-x_0)^2+(y-y_0)^2)/ σ^2} $$
Finally to normalize rho, divide rho by the mean of rho, and since rho is a density function, take the square root to transform the matrix from rho space to psi space. Remember that "psi squared" is a probability density function.
Wow. We arrived. Using computers to visualize a formula so beautiful and fundamental to physics. Now run the loop and enjoy the show!
A lot of what was explored here, like the use of meshgrids and Fourier Transforms to find solutions for natural systems, can be applied for making cool video software!
Here's the full Script very lightly adapted from @Mocz's blog post:
"""Code adapted from Philip Mocz https://levelup.gitconnected.com/create-your-own-quantum-mechanics-simulation-with-python-51e215346798"""
import numpy as np
import matplotlib.pyplot as plt
# Create Wavevectors
N = 128 # resolution
L = 1 # box length
t = 0 # time of simulation
dt = 0.0001 # timestep length
tEnd = 0.1 # length of simulation
Nt = int(np.ceil(tEnd/dt)) # number of timesteps in simulation, loops
# Create 2D radial meshgrid to represent the potential distribution, building a radial geometry into a 2D array
klin = (2.0 * np.pi / L) * np.arange(-N/2,N/2) # create array length N with increments of 2Pi / box length
kx, ky = np.meshgrid(klin, klin) # create N x N matrices where kx is N rows of klin, and ky is N columns of klin
kx = np.fft.ifftshift(kx) # shift row vals so 0 -> + values are first half of row, and - --> 0 is second half
ky = np.fft.ifftshift(ky) # shift col vals so 0 -> + values are first half of col, and - --> 0 is second half
kSq = kx**2 + ky**2 # making a radial meshgrid centered in the middle!... This will be the potential distribution (spacial force dropoff) for each point.
# Constants
G = 4000 # Gravitational Constant
# initial conditions
# Domain [0,1] x [0,1]
# Create some linear meshgrids for creating initial conditions
xlin = np.linspace(0,L, num=N+1) # Note: x=0 & x=1 are the same point!
xlin = xlin[0:N] # chop off periodic point
xx, yy = np.meshgrid(xlin, xlin) # make NxN arrays
# Intial Condition
amp = 0.01
sigma = 0.03
# Using the linear meshgrids, make rho, meshgrid containing a bunch of gaussian blobs
# based on a 2d gaussian probability density function
# this meshgrid will be our initial condition of the system
rho = 2*amp*np.exp(-((xx-0.66)**2+(yy-0.66)**2)/2/sigma**2)/(sigma**3*2*np.pi)
rho+= 2*amp*np.exp(-((xx-0.33)**2+(yy-0.33)**2)/2/sigma**2)/(sigma**3*np.sqrt(2*np.pi)**2)
rho+= 2*amp*np.exp(-((xx-0.66)**2+(yy-0.33)**2)/2/sigma**2)/(sigma**3*np.sqrt(2*np.pi)**2)
rho+= 2*amp*np.exp(-((xx-0.33)**2+(yy-0.66)**2)/2/sigma**2)/(sigma**3*np.sqrt(2*np.pi)**2)
## generate psi initial condition from normalized rho
rhobar = np.mean( rho )
rho /= rhobar
psi = np.sqrt(rho)
# prep figure
fig = plt.figure(figsize=(6,4), dpi=80)
grid = plt.GridSpec(1, 2, wspace=0.0, hspace=0.0)
ax1 = plt.subplot(grid[0,0])
ax2 = plt.subplot(grid[0,1])
outputCount = 1
### Potentials
Vhat = -np.fft.fftn(4.0*np.pi*G*(np.abs(psi)**2-1.)) / ( kSq + (kSq==0)) # Newtonian Gravity
V = np.real(np.fft.ifftn(Vhat))
# solve time dependent schodinger equation with given potential for each time point assuming hbar and m are = 1
for i in range(Nt):
# 1/2dt kick
psi = np.exp(-1j*dt/2.0*V) * psi
# drift 1dt
# multiplication in FFT space is same as Gradient/Derivative, since the laplacian must be evaluated, we transform to frequency space and multiply
psihat = np.fft.fftn(psi)
psihat = np.exp(dt*(-1j*kSq/2.)) * psihat
psi = np.fft.ifftn(psihat)
# Update the Potential
Vhat = -np.fft.fftn(4.0*np.pi*G*(np.abs(psi)**2-1.)) / ( kSq + (kSq==0)) # Newtonian Gravity
V = np.real(np.fft.ifftn(Vhat))
# 1/2dt kick
psi = np.exp(-1j*dt/2.0*V) * psi
t += dt
plt.sca(ax1)
plt.cla()
plt.imshow(np.log10(np.abs(psi)**2), cmap = 'inferno')
plt.clim(-1, 2)
ax1.get_xaxis().set_visible(False)
ax1.get_yaxis().set_visible(False)
ax1.set_aspect('equal')
plt.sca(ax2)
plt.cla()
plt.imshow(np.angle(psi), cmap = 'bwr')
plt.clim(-np.pi, np.pi)
ax2.get_xaxis().set_visible(False)
ax2.get_yaxis().set_visible(False)
ax2.set_aspect('equal')
plt.pause(0.001)
outputCount += 1