Solution of the Time-dependent Schrödinger Equation

using the Crank-Nicolson algorithm with MPI on a 2-D

regular Cartesian grid
Seminar on High Performance Computing 2
Summer term 2017

Peter Leitner & Stefan Hofmeister

Wednesday, May 10, 2017

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 1 / 13
Time-dependent Schrödinger Equation
Schrödinger Equation in 2-D for an electron in a given potential
∂ p̂ 2 p̂y2
i~|ψ(t, x)i = Ĥ(p̂, x̂) |ψ(t, x)i , Ĥ = x + + V̂ (x, y ) (1)
∂t 2m 2m
Attempt to simply discretize the differential operators,
n+1 n
ψj,l − ψj,l
∂t ψ(t, x) ≈ ,
∆t
n+1 n+1 n+1
p̂ 2 ~2 2 ~2 ψj+1,l − 2ψj,l + ψj−1
ψ(t, x) = − ∇ ψ(t, x) ≈ − + ...
2m 2m 2m (∆x)2
is – as the von Neumann stability analysis shows – unconditionally stable.
However, each solution of Schrödinger’s equation has to fulfill a continuity
equation for the probability w (t, x) d3 x = ψ ∗ ψ d3 x to find the electron
within a volume d3 x around x,
~
∂t w (t, x) + ∇ · j = 0, j = [ψ ∗ (∇ψ) − (∇ψ ∗ )ψ] (2)
2im
Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 2 / 13
From the continuity equation follows the conservation of the norm of the
wave function:
Z Z Z Z
d
d3 x w (t, x) = d3 x ∂t w (t, x) = − d3 x ∇ · j (t, x) = − dS · j (t, x)
dt
Z Z
SE
d3 x kw (0, x)k2 = 1 −→ d3 x kw (t, x)k2 = 1 (3)
R3 R3
The normalization of the wave function ψ(t, x) is however not conserved by the fully
implicit discretization as this method is not unitary ⇒ a discretization preserving the
hermiticity of the Hamiltonian is needed that maintains the normalization of the
wavefunction.
The temporal evolution of the wave function can be expressed using the unitary
time-evolution operator Û(t):

|ψ(t + ∆t, x)i ' Û(t) |ψ(t, x)i , Û(t) = exp(−iĤt/~) (4)

With the time-evolution operator one can find the explicit Forward-Time Centered-Space
n+1 n
(FTCS) scheme ψj,l = (1 − iĤ∆t)ψj,l and the implicit Backward-Time scheme
n+1 n
(1 + iĤ∆t)ψj,l = ψj,l . Neither scheme is unitary and thus violates probability
conservation.

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 3 / 13
A finite-difference approximation of exp(−iĤ∆t/~) that is unitary was
found by Cayley:
iĤ∆t
1−
exp(−iĤ∆t/~) = 2~
+ O((∆t)2 )
iĤ∆t
1+ 2~
 i   i 
(4) ⇒ 1 + Ĥ∆t |ψ(t + ∆t, x)i = 1 − Ĥ∆t |ψ(t, x)i (5)
2~ 2~
This is exactly the Crank-Nicolson scheme for the Schrödinger equation.

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 4 / 13
Discretization

n into
By substituting the 2-D discrete wave function ψ(t; x, y ) → ψj,l
Eq. (5) and considering the Hamiltonian in the form

~2 ∂ 2 ~2 ∂ 2
Ĥ = |p̂ = −i~∇| = − − + V̂ (x, y )
2m ∂x 2 2m ∂y 2

one obtains the discretized 2-D Schrödinger Equation

n+1 n+1 n+1 n+1 n+1 !
n+1 i∆t ~2 ψj+1,l n + 1 − 2ψj,l + ψj−1,l ψj,l+1 − 2ψj,l + ψj,l−1 n+1
ψj,l − 2
+ 2
− Vj,l ψj,l
2~ 2m (∆x) (∆y )
n n n n n n
!
n i∆t ~2 ψj+1,l − 2ψj,l + ψj−1,l ~2 ψj,l+1 − 2ψj,l + ψj,l−1
= ψj,l + 2
+ 2
− Vj,l ψj,l (6)
2~ 2m (∆x) 2m (∆y )

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 5 / 13
Rewrite Eq. (6) in matrix-form. Assign a short-hand notation for the prefactors of the wave
function at the next time step at various spatial grid points:
i∆t ~2
 
n+1 1 1 i∆t
ψj,l : bj,l = 1 + + + Vj,l
~ 2m (∆x)2 (∆y )2 2~
n+1 i∆t ~2 1
ψj−1,l : c=−
2~ 2m (∆x)2
n+1
ψj+1,l : a=c

n+1 i∆t ~2 1
ψj,l+1 : d =−
2~ 2m (∆x)2
n+1
ψj,l−1 : e=d
and for the RHS:
i∆t ~2
 
n 1 1 i∆t
ψj,l : fj,l = 1 − + − Vj,l
~ 2m (∆x)2 (∆y )2 2~
n i∆t ~2 1
ψj+1,l : g =
2~ 2m (∆x)2
n
ψj−1,l : h=g

n i∆t ~2 1
ψj,l+1 : k=
2~ 2m (∆y )2
n
ψj,l−1 : p=k

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 6 / 13
LHS of the discretized Schrödinger Equation

n on a single one (j, l) 7→ i:

Map the two indices j and l of the wave function ψj,l

i = 1 + (j − 1)(L − 1) + l − 1 for j = 0, 1, . . . , J; l = 0, 1, . . . , L

n → ψ n , ψ n → ψ n , . . ., ψ n → ψ n , ψ n → ψ n , . . . For simplicity (and lack of

Thereby ψ1,1 1 1,2 2 1,J J 2,1 J+1
space) shown here only for J = L = 4, though later J, L ∼ 103 will be used. It exhibits
sub-blocks of shape J × L:

 ψ n+1 
