Vectorized RootFind for interpolated function

Question

I've spent the last few days trying to vectorize my code to speed up root-finding on a grid. Apologies in advance if my code is poorly formatted. I have a fair amount of setup for the problem. It really starts from here, where I solve a large system of equations and find an interpolated function trcfnint:

Clear["`*"];
E2H[E_, c_] := 1/(c^2 - 1) (E)
E2H1 = Compile[{{E, _Real}, {c, _Real}}, 1/(c^2 - 1) (E)]
c = 1.1
E1 = 0.5
H = E2H1[E1, c]
If[c < 1 , 
  period = Abs[
    2 NIntegrate[
      1/Sqrt[2 H - 1/(c^2 - 1) (u^4/2 - u^2)], {u, -Sqrt[
         1 - Sqrt[1 + 4 (c^2 - 1) H]], 
       Sqrt[1 - Sqrt[1 + 4 (c^2 - 1) H]]}]], 
  period = Abs[
    2 NIntegrate[
      1/Sqrt[2 H - 1/(c^2 - 1) (u^4/2 - u^2)], {u, -Sqrt[
         1 + Sqrt[1 + 4 (c^2 - 1) H]], 
       Sqrt[1 + Sqrt[1 + 4 (c^2 - 1) H]]}]]];
sgtwv = NDSolveValue[{(c^2 - 1) u2''[t] + u2[t]^3 - u2[t] == 0, 
    u2[0] == 0, u2'[0] ==  Sqrt[2 H] },
   u2, {t, 0, 50}];
grid = Table[a + I b, {a, -1, 1, 0.02}, {b, -1, 1, 0.02}];
gvals = Riffle[Flatten[grid], Flatten[grid]];
A[t_] := SparseArray[
  Flatten[Table[{{2 j - 1, 2 j} -> 
      1, {2 j, 
       2 j - 1} -> -1/(c^2 - 1) (gvals[[j]]^2/(c^2 - 1) + 
         3 sgtwv[t]^2 - 1)}, {j, Length[gvals]}], 1]]
initialc = 
  Table[If[Mod[i, 4] == 1 || Mod[i, 4] == 0, 1, 0], {i, 
    2 Length[gvals]}];
ff = NDSolveValue[{q'[t] == A[t].q[t], q[0] == initialc}, {q}, {t, 0, 
     3 period}]; // AbsoluteTiming
ggg = Flatten[Through[ff[period]]];
trc = ggg[[1 ;; -1 ;; 4]] + ggg[[4 ;; -1 ;; 4]];
listint = ArrayReshape[trc, Dimensions[grid]];
trcfnint = ListInterpolation[listint, {{-1, 1}, {-1, 1}}]
evs1[a_?VectorQ, b_?VectorQ, d_?VectorQ] := 
  trcfnint[a, b] - 2 Cos[c period (a + I b)/(c^2 - 1) + d];
evs2[a_, b_, d_] := 
  trcfnint[a, b] - 2 Cos[c period (a + I b)/(c^2 - 1) + d];

My goal is to find (a,b,d) in the interpolated range such that evs2 is 0 (evs1 is just the vectorized version). This is very quick and easy when b = 0:

imAxis1 = {}; Do[
 plot1 = Plot[
   trcfnint[a, 0] - 2 Cos[c period a/(c^2 - 1) + \[Theta]], {a, 0, 1},
    Mesh -> {{0}}, MeshFunctions -> {#2 &}, 
   MeshStyle -> PointSize[Medium], PlotRange -> {{0, 1}, {0, 0}}, 
   PlotPoints -> 100]; 
 AppendTo[imAxis1, 
  Sort@Cases[Normal@plot1, Point[{x_, y_}] -> x, 
    Infinity]], {\[Theta], 0, 2 Pi + 0.1, 0.01}];
imAxis1 = Flatten[imAxis1]; imAxis1 = 
 Transpose[{ConstantArray[0, Length[imAxis1]], imAxis1}];
ListPlot[imAxis1, PlotStyle -> Directive[Red, PointSize[0.005]], 
 PlotRange -> {{-1, 1}, {-1, 1}}, 
 FrameLabel -> {Style["Re(\[Lambda])", FontSize -> 20], 
   Style["Im(\[Lambda])", FontSize -> 20]}, 
 LabelStyle -> Directive[Black], ImageSize -> Large, AspectRatio -> 1,
  Frame -> True, Axes -> False]

When b != 0, the most reliable results I have are using the non-vectorized function evs2 in the following code:

total = {};
Table[Do[
    root1 = 
     FindRoot[{Re[evs2[aa, b, dd]] == 0, 
       Im[evs2[aa, b, dd]] == 0}, {{aa, a}, {dd, theta}}, 
      MaxIterations -> 20, AccuracyGoal -> 6];
     test = Abs[evs2[aa, b, dd]] /. root1;
    If[Abs[test] < 10^(-5), AppendTo[total, {b, root1[[1]][[2]]}]];
    Break, {theta, 0, 2 Pi, 2}], {a, {0.5, 0.6}}, {b, 0, 0.18, 
    0.001}]; // AbsoluteTiming
ListPlot[total, PlotStyle -> Directive[Red, PointSize[0.005]], 
 PlotRange -> {{-1, 1}, {-1, 1}}, 
 FrameLabel -> {Style["Re(\[Lambda])", FontSize -> 20], 
   Style["Im(\[Lambda])", FontSize -> 20]}, 
 LabelStyle -> Directive[Black], ImageSize -> Large, AspectRatio -> 1,
  Frame -> True, Axes -> False]

Running this, I only take the upper half-bubble. The initial values for a are chosen because they are in a spectral gap on the a-axis and this saves a lot of time in the calculations. In general I would like the ability to quickly solve for many values of a. I have looked at various different questions on this vectorization:Findroot with a precompiled function with parameters,Jacobian for parallelised FindRoot with multiple variables,FindRoot with vector functions were the most useful. I can post my successful vectorizations, however they suffer from being slower than the unvectorized implementation. If anyone has some advice that would be much appreciated.

Updated with the beginning of the code. Apologies. – user71709 Apr 05 '20 at 08:54 — user71709, Apr 05 '20 at 08:54

Henrik Schumacher · Answer 1 · 2020-04-05T13:23:53.507

Not an actual answer to your problem, but when I tried to run your code, I observed that the solve for ff is awfully slow: It takes two minutes on my machine.

For numerical input t, the following function Abackend is 300 times faster than A[t] on my machine (and does the same of course).

Abackend = With[{n = 2 Length[gvals]},
   With[{
     a = -gvals^2/(c^2 - 1)^2 - 1./(1. - c^2),
     b = 3./(1. - c^2) ,
     pat = 
      Join[Partition[Range[n], 2], 
       Transpose[Transpose[Partition[Range[n], 2]][[{2, 1}]]]]
     },
    t \[Function] SparseArray[
      pat -> Join[ConstantArray[1. + I 0., n/2], a + b sgtwv[t]^2],
      {n, n}, 0. + 0. I
      ]
    ]
   ];

We can enforce numerical evaluation of Abackend from within NDSolveValue by using the following wrapper:

ClearAll[A];
A[t_?NumericQ] := Abackend[t];

Now, the solve for ff should take only a fifth of the time.

By the way, you seem to need only the value of ff at time period. You can obtain it directly with

ggg = NDSolveValue[
 {q'[t] == A[t].q[t], q[0] == initialc}, 
 q[period], 
 {t, 0, 3 period}
 ];

Great suggestion, thank you very much. – user71709 Apr 05 '20 at 11:33 — user71709, Apr 05 '20 at 11:33

Vectorized RootFind for interpolated function

1 Answers1