Numerical solution of the 2D time-dependent Schrödinger equation using the Chebyshev method

Colin Greenshields

Program developed and used during the period 2013-2015 as part of my PhD research at the University of Glasgow.

Global variables

(* Using $ to indicate global variables *)

$Lx=7.5; (* Length of simulation area along x axis *)
$Nx=64; (* Number of grid points along x axis *)
$Ly=7.5;
$Ny=64;
$m=1; (* Mass of particle *)
$ℏ=1;
$q=-1; (* Charge of particle *)
$δ=100;  (* Controls the accuracy of the Chebyshev expansion *)

Discretisation

(* Note that the grid is centred on the origin *)

x[i_]:=$Lx/$Nx (i-$Nx/2)  (* x coordinate corresponding to index i *)
y[j_]:=$Ly/$Ny (j-$Ny/2)
kx[i_]:=(2 π)/$Lx (i-$Nx/2) (* Spatial frequency corresponding to index i, when working in Fourier space *)
ky[j_]:=(2 π)/$Ly (j-$Ny/2)

(* Discrete Fourier transforms. RotateLeft is used since the Mathematica functions Fourier and InverseFourier require the zero frequency to be at the start of the array, rather than in the centre. *)
FT[Ψ_]:=($Lx/Sqrt[2 π]) ($Ly/Sqrt[2 π]) RotateLeft[Fourier[RotateLeft[Ψ,{$Ny/2,$Nx/2}],FourierParameters->{-1,-1}],{$Ny/2,$Nx/2}]//N
InverseFT[ΨFT_]:=(Sqrt[2 π]/$Lx) (Sqrt[2 π]/$Ly) RotateLeft[InverseFourier[RotateLeft[ΨFT,{$Ny/2,$Nx/2}],FourierParameters->{-1,-1}],{$Ny/2,$Nx/2}]//N

(* Spatial derivatives of wavefunction *)
GridLaplacian[Ψ_]:=InverseFT[Table[-kx[i]^2-ky[j]^2,{i,0,$Nx-1},{j,0,$Ny-1}] FT[Ψ]]
GridxDeriv[Ψ_]:=InverseFT[Table[I kx[i],{i,0,$Nx-1},{j,0,$Ny-1}] FT[Ψ]]
GridyDeriv[Ψ_]:=InverseFT[Table[I ky[j],{i,0,$Nx-1},{j,0,$Ny-1}] FT[Ψ]]

kxMax[Lx_,Nx_]:=π Nx/Lx (* Maximum frequency represented on grid *)
kyMax[Ly_,Ny_]:=π Ny/Ly

(* Take an analytical function f(x,y) and represent it on the grid *)
ToGrid[f_]:=Table[f[x[i],y[j]],{i,0,$Nx-1},{j,0,$Ny-1}]//N

Wavefunctions and potentials

(* Introduce cylindrical coordinates *)
ρ[x_,y_]:=Sqrt[x^2+y^2]
ϕ[x_,y_]:=If[{x,y}=={0,0},0,ArcTan[x,y]]

(* Cylindrically symmetric vector potential describing uniform B field along the z axis.  The symmetry axis of the vector potential is located at (x0,y0). *)
CylindricalA[x0_,y0_,B_]:=ToGrid[Function[{x,y},{-0.5 B ρ[x-x0,y-y0] Sin[ϕ[x-x0,y-y0]],0.5 B ρ[x-x0,y-y0] Cos[ϕ[x-x0,y-y0]],0}]]

