# Split step Fourier method to solve Schrödinger equation for moving potential

I'm trying to use the excellent Schrodinger Python class by Jake VanderPlas (https://jakevdp.github.io/blog/2012/09/05/quantum-python/) to simulate a wave packet within a moving Gaussian potential. I want to make it work with actual experimental parameters, i.e. $\hbar$, m and distances all in SI units. When running the simulation with a starting wave packet close to an eigenstate of the potential, and the potential not moving, I see an unexpected behaviour: The wave packet seems to delocalise in k-space very quickly, leaving just random noise for the k-space wave function.

Is this an insurmountable numerical error? Or have I made a mistake while switching to SI units? And while we're at it: Does anyone understand what the class is doing with the internal psi_mod_k and psi_mod_x?

Here is my code:

"""
General Numerical Solver for the 1D Time-Dependent Schrodinger's equation.

author: Jake Vanderplas
email: vanderplas@astro.washington.edu
website: http://jakevdp.github.com
Please feel free to use and modify this, but keep the above information. Thanks!
"""

import numpy as np
from matplotlib import pyplot as pl
from matplotlib import animation
from scipy.fftpack import fft,ifft

class Schrodinger(object):
"""
Class which implements a numerical solution of the time-dependent
Schrodinger equation for an arbitrary potential
"""
def __init__(self, x, psi_x0, V_x,
k0 = None, hbar=1, m=1, t0=0.0, displacement=lambda t: 0):
"""
Parameters
----------
x : array_like, float
length-N array of evenly spaced spatial coordinates
psi_x0 : array_like, complex
length-N array of the initial wave function at time t0
V_x : array_like, float
length-N array giving the potential at each x
k0 : float
the minimum value of k.  Note that, because of the workings of the
fast fourier transform, the momentum wave-number will be defined
in the range
k0 < k < 2*pi / dx
where dx = x[1]-x[0].  If you expect nonzero momentum outside this
range, you must modify the inputs accordingly.  If not specified,
k0 will be calculated such that the range is [-k0,k0]
hbar : float
value of planck's constant (default = 1)
m : float
particle mass (default = 1)
t0 : float
initial tile (default = 0)
"""
# Validation of array inputs
self.x, psi_x0, self.V_x = map(np.asarray, (x, psi_x0, V_x))
N = self.x.size
assert self.x.shape == (N,)
assert psi_x0.shape == (N,)
assert self.V_x.shape == (N,)

# Set internal parameters
self.hbar = hbar
self.m = m
self.t = t0
self.dt_ = None
self.N = len(x)
self.dx = self.x[1] - self.x[0]
self.dk = 2 * np.pi / (self.N * self.dx)
self.displacement = displacement
self.V_x_shifted = self.V_x

# set momentum scale
if k0 == None:
self.k0 = -0.5 * self.N * self.dk
else:
self.k0 = k0
self.k = self.k0 + self.dk * np.arange(self.N)

self.psi_x = psi_x0
self.compute_k_from_x()

# variables which hold steps in evolution of the
self.x_evolve_half = None
self.x_evolve = None
self.k_evolve = None

# attributes used for dynamic plotting
self.psi_x_line = None
self.psi_k_line = None
self.V_x_line = None

self.moves_right = False
self.argmax = 0

def _set_psi_x(self, psi_x):
self.psi_mod_x = (psi_x * np.exp(-1j * self.k[0] * self.x) * self.dx / np.sqrt(2 * np.pi))

def _get_psi_x(self):
return (self.psi_mod_x * np.exp(1j * self.k[0] * self.x) * np.sqrt(2 * np.pi) / self.dx)

def _set_psi_k(self, psi_k):
self.psi_mod_k = psi_k * np.exp(1j * self.x[0] * self.dk * np.arange(self.N))

def _get_psi_k(self):
return self.psi_mod_k * np.exp(-1j * self.x[0] * self.dk * np.arange(self.N))

def _get_dt(self):
return self.dt_

def _set_dt(self, dt):
if dt != self.dt_:
self.dt_ = dt
self.x_evolve_half = np.exp(-0.5 * 1j * np.roll(self.V_x, int(np.rint(1/dx*self.displacement(self.t)))) / self.hbar * dt )
self.x_evolve = self.x_evolve_half * self.x_evolve_half
self.k_evolve = np.exp(-0.5 * 1j * self.hbar / self.m * (self.k * self.k) * dt)

psi_x = property(_get_psi_x, _set_psi_x)
psi_k = property(_get_psi_k, _set_psi_k)
dt = property(_get_dt, _set_dt)

def compute_k_from_x(self):
self.psi_mod_k = fft(self.psi_mod_x)

def compute_x_from_k(self):
self.psi_mod_x = ifft(self.psi_mod_k)

def time_step(self, dt, Nsteps = 1):
"""
Perform a series of time-steps via the time-dependent
Schrodinger Equation.

Parameters
----------
dt : float
the small time interval over which to integrate
Nsteps : float, optional
the number of intervals to compute.  The total change
in time at the end of this method will be dt * Nsteps.
default is N = 1
"""
self.dt = dt
self.x_evolve_half = np.exp(-0.5 * 1j * np.roll(self.V_x, int(np.rint(1 / dx * self.displacement(self.t))))
/ self.hbar * dt)
self.x_evolve = self.x_evolve_half * self.x_evolve_half

