Improve simulation speed using Compile

Question

I have two square matrices, a hermitian matrix h and an unitary matrix v. I am interested in the following function

step[h_,v_,n_]:=Map[Tr,Dot[NestList[MatrixExp[-I h].#.MatrixExp[I h]&,v,n-1],v]];

which generates a list of n real numbers. The dimensions of the matrices are not so big so MatrixExp is still quite efficient. I want to make the above function faster by using Compile,

stepcp:=Compile[{{m,_Complex,2},{v,_Complex,2},{n,_Integer}}, Map[Tr,Dot[NestList[m.#.ConjugateTranspose[m]&,v,n-1],v]]];

The compiled version is only a few times faster than the uncompiled version. As an example

l = 128;
v = KroneckerProduct[SparseArray[{{1, 1} -> 1., {2, 2} -> -1.}],SparseArray[{i_, i_} -> 1., l]];
h := 1/2 ArrayFlatten[{{PadLeft[1/2 SparseArray[{Band[{1, 2}] -> 1., Band[{2, 1}] -> -1.}, l], {l, 2 l}]}, {PadRight[-(1/2)SparseArray[{Band[{1, 2}] -> 1., Band[{2, 1}] -> -1.}, l], {l, 2 l}]}}] + 1/2 KroneckerProduct[SparseArray[{{1, 1} -> 1., {2, 2} -> -1.}, 2],DiagonalMatrix[SparseArray[Table[RandomReal[{-1., 1.}], {j, l}]]]];

step[m_, u_, n_] := 1/(2 l) Map[Tr,Dot[NestList[MatrixExp[-I m].#.MatrixExp[I m] &, v, n - 1],v]]; 
stepcp := Compile[{{m, _Complex, 2}, {u, _Complex, 2}, {n, _Integer}},1./(2 l) Map[Tr,Dot[NestList[m.#.ConjugateTranspose[m] &, u, n - 1], u]]];

step[h, v, 64]; // AbsoluteTiming
stepcp[Normal[MatrixExp[-I h]], Normal[v], 64]; // AbsoluteTiming

My questions are:

1. How to compile this function more efficiently?
2. The matrices h and v are SparseArray, why can't I directly pass them into compiled functions, why do I have to use Normal[]?
3. The CompilePrint is shown below, is the MainEvaluate in line 1 and 29 a source of slowdown, if so can it be circumvented?

My eventual goal is to run an efficiently compiled version of step[h,v,n] 10000 times and take the Mean. Any suggestions are greatly appreciated.

Answer 1

This gives a slight speed up (~40%) for me:

stepcp = Compile[{{m, _Complex, 2}, {u, _Complex, 2}, {n, _Integer}, {l, _Real}},
  Block[{
    m2 = ConjugateTranspose@m, x = u, ut = Flatten@Transpose@u},
   Join[{1.0}, Table[x = m.x.m2; Flatten[x].ut, {i, n - 1}]/(2 l)]]]

I replaced the NestList with a more procedural approach, calculating the trace of the dot product at each step using the relation Tr[a.b] == Flatten[a].Flatten[Transpose[b]]

I also put l in as an argument - there's no real performance increase as a result but I find it neater than having a callback to the main evaluator.

Oleksandr R. points out that the run time is dominated by calls to the math library so compilation is not really helping much. In fact almost all the gain you got from the compiled version versus the original was due to not computing MatrixExp at each step. You get almost the same speed with this:

step[m_, u_, n_, l_] :=
 With[{m1 = MatrixExp[-I m], m2 = MatrixExp[I m]}, 
  1/(2 l) Map[Tr, Dot[NestList[m1.#.m2 &, v, n - 1], v]]]

Answer 2

1. Sorry, I don't know. Have you tried compiling to C? One point: you should define stepcp with Set (=) rather than SetDelayed (:=) or you recompile it every time you use it.

2. SparseArray is a specialized format of the high-level Mathematica language. It is not supported by compilation. Many operations are optimized to work on sparse arrays, saving unnecessary computation, which may be why your uncompiled method is proportionately faster than you expect. You would essentially have to reimplement sparse arrays to get the same optimizations in your compiled code.

3. Tr is not a compilable function. (Don't ask me why.) Compiled alternatives that I can think of are not as fast as Tr to begin with so the overhead of a MainEvaluate call seems to balance out. An example alternative to Tr:

   Compile[{{m, _Real, 2}}, Sum[Compile`GetElement[m, i, i], {i, Length @ m}]]