How to create an user-function?

Hello everyone,

I am trying to do a trajectory optimization using Julia. To begin to familiarize myself with the environment I have tried to implement a simpler problem such as maximizing the height of ascent of a rocket. Initially my code looked like this:

using JuMP
import Ipopt
import Plots 
using Plots

#Constants
T_c = 3.5 #Usage in thrust
h_c = 500 #Usage in drag
v_c = 620 #Usage in drag
m_c = 0.6 #Usage in drag
n = 1000 #step time
h_0 = 1 #Initial height
v_0 = 0 #Initial velocity
m_0 =  1 #Initial mass
g_0 = 1 #Gravity on the floor
const integration_rule = "Hermite simpson"

# Derivated parameters
c = 0.5*sqrt(g_0*h_0) #Thrust-to-fuel mass
m_f = 0.6 # Final mass
D_c = 0.5*v_c*m_0/g_0 #Drag scaling 
T_max = T_c*g_0*m_0 #Maximun thrust

# Create the Model

rocket = Model(Ipopt.Optimizer)
set_attribute(rocket,"print_level", 4) #Da mas informacion sobre la optimizacion 
set_attribute(rocket, "max_iter", 20000) #Define el numero maximo de iteraciones del solucionador
set_attribute(rocket, "print_timing_statistics", "yes") #Nos proporciona informacion sobre el tiempo que toma cada fase del proceso de optimizacion 

# Create the variables

@variables(rocket, begin
    v[1:n] ≥ 0 #Velocity
    h[1:n] ≥ h_0 #height
    m_f ≤ m[1:n] ≤ m_0 #Mass
    0 ≤ T[1:n] ≤ T_max  #Thrust
end)

# Initial and final conditions 

fix(v[1], v_0, force=true)
fix(h[1], h_0, force=true)
fix(m[1], m_0, force=true)
fix(m[n], m_f, force=true)

# Time step

@NLexpression(rocket, Δt, 0.2/n)

# Aditional Equations

@NLexpressions(rocket,begin
    Drag[j=1:n], D_c*(v[j]^2)*exp(-h_c*(h[j]-h_0)/h_0) 
    g[j=1:n], g_0*((h_0/h[j])^2) #Gravity
    t_f[j=1:n], Δt*n #Time of flight 
end)

# Dinamic equations

@NLexpressions(rocket, begin
    δh[j=1:n], v[j]
    δv[j=1:n], (T[j] - Drag[j])/m[j] -g[j]
    δm[j=1:n], -T[j]/c
end)

# Aditional equations and dynamic equations in collocation points 

@NLexpressions(rocket, begin
    h_cl[j=1:n-1], 0.5*(h[j]+h[j+1]) + (Δt/8)*(δh[j]-δh[j+1])
    v_cl[j=1:n-1], 0.5*(v[j]+v[j+1]) + (Δt/8)*(δv[j]-δv[j+1])
    m_cl[j=1:n-1], 0.5*(m[j]+m[j+1]) + (Δt/8)*(δm[j]-δm[j+1])
end)

@NLexpressions(rocket, begin
    T_cl[j=1:n-1], (T[j] + T[j+1])/2
    Drag_cl[j=1:n-1], (Drag[j] + Drag[j+1])/2
    g_cl[j=1:n-1], (g[j] + g[j+1])/2
end)

@NLexpressions(rocket, begin
    δh_cl[j=1:n-1], v_cl[j]
    δv_cl[j=1:n-1], (T_cl[j] - Drag_cl[j])/m_cl[j] -g_cl[j]
    δm_cl[j=1:n-1], -T_cl[j]/c
end)

# Collocation method

function CollocationMethods(x, y, z, Δt)
    if integration_rule == "Rectangular integration"
        for j in 2:n
            @NLconstraint(rocket, x[j] == x[j-1] + Δt*y[j-1])
        end     
    elseif integration_rule == "Trapezoidal integration"
        for j in 2:n
            @NLconstraint(rocket, x[j] == x[j-1] + 0.5*Δt*(y[j]+y[j-1]))
        end  
    elseif integration_rule == "Hermite simpson"
        for j in 1:n-1
            @NLconstraint(rocket, x[j+1] == x[j] + (Δt/6)*(y[j]+4*z[j]+y[j+1]))
        end
    else
        @error "Unexpected integration rule '$(integration_rule)'"
    end
end

# Use Collocation method 

CollocationMethods(h, δh, δh_cl, Δt)
CollocationMethods(v, δv, δv_cl, Δt)
CollocationMethods(m, δm, δm_cl, Δt)

