Optimizing a Parameterized ODE

Let us fit a parameterized ODE to some data. We will use the Lotka-Volterra model as an example. We will use Single Shooting to fit the parameters.

using OrdinaryDiffEq, NonlinearSolve, Plots

Let us simulate some real data from the Lotka-Volterra model.

function lotka_volterra!(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = dx = α * x - β * x * y
    du[2] = dy = -δ * y + γ * x * y
end

# Initial condition
u0 = [1.0, 1.0]

# Simulation interval and intermediary points
tspan = (0.0, 2.0)
tsteps = 0.0:0.1:10.0

# LV equation parameter. p = [α, β, δ, γ]
p = [1.5, 1.0, 3.0, 1.0]

# Setup the ODE problem, then solve
prob = ODEProblem(lotka_volterra!, u0, tspan, p)
sol = solve(prob, Tsit5(); saveat = tsteps)

# Plot the solution
using Plots
plot(sol; linewidth = 3)
Example block output

Let us now formulate the parameter estimation as a Nonlinear Least Squares Problem.

function loss_function(ode_param, data)
    sol = solve(prob, Tsit5(); p = ode_param, saveat = tsteps)
    return vec(reduce(hcat, sol.u)) .- data
end

p_init = zeros(4)

nlls_prob = NonlinearLeastSquaresProblem(loss_function, p_init, vec(reduce(hcat, sol.u)))
NonlinearLeastSquaresProblem with uType Vector{Float64}. In-place: false
u0: 4-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0

Now, we can use any NLLS solver to solve this problem.

res = solve(nlls_prob, LevenbergMarquardt(); maxiters = 1000, show_trace = Val(true),
    trace_level = TraceWithJacobianConditionNumber(25))

Algorithm: LevenbergMarquardt(
   trustregion = LevenbergMarquardtTrustRegion(β_uphill = 1.0),
   descent = GeodesicAcceleration(descent = DampedNewtonDescent(initial_damping = 1.0, damping_fn = LevenbergMarquardtDampingFunction()), finite_diff_step_geodesic = 0.1, α = 0.75)
)

----     -------------        -----------          -------
Iter     f(u) inf-norm        Step 2-norm          cond(J)
----     -------------        -----------          -------
0        5.89103198e+00       1.19656486e-309      Inf
1        5.89103198e+00       0.00000000e+00       3.46713712e+16
26       1.64552816e-11       3.62980383e-10       4.01915942e+01
51       1.06670228e-12       0.00000000e+00       4.01915942e+01
76       8.85513884e-13       6.79774583e-16       4.01915942e+01
101      8.26005930e-13       0.00000000e+00       4.01915942e+01
126      7.72715225e-13       2.60358316e-16       4.01915942e+01
Final    7.54063478e-13
----------------------
res
retcode: Stalled
u: 4-element Vector{Float64}:
 1.499999999999881
 0.9999999999995396
 2.9999999999992846
 0.9999999999997645

We can also use Trust Region methods.

res = solve(nlls_prob, TrustRegion(); maxiters = 1000, show_trace = Val(true),
    trace_level = TraceWithJacobianConditionNumber(25))

Algorithm: TrustRegion(
   trustregion = GenericTrustRegionScheme(method = RadiusUpdateSchemes.Simple),
   descent = Dogleg(newton_descent = NewtonDescent(), steepest_descent = SteepestDescent())
)

----     -------------        -----------          -------
Iter     f(u) inf-norm        Step 2-norm          cond(J)
----     -------------        -----------          -------
0        5.89103198e+00       1.38167631e-309      Inf
1        5.89103198e+00       0.00000000e+00       3.46713712e+16
Final    2.48689958e-14
----------------------
res
retcode: Success
u: 4-element Vector{Float64}:
 1.49999999999999
 0.999999999999994
 3.0000000000000657
 1.0000000000000187

Let's plot the solution.

prob2 = remake(prob; tspan = (0.0, 10.0))
sol_fit = solve(prob2, Tsit5(); p = res.u)
sol_true = solve(prob2, Tsit5(); p = p)
plot(sol_true; linewidth = 3)
plot!(sol_fit; linewidth = 3, linestyle = :dash)
Example block output