How to create an user-defined function with multiple vector as inputs

Specific Domains Optimization (Mathematical)
, ,
1

Hello, I’m new to JuMP module and also Julia.

I’m trying to perform a trajectory optimization with Hermite-Simpson direct collocation method. I’ve trying to create a function for the dynamic equations, but I’m still getting the same error to register the functions (and also splatting the vectors), here’s is my code:

function StateFunc(δ, j, X...)
   δ[1,j]= X[3,j]
   δ[2,j]= X[4,j]
   δ[3,j]=(X[6,j]*sin(X[7,j]))/X[5,j]
   δ[4,j]= -9.81 + (X[6,j]*cos(X[7,j]))/X[5,j]
   δ[5,j]= -X[6,j]/2000
  return 
end

X is a vector where each row is a state variable, and \delta would be the dynamic equation evaluated.
I want to avoid doing this (repeat each evaluation of the dynamic functions):

@NLexpression(rocket2d_hs, Δt, t_f/(n-1))
@NLexpressions(rocket2d_hs, begin
      δx[j=1:n], u[j]
      δy[j=1:n], v[j]
      δu[j=1:n], (T[j]*sin(ϵ[j]))/m[j]
      δv[j=1:n], -g_0 + (T[j]*cos(ϵ[j]))/m[j]
      δm[j=1:n], -sqrt((T[j]*sin(ϵ[j]))^2 + (T[j]*cos(ϵ[j]))^2)/c
end)
@NLexpressions(rocket2d_hs, begin
       x_c[j=1:n-1], 0.5*(x[j] + x[j+1]) + (Δt/8)*(δx[j]-δx[j+1])
       y_c[j=1:n-1], 0.5*(y[j] + y[j+1]) + (Δt/8)*(δy[j]-δy[j+1])
       u_c[j=1:n-1], 0.5*(u[j] + u[j+1]) + (Δt/8)*(δu[j]-δu[j+1])
       v_c[j=1:n-1], 0.5*(v[j] + v[j+1]) + (Δt/8)*(δv[j]-δv[j+1])
       m_c[j=1:n-1], 0.5*(m[j] + m[j+1]) + (Δt/8)*(δm[j]-δm[j+1])
end)
@NLexpressions(rocket2d_hs, begin
        T_c[j=1:n-1], (T[j] + T[j+1])/2
        ϵ_c[j=1:n-1], (ϵ[j] + ϵ[j+1])/2
end)
@NLexpressions(rocket2d_hs, begin
        δx_c[j=1:n-1], u_c[j]
        δy_c[j=1:n-1], v_c[j]
        δu_c[j=1:n-1], (T_c[j]*sin(ϵ_c[j]))/m_c[j]
        δv_c[j=1:n-1], -g_0 + (T_c[j]*cos(ϵ_c[j]))/m_c[j]
        δm_c[j=1:n-1], -sqrt((T_c[j]*sin(ϵ_c[j]))^2 + (T_c[j]*cos(ϵ_c[j]))^2)/c
end)
for j in 1:n-1
    @NLconstraints(rocket2d_hs, begin
        m[j+1] - m[j] == (Δt/6)*(δm[j] + 4*δm_c[j] + δm[j+1]) 
        v[j+1] - v[j] == (Δt/6)*(δv[j] + 4*δv_c[j] + δv[j+1])
        u[j+1] - u[j] == (Δt/6)*(δu[j] + 4*δu_c[j] + δu[j+1]) 
        y[j+1] - y[j] == (Δt/6)*(δy[j] + 4*δy_c[j] + δy[j+1])
        x[j+1] - x[j] == (Δt/6)*(δx[j] + 4*δx_c[j] + δx[j+1]) 
    end)
end
@objective(rocket2d_hs, Min, t_f)
#set_silent(rocket2d)
optimize!(rocket2d_hs)
println("Tf optimo es: ",objective_value(rocket2d_hs))

Here is my attempt of introducing the function:


@NLexpression(rocket2d_hs_2, Δt, t_f/(n-1))
register(rocket2d_hs_2, :StateFunc, 7,StateFunc; autodiff= true)
register(rocket2d_hs_2, Collocation_state, 7, Collocation_state; autodiff=true)
@objective(rocket2d_hs_2, Min, t_f)

for j in 1:n-1
    StateFunc(δ, j, X...)
    for i in 1:5
        Collocation_state(X_c, δ, i, j, X...)
    end
    StateFunc(δ_c, j, X_c...)
    @NLconstraints(rocket2d_hs_2, begin
        X[5,j+1] - X[5,j+1] == (Δt/6)*(δ[5,j] + 4*δ_c[5,j] + δ[5,j+1]) 
        X[4,j+1] - X[4,j+1] == (Δt/6)*(δ[4,j] + 4*δ_c[4,j] + δ[4,j+1])
        X[3,j+1] - X[3,j+1] == (Δt/6)*(δ[3,j] + 4*δ_c[3,j] + δ[3,j+1]) 
        X[2,j+1] - X[2,j+1] == (Δt/6)*(δ[2,j] + 4*δ_c[2,j] + δ[2,j+1])
        X[1,j+1] - X[1,j+1] == (Δt/6)*(δ[1,j] + 4*δ_c[1,j] + δ[1,j+1]) 
    end)
end

In summary, It is possible to create an user defined function with multiple arrays as inputs?

2

Hi @PatrickeTownsend, welcome to the forum.

I think you’re slightly off in how you’re you’re approaching JuMP’s user-defined functions.

With a slightly simpler example, your code looks something like this:

using JuMP
function state_f(x, y...)
    x[1] = sum(sin(y[i]) for i in 1:3)
    x[2] = sum(cos(y[i]) for i in 1:3)
end
model = Model()
register(model, :state_f, 4, state_f; autodiff = true)
@variable(model, x[1:2])
@variable(model, y[1:3])
state_f(x, y...)
@optimize(model, Min, sum(x))

This has a few things wrong with it:

  • User-defined functions must take scalars as input
  • User-defined functions must return a scalar
  • User-defined functions must be used in @NL macros

This means that you cannot write a function like StateFunc which takes a vector as the first argument and modifies in-place, and it also means that calling StateFunc and Collocation_state outside the macros is not adding a constraint to the model. It’s just evaluating the function.

Relatedly:

It is possible to create an user defined function with multiple arrays as inputs?

is “no” because you cannot have a user-defined function which takes even one array as input :smile:.

To change my example somewhat, here’s something that you could write in JuMP:

using JuMP
state_f1(y...) = sum(sin(y[i]) for i in 1:3)
state_f2(y...) = sum(cos(y[i]) for i in 1:3)
model = Model()
register(model, :state_f1, 3, state_f; autodiff = true)
register(model, :state_f2, 3, state_f; autodiff = true)
@variable(model, x[1:2])
@variable(model, y[1:3])
@NLconstraint(model, x[1] == state_f1(y...))
@NLconstraint(model, x[2] == state_f2(y...))
@optimize(model, Min, sum(x))

Now writing out state_f1, state_f2 etc can be painful. So you can write functions which return a vector as output, but then you need to follow this tutorial: Tips and tricks · JuMP

1 Like