How to be more efficient/faster with this MatrixExp-vector multiplication?

Question

I am looking for expertise that helps to improve speed of the code below. First, a little bit of background:

There is some system of differential equations $\dot{\vec{m}}(t)=S(t)\vec{m}(t)$. In order to solve this without using NDSolve/NDSolveValue for this one can just proceed as $$\vec{m}(T)=\left(\prod\limits_{j=N}^1 e^{S(j\,dt)}\right) \cdot\vec{m}(0)$$ where $dt$ is length of one timestep and $N=T/dt$ the number thereof. Obviously, this method involves the matrix exponential of $S$ which can be time consuming, especially since $N$ should be on the order of $10^3$ for my purposes to achieve reasonable accuracy.

Let me define the system in Mathematica code (the important variable is mat which corresponds to $S(t)$ - the tiny rest is needed to construct a sample matrix that is almost of same dimension and "sparsity" as the ones I am actually dealing with)

ClearAll[init, listDiag, listOffDiag, valOffDiag, tmp, mat, col, squ];

col[mat_?MatrixQ] := Flatten[Transpose[mat]]; (* stack columns of a matrix *)

squ[list_?VectorQ] := Transpose[ArrayReshape[list, {Sqrt[Length@list], Sqrt[Length@list]}]]; (* transform stacked column form into square matrix again *)

(* preliminary definitions of initial condition and `tmp` that is needed to construct `mat` *)
init = ConstantArray[0, {36, 36}];
init[[8, 8]] = 1;

listOffDiag = {{13, 1}, {13, 7}, {14, 2}, {14, 8}, {15, 3}, {15, 9}, {16, 4}, {16, 10}, {17, 5}, {17, 11}, {18, 6}, {18, 12}, {19, 1}, {19, 7}, {20, 2}, {20, 8}, {21, 3}, {21, 9}, {22, 4}, {22,10}, {23, 5}, {23, 11}, {24, 6}, {24, 12}, {25, 1}, {25, 7}, {26,2}, {26, 8}, {27, 3}, {27, 9}, {28, 4}, {28, 10}, {29, 5}, {29, 11}, {30, 6}, {30, 12}};

listDiag = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}, {6, 6}, {7, 7}, {9, 9}, {10, 10}, {11, 11}, {13, 13}, {14, 14}, {15, 15}, {16,16}, {17, 17}, {18, 18}, {19, 19}, {20, 20}, {21, 21}, {22, 22}, {23, 23}, {24, 24}, {25, 25}, {26, 26}, {27, 27}, {28, 28}, {29, 29}, {30, 30}, {31, 31}, {33, 33}, {34, 34}, {35, 35}};

valOffDiag = t*RandomReal[{-5, 5}, Length@listOffDiag];
tmp = SparseArray[Join[Thread[Rule[listDiag, RandomReal[{-100, 100},Length@listDiag]]],Thread[Rule[Table[{30 + i, 30 + i}, {i, 1, 6}],ConstantArray[0, 6]]]]] + SparseArray[Join[Thread[Rule[listOffDiag,valOffDiag]], {{36, 36} -> 0}]] + Transpose@SparseArray[Join[Thread[Rule[listOffDiag, valOffDiag]],{{36, 36}->0}]];

mat = KroneckerProduct[tmp, IdentityMatrix[36]] + KroneckerProduct[IdentityMatrix[36], Transpose@tmp];

Now here are my two approaches. Observing AbsoluteTiming of

MatrixExp[-I*SparseArray[ArrayRules[mat] /. t -> 3,Dimensions[mat]]].col[init]; // AbsoluteTiming
MatrixExp[-I*SparseArray[ArrayRules[mat] /. t -> 3, Dimensions[mat]],col[init]]; // AbsoluteTiming

yields an order of magnitude improvement of the latter over the former (0.254 vs 0.027) on my machine. So instead of first computing the matrix product surrounded by the parentheses in the system of ODEs above, it should be faster to use MatrixExp[matrix,vector] sequentially.

evol1[mat_, initial_, ti_, tf_] := Module[
   {dt = (tf - ti)/10, res, d = Dimensions[mat][[1]]},
   res = Prepend[Table[MatrixExp[-I*SparseArray[ArrayRules[mat]/.t->i, Dimensions[mat]]], {i,ti, tf, dt}], col[initial]];
   Return[squ[Apply[Dot, Reverse[res]]]]];

evol2[mat_, initial_, ti_, tf_] := Module[
   {dt = (tf - ti)/10, res = col[initial]},
   Do[res = MatrixExp[-I*SparseArray[ArrayRules[mat]/.t->i,Dimensions[mat]],res], {i, ti, tf, dt}];
   Return[squ[res]]];

evol1 is the straightforward method that computes the matrix product of all matrix exponentials and then applies it to the initial vector. evol2 makes use of MatrixExp[matrix,vector]. Comparing speed and results:

res1 = evol1[mat, init, 0.01, 10]; // AbsoluteTiming
(* {2.634993, Null} *)
res2 = evol2[mat, init, 0.01, 10]; // AbsoluteTiming
(* {0.302688, Null} *)
Chop[res1 - res2] == ConstantArray[0, {36, 36}]
(* True *)

I feel like the Do in evol2 is all but efficient but I do not get any idea about how to replace it. Is there any possibility to increase speed - not caring about memory usage?

Edit I am sorry for the initial confusion arising from a copy-paste error that affected the results.

Answer 1

