Optimal control of Algal-Bacterial consortium in Julia using OptimalControl.jl

ab_docp_ctdirect.jl

# Standalone script to calculate optimal control of an algal-bacterial consortium system
# using the direct method, using Julia's CTDirect.jl package from https://control-toolbox.org/.

import Pkg
Pkg.Registry.add(Pkg.RegistrySpec(url = "https://github.com/control-toolbox/ct-registry.git"))
Pkg.add.(["OptimalControl", "Plots", "LaTeXStrings"])
using OptimalControl, Plots, LaTeXStrings

# System parameters
s_in = 0.5
β = 23e-3
γ = 0.44
dmax = 1.5

ϕmax = 6.48; ks = 0.09;
ρmax = 27.3e-3; kv = 0.57e-3;
μmax = 1.0211; qmin = 2.7628e-3;
ϕ(s) = ϕmax * s / (ks + s)
ρ(v) = ρmax * v / (kv + v)
μ(q) = μmax * (1 - qmin / q)

# Initialize model
N = 5000    # time steps
t0 = 0; tf = 20
x0 = [0.1629, 0.0487, 0.0003, 0.0177, 0.035, 0]
@def docp begin
    t ∈ [t0, tf], time
    x ∈ R⁶, state
    u ∈ R², control
    
    x(t0) == x0
    x₁(t) ≥ 0
    x₂(t) ≥ 0
    x₃(t) ≥ 0
    x₄(t) ≥ qmin
    x₅(t) ≥ 0
    0 ≤ u₁(t) ≤ 1
    0 ≤ u₂(t) ≤ dmax
    
    ẋ(t) == [
        u₂(t)*(s_in - x₁(t)) - ϕ(x₁(t))*x₂(t)/γ,
        ((1 - u₁(t))*ϕ(x₁(t)) - u₂(t))*x₂(t),
        u₁(t)*β*ϕ(x₁(t))*x₂(t) - ρ(x₃(t))*x₅(t) - u₂(t)*x₃(t),
        ρ(x₃(t)) - μ(x₄(t))*x₄(t),
        (μ(x₄(t)) - u₂(t))*x₅(t),
        u₂(t)*x₅(t),
    ]
    x₆(tf) → max
end

# Optimization --------------------------------------------------------------

# Solution data structure
struct OCSolution
    t::Vector{Real}
    x::Matrix{Real}
    λ::Matrix{Real}
    u::Matrix{Real}
    obj::Real
end

@timev begin
    # solve DOCP
    sol = solve(docp, :direct, :ipopt; grid_size=N, tol=1e-12, mu_strategy="adaptive", print_level=4)

    # Retrieve OCP solution data
    t = sol.times
    x = reduce(vcat, transpose.(sol.state.(t)))
    λ = reduce(vcat, transpose.(sol.costate.(t)))
    u = reduce(vcat, transpose.(sol.control.(t)))
    sol = OCSolution(t, x, λ, u, sol.objective)
end

# Plots ---------------------------------------------------------------------

latexify(tab) = latexstring.(L"$".* tab .* L"$")
x = ["s", "e", "v", "q", "c"]
u = [raw"\alpha", "d"]
ind = [1, 3, 4, 2, 5]
l = @layout [° ° °; ° °]

function plot_state(sol::OCSolution)
    plot(sol.t, sol.x[:, ind], label=latexify(reshape(x[ind], (1,5))), layout=l)
    plot!(fontfamily="sans-serif")
    xlabel!("Time")
end

function plot_costate(sol::OCSolution)
    plot(sol.t, sol.λ[:, ind], label=latexify(reshape("\\lambda_".* x[ind], (1,5))), layout=l)
    xlabel!("Time")
    plot!(fontfamily="sans-serif")
end

function plot_control(sol::OCSolution)
    p1 = plot(sol.t, sol.u[:,1], label=latexstring(u[1]))
    p2 = plot(sol.t, sol.u[:,2], label=latexstring(u[2]))
    plot(p1, p2, layout=(2, 1))
end

function plot_objective(sol::OCSolution)
    plot(sol.t, sol.x[:,end], fillrange=zeros(length(sol.t)), alpha=.5, legend=false)
    mid = Integer(ceil(length(sol.t) * 0.75))
    xpos = sol.t[mid]
    ypos = sol.x[mid,end] / 2
    plot!(annotations=(xpos, ypos, Plots.text(round(sol.obj; digits=3), :hcenter)))
end

# save plots
mkpath("plots/")
fname(s) = "plots/ctdirect_" * string(N) * "_" * s * ".pdf"
function save_plots(sol)
    savefig(plot_state(sol), fname("state"))
    savefig(plot_costate(sol), fname("costate"))
    savefig(plot_control(sol), fname("control"))
    savefig(plot_objective(sol), fname("objective"))
end

save_plots(sol)