# How to speed up this TDSE code?

I'm trying to build a 1D time-dependent Schrodinger equation solver, using the Crank-Nicholson method. And I think the code is kind of working now, but the speed is still slow for me.

Here is the code that calculates the wave function $\psi(x,t)$ at certain time given initial wave function $\psi(x,0)$ and a position dependent potential $V(x)$.

TDSE1[V_Function, Ct0_, {xmin_, xmax_}, {tmin_, tmax_, dt_}] :=
Module[{I0, Tmtx, Vmtx, Hmtx, Ct, Cls, dx, tls, xls,
N0Grid = Length@Ct0},
dx = (xmax - xmin)/(N0Grid - 1);
xls = Range[xmin, xmax, dx];
tls = Range[tmin + 0.5 dt, tmax, dt];
Tmtx = -(1/(2 (dx)^2))
SparseArray[{{i_, i_} -> -2, {i_, j_} /; Abs[i - j] == 1 -> 1}, {N0Grid, N0Grid}];(* kinetic energy *)
I0 = SparseArray[{{i_, i_} -> 1.}, {N0Grid, N0Grid}]; (*identity matrix*)
Ct = Ct0;
Vmtx = SparseArray[Band[{1, 1}] -> V /@ xls];(*potential energy*)
Hmtx = Tmtx + Vmtx;(*total Hamiltonian*)
Cls = Table[
Ct = LinearSolve[(I0 + I/2. (Hmtx)*dt), (I0 - I/2. (Hmtx)*dt).Ct]
, {t, tls}];
Cls
]


and here is an example of calculating a wave function using this solver

Ct0 = Array[Exp[-#^2] &, 400, {-10, 10}] // N;(*initial wave function*)
Cls1 = TDSE1[Function[{x}, 0.5 x^2], Ct0, {-10, 10}, {0., 50., 0.01}]; // AbsoluteTiming
(*{1.329230,Null}*)

ListPlot[Re[Cls1[[1 ;; 1000 ;; 100]]], Joined -> True, PlotRange -> All]


And I'm pretty satisfied with this speed.

However, I'm more interested in a time-dependent potential, ie, $V=V(x,t)$, than a time independent one. So I changed the code a little bit so that at every time step, it recalculates the potential Vmtx at that time.

TDSE2[V_Function, Ct0_, {xmin_, xmax_}, {tmin_, tmax_, dt_}] :=
Module[{I0, Tmtx, Vmtx, Hmtx, Ct, Cls, dx, tls, xls,
N0Grid = Length@Ct0},
dx = (xmax - xmin)/(N0Grid - 1);
xls = Range[xmin, xmax, dx];
tls = Range[tmin + 0.5 dt, tmax, dt];
Tmtx = -(1/(2 (dx)^2))
SparseArray[{{i_, i_} -> -2, {i_, j_} /; Abs[i - j] == 1 -> 1}, {N0Grid, N0Grid}];
I0 = SparseArray[{{i_, i_} -> 1.}, {N0Grid, N0Grid}];
Ct = Ct0;
Cls = Table[
Vmtx = SparseArray[Band[{1, 1}] -> (V[t] /@ xls)];
Hmtx = Tmtx + Vmtx;
Ct = LinearSolve[(I0 + I/2. (Hmtx)*dt), (I0 - I/2. (Hmtx)*dt).Ct]
, {t, tls}];
Cls
]


and the function call is also changed slitly

Cls2 = TDSE2[Function[{t}, Function[{x}, 0.5 x^2]], Ct0, {-10, 10}, {0., 50., 0.01}]; // AbsoluteTiming
(*{9.446716, Null}*)


but now it is about 7X slower than before!

I tried to compile the function inside the module, replacing the Vmtx line by something like

Vmtx = SparseArray[Band[{1, 1}] -> (VCpf[t, xls])];


where VCpf is a compiled function defined by

VCpf = Compile[{{t, _Real}, {x, _Real}}, V[t][x], RuntimeAttributes -> {Listable}]


and that makes the problem even worse, with about 20X slower than TDSE1.

So my question is, is it possible to speed up TDSE2 to get about the same speed as TDSE1 or even faster?

-
Your Ct0 definition won't work for <V9, worth to add. –  Kuba Mar 29 '14 at 8:20
@Kuba I'm curious why it won't work for <v9? –  xslittlegrass Mar 29 '14 at 18:32
@xslittlegrass To make it work on version <9, you could replace the definition of Ct0 by this: Ct0 = Table[Exp[-x^2], {x, -10, 10, 20/399}] // N; (the construct where Array fills slots in a range {a,b} given by the third argument doesn't exist in earlier versions). –  Jens Mar 29 '14 at 18:36

You have 2 sources of inefficiency in your code. One is that you don't use the machine-precision for your dx (and therefore xls), and another one is that Band is not the fastest way to build a SparseArray object, and in the time-dependent case you have to build a new one for every time point.

Here is the code which is on the same level of performance as your time-independent code:

ClearAll[TDSE2];
TDSE2[V_Function, Ct0_, {xmin_, xmax_}, {tmin_, tmax_, dt_}] :=
Module[{I0, Tmtx, Vmtx, Hmtx, Ct, Cls, dx, tls, xls,
N0Grid = Length@Ct0, len, ic, jr},
dx = N@(xmax - xmin)/(N0Grid - 1);
xls = Range[xmin, xmax, dx];
tls = Range[tmin + 0.5 dt, tmax, dt];
Tmtx =
-(1/(2 (dx)^2)) SparseArray[{{i_, i_} -> -2, {i_, j_} /; Abs[i - j] == 1 -> 1}, {N0Grid, N0Grid}];
I0 = SparseArray[{{i_, i_} -> 1.}, {N0Grid, N0Grid}];
Ct = Ct0;
len = Length[xls];
ic = Prepend[Range[len], 0];
jr = Transpose[{Range[len]}];
Cls =
Table[
Vmtx =
SparseArray @@ {Automatic, {len, len}, 0, {1, {ic, jr}, V[t] /@ xls}};
Hmtx = Tmtx + Vmtx;
Ct = LinearSolve[(I0 + I/2. (Hmtx)*dt), (I0 - I/2. (Hmtx)*dt).Ct]
,
{t, tls}
];
Cls
]


You can see that I used dx = N@(xmax - xmin)/(N0Grid - 1);, and also the lower-level SparseArray - building code

len = Length[xls];
ic = Prepend[Range[len], 0];
jr = Transpose[{Range[len]}];


and

SparseArray @@ {Automatic, {len, len}, 0, {1, {ic, jr}, V[t] /@ xls}};


See this recent answer for an explanation of this code.

EDIT

As Jens pointed out in comments, an easier option than the low-level SparseArray construction in this particular case would be to use

Vmtx = DiagonalMatrix[SparseArray[(V[t] /@ xls)]];


in place of the above 4 lines of code. This gives comparable performance, without the need to deal with low-level details.

-
Thanks a lot! I didn't know that SparseArray object can be directly constructed. –  xslittlegrass Mar 29 '14 at 18:27
@xslittlegrass Welcome :) –  Leonid Shifrin Mar 29 '14 at 18:36
@LeonidShifrin It looks like you can get the same speed by replacing your SparseArray construct with the simpler Vmtx=DiagonalMatrix[SparseArray[(V[t]/@xls)]] –  Jens Mar 29 '14 at 20:12
@Jens Thanks, I was thinking in this direction too, but did not follow this to the end. My tests show that it's a bit slower, but I will edit your suggestion in. –  Leonid Shifrin Mar 29 '14 at 20:15