Performance problem with deeply nested nonlinear expressions #4024
Description
Hello, I am interested in switching from CasADi tu JuMP for optimal control (the system is quaternion kinematics, so not a lot of variables). My MATLAB CasADi model does multistep RK4 integration for multiple shooting (4 steps/interval, 10 intervals) and builds and solves the problem in < 2s. The equivalent JuMP code builds the model for many minutes, using up to 10GB of RAM, and is basically unusable. This is the JuMP code I wrote:
using ReferenceFrameRotations, DifferentialEquations, Plots
##
function qdot(X, u)
collect(dquat(Quaternion(X), u))
end
function RK4(f, X, u, dt)
k1 = dt*f(X , u)
k2 = dt*f(X+k1/2, u)
k3 = dt*f(X+k2/2, u)
k4 = dt*f(X+k3 , u)
X + (k1+2k2+2k3+k4)/6
end
model = Model(Ipopt.Optimizer)
DCMf = [
-1 0 0
0 0 1
0 1 0
]
# two quaternions represent the same orientation
qf = -convert(Quaternion, DCM(DCMf)) |> collect
N = 2
T = 1
dt = T / (N-1)
tabq = @variable(model, [1:4, 1:N], base_name="q")
tabu = @variable(model, [1:3, 1:N-1], base_name="ω_b")
cost = 0
for i = 1:N-1
F = [tabq[:, i]; cost]
for M = 1:4
F = RK4((X, u) -> [qdot(X[1:4], u); 1/2 * u' * u], F, tabu[:, i], dt)
end
@constraint(model, tabq[:, i+1] == F[1:4])
cost = F[end]
end
#traj constraints
@constraint(model, tabq[:, 1] == q0)
@constraint(model, tabq[:, N] == qf)
# @constraint(model, c[i=1:N], tabq[:, i]' * tabq[:, i] == 1)
@objective(model, Min, cost)
#start values
set_start_value.(tabq, q0)I set N = 2 to prove it's slow even for a single step (similar problems happen with single-step high order RK schemes, such as RK8, which leads me to believe JuMP can't handle big expressions really well). For comparison, if I build a single-step model with 50 steps (approx. 4steps/interval x 10 intervals for the working CasADi model) like this:
model = Model(Ipopt.Optimizer)
DCMf = [
-1 0 0
0 0 1
0 1 0
]
# two quaternions represent the same orientation
qf = -convert(Quaternion, DCM(DCMf)) |> collect
N = 50
T = 1
dt = T / (N-1)
tabq = @variable(model, [1:4, 1:N], base_name="q")
tabu = @variable(model, [1:3, 1:N-1], base_name="ω_b")
cost = 0
for i = 1:N-1
#can't do multistep integration
F = RK4((X, u) -> [qdot(X[1:4], u); 1/2 * u' * u], [tabq[:, i]; cost], tabu[:, i], dt)
@constraint(model, tabq[:, i+1] == F[1:4])
cost = F[end]
end
#traj constraints
@constraint(model, tabq[:, 1] == q0)
@constraint(model, tabq[:, N] == qf)
# @constraint(model, c[i=1:N], tabq[:, i]' * tabq[:, i] == 1)
@objective(model, Min, cost)
#start values
set_start_value.(tabq, q0)
model
##
optimize!(model)It then works reasonably fast. I don't understand why JuMP can't handle the first code; in my view, this is a big limitation. Is there a particular way of coding in JuMP I'm missing, or is this truly JuMP's fault?