Jacobian of ParametricNDSolve and FindRoot for the Three Body Problem

Question

The main problem

I took the time to reformulate the question in a more appealing and concise way.

I want to find bifurcations of solutions to the three body problem. In order to do that, I define the symplectic matrix, the Gravitational Potential and the Hamiltonian:

(* Symplectic matrix *)
Clear[mJ]
mJ = ArrayFlatten[{{0, IdentityMatrix[9]}, {-IdentityMatrix[9], 0}}];

(* Gravitational potential *)
Clear[V]
V[x_, m_] := -((#1 - #2).(#1 - #2))^(-1/2) & @@@ Subsets[x, {2}].
    (#1 #2 & @@@ Subsets[m, {2}]);

(* Hamiltonian *)
Clear[H]
H[x_, p_, m_] := Total[p^2/(2 m), 2] + V[x,m];

The main problem is to find a periodic solution of

$$\label{E:HS}\pmatrix{\dot{x}\\\dot{p}} = J \nabla H,\tag{1}$$

where $x$ is the position vector and $p$ is the momentum vector, $J$ the symplectic matrix, and $H$ the Hamiltonian, and then chase it as one of the masses varies (the bifurcation parameter).

A remarkable solution for this system is the Chenciner-Montgomery figure-8 solution:

From this solution I'll be chasing the bifurcation.

Given the Hamiltonian structure of \eqref{E:HS}, a better numerical approach is to use its conserved quantities to extend (unfold) it in a proper way. In this case, the conserved quantities are the linear and angular momenta, along with the Hamiltonian itself. Accordingly, I define the Linear and Angular Momenta:

(* Linear momentum *)
P[p_] := Total@p;

(* Angular momentum *)
L[x_, p_] := Total@MapThread[Cross[#1, #2] &, {x, p}];

Next, I define the variables:

(* Variables *)
x = {{x1, y1, z1}, {x2, y2, z2}, {x3, y3, z3}};
p = {{px1, py1, pz1}, {px2, py2, pz2}, {px3, py3, pz3}};

(* Initial conditions *)
xI = {{xI1, yI1, zI1}, {xI2, yI2, zI2}, {xI3, yI3, zI3}};
pI = {{pxI1, pyI1, pzI1}, {pxI2, pyI2, pzI2}, {pxI3, pyI3, pzI3}};

(* Dummy variables *)
xD = {{x1D, y1D, z1D}, {x2D, y2D, z2D}, {x3D, y3D, z3D}};
pD = {{px1D, py1D, pz1D}, {px2D, py2D, pz2D}, {px3D, py3D, pz3D}};

The unfolded system is (Doedel, et.al.):

$$\begin{multline}\label{E:US} \pmatrix{\dot{x}\\ \dot{p}} = T \nabla H + \alpha \nabla H + \beta_1 \nabla P_x + \beta_2 \nabla P_y \\ + \beta_3 \nabla P_z + \lambda_1 \nabla L_x + \lambda_2 \nabla L_y + \lambda_3 \nabla L_z, \tag{2}\end{multline}$$

where we have rescaled time (in order to work in $0 \le t < 1$). Here, $T$ is the period, and $\alpha$, $\beta_i$, and $\lambda_i$ are the unfolding parameters.

The main purpose is to make the system dissipative in the direction of the central manifold (pretty clever, if you ask me). It can be shown that the solution of the extended system is a solution of the original system when all the folding parameter are equal to zero (Muñoz-Almaraz, et.al.).

So:

(* Unfolded Hamiltonian system *)
threeBSys = MapThread[#1 == #2 &, {D[Through[(Flatten@{x, p})[t]], t],
    (T mJ + a IdentityMatrix[18]).Grad[H[x, p, {m1, m2, m3}], 
    Flatten@{x, p}] + Total@Grad[MapThread[#1 #2 &, {{b1, b2, b3}, P[p]}], 
    Flatten@{x, p}] + Total@Grad[MapThread[#1 #2 &, {{L1, L2, L3}, L[x, p]}], 
    Flatten@{x, p}] /. Map[# -> #[t] &, Flatten@{x, p}]}];

(* Initial conditions equations *)
iC = MapThread[#1 == #2 &, {Through[(Flatten@{x, p})[0]], Flatten@{xI, pI}}];

In order to study the bifurcation of the figure-8 solution to the 3-Body problem, I start with the appropriate coordinates:

(* Initial conditions for the figure 8 *)
{x0, p0} = {{{-1, 0, 0}, {1, 0, 0}, {0, 0, 0}}, {{0.347111, 0.532728, 0},
    {0.347111, 0.532728, 0}, -2 {0.347111, 0.532728, 0}}};
T0 = 6.324449;

And then I define the flow:

Clear[flow1, flowT]
(* Flow at t=1 *)
flow1 = ParametricNDSolveValue[Flatten@{threeBSys, iC},
    Through[(Flatten@{x, p})[1]], {t, 0, 1}, 
    Flatten@{xI, pI, T, m1, m2, m3, a, b1, b2, b3, L1, L2, L3}, 
    MaxSteps -> Infinity, PrecisionGoal -> 12, AccuracyGoal -> 12, 
    Method -> {"ExplicitRungeKutta"}];

(* Flow for 0<t<1 *)
flowT = ParametricNDSolveValue[Flatten@{threeBSys, iC}, 
    Through[(Flatten@{x, p})[t]], {t, 0, 1}, 
    Flatten@{xI, pI, T, m1, m2, m3, a, b1, b2, b3, L1, L2, L3}, 
    Method -> {"DAEInitialization" -> {"Collocation"}}];

Now I'm almost ready to chase the bifurcation with respect to mass m1. As seven extra folding parameters were added to the system, seven more equations must be added to close it. This are, of course, the conserved quantities. If $x_0$ is the starting point (initial condition), and $x_1$ is the point at time $t=1$, then, if $\varphi_t$ is the flow at time $t$,

$$\begin{align} x_1 - \varphi_1(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla H(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla P_x(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla P_y(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla P_z(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla L_x(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla L_y(x_0) &= 0,\\ (x_1 - x_0)^T J \nabla L_z(x_0) &= 0, \end{align}$$ where the first condition is the periodicity condition, the second is the conserved Hamiltonian condition, the third, fourth, and fifth conditions are the conserved linear momenta, and the sixth, seventh, and eighth conditions are the conserved angular momenta.

Now that I've closed the system, I'm ready to chase the bifurcation, but, first, I want to refine my initial point (where all masses are equal to $1$):

(* Fig-8 conditions refinement *)
{x0B, p0B, T0B, m1B, m2B, m3B, aB, b1B, b2B, b3B, L1B, L2B, L3B} =
 {xI, pI, T0, 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.} /. 
 FindRoot[Print[NumberForm[Norm@Flatten@{flow1 @@ 
 Flatten@{xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3} - Flatten@{xI, pI}}, 20]];
 Flatten@{flow1 @@ Flatten@{xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3} - 
 Flatten@{xI, pI}, Flatten[{xI, pI} - {x0, p0}].mJ.(Grad[H[x, p, {1., 1., 1.}], 
 Flatten@{x, p}] /. MapThread[#1 -> #2&, {Flatten@{x, p}, Flatten@{x0, p0}}]), Total@(xI-x0),
 Flatten[{xI, pI} - {x0, p0}].mJ.# & /@ (Grad[L[x, p], Flatten@{x, p}] /.
 MapThread[#1 -> #2 &, {Flatten@{x, p}, Flatten@{x0, p0}}])},
 MapThread[{#1, #2} &, {Flatten@{xI, pI, a, b1, b2, b3, L1, L2, L3}, 
 Flatten@{x0, p0, 0., 0., 0., 0., 0., 0., 0.}}], Evaluated -> False];

The Print command allows me to see how the periodicity condition is behaving. This quickly yields the output

0.003405540054764908
0.003405540054764908
0.0034066427736885
0.003405540055840062
0.00340554005583874
0.003405540055840691
0.003405540054936841
0.003405540055896425
0.003405657294778178

FindRoot::jsing: Encountered a singular Jacobian at the point 
{xI1,yI1,zI1,xI2,yI2,zI2,xI3,yI3,zI3,pxI1,pyI1,pzI1,pxI2,pyI2,pzI2,pxI3,pyI3,pzI3,a,b1,b2,b3,L1,L2,L3} = 
{-1.,0.,0.,1.,0.,0.,0.,0.,0.,0.347111,0.532728,0.,0.347111,0.532728,0.,-0.694222,-1.06546,0.,0.,0.,0.,0.,0.,0.,0.}.
Try perturbing the initial point(s).

The option Jacobian -> "FiniteDifferences" doesn't do anything to help.

Very clumsily, I think, I define the Jacobian matrix using a wrapper for flow1:

(* Flow Wrapper *)
Clear[fWrap, solp]
solp = ParametricNDSolve[Flatten@{threeBSys, iC}, Through[(Flatten@{x, p})[1]], {t, 0, 1},
    Flatten@{xI, pI, T, m1, m2, m3, a, b1, b2, b3, L1, L2, L3}, 
    MaxSteps -> Infinity, Method -> "Adams", PrecisionGoal -> 12, 
    AccuracyGoal -> 12];
fWrap[xEv__, pEv__, T_, m1_, m2_, m3_, aEv_, b1Ev_, b2Ev_, b3Ev_, L1Ev_, L2Ev_, L3Ev_] :=
    (# @@ Flatten@{xEv, pEv, T, m1, m2, m3, aEv, b1Ev, b2Ev, b3Ev, L1Ev, L2Ev, L3Ev} /.
    solp) & /@ ((#@1) & /@ Flatten@{x, p}) - Flatten@{xEv, pEv};

(* Jacobian Matrix *)
Clear[J]
J[xEv__, pEv__, T_, m1_, m2_, m3_, aEv_, b1Ev_, b2Ev_, b3Ev_, L1Ev_, L2Ev_, L3Ev_] :=
    D[Flatten@{ fWrap[xD, pD, T, m1, m2, m3, aD, b1D, b2D, b3D, L1D, L2D, L3D],
    Flatten[{xD, pD} - {xEv, pEv}].mJ.(Grad[H[xD, pD, {m1, m2, m3}], Flatten@{xD, pD}] /. 
    MapThread[#1 -> #2 &, {Flatten@{xD, pD}, Flatten@{xEv, pEv}}]),
    Total@(xD-xEv),Flatten[{xD, pD} - {xEv, pEv}].mJ.#&/@(Grad[L[xD, pD],Flatten@{xD, pD}]/. 
    MapThread[#1 -> #2 &, {Flatten@{xD, pD}, Flatten@{xEv, pEv}}])},
    {Flatten@{xD, pD, aD, b1D, b2D, b3D, L1D, L2D, L3D}}] /.
    MapThread[#1 -> #2 &, {Flatten@{xD, pD, aD, b1D, b2D, b3D, L1D, L2D, L3D},
    Flatten@Flatten@{xEv, pEv, aEv, b1Ev, b2Ev, b3Ev, L1Ev, L2Ev, L3Ev}}]

(I do this to explicitly show the dimensions of flow1 to J, otherwise I get mismatching dimensions errors.)

A brief test shows how inefficient this is:

In[1]:= Timing[J[x0B, p0B, T0, 1., 1., 1., aB, b1B, b2B, b3B, L1B, L2B, L3B]]
Out[1]:= {112.516, <<..>>}

But when I add the option Jacobian -> J[xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3]] to FindRoot, I end up with a very nice refinement:

In[2]:= Timing[FindRoot[<<..>>,
    Jacobian -> J[xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3]];

0.003405540054764908
9.17370958132291*10^(-6)
5.906750301560982*10^(-7)
1.535443284056098*10^(-11)
7.74379558359556*10^(-15)

Out[2]:= {826.234, {{{-0.999846, -5.68972*10^-6, 0.}, {0.999846, 5.68972*10^-6, 0.},
    {2.14173*10^-12, -3.2671*10^-12, 0.}}, {{0.347141, 0.532768, 0.},
    {0.347141, 0.532768, 0.}, {-0.694281, -1.06554, 0.}}, 6.32445,
    1., 1., 1., 0., 0., 0., 0., 0., 0., 0.}}

Question 1

Is there a faster way to implement the Jacobian? I've tried to use ND with little succes

Question 2

How can I sow the Jacobian computed by FindRoot?

Question 3

In general, how can I improve the overall performance of the code? (I'm not sure if I'm clogging the code somewhere with bad practice)

---

Update

Reading How to work with Experimental`NumericalFunction?, I've been able to define a remarkably faster Jacobian:

(* Numerical Jacobian *)
Clear[JN]
JN[xEv_?(MatrixQ[#, NumericQ] &), pEv_?(MatrixQ[#, NumericQ] &), 
    TEv_?NumericQ, m1Ev_?NumericQ, m2Ev_?NumericQ, m3Ev_?NumericQ, 
    aEv_?NumericQ, b1Ev_?NumericQ, b2Ev_?NumericQ, b3Ev_?NumericQ, 
    L1Ev_?NumericQ, L2Ev_?NumericQ, L3Ev_?NumericQ] := Module[{fluxWrap, smallJ, ppDim},
    ppDim = Length[Flatten@{xEv, pEv}];
    fluxWrap = Experimental`CreateNumericalFunction[
    Flatten@{xD, pD, aD, b1D, b2D, b3D, L1D, L2D, L3D}, 
    ParametricNDSolveValue[Flatten@{threeBSys, iC}, Through[(Flatten@{x, p})[1]], {t, 0, 1}, 
    Flatten@{xI, pI, TI, m1I, m2I, m3I, aI, b1I, b2I, b3I, L1I, L2I, L3I},
    MaxSteps -> Infinity, PrecisionGoal -> 12, AccuracyGoal -> 12,
    Method -> {"ExplicitRungeKutta"}] @@ (Flatten@{xD, pD, TEv, 
    m1Ev, m2Ev, m3Ev, aD, b1D, b2D, b3D, L1D, L2D, L3D}), {ppDim},
    Jacobian -> "FiniteDifference"];
    smallJ = D[Flatten@{Flatten[{xD, pD} - {xEv, pEv}].mJ.(Grad[H[xD, pD,{m1Ev, m2Ev, m3Ev}],
    Flatten@{xD, pD}] /. MapThread[#1 -> #2 &, {Flatten@{xD, pD}, Flatten@{xEv, pEv}}]),
    Total@(xD-xEv), Flatten[{xD, pD} - {xEv, pEv}].mJ.# & /@ (Grad[L[xD, pD], 
    Flatten@{xD, pD}] /. MapThread[#1 -> #2 &, {Flatten@{xD, pD}, Flatten@{xEv, pEv}}])},
    {Flatten@{xD, pD, aD, b1D, b2D, b3D, L1D, L2D, L3D}}] /. 
    MapThread[#1 -> #2 &, {Flatten@{xD, pD, aD, b1D, b2D, b3D, L1D, L2D, L3D}, 
    Flatten@Flatten@{xEv, pEv, aEv, b1Ev, b2Ev, b3Ev, L1Ev, L2Ev, 
     L3Ev}}];
    Join[fluxWrap["Jacobian"[Flatten@{xEv, pEv, aEv, b1Ev, b2Ev, b3Ev, L1Ev, L2Ev, L3Ev}]] -
    Join[IdentityMatrix[ppDim], ConstantArray[0, {ppDim, 7}], 2], smallJ]
    ]

(The function calculates the Jacobian of the flow at $t=1$, subtracts an identity matrix (due to the fixed point), and joins the corresponding part of the conserved equations.)

In[3]:= Timing[JN @@ {x0B, p0B, T0B, m1B, m2B, m3B, aB, b1B, b2B, b3B, L1B, L2B, L3B}]
Out[3]:= {0.359375, <<...>>}

The problem now lies in that I've been unable to pass this Jacobian to FindRoot:

In[4]:= FindRoot[<<..>>, Jacobian -> JN[xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3]]

FindRoot::njnum: The Jacobian is not a matrix of numbers at {xI1,yI1,zI1,xI2,yI2,zI2,xI3,yI3,zI3,pxI1,pyI1,pzI1,pxI2,pyI2,pzI2,pxI3,pyI3,pzI3,a,b1,b2,b3,L1,L2,L3} = {-1.,0.,0.,1.,0.,0.,0.,0.,0.,0.347111,0.532728,0.,0.347111,0.532728,0.,-0.694222,-1.06546,0.,0.,0.,0.,0.,0.,0.,0.}.

At this stage, I'm able to implement a Newton method, but I'd really like to take advantage of FindRoot.

Here are some nice near collision continuation orbits:

I guess at this point, sowing the Jacobian is no longer a priority, much as it is to correctly pass it to FindRoot, and the overall improvement of the code.

Thanks for paying attention.

Answer 1

There's no need to struggle with Jacobian, all we need to do is to take away Evaluated -> False and modify the equations accordingly.

It's not clearly mentioned, but the document of Jacobian (it doesn't have a separate page, but is included in the document of FindRoot) suggests that, when calcuating Jacobian matrix, FindRoot performs symbolic differentiation if the model can be symbolically differentiated, and use finite difference if the model is a blackbox. ParametricFunction can be "symbolicly" differentiated, but it seems that, if Evaluated -> False is set, any function inside FindRoot will be treated as a black box.

Remark

1. I think this is the reason why Jacobian -> "FiniteDifferences" doesn't help in this case: the warning FindRoot::jsing is exactly the result of calculating Jacobian matrix using finite difference.

2. I know Evaluated -> False is necessary for the FindRoot[Print…… trick, but now we have EvaluationMonitor, right?

Still, simply taking away Evaluated -> False will only result in warning FindRoot::nveq, because Flatten@{xI, pI} is an explicit list while flow1 @@ Flatten@{xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3} is an implicit one, what happens here is the same as:

(* "lst" is supposed to become a list later *)
lst + {a, b, c}
(* {a + lst, b + lst, c + lst} *)

so we need to stop the auto-threading, with the _?NumericQ trick.

The following is the modified code. I've taken away those AccracyGoal because of the reason mentioned here:

Is manual adjustment of AccuracyGoal and PrecisionGoal useless?

flow1 = ParametricNDSolveValue[{threeBSys, iC},
      (Flatten@{x, p})[1] // Through, {t, 0, 1}, 
      Flatten@{xI, pI, T, m1, m2, m3, a, b1, b2, b3, L1, L2, L3}, 
      MaxSteps -> Infinity]

(helper[##] = Flatten@{xI, pI}) & @@ ((Pattern[#, _])?NumericQ & /@ Flatten@{xI, pI});
helper // DownValues
(* {HoldPattern[
   helper[xI1_?NumericQ, yI1_?NumericQ, zI1_?NumericQ, xI2_?NumericQ, yI2_?NumericQ, 
    zI2_?NumericQ, xI3_?NumericQ, yI3_?NumericQ, zI3_?NumericQ, pxI1_?NumericQ, 
    pyI1_?NumericQ, pzI1_?NumericQ, pxI2_?NumericQ, pyI2_?NumericQ, pzI2_?NumericQ, 
    pxI3_?NumericQ, pyI3_?NumericQ, pzI3_?NumericQ]] :> {xI1, yI1, zI1, xI2, yI2, zI2, 
   xI3, yI3, zI3, pxI1, pyI1, pzI1, pxI2, pyI2, pzI2, pxI3, pyI3, pzI3}} *)

test = FindRoot[
      Flatten@{flow1 @@ Flatten@{xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3} - 
        helper @@ Flatten@{xI, pI} , 

       Flatten[{xI, pI} - {x0, p0}].mJ.(Grad[H[x, p, {1., 1., 1.}], 
            Flatten@{x, p}] /. 
          MapThread[#1 -> #2 &, {Flatten@{x, p}, Flatten@{x0, p0}}]), Total@(xI - x0),
        Flatten[{xI, pI} - {x0, p0}].mJ.# & /@ (Grad[L[x, p], Flatten@{x, p}] /.
           MapThread[#1 -> #2 &, {Flatten@{x, p}, Flatten@{x0, p0}}])},
      MapThread[{#1, #2} &, {Flatten@{xI, pI, a, b1, b2, b3, L1, L2, L3}, 
        Flatten@{x0, p0, 0., 0., 0., 0., 0., 0., 0.}}]]; // AbsoluteTiming
(* {0.296140, Null} *)

Norm[flow1 @@ Flatten@{xI, pI, T0, 1., 1., 1., a, b1, b2, b3, L1, L2, L3} - 
   helper @@ Flatten@{xI, pI}] /. test 

(* 1.55795*10^-13 *)