if Nsteps > 0:
self.psi_mod_x *= self.x_evolve_half

for i in range(Nsteps - 1):
self.compute_k_from_x()
self.psi_mod_k *= self.k_evolve
self.compute_x_from_k()
self.psi_mod_x *= self.x_evolve

self.compute_k_from_x()
self.psi_mod_k *= self.k_evolve

self.compute_x_from_k()
self.psi_mod_x *= self.x_evolve_half

self.compute_k_from_x()

self.V_x_shifted = np.roll(self.V_x, int(np.rint(1/dx*self.displacement(self.t))))

self.t += dt * Nsteps

######################################################################
# Utility functions for running the animation

def gaussian(x, x0, sigma, height):
return height * np.exp(-(x-x0)**2/(2*sigma**2))

def harmonic(x, x0, omega, depth):
norm = 0.5*m*omega**2
x = np.asarray(x)
y = norm*(x - x0)** 2 - depth
return y

######################################################################
# Create the animation

# specify time steps and duration
dt = 0.0001
N_steps = 50
t_max = 120
frames = int(t_max / float(N_steps * dt))

# specify constants
hbar = 1.05*10**(-34)   # planck's constant
m = 87*1.66*10**(-27)      # particle mass
kB = 1.38*10**(-23)
uK = 10**-6
mm = 10**-3

# specify range in x coordinate
N = 2 ** 14 # powers of two work most efficiently with FFT
dx = 0.05*mm # width of each simulation point
x = dx * (np.arange(N) - 0.5 * N) # the x values are saved in m

# POTENTIAL
V0 = 150*kB*uK  # potential depth (kB*uK)
x0 = 0*mm # position of trap minimum at t=0
omega=3.7*2*np.pi# trap frequency in Hz
V_x = gaussian(x, x0, 4*mm, -V0)

# now generate an initial WAVEFUNCTION that fits well into the potential given by V0, omega and m:
psi_abs = (m*omega/(np.pi*hbar))**0.25
sigma = 3*mm
offset = 0 # offset of wave packet position from trap minimum
psi_x0 = psi_abs*np.exp(-(x-x0-offset)**2/(2*sigma**2))

d = 300*mm # transport length
t_f = 1 # transport duration in s
def stay(t): return 0

# define the Schrodinger object which performs the calculations
S = Schrodinger(x=x,
psi_x0=psi_x0,
V_x=V_x,
hbar=hbar,
m=m,
#k0=-50000,
displacement=stay)

######################################################################
# Set up plot
fig = pl.figure(figsize=(17, 6), dpi=100)

# POTENTIAL
ymin = -V0
ymax = V0
xlim = (-10, 350)
ylim=(ymin - 0.2 * (ymax - ymin),
ymax + 0.2 * (ymax - ymin)))
V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')

ax1.legend(prop=dict(size=12))
ax1.set_xlabel('$x (mm)$')
ax1.set_ylabel(r'$V(x)$')

# WAVEFUNCTION
ax2 = ax1.twinx()
ax2.set_ylim((-psi_abs, psi_abs))
psi_x_line, = ax2.plot([], [], c='r', label=r'$|\psi(x)|$')
ax2.set_ylabel(r'$|\psi(x)|$', color='r')
ax2.tick_params('y', colors='r')

# WAVEFUNCTION IN K-SPACE
ymin = abs(S.psi_k).min()
ymax = abs(S.psi_k).max()
ax3 = fig.add_subplot(212, ylim=(ymin - 0.2 * (ymax - ymin), ymax + 0.2 * (ymax - ymin)))
psi_k_line, = ax3.plot([], [], c='r', label=r'$|\psi(k)|$')

mV_line = ax3.axvline(np.sqrt(2 * V0) / hbar, c='k', ls='--', label=r'$\sqrt{2mV_0}$')
ax3.legend(prop=dict(size=12))
ax3.set_xlim((S.k0, S.k0+S.dk*len(S.k)))
ax3.set_xlabel('$k$')
ax3.set_ylabel(r'$|\psi(k)|$')

V_x_line.set_data(S.x, S.V_x)

######################################################################
# Animate plot
def init():
psi_x_line.set_data([], [])
V_x_line.set_data([], [])

psi_k_line.set_data([], [])
return (psi_x_line, V_x_line, psi_k_line)

def animate(i):
S.time_step(dt, N_steps)
psi_x_line.set_data(S.x/mm, abs(S.psi_x)) # for plotting, convert m into mm
V_x_line.set_data(S.x/mm, S.V_x_shifted)

psi_k_line.set_data(S.k, abs(S.psi_k))
ax1.set_title("t = %.2f s" % S.t)
ax1.set_xlim((-50, 350))
return (psi_x_line, V_x_line, psi_k_line)

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=frames, interval=30, blit=False)

pl.show()


## migrated from physics.stackexchange.comAug 6 '18 at 8:54

This question came from our site for active researchers, academics and students of physics.

• double precision arithmetic has rounding error that is 10^7 or more times larger than the constants you want to use. your simulation will always be noise – smh Aug 6 '18 at 13:37