(* Laguerre-Gaussian (LG) wavefunction.  Here n and ℓ are the radial and azimuthal indices respectively, while ρ0 is the width parameter.  The wavefunction is centred on (x0,y0) in real space and at (kx0,ky0) in Fourier space. *)
Nnℓ[n_,ℓ_]:=Sqrt[(2 n!)/(π (n+Abs[ℓ])!)] (* Normalisation constant *)
LG[x0_,y0_,kx0_,ky0_,n_,ℓ_,ρ0_]:=Nnℓ[n,ℓ] Table[1/ρ0 If[ℓ≠0,((ρ[x[i]-x0,y[j]-y0] Sqrt[2])/ρ0)^Abs[ℓ],1] Exp[-(ρ[x[i]-x0,y[j]-y0]^2/ρ0^2)] LaguerreL[n,Abs[ℓ],2 ρ[x[i]-x0,y[j]-y0]^2/ρ0^2] Exp[I ℓ ϕ[x[i]-x0,y[j]-y0]] Exp[I kx0 (x[i]-x0)] Exp[I ky0 (y[j]-y0)],{i,0,$Nx-1},{j,0,$Ny-1}]//N

B0[ρ0_]:=(4 $ℏ)/(Abs[$q] ρ0^2) (* Strength of uniform B field for which an LG wavefunction with width ρ0 is a Landau state *)

ρB[B_]:=2 Sqrt[$ℏ/Abs[$q B]] (* Width that an LG wavefunction must have to be a Landau state of a field with strength B *)

Time-evolution

(* Calculation of time-evolution using the Chebyshev method.  See manuscript for description of the algorithm and definition of all variables.  The function GeneralCheby works for any Hamiltonian with the form of ℋ defined below, while ChebyTDSEuniformB simplifies this for use with the CylindricalA vector potential. *)

a[S2_,S1x_,S1y_,S0_]:=(S2 (-kxMax[$Lx,$Nx]^2-kyMax[$Ly,$Ny]^2)+2 (Max[S1x] kxMax[$Lx,$Nx]+Max[S1y] kyMax[$Ly,$Ny])+Max[S0]-Min[S0])/2
b[S2_,S1x_,S1y_,S0_]:=(S2 (-kxMax[$Lx,$Nx]^2-kyMax[$Ly,$Ny]^2)+Max[S0]+Min[S0])/2
ℋ[Ψ_,S2_,S1x_,S1y_,S0_]:=S2 GridLaplacian[Ψ]+I S1x GridxDeriv[Ψ]+I S1y GridyDeriv[Ψ]+S0 Ψ
(OverTilde[ℋ])[Ψ_,S2_,S1x_,S1y_,S0_]:=1/a[S2,S1x,S1y,S0] (ℋ[Ψ,S2,S1x,S1y,S0]-b[S2,S1x,S1y,S0] Ψ)
M[S2_,S1x_,S1y_,S0_,η_]:=Ceiling[a[S2,S1x,S1y,S0] η/$ℏ]+$δ

