Repository navigation
Performance problem with deeply nested nonlinear expressions #3065
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?
Activity
blegat commentedon Jul 7, 2025
Hi @pekpuglia
Good question. The issue here is that we expand everything into an expression graph at the level of elementary operations:
*,+, ... But that graph is way to big.So you would want adding more allowed operators to this graph such as
RK4etc... Let's even add the whole multistep integration as one operator.One limitation though is that JuMP currently only support scalar inputs and scalar outputs so we have to scalarize like so:
Since we created separate functions for each output to make it scalar output, there might be duplicated computation, see this tutorial to avoid it.
pekpuglia commentedon Jul 7, 2025
Thank you! will try this in a bit. However, I still don't understand wht is the difference between JuMP and CasADi that allows CasADi to trace multi-step integration and not JuMP. I believe CasADi also expands expressions, as can be verified interactively in MATLAB.
blegat commentedon Jul 7, 2025
Can you post the expression graph you obtain with CasADi (so
F) after one or two steps ?odow commentedon Jul 7, 2025
I'm tempted to close this issue as out-of-scope. JuMP is not built for this sort of modeling approach. With large nested expressions.
The solution is to introduce temporary variables to represent intermediate expressions:
Alternatively, you might want to look into https://github.com/infiniteopt/InfiniteOpt.jl
pekpuglia commentedon Jul 8, 2025
I just ran the equivalent CasADi code (using a
casadi.Functionfor multistep RK4, which I'm not sure is important) and it builds this expression:Which apparently builds sub-expressions for the multi-step integration.
So, in practice, there are feasible ways of doing this with JuMP, and I can go on working on my problem (thanks!). It just feels frustrating/limiting that, even as Julia's main optimization/optimal control package, JuMP needs workarounds (such as operator registering) where CasADi just works. The memoization approach is somewhat cumbersome and I think there would be interest from the community in making this more streamlined. Maybe from the point of view of the developers this is out-of-scope, but as a Julia user, I would reach for JuMP as a drop-in replacement (meaning, the same mathematical algorithms should work in JuMP as easily as in CasADi).
blegat commentedon Jul 8, 2025
Thanks, it seems Casadi is automatically creating the intermediate variable as @odow mentioned.
Actually, building the JuMP functions is quite fast. The fact that they have common sub-expressions is not an issue because they are shared as reference, not copied.
So as long as you don't print them you are fine.
The bottleneck is
moi_functionbecause, it's being called many times on the sameGenericNonlinearExpr. If we had a dict that would check the ref and check whether we already built theMOI.ScalarNonlinearFunctionfor it then it would speed up the generation of the MOI model but the AD will still be slow because it will not exploit the fact that there are common sub-expressions. So just improving the build time without creating sub-expressions is not helpful.We discussed in jump-dev/JuMP.jl#3738 for a syntax for the user to explictly add sub-expressions.
Actually, maybe we should just automatically create sub-expressions when we have several times a reference to the same
GenericNonlinearExpr.This seems to be what CasADi is doing and as mentioned by @pekpuglia , this will help for the problems users write naturally.
odow commentedon Jul 8, 2025
I explicitly considered this case when designing the data structure, and we chose not to handle it. If you're solving optimal control type problems, CasADi is an excellent tool for the job. JuMP doesn't have to be great at everything.
This is very much "it depends." We don't detect equivalent subexpressions, and subexpressions are mutable, so they don't have object equality.
The problem is that this is not easily decidable.
We have chosen to force people to explicitly implement their own subexpressions, like I did above. If JuMP got into that game it opens a massive can of worms.
As one example: we couldn't detect that two subexpressions are equivalent and replace them by one shared reference because the user may later modify one of them.
[-]JuMP can't handle multi-step Runge Kutta for optimal control[/-][+]Performance problem with deeply nested nonlinear expressions[/+]pekpuglia commentedon Jul 28, 2025
As a follow-up question, since I'm not completely sure from reading the docs: going the operator registering way disables the use of the Hessian?
odow commentedon Jul 28, 2025
If you don't provide a Hessian operator, yes. But you can provide a Hessian callback for the operator (see the docs).
odow commentedon Jun 29, 2026
I'm going to move this to MOI. I don't think there's much to do here from the JuMP perspective. We push objectid-equal subexpressions to MOI. Taking advantage of that requires changes in MOI.
blegat commentedon Jun 30, 2026
Yes, what's needed in JuMP has now landed in jump-dev/JuMP.jl#4032, I think the next step is #2788 so that this works if the user uses
direct_modeland notModel.