How to speed up by Compile?

Question

I need to NDSolve a system many times by scanning some of its parameters and do some matrix calculation with the (discretized) solutions. The example is the following code. The system eqs looks formidable, but the concrete form doesn't matter.

Because the computation is quite heavy (I actually need much denser meshes for scanning the parameters, I mean the lists I defined in the middle), I was wondering if any further speed improvement possible. Probably by properly Compile the program? I've heard of great gain by compile to C or machine code. But I somehow got lost in the related documentation. Any suggestion would be appreciated.

Note that the NDSolve part sometimes takes less or comparable time than the matrix calculations followed, depending on how dense we scan particular parameters.

eqs = {I Derivative[1][B1][t] == 
    Abs[B1[t]]^2 (m B1[t] + B3[t] (x - I y + Cos[t] - I Sin[t])) + 
     Abs[B2[t]]^2 (m B1[t] + 
        B3[t] (x - I y + Cos[t] - I Sin[t])) + (B1[t] Conjugate[
          B3[t]] + B2[t] Conjugate[B4[t]]) (-m B3[t] + 
        B1[t] (x + I y + Cos[t] + I Sin[t])), 
   I Derivative[1][B2][t] == 
    Abs[B1[t]]^2 (m B2[t] + B4[t] (x - I y + Cos[t] - I Sin[t])) + 
     Abs[B2[t]]^2 (m B2[t] + 
        B4[t] (x - I y + Cos[t] - I Sin[t])) + (B1[t] Conjugate[
          B3[t]] + B2[t] Conjugate[B4[t]]) (-m B4[t] + 
        B2[t] (x + I y + Cos[t] + I Sin[t])), 
   I Derivative[1][B3][t] == 
    B1[t] (m (B3[t] Conjugate[B1[t]] + 
           B4[t] Conjugate[B2[t]]) + (Abs[B3[t]]^2 + 
           Abs[B4[t]]^2) (x + I y + Cos[t] + I Sin[t])) + 
     B3[t] (-m (Abs[B3[t]]^2 + 
           Abs[B4[t]]^2) + (B3[t] Conjugate[B1[t]] + 
           B4[t] Conjugate[B2[t]]) (x + Cos[t] - I (y + Sin[t]))), 
   I Derivative[1][B4][t] == 
    B2[t] (m (B3[t] Conjugate[B1[t]] + 
           B4[t] Conjugate[B2[t]]) + (Abs[B3[t]]^2 + 
           Abs[B4[t]]^2) (x + I y + Cos[t] + I Sin[t])) + 
     B4[t] (-m (Abs[B3[t]]^2 + 
           Abs[B4[t]]^2) + (B3[t] Conjugate[B1[t]] + 
           B4[t] Conjugate[B2[t]]) (x + Cos[t] - I (y + Sin[t])))};