b1,1 d 0 a 0 0 0 0 0

1
 e b1,2 d 0 a 0 0 0 0  ψ2n+1 
 ψ3n+1 

 0 e b1,3 0 0 a 0 0 0


 n+1 

 c 0 0 b2,1 d 0 a 0 0  ψ4 
 
  n+1 
 0 c 0 e b2,2 d 0 a 0  ψ5  = r

  n+1 
 0 0 c 0 e b2,3 0 0 a  ψ6 

  n+1 
 0 0 0 c 0 0 b3,1 d 0

 ψ7 
 0 0 0 0 c 0 e b3,2 d   n+1 
ψ8
0 0 0 0 0 c 0 e b3,3 ψ9n+1

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 7 / 13
RHS of the discretized Schrödinger Equation

Vectorization of the RHS is not necessary when working with Fortran

exclusively but essential when using the numpy module with Python:
  n
f1,1 k 0 g 0 0 0 0 0 ψ1

 p f1,2 k 0 g 0 0 0 0  ψ2n 
   n
 0 p f1,3 0 0 g 0 0 0  ψ3n 
 

 h
 0 0 f2,1 k 0 g 0 0  ψ4n 
 
r = 0
 h 0 p f2,2 k 0 g 0 

ψ5 

 0 0 h 0 p f2,3 0 0 g  ψ6 
  n
 
 0 0 0 h 0 0 f3,1 k 0  n
  ψ7 
 0 0 0 0 h 0 p f3,2 k  ψ8n 
0 0 0 0 0 h 0 p f33 ψ9n

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 8 / 13
After vectorization, the discretized Schrödinger Equation simplifies to

Âψ n+1 = r = B̂ψ n

Pseudo-code:

for n in range(N):

Compute B̂ψ n = r as a matrix-vector product

Solve Âψ n+1 = r with LU decomposition
This yields ψ n+1
Map the rolling index i back to indices j, l: ψin+1 → ψj,l
n+1

Check that the wave packet remains normalized:

Z J,L
2 2
X n 2 !
d x kψ(t, x)k ≈ ∆x∆y ψj,l =1
∆2 j,l=1

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 9 / 13
Setting up the matrices in Python
# Build the main diagonal
b_main_diag = [] # coefficient matrix of LHS
f_main_diag = [] # coefficient matrix of RHS
for j in range(J-1):
for l in range(J-1):
b_main_diag.append(b[j,l]); f_main_diag.append(f[j,l])

e_minor_diag = np.ones((len(b_main_diag)-1,), dtype=complex)*e

d_minor_diag = np.ones((len(b_main_diag)-1,), dtype=complex)*d
a_minor_diag = np.ones((len(b_main_diag)-3,), dtype=complex)*a
c_minor_diag = np.ones((len(b_main_diag)-3,), dtype=complex)*c
k_minor_diag = np.ones((len(f_main_diag)-1,), dtype=complex)*k
p_minor_diag = np.ones((len(f_main_diag)-1,), dtype=complex)*p
g_minor_diag = np.ones((len(f_main_diag)-3,), dtype=complex)*g
h_minor_diag = np.ones((len(f_main_diag)-3,), dtype=complex)*h

# Build the coefficient matrices for A*psi^(n+1) = B*psi^n

A = np.diag(b_main_diag) + np.diag(e_minor_diag,-1) + np.diag(d_minor_diag,1) \
+ np.diag(c_minor_diag,-3) + np.diag(a_minor_diag,3)
B = np.diag(f_main_diag) + np.diag(p_minor_diag,-1) + np.diag(k_minor_diag,1) \
+ np.diag(h_minor_diag,-3) + np.diag(g_minor_diag,3)

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 10 / 13
Setting-up the RHS in Python

for n in range(Nt):
if n == 0:
psi_n_mat = psi_0
# Rewrite matrix psi_jl as a vector psi_i
psi_n_vec = np.empty(((J-1)**2,), dtype=complex)
for j in range(J-1):
for l in range(J-1):
i = ind(j,l)
psi_n_vec[i] = psi_n_mat[j,l]
else:
psi_n_vec = psi_np1_vec

RHS = np.dot(B,psi_n_vec)
r = RHS

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 11 / 13
LU decomposition; finally to be parallelized as Fortran code
# LU decomposition:
P, L, U = scipy.linalg.lu(A)
LU = L + U
nA = np.size(A, 0)
y = np.zeros((nA,), dtype=complex)
x = np.zeros((nA,), dtype=complex)

r = np.dot(np.linalg.inv(P),r)

for i in range(1, nA+1):

s = r[i-1]
for j in range(1, i):
s = s - LU[i-1,j-1]*y[j-1]
y[i-1] = s

for i in range(nA, 0, -1):

s = y[i-1]
for j in range(i+1, nA+1):
s = s - LU[i-1,j-1]*x[j-1]
x[i-1] = s/LU[i-1,i-1]

psi_np1_vec = x
Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 12 / 13
Parallelizing the code

The following strategies for code parallelization systems are envisaged:

Application of Distributed SuperLU (SuperLU_DIST) of the SuperLU
library
Crank-Nicolson predictor-corrector (CNPC) method
Decomposition of spatial domain Ω in subdomains Ωi which can be
assigned to their own processor
Forward Euler to predict unknown values on the interface Γ between
subdomains Ωi ⊂ Ω
Solving the problem independently on the subdomains Ωi with
Backward Euler using the predicted values ψΓ∗ at tn+1
Correction of interface values ψΓn+1 with Backward Euler using interior
values

Peter Leitner & Stefan Hofmeister Crank-Nicolson using MPI Wednesday, May 10, 2017 13 / 13