Writting runge-kutta stages as an array of expression

Specific Domains Optimization (Mathematical)
, ,
1

Hi, I am writting a trajectory optimization scheme and for very specific reasons I want to constrain the change from state x_j to x_{j+1} to be a step of a runge-kutta method (I know multi-step methods would be more suitable here). Just to test things I am doing some optimizations on a pendulum using Heun’s method. Bellow is a rather verbose version where stages are written as individual expressions.

using JuMP
using NLopt

# Model initialization
model = Model(NLopt.Optimizer)
set_attribute(model, "algorithm", :LD_SLSQP) # LN_COBYLA

const n = 80                  # number of time-steps

@variables(model, begin
    x[1:n+1, 1:2]             # states
    u[1:n]                    # inputs
    Δt[1:n]                   # time-steps
end)

# System's ODE
const g = 9.81 # ms⁻² - gravity
const m = 1    # kg - mass
const l = 1    # m - length

f(θ, ω, u) = [ω, -g / l * sin(θ) + 1 / (m * l^2) * u]

# Integration scheme
A = [
    0 0
    1 0
]
b = [0.5, 0.5]

@NLexpression(model, k1[j=1:n, i=1:2],
    f(
        x[j, 1],
        x[j, 2],
        u[j]
    )[i]
)
@NLexpression(model, k2[j=1:n, i=1:2],
    f(
        x[j, 1] + Δt[j] * A[2, 1] * k1[j, 1],
        x[j, 2] + Δt[j] * A[2, 1] * k1[j, 2],
        u[j]
    )[i]
)
@NLexpression(model, xₕ[j=2:n+1, i=1:2], x[j-1, i] + Δt[j-1] * (b[1] * k1[j-1, i] + b[2] * k2[j-1, i])) # Heun's method

## Dynamics' constraint
@NLconstraint(model, [j = 2:n+1, i = 1:2], x[j, i] == xₕ[j, i])

At the moment I am trying to write all stages k_l as a single expression so that I can create them programatically for higher order methods.
Bellow is what I have come up with so far as a modification

@NLexpression(model, k[j=1:n, i=1:2, l=1:2],
    f(
        x[j, 1] + Δt[j] * sum(A[l,m] * k[j,1,m] for m in 1:l-1; init=0.0),
        x[j, 2] + Δt[j] * sum(A[l,m] * k[j,2,m] for m in 1:l-1; init=0.0),
        u[j]
    )[i]
)

@NLexpression(model, xₕ[j=2:n+1, i=1:2], x[j-1, i] + Δt[j-1] * sum(b[l] * k[j-1, i, l] for l in 1:2)) # Heun's method

but when defining the expression I get the error

`TypeError: in typeassert, expected NonlinearExpr, got a value of type VariableRef`

Does anyone have some suggestions on what I should try next?

2

Rubber ducky method to the rescue

## stages
K = Matrix{Vector}(undef, n, s)

for j in 1:n, l in 1:s
    K[j, l] = f(x[j, :], u[j]) + Δt[j] * sum(A[l, m] * K[j, m][:] for m in 1:l-1; init=zeros(2))
end

@NLexpression(model, k[j=1:n, l=1:s, i=1:r], K[j, l][i]) # step,order,state

## linear combination
@NLexpression(model, xₕ[j=2:n+1, i=1:r], x[j-1, i] + Δt[j-1] * sum(bᵣₖ₂[l] * k[j-1, l, i] for l in 1:s)) # Heun's method
3

The error is because you are attempting to combine the new NonlinearExpr syntax in JuMP with the old @NL expression syntax.

See ANN: JuMP v1.15 is released

Here’s how I would write your model

using JuMP, Ipopt
const n = 80
const g = 9.81
const m = 1
const l = 1
f(θ, ω, u) = [ω, -g / l * sin(θ) + 1 / (m * l^2) * u]
A = [0 0; 1 0]
b = [0.5, 0.5]
model = Model(Ipopt.Optimizer)
@variables(model, begin
    x[1:n+1, 1:2]
    u[1:n]
    Δt[1:n]
end)
@expressions(model, begin
    k1[j in 1:n], f(x[j,1], x[j,2], u[j])
    y[j in 1:n, i in 1:2], x[j,i] + Δt[j] * A[2,1] * k1[j][i]
    k2[j in 1:n], f(y[j,1], y[j,2], u[j])
    xₕ[j in 2:n+1, i in 1:2],
        x[j-1, i] + Δt[j-1] * (b[1] * k1[j-1][i] + b[2] * k2[j-1][i])
end)
@constraint(model, [j in 2:n+1, i in 1:2], x[j,i] == xₕ[j,i])