JuMP: User defined functions with vector inputs and outputs

Specific Domains Optimization (Mathematical)
1

In the JuMP documentation it is mentioned that functions with vector inputs and outputs require a workaround. However the following example works:

MWE 
using JuMP
import Ipopt

mwe = Model(Ipopt.Optimizer)

@variable(mwe, x[1:5])
@variable(mwe, y[1:5])

xref = ones(5)*10
yref = ones(5)*20

function obj(x, y, xref, yref)
    sum = zero(eltype(x))
    for i=1:length(xref)
        sum = sum + (x[i] - xref[i])^2 + (y[i] - yref[i])^2
    end

    return [sum, sum]
end

@objective(mwe, Min, obj(x[1:5], y[1:5], xref, yref)[1])

optimize!(mwe)

println("Value of x and y : ", [value.(x) value.(y)])

It is not clear from the docs if user defined functions with vector inputs and outputs are supported?

I could even write an RK4 scheme for direct optimal control nicely:

RK4.jl 
using JuMP
import Ipopt
using GLMakie

dubins = Model(Ipopt.Optimizer)

vm = 1.0   # Max forward speed
um = 1.0   # Max turning speed
ns = 3     # Number of states
nu = 1     # Number of inputs
n = 200    # Time steps

# System dynamics
function f(x, u)
    return [vm*cos(x[3]), vm*sin(x[3]), u]
end

# Objective Function
function obj_f(x, u, Δt)
    return Δt
end

# Decision variables
@variables(dubins, begin
    0 <= Δt <= 0.05  # Time step
    # State variables
    x[1:3,1:n]               # Velocity
    # Control variables
    -um ≤ u[1:1,1:n] ≤ um    # Thrust
end)

# Objective
@objective(dubins, Min, obj_f(x[1:3,1:n], u[1,1:n], Δt))
#Initial conditions
@constraint(dubins, x[1:3, 1] .== [0, 0, -π/2 ] )
# Final conditions
@constraint(dubins, x[1:3, n] == [5.0, 5.0, π/2 + π/4])

# Dynamic constraints: RK4
for j in 1:n-1
    k1 = f(x[1:3,j], u[j])
    k2 = f(x[1:3,j] + Δt*k1/2, u[j])
    k3 = f(x[1:3,j] + Δt*k2/2, u[j])
    k4 = f(x[1:3,j] + Δt*k3, u[j])
    @constraint(dubins, x[1:3, j+1] == x[1:3,j] + Δt*(k1 + 2*k2 + 2*k3 + k4)/6)
end

# Solve for the control and state
println("Solving...")
optimize!(dubins)
solution_summary(dubins)

# Display results
println("Min time: ", objective_value(dubins)*n)

fig = Figure()
ax = Axis(fig[1,1], autolimitaspect = 1)
lines!(ax, value.(x[1,1:n]), value.(x[2, 1:n]))
fig
2

The restriction to user-defined functions with vector outputs applies only to functions that you registered with @operator.

In your examples, you are using the function-tracing feature, which builds a NonlinearExpr object, and so everything works.

I agree that this isn’t clear in the documentation, so I need to improve it.

1 Like
3

Thanks a lot for the clarification. What are the limitations of the function-tracing approach? And when is it necessary to use @operator for function registration?

If it may be of help for documentation: most of the examples in the Nonlinear section add constraint equations one at a time. Hence, it creates an impression that the vectors are not supported. Maybe it would be helpful to add an example which uses function-tracing?

4

It needs to use only operators for which operator overloading has been implemented. Your obj example uses only +, - and ^.

It won’t work for something like LinearAlgebra.det(x' * x) or call_to_c_library(x).

In general, using function tracing is preferred to registering an operator, except for when the function produces a very large expression.

If it may be of help for documentation: most of the examples in the Nonlinear section add constraint equations one at a time. Hence, it creates an impression that the vectors are not supported. Maybe it would be helpful to add an example which uses function-tracing?

Yip. And in addition, I think I need to add a section explaining what functions tracing is, and what the benefits and limitations are :smile:.

My initial hope was that things like obj(...) would just work (it does), and people wouldn’t get confused, but I guess they do!

2 Likes
5

I’ve opened a PR to improve the documentation: [docs] add section on function tracing by odow · Pull Request #3570 · jump-dev/JuMP.jl · GitHub. Let me know if you have any comments.

1 Like
6

Thanks a lot. I will try out some examples from optimal control with this approach and report if I face any issues.

1 Like