GeneralCheby[S2_,S1x_,S1y_,S0_,Δt_,ΨInitial_]:=Module[{TqΨ,TqMinus2Ψ,TqMinus1Ψ,sum},TqMinus2Ψ=ΨInitial;
TqMinus1Ψ=(OverTilde[ℋ])[ΨInitial,S2,S1x,S1y,S0];
sum=BesselJ[0,a[S2,S1x,S1y,S0] Δt/$ℏ] TqMinus2Ψ-2 I BesselJ[1,a[S2,S1x,S1y,S0] Δt/$ℏ] TqMinus1Ψ;
Do[TqΨ=2 (OverTilde[ℋ])[TqMinus1Ψ,S2,S1x,S1y,S0]-TqMinus2Ψ;
sum=sum+2 (-I)^q BesselJ[q,a[S2,S1x,S1y,S0] Δt/$ℏ] TqΨ;
TqMinus2Ψ=TqMinus1Ψ;
TqMinus1Ψ=TqΨ;,{q,2,M[S2,S1x,S1y,S0,Δt]}];
Exp[-I b[S2,S1x,S1y,S0] Δt] sum//N]

ChebyTDSEuniformB[ΨInitial_,B_,Δt_]:=Module[{Ax,Ay,Az,S2,S1x,S1y,S0},Ax=CylindricalA[0,0,B][[All,All,1]];
Ay=CylindricalA[0,0,B][[All,All,2]];
Az=CylindricalA[0,0,B][[All,All,3]];
S2=-(($ℏ)^2)/(2$m);
S1x=$q $ℏ Ax/$m;
S1y=$q $ℏ  Ay/$m;
S0=($q^2)(Ax^2+Ay^2)/(2 $m);
GeneralCheby[S2,S1x,S1y,S0,Δt,ΨInitial]]

Calculation of observables

(* Action of operator for x-component of kinetic momentum on a wavefunction Ψ *)
pkinxΨ[Ψ_,A_]:=Module[{Ax},Ax=A[[All,All,1]]; 
-I $ℏ GridxDeriv[Ψ]-$q Ax Ψ]

pkinyΨ[Ψ_,A_]:=Module[{Ay},Ay=A[[All,All,2]];
-I $ℏ GridyDeriv[Ψ]-$q Ay Ψ]

(* x and y components of current density *)
jx[Ψ_,A_]:=Re[Conjugate[Ψ] pkinxΨ[Ψ,A]]/$m 
jy[Ψ_,A_]:=Re[Conjugate[Ψ] pkinyΨ[Ψ,A]]/$m

(* Combine two arrays into an array of 2D vectors, in the form required by ListVectorPlot *)
Create2DVectors[vxArray_,vyArray_]:=MapThread[Function[{vx,vy},{vx,vy}],{vxArray,vyArray},2]

jxy[Ψ_,A_]:=Create2DVectors[jx[Ψ,A],jy[Ψ,A]] 

Absjxy[Ψ_,A_]:=Sqrt[jx[Ψ,A]^2+jy[Ψ,A]^2] (* Absolute value of the current density within the x-y plane *)

(* Action of operator for z-component of kinetic orbital angular momentum on the wavefunction Ψ. This is calculated for a rotation axis located at (x0,y0). *)
LkinzPsi[Ψ_,x0_,y0_,A_]:=Module[{Ax,Ay},Ax=A[[All,All,1]];
Ay=A[[All,All,2]];-I $ℏ (ToGrid[Function[{x,y},x-x0]] GridyDeriv[Ψ]-ToGrid[Function[{x,y},y-y0]] GridxDeriv[Ψ])-$q (ToGrid[Function[{x,y},x-x0]] Ay-ToGrid[Function[{x,y},y-y0]] Ax) Ψ]

(* Expectation value of z-component of kinetic orbital angular momentum. *)
ExpectationLkinz[Ψ_,x0_,y0_,A_]:=Re[Total[Conjugate[Ψ] LkinzPsi[Ψ,x0,y0,A],2]/Total[Abs[Ψ]^2,2]]

(* Expectation values of x and y coordinates. *)
Expectationx[Ψ_]:=Total[ToGrid[Function[{x,y},x]] Abs[Ψ]^2,2]/Total[Abs[Ψ]^2,2]
Expectationy[Ψ_]:=Total[ToGrid[Function[{x,y},y]] Abs[Ψ]^2,2]/Total[Abs[Ψ]^2,2]

(* Vorticity of the current density, defined here as half of the curl of the current density. *)
Ωz[Ψ_,A_]:= 0.5 Re[GridxDeriv[jy[Ψ,A]]-GridyDeriv[jx[Ψ,A]]]

Plotting

ωL[B_]:=-$q B/(2 $m) (* Larmor frequency corresponding to field strength B *)
T[B_]:=2Pi/(2Abs[ωL[B]]) (* Cyclotron period *)

(* Radius of classical cyclotron orbit corresponding to field strength B and transverse momentum pC *)
σ[B_,pC_]:=(Abs[pC]/(2 $m Abs[ωL[B]]))

(* When use ListDensityPlot need to transpose the array first.  It also appears better centred if shifted by one grid point in each direction. *)
ProbDensForPlot[Ψ_]:=((Abs[Transpose[RotateLeft[Ψ,{1,1}]]])^2) 

(* Plotting probability density in x-y plane at a given time.  The colour scale runs from 0 to intensityScaleMax. xLength and yLength specify the plot area, which may be smaller than the simulation area.  yOffest allows the centre of the plot area to be shifted along the y-axis. *)
PlotProbDens[Ψ_,probDensScaleMax_,xLength_,yLength_,yOffset_,opts___]:=ListDensityPlot[ProbDensForPlot[Ψ],opts,PlotRange->{{-xLength/2,xLength/2},{-yLength/2+yOffset,yLength/2+yOffset},{0,probDensScaleMax}},ColorFunction->(Directive[Opacity[Rescale[#,{0,1},{0,((Max[((Abs[Ψ])^2)])/probDensScaleMax)}]],Lighter[Blue,1/3]]&),ColorFunctionScaling->True,DataRange->{{-$Lx/2,$Lx/2},{-$Ly/2,$Ly/2}},AspectRatio->yLength/xLength,Frame->False,Background->None]

(* Similar to PlotProbDens above *)
PlotCurrentDens[Ψ_,A_,currentScaleMax_,xLength_,yLength_,yOffset_]:=ListVectorPlot[jxy[Ψ,A],DataRange->{{-$Lx/2,$Lx/2},{-$Ly/2,$Ly/2}},AspectRatio->yLength/xLength,VectorScale->{0.06,Scaled[1.8],None},Frame->False,VectorColorFunction->(Directive[Opacity[Rescale[#5,{0,1},{0,((Max[Absjxy[Ψ,A]])/currentScaleMax)}]],Darker[Red,1]]&),VectorColorFunctionScaling->True,VectorPoints->16,PlotRange->{{-xLength/2,xLength/2},{-yLength/2+yOffset,yLength/2+yOffset}},Background->None]

(* Plotting evolution of probability density and probability current density.  Accepts as an argument a list containing the wavefunction at a number of different time steps, which can be generated by applying ChebyTDSEuniformB iteratively using NestList. *)
Surface[ΨList_,step_,xLength_,yLength_,yOffset_]:=Polygon[{{((step-1)/(Length[ΨList]-1)),-xLength/2,-yLength/2+yOffset},{((step-1)/(Length[ΨList]-1)),xLength/2,-yLength/2+yOffset},{((step-1)/(Length[ΨList]-1)),xLength/2,yLength/2+yOffset},{((step-1)/(Length[ΨList]-1)),-xLength/2,yLength/2+yOffset}},VertexTextureCoordinates->{{0,0},{1,0},{1,1},{0,1}}]
PlotEvolution[ΨList_,A_,probDensScaleMax_,currentScaleMax_,xLength_,yLength_,yOffset_]:=Show[Map[Function[step,Graphics3D[Style[{Texture[ImageData[Show[PlotProbDens[ΨList[[step]],probDensScaleMax,xLength,yLength,yOffset],PlotCurrentDens[ΨList[[step]],A,currentScaleMax,xLength,yLength,yOffset],Background->None]]],EdgeForm[Gray],Surface[ΨList,step,xLength,yLength,yOffset]},Lighting->{{"Ambient",White}}]]],Range[Length[ΨList]]],Graphics3D[{Gray,Line[{{0,0,0},{1,0,0}}]}],BoxRatios->{3,1,yLength/xLength},Boxed->False,ViewPoint->{1.9,-2,0.7},ImageSize->72*3.125,Background->None,Axes->True,AxesEdge->{{-1,-1},Automatic,Automatic},AxesLabel->{"(Subscript[ω, c])t","x/(Subscript[ρ, B])","y/(Subscript[ρ, B])"},Ticks->{{{0,0},{0.5,"π"},{1,2π}},{-1,0,1},{-2,-1,0,1,2}}]