Nonlinear Modeling: user-defined functions with vector inputs *and* output

[ferrolho]

Consider the following user-defined functions with vector inputs:

function my_dynamics(x::Vararg{T}) where T
    q = [i for i in x[0*dof+1:1*dof]]
    v = [i for i in x[1*dof+1:2*dof]]
    τ = [i for i in x[2*dof+1:3*dof]]

    state = statecache[T]
    set_configuration!(state, q)
    set_velocity!(state, v)

    result = dynamicsresultcache[T]
    dynamics!(result, state, τ)

    euler_q_next, euler_v_next = EulerStep(q, v, result.v̇, sim_Δt)
end

register(model, :my_dynamics, 3*dof, my_dynamics, autodiff=true)

This function does not return a single scalar; in fact, it returns a tuple comprised of two arrays.

Ideally, I'd like to be able to define some set of constraints with the following (non-working) syntax:

for i in 1:sim_num_knots-1
    q = configurations[i]
    v = velocities[i]
    τ = torques[i]

    q_next = configurations[i+1]
    v_next = velocities[i+1]
    
    @NLconstraint(model, my_dynamics(q..., v..., τ...) == (q_next, v_next))
end

Now, I am aware of the current NL interface limitations and know this is not currently possible. Moreover, I have been trying to find an example of such a situation and some workaround with the methods currently available in the NL interface but without success...

For that matter, my current workaround has ended up resembling a horrible, horrible abomination:

function my_dynamics(x::Vararg{T}) where T
    q = [i for i in x[0*dof+1:1*dof]]
    v = [i for i in x[1*dof+1:2*dof]]
    τ = [i for i in x[2*dof+1:3*dof]]

    state = statecache[T]
    set_configuration!(state, q)
    set_velocity!(state, v)

    result = dynamicsresultcache[T]
    dynamics!(result, state, τ)

    euler_q_next, euler_v_next = EulerStep(q, v, result.v̇, sim_Δt)

    vcat(euler_q_next, euler_v_next)
end

register(model, :my_dynamics, 3*dof, my_dynamics, autodiff=true)

my_dynamics1(x...) = my_dynamics(x...)[1]
my_dynamics2(x...) = my_dynamics(x...)[2]
my_dynamics3(x...) = my_dynamics(x...)[3]
my_dynamics4(x...) = my_dynamics(x...)[4]
my_dynamics5(x...) = my_dynamics(x...)[5]
my_dynamics6(x...) = my_dynamics(x...)[6]
my_dynamics7(x...) = my_dynamics(x...)[7]

my_dynamics11(x...) = my_dynamics(x...)[7+1]
my_dynamics12(x...) = my_dynamics(x...)[7+2]
my_dynamics13(x...) = my_dynamics(x...)[7+3]
my_dynamics14(x...) = my_dynamics(x...)[7+4]
my_dynamics15(x...) = my_dynamics(x...)[7+5]
my_dynamics16(x...) = my_dynamics(x...)[7+6]
my_dynamics17(x...) = my_dynamics(x...)[7+7]

register(model, :my_dynamics1, 3*dof, my_dynamics1, autodiff=true)
register(model, :my_dynamics2, 3*dof, my_dynamics2, autodiff=true)
register(model, :my_dynamics3, 3*dof, my_dynamics3, autodiff=true)
register(model, :my_dynamics4, 3*dof, my_dynamics4, autodiff=true)
register(model, :my_dynamics5, 3*dof, my_dynamics5, autodiff=true)
register(model, :my_dynamics6, 3*dof, my_dynamics6, autodiff=true)
register(model, :my_dynamics7, 3*dof, my_dynamics7, autodiff=true)

register(model, :my_dynamics11, 3*dof, my_dynamics11, autodiff=true)
register(model, :my_dynamics12, 3*dof, my_dynamics12, autodiff=true)
register(model, :my_dynamics13, 3*dof, my_dynamics13, autodiff=true)
register(model, :my_dynamics14, 3*dof, my_dynamics14, autodiff=true)
register(model, :my_dynamics15, 3*dof, my_dynamics15, autodiff=true)
register(model, :my_dynamics16, 3*dof, my_dynamics16, autodiff=true)
register(model, :my_dynamics17, 3*dof, my_dynamics17, autodiff=true)

for i in 1:sim_num_knots-1
    q = configurations[i]
    v = velocities[i]
    τ = torques[i]

    q_next = configurations[i+1]
    v_next = velocities[i+1]
    
    @NLconstraint(model, my_dynamics1(q..., v..., τ...) == q_next[1])
    @NLconstraint(model, my_dynamics2(q..., v..., τ...) == q_next[2])
    @NLconstraint(model, my_dynamics3(q..., v..., τ...) == q_next[3])
    @NLconstraint(model, my_dynamics4(q..., v..., τ...) == q_next[4])
    @NLconstraint(model, my_dynamics5(q..., v..., τ...) == q_next[5])
    @NLconstraint(model, my_dynamics6(q..., v..., τ...) == q_next[6])
    @NLconstraint(model, my_dynamics7(q..., v..., τ...) == q_next[7])

    @NLconstraint(model, my_dynamics11(q..., v..., τ...) == v_next[1])
    @NLconstraint(model, my_dynamics12(q..., v..., τ...) == v_next[2])
    @NLconstraint(model, my_dynamics13(q..., v..., τ...) == v_next[3])
    @NLconstraint(model, my_dynamics14(q..., v..., τ...) == v_next[4])
    @NLconstraint(model, my_dynamics15(q..., v..., τ...) == v_next[5])
    @NLconstraint(model, my_dynamics16(q..., v..., τ...) == v_next[6])
    @NLconstraint(model, my_dynamics17(q..., v..., τ...) == v_next[7])
