I managed to extract a MWE, which should be readable enough.
import Symbolics
using JuMP, Ipopt, LinearAlgebra
using Plots
function vech(L::AbstractMatrix{T}) where {T}
n = LinearAlgebra.checksquare(L)
c = zeros(T, n * (n + 1) Ć· 2)
k = 0
for j = 1:n, i = j:n
k += 1
c[k] = L[i, j]
end
return c # half-vectorization of L
end
function unvech(c::AbstractVector{T}) where {T}
ns2 = length(c)
n = Int((sqrt(1 + 8ns2) - 1) / 2)
L = zeros(T, n, n)
k = 0
for j = 1:n, i = j:n
k = k + 1
L[i, j] = c[k]
end
return L # lower-triangular matrix built from vector c
end
function dynamics(x, u, t, p)
g, Ļ = p # unpack parameters
f = [x[2], -u[1] * x[1] + g * sin(x[1])] # drift
G = [0, Ļ] # diffusion
return f, G
end
function solve_ocp()
# Dimensions
nx = 2
nu = 1
nc = (nx * (nx + 1)) Ć· 2
nz = nx + nc
# General settings
Ļ = 0.1
g = 1
T = 2.0 # Number of seconds
N = 201 # Number of time steps
dt = T / (N - 1)
# Parameters tuple
p = (g, Ļ)
#------ Symbolic calculations
Symbolics.@variables t
Symbolics.@variables x[1:nx]
x = Symbolics.scalarize(x)
Symbolics.@variables c[1:nc]
c = Symbolics.scalarize(c)
Symbolics.@variables u[1:nu]
u = Symbolics.scalarize(u)
L = unvech(c)
P = L * L'
f, G = dynamics(x, u, t, p)
# Computation of symbolic Jacobian matrix
A = Symbolics.jacobian(f, x)
dP = A * P + P * A' + G * G'
invL = inv(L)
H = invL * dP * invL'
dx = f
dc = vech(L * (LowerTriangular(H) - Diagonal(H) / 2))
z = vcat(x, c)
dz = vcat(dx, dc)
F_expr = Symbolics.build_function(dz, z, u, t; expression=Val{false}, nanmath=false)[1]
# initial state
zini = [0, 0, 1e-3, 0, 1e-3]
# Set Bounds
umin = [1e-4]
umax = [10]
zmin = [-10, -10, 1e-4, -10, 1e-4]
zmax = [10, 10, 10, 10, 10]
#------ Build JuMP model
model = Model(Ipopt.Optimizer)
# Create JuMP variables
@variables(model, begin
zmin[i] ā¤ z[i=1:nz, j=1:N] ā¤ zmax[i]
umin[i] ā¤ u[i=1:nu, j=1:N] ā¤ umax[i]
end)
# Create useful expressions
@expression(model, t[j=1:N], (j - 1) * dt)
@expression(model, P[j=1:N], unvech(1.0z[nx+1:end, j]) * unvech(1.0z[nx+1:end, j])')
# /!\ the line below can be time-consuming with bigger models /!\
@expression(model, F[i=1:nz, j=1:N], F_expr(z[:, j], u[:, j], t[j])[i])
# Set constraints
@constraint(model, fĢ_const[j=1:N-1], z[:, j+1] - z[:, j] - dt * 0.5 * (F[:, j] + F[:, j+1]) .== 0)
@constraint(model, xend_const[i=1:nx], z[i, end] == 0.0)
# Fix initial state
fix.(z[:, 1], zini; force=true)
# Generate and set initial guess
zguess = repeat(zini, 1, N)
uguess = repeat(ones(nu), 1, N)
set_start_value.(z, zguess)
set_start_value.(u, uguess)
# Objective
R = diagm(ones(nu))
Q = 1e4 * diagm(ones(nx))
@expression(model, cost[j=1:N], z[1:nx, j]' * Q * z[1:nx, j] + tr(Q * P[j]) + u[:, j]' * R * u[:, j])
@objective(model, Min, sum(dt * 0.5(cost[j] + cost[j+1]) for j = 1:N-1))
# Optimize!
optimize!(model)
# Extraction of the optimal control
uopt = Array(value.(u))
return uopt
end
uopt = solve_ocp()
plot(uopt')
I emphasized the line taking more time with larger modelsā¦
Iām not sure whether I can avoid the symbolic Jacobian calculation (or whether I can replace it with some AD solution, but I donāt know how to do that in combination with JuMP).
Thank you for your advices.
