Lotka Volterra fishing problem (Julia Neural Network solve)

This is an implementation of the Lotka Volterra fishing problem using Julia and a neural network by Julius Martensen, see https://github.com/AlCap23/IMPRS2021

The control function was relaxed, i.e. $u(t) \in [0,1]$.

using LinearAlgebra
using Plots
 
using OrdinaryDiffEq
using Flux
using DiffEqFlux
 
# Example for optimal control
 
controller = Chain(Dense(1, 10, relu), Dense(10,10, relu), Dense(10,10, relu), Dense(10, 1, σ))
 
pnn, f = Flux.destructure(controller)
 
function lotka!(du, u, p, t)
    c = f(p)([t])[1]
    du[1] = (1-0.4*c)*u[1] - u[1]*u[2]
    du[2] = u[1]*u[2] - (1+0.2*c)*u[2]
    du[3] = (u[1] - 1)^2 + (u[2] -1)^2
end
 
u0 = Float32[0.5; 0.7; 0.0]
tspan = (0.0f0, 12.0f0)
Δt = 12.0f0/100
 
prob = ODEProblem(lotka!, u0, tspan, pnn)
solution = solve(prob, Tsit5(), saveat = Δt)
length(solution)
 
function loss(p)
    s_ = solve(prob, Tsit5(), p = p)
    l = sum(abs, s_[end,end])
    return l, s_
end
 
loss(pnn)
 
 
ps = typeof(solution)[]
ls = eltype(p_init)[]
 
callback = function (p, l, pred)
    display(l)
    push!(ps, pred)
    push!(ls, l)
    l <= 1.35 && return true
    return false
end
 
# Steer away from initial position and try to avoid local minima
res_1 = DiffEqFlux.sciml_train(loss, pnn, cb = callback, maxiters = 100)
 
anm = @animate for i in 1:1:length(ps)
    pl_ = plot(ps[i], ylim = (0, 4), title = "Iteration $i : $(ls[i])")
    plot!(ps[i].t, 1.34408f0*ones(length(ps[i].t)), color = :black, style = :dash, label = "Optimum")
    pl_
end
 
gif(anm, joinpath(pwd(), "figures", "lotka_volterra_optimal_control.gif"), fps = 10)
 
# Plot the control value
s_ = solve(prob, Tsit5(), p = res_1.u, saveat = Δt)
us = f(res_1.u)(permutedims(s_.t))
 
p1 = plot(s_, legend = :bottomright, ylabel = "u(t)")
plot!(s_.t, 1.34408f0*ones(length(s_.t)), color = :black, style = :dash, label = "Optimum")
 
plot(
    p1,
    plot(s_.t, us', xlabel = "t", ylabel = "w(t)", label = nothing, xticks = 0:2:12),
    layout = (2,1), link = :x
)
 
savefig(joinpath(pwd(), "figures", "lotka_volterra_optimal_control_input.png"))
 
plot(ls, title = "Objective", xlabel = "Iterations", label = nothing)
savefig(joinpath(pwd(), "figures", "lotka_volterra_optimal_control_objective.png"))

• Reference solution plot
• 

  Animated solution process