How to speed up this TDSE code?

Question

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?

Answer 1

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.