Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

bugfix for midpoint #284

Merged
merged 1 commit into from
Sep 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ ADNLPModels = "0.8"
CTBase = "0.13"
DocStringExtensions = "0.9"
HSL = "0.4"
JLD2 = "0.5"
JLD2 = "0.4" # 0.5 has ugly warnings for functions
JSON3 = "1"
MadNLP = "0.8"
NLPModels = "0.21"
Expand Down
2 changes: 1 addition & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ CTDirect = "0.12"
Documenter = "1"
DocumenterMermaid = "0.1"
HSL = "0.4"
JLD2 = "0.5"
JLD2 = "0.4" # 0.5 has ugly warnings for functions
JSON3 = "1"
NLPModelsIpopt = "0.10"
Plots = "1"
35 changes: 23 additions & 12 deletions src/midpoint.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,17 @@ with the convention u([t_i,t_i+1[) = U_i and u(tf) = U_N-1
=#

# +++ TODO: use args

# NB. could be defined as a generic IRK
struct Midpoint <: Discretization

stage::Int
additional_controls::Int
Midpoint() = new(1, 0)
butcher_a::Matrix{Float64}
butcher_b::Vector{Float64}
butcher_c::Vector{Float64}
info::String

Midpoint() = new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
end

"""
Expand Down Expand Up @@ -40,7 +46,7 @@ function get_variables_at_time_step(xu, docp::DOCP{Midpoint}, i)
if i < N
offset_u = offset
else
offset_u = (nx * 2 + m) * (i - 1)
offset_u = (nx * (1 + docp.discretization.stage) + m) * (i - 1)
end
if m == 1
ui = xu[offset_u + nx + 1]
Expand All @@ -66,15 +72,15 @@ function get_NLP_variables_at_time_step(xu, docp, i, disc::Midpoint)
nx = docp.dim_NLP_x
m = docp.dim_NLP_u
N = docp.dim_NLP_steps
offset = (nx * 2 + m) * i
offset = (nx * (1 + docp.discretization.stage) + m) * i

# state
xi = xu[(offset + 1):(offset + nx)]
# control
if i < N
offset_u = offset
else
offset_u = (nx * 2 + m) * (i - 1)
offset_u = (nx * (1 + docp.discretization.stage) + m) * (i - 1)
end
ui = xu[(offset_u + nx + 1):(offset_u + nx + m)]
# stage
Expand All @@ -92,7 +98,7 @@ function set_variables_at_time_step!(xu, x_init, u_init, docp, i, disc::Midpoint
n = docp.dim_OCP_x
m = docp.dim_NLP_u
N = docp.dim_NLP_steps
offset = (nx * 2 + m) * i
offset = (nx * (1 + docp.discretization.stage) + m) * i

# NB. only set the actual state variables from the OCP
# - skip the possible additional state for lagrange cost
Expand Down Expand Up @@ -155,9 +161,9 @@ $(TYPEDSIGNATURES)
Set the constraints corresponding to the state equation
"""
function setStateEquation!(docp::DOCP{Midpoint}, c, index::Int, args, v, i)

ocp = docp.ocp

# +++ later use butcher table in struct ?
disc = docp.discretization

# variables
ti = args.time
Expand All @@ -171,16 +177,21 @@ function setStateEquation!(docp::DOCP{Midpoint}, c, index::Int, args, v, i)
hi = tip1 - ti

# midpoint rule
@. c[index:(index + docp.dim_OCP_x - 1)] = xip1 - (xi + hi * ki[1:(docp.dim_OCP_x)])
h_sum_bk = hi * disc.butcher_b[1] * ki[1:docp.dim_NLP_x]
c[index:(index + docp.dim_OCP_x - 1)] .= xip1 .- (xi .+ h_sum_bk[1:docp.dim_OCP_x])
# +++ just define extended dynamics !
if docp.has_lagrange
c[index + docp.dim_OCP_x] = xlip1 - (xli + hi * ki[end])
c[index + docp.dim_OCP_x] = xlip1 - (xli + h_sum_bk[end])
end
index += docp.dim_NLP_x

# stage equation at mid-step
t_s = 0.5 * (ti + tip1)
x_s = 0.5 * (xi + xip1)
t_s = ti + hi * disc.butcher_c[1]
if docp.dim_OCP_x == 1
x_s = xi + hi * disc.butcher_a[1][1] * ki[1] #FFS
else
x_s = xi .+ hi .* (disc.butcher_a[1][1] .* ki[1:docp.dim_OCP_x])
end
c[index:(index + docp.dim_OCP_x - 1)] .= ki[1:(docp.dim_OCP_x)] .- ocp.dynamics(t_s, x_s, ui, v)
# +++ just define extended dynamics !
if docp.has_lagrange
Expand Down
9 changes: 6 additions & 3 deletions src/trapeze.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,14 @@ Internal layout for NLP variables:
[X_0,U_0, X_1,U_1, .., X_N,U_N, V]
=#

# NB. could be defined as a generic IRK
struct Trapeze <: Discretization

stage::Int
additional_controls::Int
# add control at tf
Trapeze() = new(0, 1)
additional_controls::Int # add control at tf
info::String

Trapeze() = new(0, 1, "Implicit Trapeze aka Crank-Nicolson, 2nd order, A-stable")
end

"""
Expand Down
2 changes: 1 addition & 1 deletion test/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
[compat]
CTBase = "0.13"
HSL = "0.4"
JLD2 = "0.5"
JLD2 = "0.4" # 0.5 has ugly warnings for functions
JSON3 = "1"
MadNLP = "0.8"
NLPModelsIpopt = "0.10"
Expand Down
Loading