Jacobian for parallelised FindRoot with multiple variables

Question

Michael E2 explains how to use parallel processing to speed up multiple FindRoot calculation with the same function in https://mathematica.stackexchange.com/a/163295/40271.

An example of some code which does this is

f[x_] := x^x + (4 - x)^(4 - x) - 20;
df0[x_?VectorQ] = D[f[x], x];
df[x_?VectorQ] := DiagonalMatrix@SparseArray@df0[x];
params = RandomReal[{2, 3}, 10];
FindRoot[f[x], {x, params}, Jacobian :> df[x]]

I would now like to extend this to work in more than one variable. A syntax to do this is quite simple;

f[x_, y_] := {x^x + y^y - 20, x + y - 4}
inits = RandomReal[{2, 3}, {2, 10}];
FindRoot[f[x,y], {{x, inits[[1]]}, {y, inits[[2]]}}]

However I can't find the right syntax to add in an explicit Jacobian. My example code is

f[x_, y_] := {x^x + y^y - 20, x + y - 4}
df0[x_, y_] = Outer[D, f[x, y], {x,y}];
df[x_, y_] := df0[x, y]
inits = RandomReal[{2, 3}, {2, 10}];
Block[{x, y}, 
  FindRoot[f[x, y], {{x, inits[[1]]}, {y, inits[[2]]}}, 
   Jacobian :> df[x, y]]]

And I get an error that The Jacobian is not a matrix of numbers at {x,y} = inits. So I try instead

f[x_, y_] := {x^x + y^y - 20, x + y - 4}
df0[x_, y_] = Outer[D, f[x, y], {x,y}];
SetAttributes[df, Listable];
df[x_, y_] := df0[x, y]
inits = RandomReal[{2, 3}, {2, 10}];
Block[{x, y}, 
  FindRoot[f[x, y], {{x, inits[[1]]}, {y, inits[[2]]}}, 
   Jacobian :> df[x, y]]]

Which produces the same error, however in this case if I evaluate df@@inits, I do indeed get a list of matrices of numbers. So I am confused about this.

I have tried also with explicit ?VectorQ, like so;

f[x_?VectorQ, y_?VectorQ] := {x^x + y^y - 20, x + y - 4}
df0[x_?VectorQ, y_?VectorQ] = 
  Outer[D, {x^x + y^y - 20, x + y - 4}, {x, y}];
df[x_?VectorQ, y_?VectorQ] := df0[x, y];
inits = RandomReal[{2, 3}, {2, 10}];
FindRoot[f[x, y], {{x, inits[[1]]}, {y, inits[[2]]}}, 
 Jacobian :> df[x, y]]

And this doesn't work either.

Who knows how to do this?

Answer 1

The major problem is to have a fast way to obtain a block diagonal matrix. Moreover, there is a problem with the vectorization of df. In the following I skip the vectorization part (for complicated functions in several variables, vectorization is not always more performant).

xx = Table[Compile`GetElement[x, i], {i, 1, 2}];
f = x \[Function] Evaluate[N[{xx[[1]]^xx[[1]] + xx[[2]]^xx[[2]] - 20, xx[[1]] + xx[[2]] - 4}]];
Df = x \[Function] Evaluate[N[D[f[xx], {xx, 1}]]];

Regarding the block diagonal matrix, I compute its column indices and row pointers by hand and store them for later use in DF.

m = Length[xx];
n = 10000;    
ci = Flatten[
   Transpose[{Partition[Partition[Range[m n], 1], m]}[[ConstantArray[ 1, m]]]], 
  2];
rp = Range[0, m m n, m];

F and DF are sort of wrapper functions that constitute the high dimensional system. I use BlockMap in order to apply f and Df to flat list of length 2 n. This is actually where the block diagonal matrix is produced.

F = X \[Function] Flatten@(BlockMap[f, X, m]);
DF = X \[Function] SparseArray @@ {Automatic, {m n, m n}, 
     0., {1, {rp, ci}, Flatten[BlockMap[Df, X, m]]}};

For speeding up even further (and as substitute for vectorization), I also compile f and Df and create wrapper functions for them.

With[{code = f[xx]},
  cf = Compile[{{x, _Real, 1}},
    code,
    CompilationTarget -> "C",
    RuntimeAttributes -> {Listable},
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
  ];
With[{code = Df[xx]},
  cDf = Compile[{{x, _Real, 1}},
    code,
    CompilationTarget -> "C",
    RuntimeAttributes -> {Listable},
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
  ];
cF = X \[Function] cf[Partition[Abs[X], m]];
cDF = X \[Function] SparseArray @@ {Automatic, {m n, m n}, 0., {1, {rp, ci}, Flatten[cDf[Partition[Abs[X], m]]]}};

Now, we have at least 4 ways to solve the systems; let's try them out and compare! (I use Quiet to suppress some errors that get produced in the compiled function, probalby during the line search...)

X = X0 = RandomReal[{2, 3}, {n m}];

Y = Map[x \[Function] FindRoot[f, {{x}}, Jacobian -> Df], 
    Partition[X0, m]]; // AbsoluteTiming // First
cY = Quiet[Map[x \[Function] FindRoot[cf, {{x}}, Jacobian -> cDf], 
     Partition[X0, m]]]; // AbsoluteTiming // First
X = FindRoot[F, {X0}, Jacobian -> DF][[1]]; // AbsoluteTiming // First
cX = Quiet[FindRoot[cF, {X0}, Jacobian -> cDF][[1]]]; // AbsoluteTiming // First

3.29515

2.45185

1.19915

0.353124

Residuals:

Max[Abs[F[Flatten[Y]]]]
Max[Abs[F[Flatten[cY]]]]
Max[Abs[F[X]]]
Max[Abs[F[cX]]]

8.52651*10^-14

8.52651*10^-14

3.55271*10^-15

3.55271*10^-15

And the differences are:

Max[Abs[Flatten[X] - Flatten[Y]]]
Max[Abs[Flatten[X] - Flatten[cY]]]
Max[Abs[Flatten[X] - Flatten[cX]]]

2.88658*10^-15

2.88658*10^-15

6.66134*10^-16