1
$\begingroup$

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}];
$\endgroup$
3
  • 3
    $\begingroup$ 1. Given the system isn't that large, I doubt if Compile will help. 2. Why not ParametricNDSolve? $\endgroup$
    – xzczd
    Commented Feb 14, 2021 at 11:52
  • $\begingroup$ @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. $\endgroup$
    – xiaohuamao
    Commented Feb 14, 2021 at 12:36
  • 1
    $\begingroup$ I presume you are already aware of this list? $\endgroup$ Commented Feb 14, 2021 at 12:52

1 Answer 1

2
$\begingroup$
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.

$\endgroup$
7
  • $\begingroup$ 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? $\endgroup$
    – xiaohuamao
    Commented Feb 15, 2021 at 14:20
  • $\begingroup$ @xiaohuamao Oh RepeatedTiming is missleading in this case, because ParametricNDSolveValue caches the previous results… Timing in answer revised. $\endgroup$
    – xzczd
    Commented Feb 15, 2021 at 14:43
  • $\begingroup$ @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[…], …]. ) $\endgroup$
    – xzczd
    Commented Feb 16, 2021 at 4:14
  • $\begingroup$ 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. $\endgroup$
    – xiaohuamao
    Commented Feb 16, 2021 at 8:13
  • $\begingroup$ @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}}}] $\endgroup$
    – xzczd
    Commented Feb 16, 2021 at 8:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.