@objective(rocket, Max, h[n])

optimize!(rocket)

println( "La altitud maxima es: ", objective_value(rocket))

# Plot results

altitud = value.(h)
masa = value.(m)
velocidad = value.(v)
empuje = value.(T)
time = [i * value(Δt) for i in 1:n]

p1 = plot(time, altitud, xlabel="Tiempo", ylabel="Altura", label="Altura del cohete", legend=:bottomright)
p2 = plot(time, masa, xlabel="Tiempo", ylabel="Masa", label="Masa del cohete", legend=:bottomright)
p3 = plot(time, velocidad, xlabel="Tiempo", ylabel="Velocidad", label="Velocidad del cohete", legend=:bottomright)
p4 = plot(time, empuje, xlabel="Tiempo", ylabel="Empuje", label="Empuje del cohete", legend=:bottomright)
plot(p1,p2,p3,p4,layout = (2,2),legend = false)

This code works but it would be useful for the implementation of Hermite simpson to have a function that picks up the dynamic system and that you can evaluate for placement points. This was the attempt made but it didn’t work, how could I do to define it and make it work?

using JuMP
import Ipopt
using Plots

# Constants
T_c = 3.5  # Usage in thrust
h_c = 500  # Usage in drag
v_c = 620  # Usage in drag
m_c = 0.6  # Usage in drag
n = 1000   # Step time
h_0 = 1    # Initial height
v_0 = 0    # Initial velocity
m_0 = 1    # Initial mass
g_0 = 1    # Gravity on the floor
ns = 3     # Number of stages

# Derived parameters
c = 0.5 * sqrt(g_0 * h_0)  # Thrust-to-fuel mass
m_f = 0.6                  # Final mass
D_c = 0.5 * v_c * m_0 / g_0  # Drag scaling 
T_max = T_c * g_0 * m_0     # Maximum thrust
const integration_rule = "Rectangular integration"

# Functions for drag and gravity
drag(x, j) = D_c * (x[2, j]^2) * exp(-h_c * ((x[1, j] - h_0) / h_0))
grav(x, j) = g_0 * (h_0 / x[1, j]^2)

# Dynamic system function
state_f1(x, u, j) = x[2, j]
state_f2(x, u, j) = (u[1, j] - drag(x, j)) / x[3, j] - grav(x, j)
state_f3(x, u, j) = u[1, j] / c

model = Model(Ipopt.Optimizer)
set_optimizer_attribute(model, "print_level", 4)  # Gives more information about optimization
set_optimizer_attribute(model, "max_iter", 20000) # Defines the maximum number of solver iterations
set_optimizer_attribute(model, "print_timing_statistics", "yes") # Provides information about the time taken for each phase of the optimization process
register(model, :state_f1, 4, state_f1, autodiff = true)
register(model, :state_f2, 4, state_f2, autodiff = true)
register(model, :state_f3, 4, state_f3, autodiff = true)
register(model, :drag, 3, drag, autodiff = true)
register(model, :grav, 3, grav, autodiff = true)

# Variables
@variables(model, begin
    # State
    x[1:ns, 1:n]
    # Control
    u[1, 1:n]
end)

# Initial and final conditions
fix(x[2, 1], v_0, force=true)
fix(x[1, 1], h_0, force=true)
fix(x[3, 1], m_0, force=true)
fix(x[3, n], m_f, force=true)

# Collocation method
function CollocationMethods(x, u, Δt)
    if integration_rule == "Rectangular integration"
        for j in 2:n
            f1 = state_f1(x, u, j)
            f2 = state_f2(x, u, j)
            f3 = state_f3(x, u, j)
            @NLconstraints(model, begin
                x[1, j] == x[1, j-1] + Δt * f1
                x[2, j] == x[2, j-1] + Δt * f2
                x[3, j] == x[3, j-1] + Δt * f3
            end)  
        end     
    else
        @error "Unexpected integration rule '$(integration_rule)'"
    end
end

Δt = 0.2 / n
CollocationMethods(x, u, Δt)

# Objective
@objective(model, Max, x[3, n])

# Optimization
optimize!(model)

# Printing the maximum altitude
println("La altitud máxima es: ", objective_value(model))

Thank you so much!

You appear to have looked at an older tutorial for JuMP.

Take a look at the current tutorial: Rocket Control · JuMP

You don’t need to use @NLconstraint anymore, and we support constructing nonlinear expressions with function tracing.