Skip to content

Commit

Permalink
midpoint allocs ok
Browse files Browse the repository at this point in the history
  • Loading branch information
@PierreMartinon
PierreMartinon committed Nov 15, 2024
1 parent 34a3d2e commit f1aea78
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 61 deletions.
8 changes: 4 additions & 4 deletions benchmark/prof.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ using Profile
using PProf
using JET

include("../test/problems/goddard.jl")
#include("../test/problems/double_integrator.jl")
#include("../test/problems/goddard.jl")
include("../test/problems/simple_integrator.jl")

# local version of mayer cost
function local_mayer(obj, x0, xf, v)
Expand All @@ -20,8 +20,8 @@ end

function init(;grid_size, disc_method)
#prob = goddard_all()
prob = goddard()
#prob = double_integrator_mintf()
#prob = goddard()
prob = simple_integrator()
ocp = prob[:ocp]
docp = CTDirect.DOCP(ocp, grid_size=grid_size, time_grid=CTDirect.__time_grid(), disc_method=string(disc_method))
xu = CTDirect.DOCP_initial_guess(docp)
Expand Down
113 changes: 62 additions & 51 deletions src/midpoint.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,20 @@ Internal layout for NLP variables:
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
# NB. could be defined as a generic IRK in future IRK.jl for comparison purposes
#butcher_a::Matrix{Float64}
#butcher_b::Vector{Float64}
#butcher_c::Vector{Float64}
#return new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
struct Midpoint <: Discretization

stage::Int
additional_controls::Int
butcher_a::Matrix{Float64}
butcher_b::Vector{Float64}
butcher_c::Vector{Float64}
additional_controls::Int # add control at tf
info::String

# constructor
function Midpoint(dim_NLP_x, dim_NLP_u, dim_NLP_steps)
return new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
# constructor
function Midpoint()
return new(1, 0, "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
end
end

Expand Down Expand Up @@ -61,48 +61,70 @@ end


function get_stagevars_at_time_step(xu, docp::DOCP{Midpoint}, i)
nx = docp.dim_NLP_x
m = docp.dim_NLP_u
N = docp.dim_NLP_steps
offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x + docp.dim_NLP_u
# retrieve vector stage variable (except at final time)
if i < N+1
return @view xu[(offset + nx + m + 1):(offset + nx + m + nx) ]
if i < docp.dim_NLP_steps+1
return @view xu[(offset + 1):(offset + docp.dim_NLP_x)]
else
return nothing
return @view xu[1:docp.dim_NLP_x] # unused but keep same type !
end
end


function set_state_at_time_step!(xu, x_init, docp::DOCP{Midpoint}, i)
nx = docp.dim_NLP_x
n = docp.dim_OCP_x
m = docp.dim_NLP_u
offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u)
# initialize only actual state variables from OCP (not lagrange state)
if !isnothing(x_init)
xu[(offset + 1):(offset + n)] .= x_init
xu[(offset + 1):(offset + docp.dim_OCP_x)] .= x_init
end
end


function set_control_at_time_step!(xu, u_init, docp::DOCP{Midpoint}, i)
nx = docp.dim_NLP_x
m = docp.dim_NLP_u
N = docp.dim_NLP_steps
if i < N+1
offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
if i < docp.dim_NLP_steps+1
offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x
else
offset = (nx*(1+docp.discretization.stage) + m) * (i-2)
offset = (i-2) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x
end
if !isnothing(u_init)
xu[(offset + nx + 1):(offset + nx + m)] .= u_init
xu[(offset + 1):(offset + docp.dim_NLP_u)] .= u_init
end
end


function setWorkArray(docp::DOCP{Midpoint}, xu, time_grid, v)
return nothing
# use work array to store all dynamics + lagrange costs
work = similar(xu, docp.dim_NLP_x * (docp.dim_NLP_steps))
if docp.dim_OCP_x > 1
xs = similar(xu, docp.dim_OCP_x)
end

# loop over time steps
for i = 1:docp.dim_NLP_steps
offset = (i-1) * docp.dim_NLP_x
ti = time_grid[i]
xi = get_OCP_state_at_time_step(xu, docp, i)
ui = get_OCP_control_at_time_step(xu, docp, i)
tip1 = time_grid[i+1]
xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
ts = 0.5 * (ti + tip1)
# add a new getter ?
if docp.dim_OCP_x == 1
xs = 0.5 * (xi + xip1)
else
@. xs = 0.5 * (xi + xip1) # not ok for dim 1...
end
# OCP dynamics
docp.ocp.dynamics((@view work[offset+1:offset+docp.dim_OCP_x]), ts, xs, ui, v)
# lagrange cost
if docp.is_lagrange
docp.ocp.lagrange((@view work[offset+docp.dim_NLP_x:offset+docp.dim_NLP_x]), ts, xs, ui, v)
end
end
return work
end


"""
$(TYPEDSIGNATURES)
Expand All @@ -111,15 +133,13 @@ Convention: 1 <= i <= dim_NLP_steps (+1)
"""
function setConstraintBlock!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work)

disc = docp.discretization