end

Surely, there must be a better way...

As such, I am opening this issue here, with the hope that someone, somewhere, has been through this and can perhaps give me some pointers to some minimum working examples, or some suggestions on how to tackle this and improve the snippet above...

Thanks in advance, and apologies for this terrible code.

P.s. I was unsure whether I should have opened this issue as a "Bug report" or as a "Feature request". I hope it isn't a big deal in the case I chose wrongly.

[odow]

The short answer is that there isn't (currently) a good way to do this. JuMP's automatic differentiation library was written before much of Julia's AD tooling, and so is very specific to how JuMP likes things.

We have a Google Summer of Code project on integrating other AD-backends that may enable things like this.

[mlubin]

I'm going to close this issue given that we have no plans to implement this feature. It would require a complete rewrite of the AD system that JuMP uses. IMO the best bet is support the effort to get CasADi running in Julia: casadi/casadi#1105.

[ChrisRackauckas]

I'm going to close this issue given that we have no plans to implement this feature. It would require a complete rewrite of the AD system that JuMP uses. IMO the best bet is support the effort to get CasADi running in Julia: casadi/casadi#1105.

I think we are actually pretty far along with that through SparsityDetection.jl and SparseDiffTools.jl. Doing sparse AD and coloring directly from descriptions of a Julia function is now a thing. How to make JuMP utilize this is a different story, but I think a lot of what JuMP does doesn't necessarily need the DSL anymore with the newest iteration of Cassette-based tooling.

[mlubin]

@ChrisRackauckas Agreed, I'm totally supportive of seeing how far Cassette-based optimization modeling can go. Maybe you need the DSL aspect only for expressing constraints (i.e., changing the meaning of relation operators).

[bienpierre]

I would like to know if user-defined function with vector output is still restricted due to AD? regards

[ChrisRackauckas]

With ModelingToolkit you can automatically generate symbolic code from numeric code, so this kind of this is supported:

rosenbrock(x,p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)
p  = [1.0,100.0]
prob = OptimizationProblem(rosenbrock,x0,p)
sys = modelingtoolkitize(prob) # symbolicitize me captain!
prob = OptimizationProblem(sys,x0,p,grad=true,hess=true)

sol = solve(prob,NelderMead())
@test sol.minimum < 1e-8

sol = solve(prob,BFGS())
@test sol.minimum < 1e-8

sol = solve(prob,Newton())
@test sol.minimum < 1e-8

And for the OP, ModelingToolkit's ControlSystem can generate the optimization problem for the discretized control system:

  
using ModelingToolkit

@variables t x(t) v(t) u(t)
@parameters p
@derivatives D'~t

loss = (4-x)^2 + 2v^2 + u^2
eqs = [
    D(x) ~ v
    D(v) ~ p*u^3
]

sys = ControlSystem(loss,eqs,t,[x,v],[u],[p])
dt = 0.1
tspan = (0.0,1.0)
sys = runge_kutta_discretize(sys,dt,tspan)

u0 = rand(112) # guess for the state values
prob = OptimizationProblem(sys,u0,[0.1],grad=true)

I think the way forward here is to use this kind of tool for nonlinear things, but YMMV. We could link this to JuMP, i.e. generate JuMP symbolic code from this symbolic code, but I think directly connecting with MOI through GalacticOptim might have some advantages here, so it's what we're trying first (though outputting to JuMP in the future might be nice: automatically detect whether a numerical code is quadratic and then use quadratic optimization for example).

[bienpierre]

Thanks Chris for your answer. I am trying to reproduce the two examples, however I can't. I have an error with OptimizationProblem(). In documentation it is typed with sys::OptimizationSystem:

regards

[ChrisRackauckas]

Oh, this feature is brand new and I had forgot to tag it. It'll be released in ModelingToolkit v3.21 later today.

[hozdenoren]

@odow's talk at JuliaCon2021 included some remarks about revamping the Automatic-Differentiation engine and Nonlinear Interface. Would that mean vector inputs/outputs can be on the agenda as well?

When I had tried using splatting workaround before, it ended up being easier to move the entire problem to Matlab just to use fmincon...

[odow]

See https://jump.dev/JuMP.jl/stable/tutorials/Nonlinear%20programs/nlp_tricks/#User-defined-functions-with-vector-outputs.

Would that mean vector inputs/outputs can be on the agenda as well?

Yes. This is on our radar.

[adubredu]

@ferrolho Check out this elegant work-around by @tkoolen : https://github.com/tkoolen/RigidBodyDynamicsSandbox/blob/master/src/jump_utils.jl

With this macro, you can add nonlinear constraints with vector inputs and outputs as follows:

@VectorNLconstraint(model, my_dynamics(x) == qnext)

[ferrolho]

Thank you for sharing that, @adubredu. Twan's workaround seems along the lines of what is described in Tips and tricks > User-defined functions with vector outputs. I'm a bit sad I didn't find that a few years ago!