ic = {B1[-tmax] == 1, B2[-tmax] == 0, B3[-tmax] == 0, B4[-tmax] == 1};
var = {B1, B2, B3, B4};
tmax = 25; m = 0.4;
tlist = Table[i, {i, 10, 30, 5}];
ttmax = tmax; ttmesh = ttmax/10; ttlist = Table[tt, {tt, -ttmax, ttmax, ttmesh}];
rmax = 2.0; rmesh = 0.1; rlist = Table[r, {r, -rmax, rmax, rmesh}];
omax = 1.0; omesh = omax/60; olist = Table[o, {o, -omax, omax, omesh}];
smat = Table[Exp[-(t1 - myt)^2/5 + I o t1], {o, olist}, {myt, tlist}, {t1, ttlist}];
G0 = {{x^2, y}, {-y, x y}};
data = ParallelTable[
   Module[{sol = Last[#] & /@ (First@NDSolve[{eqs, ic}, var, {t, -tmax, tmax}]), Bsol},
    Bsol = Table[Map[#@t1 &, sol, {1}], {t1, ttlist}]; 
    Map[#.G0.ConjugateTranspose[#] &@(ttmesh Partition[#1.Bsol, 2]) &, smat, {2}]], {x, rlist}, {y, rlist}];

@xzczd 1. I actually need much denser mesh for scanning the parameters, I mean the lists I defined. 2. Thanks for point that out. Maybe that helps. But the NDSolve part actually can take much less time than the matrix calculations followed, depending on how dense we scan particular parameters. So I was wondering if any Compile can give some overall improvement. — xiaohuamao, Feb 14 '21 at 12:36

xzczd · Accepted Answer · 2021-02-15T14:42:02.890

2

tst = 
   Block[{x = rlist[[1]], y = rlist[[1]]}, 
    Module[{sol = (Last[#1] &) /@ First[NDSolve[{eqs, ic}, var, {t, -tmax, tmax}]], 
      Bsol}, Bsol = Table[Map[#1[t1] &, sol, {1}], {t1, ttlist}]; 
     Map[(#1 . G0 . ConjugateTranspose[#1] &)[ttmesh Partition[#1 . Bsol, 2]] &, 
      smat, {2}]]]; // AbsoluteTiming
(* {0.0249631, Null} *)

pnd = ParametricNDSolveValue[{eqs, ic}, 
   ArrayReshape[Outer[#2[#] &, ttlist, var], {Length@ttlist, 2, 2}], {t, -tmax, 
    tmax}, {x, y}];

cmap = With[{ttmesh = ttmesh, G0 = G0, smat = smat}, 
   Compile[{{x, _Real}, {y, _Real}, {Bsol, _Complex, 3}(*,{smat,_Complex,3}*)}, 
    With[{mat = ttmesh smat . Bsol}, 
     Module[{m, n}, {m, n} = Dimensions[mat][[;; 2]]; 
      Table[(mat[[i, j]] . G0 . ConjugateTranspose[mat[[i, j]]]), {i, m}, {j, n}]]], 
    CompilationTarget -> C]];

tst3 = Block[{x = rlist[[1]], y = rlist[[1]]}, 
    With[{Bsol = pnd[x, y]}, cmap[x, y, Bsol(*,smat*)]]]; // AbsoluteTiming
(* {0.0196469, Null} *)

tst3 - tst // Abs // Max
(* 2.8917*10^-14 *)

If you don't have a C compiler installed, take the CompilationTarget -> C away. The corresponding timing will be a bit slower, of course.

edited Feb 15 '21 at 14:42

answered Feb 15 '21 at 07:04

xzczd

65,995
9
163
468

Thank you for the answer! This indeed speeds it up. But when I use it inside ParallelTable or Table as the end of my original post, it's only 1/3 or 1/2 the original time cost, surprisingly far worse than the test you showed. Is it something wrong or inefficient in such a usage? – xiaohuamao Feb 15 '21 at 14:20
@xiaohuamao Oh RepeatedTiming is missleading in this case, because ParametricNDSolveValue caches the previous results… Timing in answer revised. – xzczd Feb 15 '21 at 14:43
@xiaohuamao BTW, it's worth pointing out differential equation solving is much more time-consuming than subsequent matrix calculation, at least for this toy example. If in real cases the matrix calculation is the bottleneck, you may consider compiling the code further. (Something like Compile[…, Table[cmap[…], …]. ) – xzczd Feb 16 '21 at 04:14
Thanks. Most of the calculation is what I show here -- just put it inside Table or ParallelTable finally. But instead of that, will it be further beneficial to use the option RuntimeAttributes -> Listable, Parallelization -> True and make cmap only take [x,y] as the argument? I tried but couldn't figure out how to deal with the input from pnd. – xiaohuamao Feb 16 '21 at 08:13
@xiaohuamao I doubt if RuntimeAttributes -> Listable, Parallelization -> True will help, because you've already used ParallelTable. (Notice Compile[…, CompilationTarget -> C] won't generate parallelized C code, parallelization happens only in Mathematica kernel. ) Anyway you can try it out. You don't need any special treatment for the input from pnd, just observe the behavior of the following sample: cf = Compile[{{a, _Real, 2}}, a[[1, 2]], RuntimeAttributes -> {Listable}]; cf[{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}] – xzczd Feb 16 '21 at 08:29
Sorry, I actually stupidly got stuck in reducing With[{Bsol = pnd[x, y]}, cmap[x, y, Bsol]] directly into the definition of cmap, i.e., getting something like cmap[x, y]. Do we need to Table the whole pnd data to be used before making cmap? – xiaohuamao Feb 16 '21 at 09:14
@xiao If I've understood what you mean correctly, yes. First generate a list of Bsol i.e. Bsollist=(*Parallel*)Table[pnd[x,y],{x,rlist},{y,rlist}], then something like cmaplistable= With[{ttmesh=ttmesh,G0=G0},Compile[{{x,_Real},{y,_Real},{Bsol, _Complex,3}, {smat, _Complex, 3}},With[{mat = ttmesh smat.Bsol}, Module[{m, n}, {m,n} = Dimensions[mat][[;; 2]];Table[(mat[[i, j]].G0.ConjugateTranspose[mat[[i,j]]]), {i, m}, {j,n}]]],CompilationTarget->C,RuntimeAttributes -> {Listable}]]; data5=cmaplistable[#\[Transpose],#,Bsollist,smat] &@ConstantArray[rlist, Length@rlist]; // AbsoluteTiming – xzczd Feb 17 '21 at 03:51

How to speed up by Compile?

1 Answers1