Let me first start with a cleaned up version of your evol2:

Clear@evol3
evol3[mat_, initial_, ti_, tf_] := 
  Module[{dt = (tf - ti)/10, res = col[initial]}, 
   Do[res = MatrixExp[-I*mat, res], {t, ti, tf, dt}];
   squ[res]];

Note that you don't need the complicated way to substitute numerical values for t, as Do uses Block internally which does this for you. (This results in a speedup of factor ~2 for me.) Also, you don't need Return here. Otherwise this is a straightforward and perfectly fine way to implement this. Note that one of the bottlenecks of evol1 is almost certainly the prepending of new results to the list.

As Marius already pointed out, Fold is a functional alternative to the Do loop. My stab at it looks like this:

Clear@evol4
evol4[mat_, initial_, ti_, tf_] := 
  Module[{dt = (tf - ti)/10, init = col[initial], func},
   func[res_, s_] := Block[{t = s}, MatrixExp[-I*mat, res]];
   squ@Fold[func, init, Range[ti, tf, dt]]];

Note that there are many different (and maybe nicer) ways to define func, but I don't want to complicate things any more right now. Speedwise evol3 and evol4 are very similar. I'm pretty sure there are faster ways to compute this.

Answer 2

Your approach is certainly faster than using NDSolve, but I think there is a mistake in the implementation, and the accuracy leaves much to be desired. Your ODE is:

$$m'(t) = -i S(t) . m(t)$$

This ODE can be discretized as follows:

$$m'(t) \Delta t \approx m(t + \Delta t) - m(t) \ \approx - i S(t) . m(t) \Delta t$$

which can be rewritten as:

$$m(t + \Delta t)\approx m(t)-i S(t) \Delta t\, .m(t)$$

or

$$m(t+\Delta t)\approx (I-i S(t) \Delta t) \, .m(t)$$

Here $I$ is the identity matrix. Using the exponential to approximate the RHS, we obtain:

$$m(t+\Delta t)\approx e^{-i S(t) \Delta t}.m(t)$$

which differs from yours with the presence of $\Delta t$ in the exponent. I will rewrite your evol2 function to include this $\Delta t$ term, and to make it a function argument (instead of always using $(t_f-t_i)/10$). Also, I won't bother reshaping the vector into a matrix:

evol[mat_, initial_, {t_, ti_, tf_, dt_}] := Module[{res=col[initial]},
    Do[res = MatrixExp[-I mat dt, res], {t, ti, tf, dt}];
    res
]

Now, let's solve your ODE using NDSolveValue:

res = NDSolveValue[{m'[t] == -I mat . m[t], m[0] == col[init]}, m, {t, 0, 10}]; //AbsoluteTiming

{10.126, Null}

Let's compare the two approaches at time $t=1$:

discrete1 = evol[mat, init, {t, 0, 1, .1}]; //AbsoluteTiming
discrete2 = evol[mat, init, {t, 0, 1, .01}]; //AbsoluteTiming
discrete3 = evol[mat, init, {t, 0, 1, .001}]; //AbsoluteTiming
MinMax @ ReIm[res[1] - discrete1]
MinMax @ ReIm[res[1] - discrete2]
MinMax @ ReIm[res[1] - discrete3]

{0.042896, Null}

{0.336994, Null}

{3.2513, Null}

{-0.0549661, 0.11708}

{-0.00586709, 0.0119637}

{-0.000589242, 0.00119601}

As expected, the discrete approach timing increases linearly with the number of steps, and the accuracy also increases linearly. However, even with a $\Delta t$ of .001 the error is rather large. The situation is much worse at time $t=10$:

discrete1 = evol[mat, init, {t, 0, 10, .1}]; //AbsoluteTiming
discrete2 = evol[mat, init, {t, 0, 10, .01}]; //AbsoluteTiming
discrete3 = evol[mat, init, {t, 0, 10, .001}]; //AbsoluteTiming
MinMax @ ReIm[res[10] - discrete1]
MinMax @ ReIm[res[10] - discrete2]
MinMax @ ReIm[res[10] - discrete3]

{0.347939, Null}

{3.38722, Null}

{34.2144, Null}

{-0.273781, 0.294955}

{-0.0395226, 0.0314886}

{-0.00394035, 0.0032858}

I think it's much better to just use NDSolve and rely on the error control that is built-in to the function.

Timing update

I forgot to mention that my version of evol is essentially the same as @sebhofer's (my previous edit erroneously asserted that evol was faster):

r1 = squ @ evol[mat, init, {t, 0, 10, 1}]; //RepeatedTiming
r2 = evol2[mat,init, 0, 10]; //RepeatedTiming
r3 = evol3[mat, init, 0, 10]; //RepeatedTiming

r1 === r2 === r3

{0.037, Null}

{0.16, Null}

{0.038, Null}

True

Answer 3

Let me at least show you how to avoid the Do loop while I think about other possible improvements.

The idiomatic way is to use Fold. Here is my version:

evol3[mat_, initial_, ti_, tf_] := 
  Module[{dt = (tf - ti)/10, res = col[initial], 
    ar = ArrayRules[mat]},
  squ[Fold[
  MatrixExp[-I*SparseArray[ar /. t -> #2, Dimensions[mat]], #1] &, 
  res, Range[ti, tf, dt]]]
];

Notice that I don't compute the ArrayRules every time, since it can be done once. But this still only gives a minor improvement over your evol2, from 0.23 to 0.21 on my machine.