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.

Answer 1

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}
 ];