Second Order ODEs and Energy Preservation

In this tutorial we consider an energy-preserving, physical dynamical system, given by a second-order ODE.

TL;DR:

1. To efficiently solve second-order ODEs, just define the problem as a SecondOrderODEProblem.
2. To preserve constant quantities, use the ManifoldUpdate callback; same syntax as DiffEqCallback.jl's ManifoldProjection.

Simulating the Hénon-Heiles system

The Hénon-Heiles model describes the motion of a star around a galactic center, restricted to a plane. It is given by a second-order ODE

\[\begin{aligned} \ddot{x} &= - x - 2 x y \\ \ddot{y} &= y^2 - y - x^2. \end{aligned}\]

Our goal is to numerically simulate this system on a time span $t \in [0, T]$, starting with initial values $x(0)=0$, $y(0) = 0.1$, $\dot{x}(0) = 0.5$, $\dot{y}(0) = 0$.

Transforming the problem into a first-order ODE

A very common approach is to first transform the problem into a first-order ODE by introducing a new variable

\[u = [dx,dy,x,y],\]

to obtain

\[\begin{aligned} \dot{u}_1(t) &= - u_3 - 2 u_3 u_4 \\ \dot{u}_2(t) &= u_4^2 - u_4 - u_4^2 \\ \dot{u}_3(t) &= u_1 \\ \dot{u}_4(t) &= u_2. \end{aligned}\]

This first-order ODE can then be solved using any conventional ODE solver - including our EK1:

using ProbNumDiffEq, Plots

function Hénon_Heiles(du, u, p, t)
    du[1] = -u[3] - 2 * u[3] * u[4]
    du[2] = u[4]^2 - u[4] - u[3]^2
    du[3] = u[1]
    du[4] = u[2]
end
u0, du0 = [0.0, 0.1], [0.5, 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(Hénon_Heiles, [du0; u0], tspan)
sol = solve(prob, EK1());
plot(sol, idxs=(3, 4)) # where `idxs=(3,4)` is used to plot x agains y

Solving the second-order ODE directly

Instead of first transforming the problem, we can also solve it directly as a second-order ODE, by defining it as a SecondOrderODEProblem.

function Hénon_Heiles2(ddu, du, u, p, t)
    ddu[1] = -u[1] - 2 * u[1] * u[2]
    ddu[2] = u[2]^2 - u[2] - u[1]^2
end
prob2 = SecondOrderODEProblem(Hénon_Heiles2, du0, u0, tspan)
sol2 = solve(prob2, EK1());
plot(sol2, idxs=(3, 4))

Benchmark: Solving second order ODEs is faster

Solving second-order ODEs is not just a matter of convenience - in fact, SciMLBase's SecondOrderODEProblem is neatly designed in such a way that all the classic solvers from OrdinaryDiffEq.jl can handle it by solving the corresponding first-order ODE. But, transforming the ODE to first order increases the dimensionality of the problem, and comes therefore at increased computational cost; this also motivates classic specialized solvers for second-order ODEs.

The probabilistic numerical solvers from ProbNumDiffEq.jl have the same internal state representation for first and second order ODEs; all that changes is the measurement model [1]. As a result, we can use the EK1 both for first and second order ODEs, but it automatically specializes on the latter to provide a 2x performance boost:

julia> using BenchmarkTools

julia> @btime solve(prob, EK1(order=3), adaptive=false, dt=1e-2);
  317.336 ms (140561 allocations: 140.41 MiB)

julia> @btime solve(prob2, EK1(order=3), adaptive=false, dt=1e-2);
  136.654 ms (180647 allocations: 47.95 MiB)

Energy preservation

In addition to the ODE given above, we know that the solution of the Hénon-Heiles model has to preserve energy over time. The total energy can be expressed as the sum of the potential and kinetic energies, given by

\[\begin{aligned} \operatorname{PotentialEnergy}(x,y) &= \frac{1}{2} \left( x^2 + y^2 + 2 x^2 y - \frac{2y^3}{3} \right), \\ \operatorname{KineticEnergy}(\dot{x}, \dot{y}) &= \frac{1}{2} \left( \dot{x}^2 + \dot{y}^2 \right). \end{aligned}\]

In code:

PotentialEnergy(x, y) = 1 // 2 * (x^2 + y^2 + 2x^2 * y - 2 // 3 * y^3)
KineticEnergy(dx, dy) = 1 // 2 * (dx^2 + dy^2)
E(dx, dy, x, y) = PotentialEnergy(x, y) + KineticEnergy(dx, dy)
E(u) = E(u...); # convenient shorthand

E (generic function with 2 methods)

Let's have a look at how the total energy changes over time when we numerically simulate the Hénon-Heiles model over a long period of time: Standard solve

longprob = remake(prob2, tspan=(0.0, 1e3))
longsol = solve(longprob, EK1(smooth=false), dense=false)
plot(longsol.t, E.(longsol.u))

It visibly loses energy over time, from an initial 0.12967 to a final 0.12899. Let's fix this to get a physically more meaningful solution.

Energy preservation with the ManifoldUpdate callback

In the language of ODE filters, preserving energy over time amounts to just another measurement model [1]. The most convenient way of updating on this additional zero measurement with ProbNumDiffEq.jl is with the ManifoldUpdate callback.

To do so, first define a (vector-valued) residual function, here chosen to be the difference between the current energy and the initial energy, and build a ManifoldUpdate callback

residual(u) = [E(u) - E(du0..., u0...)]
cb = ManifoldUpdate(residual)

DiscreteCallback{ProbNumDiffEq.var"#condition#140", ProbNumDiffEq.var"#affect!#141"{Int64, Float64, Float64, typeof(Main.residual)}, typeof(SciMLBase.INITIALIZE_DEFAULT), typeof(SciMLBase.FINALIZE_DEFAULT), Nothing}(ProbNumDiffEq.var"#condition#140"(), ProbNumDiffEq.var"#affect!#141"{Int64, Float64, Float64, typeof(Main.residual)}(100, 1.0e-25, 1.0e-15, Main.residual), SciMLBase.INITIALIZE_DEFAULT, SciMLBase.FINALIZE_DEFAULT, Bool[1, 1], nothing)

Then, solve the ODE with this callback

longsol_preserving = solve(longprob, EK1(smooth=false), dense=false, callback=cb)
plot(longsol.t, E.(longsol.u))
plot!(longsol_preserving.t, E.(longsol_preserving.u))

Voilà! With the ManifoldUpdate callback we could preserve the energy over time and obtain a more truthful probabilistic numerical long-term simulation of the Hénon-Heiles model.

References