# offset for previous steps
offset = (i-1)*(docp.dim_NLP_x * (1+docp.discretization.stage) + docp.dim_path_cons)

# variables
# 0. variables
ti = time_grid[i]
xi = get_OCP_state_at_time_step(xu, docp, i)
ui = get_OCP_control_at_time_step(xu, docp, i)
ui = get_OCP_control_at_time_step(xu, docp, i)

# 1. state equation
if i <= docp.dim_NLP_steps
Expand All @@ -128,42 +148,33 @@ function setConstraintBlock!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work)
tip1 = time_grid[i+1]
xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
hi = tip1 - ti
offset_dyn_i = (i-1)*docp.dim_NLP_x

# midpoint rule
h_sum_bk = hi * disc.butcher_b[1] * ki
if docp.dim_OCP_x == 1
c[offset+1] = xip1 - (xi + h_sum_bk[1])
else
c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + h_sum_bk[1:docp.dim_OCP_x])
end
@views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + hi * ki[1:docp.dim_OCP_x])

if docp.is_lagrange
c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + h_sum_bk[docp.dim_NLP_x])
c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + hi * ki[docp.dim_NLP_x])
end
offset += docp.dim_NLP_x

# stage equation at mid-step
ts = ti + hi * disc.butcher_c[1]
#xs = xi + hi * (disc.butcher_a[1][1] * ki)
xs = 0.5 * (xi + xip1) #compare bench
docp.ocp.dynamics((@view c[offset+1:offset+docp.dim_OCP_x]), ts, xs, ui, v)
@views c[offset+1:offset+docp.dim_OCP_x] = -c[offset+1:offset+docp.dim_OCP_x] + ki[1:docp.dim_OCP_x]
@views @. c[offset+1:offset+docp.dim_OCP_x] = ki[1:docp.dim_OCP_x] - work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x]
if docp.is_lagrange
docp.ocp.lagrange((@view c[offset+docp.dim_NLP_x:offset+docp.dim_NLP_x]), ts, xs, ui, v)
@views c[offset+docp.dim_NLP_x] = -c[offset+docp.dim_NLP_x] + ki[docp.dim_NLP_x]
c[offset+docp.dim_NLP_x] = ki[docp.dim_NLP_x] - work[offset_dyn_i+docp.dim_NLP_x]
end
offset += docp.dim_NLP_x
end

# 2. path constraints
# Notes on allocations:.= seems similar
if docp.dim_u_cons > 0
docp.control_constraints[2]((@view c[offset+1:offset+docp.dim_u_cons]),ti, docp._u(ui), v)
docp.control_constraints[2]((@view c[offset+1:offset+docp.dim_u_cons]),ti, ui, v)
end
if docp.dim_x_cons > 0
docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, docp._x(xi), v)
docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, xi, v)
end
if docp.dim_mixed_cons > 0
docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, docp._x(xi), docp._u(ui), v)
docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, xi, ui, v)
end

end
4 changes: 2 additions & 2 deletions src/problem.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,9 +180,9 @@ struct DOCP{T <: Discretization, X <: ScalVect, U <: ScalVect, V <: ScalVect}

# parameter: discretization method
if disc_method == "midpoint"
discretization = CTDirect.Midpoint(dim_NLP_x, dim_NLP_u, dim_NLP_steps)
discretization = CTDirect.Midpoint()
elseif disc_method == "trapeze"
discretization = CTDirect.Trapeze(dim_NLP_x, dim_NLP_u)
discretization = CTDirect.Trapeze()
else
error("Unknown discretization method:", disc_method)
end
Expand Down
8 changes: 4 additions & 4 deletions src/trapeze.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ struct Trapeze <: Discretization
info::String

# constructor
function Trapeze(dim_NLP_x, dim_NLP_u)
function Trapeze()
return new(0, 1, "Implicit Trapeze aka Crank-Nicolson, 2nd order, A-stable")
end
end
Expand Down Expand Up @@ -113,15 +113,15 @@ function setConstraintBlock!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work)
# more variables
tip1 = time_grid[i+1]
xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
half_hi = 0.5 * (tip1 - ti)
offset_dyn_i = (i-1)*docp.dim_NLP_x
offset_dyn_ip1 = i*docp.dim_NLP_x

# trapeze rule (no allocations ^^)
halfstep = 0.5 * (tip1 - ti)
@views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + halfstep * (work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x] + work[offset_dyn_ip1+1:offset_dyn_ip1+docp.dim_OCP_x]))
@views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + half_hi * (work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x] + work[offset_dyn_ip1+1:offset_dyn_ip1+docp.dim_OCP_x]))

if docp.is_lagrange
c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + 0.5 * (tip1 - ti) * (work[offset_dyn_i+docp.dim_NLP_x] + work[offset_dyn_ip1+docp.dim_NLP_x]))
c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + half_hi * (work[offset_dyn_i+docp.dim_NLP_x] + work[offset_dyn_ip1+docp.dim_NLP_x]))
end
offset += docp.dim_NLP_x
end
Expand Down

0 comments on commit f1aea78

Please sign in to comment.