From c45c8678685006dd7f83ea9cfcc0da1ec9f30885 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Mon, 23 Feb 2026 12:12:34 +0100 Subject: [PATCH 01/21] introduce integral function in discretization API (generalize former runningCost function) --- src/CTDirect.jl | 11 +-- src/direct_shooting.jl | 136 ++++++++++++++++++++++++++++++++++ src/{disc => ode}/common.jl | 2 +- src/{disc => ode}/euler.jl | 0 src/{disc => ode}/irk.jl | 0 src/{disc => ode}/midpoint.jl | 9 +-- src/{disc => ode}/trapeze.jl | 0 test/manual_test.jl | 4 +- 8 files changed, 149 insertions(+), 13 deletions(-) create mode 100644 src/direct_shooting.jl rename src/{disc => ode}/common.jl (99%) rename src/{disc => ode}/euler.jl (100%) rename src/{disc => ode}/irk.jl (100%) rename src/{disc => ode}/midpoint.jl (97%) rename src/{disc => ode}/trapeze.jl (100%) diff --git a/src/CTDirect.jl b/src/CTDirect.jl index 027303f6..c4371891 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -63,10 +63,11 @@ include("collocation.jl") include("collocation_core.jl") include("collocation_variables.jl") include("collocation_functions.jl") -include("disc/common.jl") -include("disc/euler.jl") -include("disc/irk.jl") -include("disc/midpoint.jl") -include("disc/trapeze.jl") +include("ode/common.jl") +include("ode/euler.jl") +include("ode/irk.jl") +include("ode/midpoint.jl") +include("ode/trapeze.jl") +# ode/variable_step end diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl new file mode 100644 index 00000000..7cf8841c --- /dev/null +++ b/src/direct_shooting.jl @@ -0,0 +1,136 @@ +# --------------------------------------------------------------------------- +# Implementation of Direct shooting discretizer +# --------------------------------------------------------------------------- +struct DirectShooting <: AbstractDiscretizer + options::Strategies.StrategyOptions +end + +# useful for OptimalControl +Strategies.id(::Type{<:DirectShooting}) = :direct_shooting + +# default options +__direct_shooting_grid_size()::Int = 100 +__direct_shooting_scheme()::Symbol = :midpoint # later use variable step ode solver +__direct_shooting_scheme_steps() = 10 + +function Strategies.metadata(::Type{<:Collocation}) + return Strategies.StrategyMetadata( + Options.OptionDefinition( + name = :grid_size, + type = Int, + default = __direct_shooting_grid_size(), + description = "Number of time steps for the direct shooting grid", + ), + Options.OptionDefinition( + name = :scheme, + type = Symbol, + default = __direct_shooting_scheme(), + description = "Time integration scheme (e.g., :midpoint, :trapeze)", + ), + ) +end + +# constructor: kwargs contains the options values +function DirectShooting(; mode::Symbol = :strict, kwargs...) + opts = Strategies.build_strategy_options(DirectShooting; mode = mode, kwargs...) + return DirectShooting(opts) +end + +Strategies.options(c::DirectShooting) = c.options + + +# ========================================================================================== +# Build discretizer API (return sets of model/solution builders) +# ========================================================================================== +function (discretizer::DirectShooting)(ocp::AbstractModel) + + # common parts for builders + docp = get_docp(discretizer, ocp) ? + exa_getter = nothing # will be set in build_exa_model + + # ========================================================================================== + # The needed builders for the construction of the final DiscretizedModel + # ========================================================================================== + + # NLP builder for ADNLPModels + function build_adnlp_model( + initial_guess::CTModels.AbstractInitialGuess; + backend, + kwargs... + )::ADNLPModels.ADNLPModel + + # functions for objective and constraints + f = x -> CTDirect.DirectShooting_objective(x, docp) + c! = (c, x) -> CTDirect.DirectShooting_constraints!(c, x, docp) + + # build initial guess + init = get_docp_initial_guess(:adnlp, docp, initial_guess) + + # unused backends (option excluded_backend = [:jprod_backend, :jtprod_backend, :hprod_backend, :ghjvprod_backend] does not seem to work) + unused_backends = ( + hprod_backend=ADNLPModels.EmptyADbackend, + jtprod_backend=ADNLPModels.EmptyADbackend, + jprod_backend=ADNLPModels.EmptyADbackend, + ghjvprod_backend=ADNLPModels.EmptyADbackend, + ) + + + # use backend preset + backend_options = (backend=backend,) + + # build NLP + nlp = ADNLPModel!( + f, + init, + docp.bounds.var_l, + docp.bounds.var_u, + c!, + docp.bounds.con_l, + docp.bounds.con_u; + minimize=(!docp.flags.max), + backend_options..., + unused_backends..., + kwargs..., + ) + + return nlp + end + + # Solution builder for ADNLPModels + function build_adnlp_solution(nlp_solution::SolverCore.AbstractExecutionStats) + + # retrieve data from NLP solver + minimize = !docp.flags.max + objective, iterations, constraints_violation, message, status, successful = CTSolvers.extract_solver_infos(nlp_solution, minimize) + + # retrieve time grid + T = get_time_grid(nlp_solution.solution, docp) + + # build OCP solution from NLP solution + sol = CTDirect.build_OCP_solution(docp, nlp_solution, T, + objective, iterations, constraints_violation, message, status, successful) + + return sol + end + + # NLP builder for ExaModels + function build_exa_model( + ::Type{BaseType}, + initial_guess::CTModels.AbstractInitialGuess; + backend + )::ExaModels.ExaModel where {BaseType<:AbstractFloat} + end + + # Solution builder for ExaModels + function build_exa_solution(nlp_solution::SolverCore.AbstractExecutionStats) + end + + #NB. it would be better to return builders as model/solution pairs since they are linked + return CTSolvers.DiscretizedModel( + ocp, + CTSolvers.ADNLPModelBuilder(build_adnlp_model), + CTSolvers.ExaModelBuilder(build_exa_model), + CTSolvers.ADNLPSolutionBuilder(build_adnlp_solution), + CTSolvers.ExaSolutionBuilder(build_exa_solution), + ) +end diff --git a/src/disc/common.jl b/src/ode/common.jl similarity index 99% rename from src/disc/common.jl rename to src/ode/common.jl index f049736a..1ea4c914 100644 --- a/src/disc/common.jl +++ b/src/ode/common.jl @@ -214,7 +214,7 @@ $(TYPEDSIGNATURES) Compute the running cost (must be implemented for each discretization scheme) """ function runningCost(docp::DOCP{D}, xu, v, time_grid) where {(D<:Discretization)} - error("running_cost not implemented for discretization ", D) + return integral(docp, xu, v, time_grid, CTModels.lagrange(ocp_model(docp))) end """ diff --git a/src/disc/euler.jl b/src/ode/euler.jl similarity index 100% rename from src/disc/euler.jl rename to src/ode/euler.jl diff --git a/src/disc/irk.jl b/src/ode/irk.jl similarity index 100% rename from src/disc/irk.jl rename to src/ode/irk.jl diff --git a/src/disc/midpoint.jl b/src/ode/midpoint.jl similarity index 97% rename from src/disc/midpoint.jl rename to src/ode/midpoint.jl index 3f0bac1a..14bbd773 100644 --- a/src/disc/midpoint.jl +++ b/src/ode/midpoint.jl @@ -77,9 +77,8 @@ $(TYPEDSIGNATURES) Compute the running cost """ -function runningCost(docp::DOCP{Midpoint}, xu, v, time_grid) - obj_lagrange = 0.0 - ocp = ocp_model(docp) +function integral(docp::DOCP{Midpoint}, xu, v, time_grid, f) + value = 0.0 # loop over time steps for i in 1:docp.time.steps @@ -92,10 +91,10 @@ function runningCost(docp::DOCP{Midpoint}, xu, v, time_grid) get_OCP_state_at_time_step(xu, docp, i+1) ) ui = get_OCP_control_at_time_step(xu, docp, i) - obj_lagrange = obj_lagrange + hi * CTModels.lagrange(ocp)(ts, xs, ui, v) + value += hi * f(ts, xs, ui, v) end - return obj_lagrange + return value end """ diff --git a/src/disc/trapeze.jl b/src/ode/trapeze.jl similarity index 100% rename from src/disc/trapeze.jl rename to src/ode/trapeze.jl diff --git a/test/manual_test.jl b/test/manual_test.jl index 3d095c0d..90afac5e 100644 --- a/test/manual_test.jl +++ b/test/manual_test.jl @@ -7,6 +7,6 @@ include("./problems/beam.jl") include("./problems/goddard.jl") include("./problems/double_integrator.jl") -sol = solve_problem(goddard(); display=true) -sol1 = solve_problem(goddard2(); display=true, modeler=:exa, solver=:madnlp, scheme=:trapeze) +sol = solve_problem(double_integrator_minenergy(); display=true) +sol1 = solve_problem(double_integrator_minenergy2(); display=true, modeler=:exa, solver=:madnlp) From 9eb5444b69d68585e23129802de90de792c3b20d Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Mon, 23 Feb 2026 12:18:08 +0100 Subject: [PATCH 02/21] integral function for euler, trapeze and IRK --- src/ode/euler.jl | 9 ++++----- src/ode/irk.jl | 13 ++++++------- src/ode/trapeze.jl | 19 +++++++++---------- 3 files changed, 19 insertions(+), 22 deletions(-) diff --git a/src/ode/euler.jl b/src/ode/euler.jl index b8d1787c..046232bd 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -118,11 +118,10 @@ $(TYPEDSIGNATURES) Compute the running cost """ -function runningCost(docp::DOCP{Euler}, xu, v, time_grid) - ocp = ocp_model(docp) +function integral(docp::DOCP{Euler}, xu, v, time_grid, f) disc = disc_model(docp) dims = docp.dims - obj_lagrange = 0.0 + value = 0.0 # loop over time steps for i in 1:docp.time.steps @@ -137,10 +136,10 @@ function runningCost(docp::DOCP{Euler}, xu, v, time_grid) xi = get_OCP_state_at_time_step(xu, docp, index) ui = get_OCP_control_at_time_step(xu, docp, index) hi = time_grid[i + 1] - time_grid[i] - obj_lagrange = obj_lagrange + hi * CTModels.lagrange(ocp)(ti, xi, ui, v) + value += hi * f(ti, xi, ui, v) end - return obj_lagrange + return value end """ diff --git a/src/ode/irk.jl b/src/ode/irk.jl index 1aedbee3..1ccd25dc 100644 --- a/src/ode/irk.jl +++ b/src/ode/irk.jl @@ -206,11 +206,10 @@ $(TYPEDSIGNATURES) Compute the running cost """ -function runningCost(docp::DOCP{<: GenericIRK}, xu, v, time_grid) - ocp = ocp_model(docp) +function integral(docp::DOCP{<: GenericIRK}, xu, v, time_grid, f) disc = disc_model(docp) dims = docp.dims - obj_lagrange = 0.0 + value = 0.0 # work array layout: [x_ij ; sum_bk] work_xij = similar(xu, dims.NLP_x) @@ -243,19 +242,19 @@ function runningCost(docp::DOCP{<: GenericIRK}, xu, v, time_grid) # split to avoid dual tag ordering error in AD if j == 1 work_sumbk[1] = - disc.butcher_b[j] * CTModels.lagrange(ocp)(tij, work_xij, ui, v) + disc.butcher_b[j] * f(tij, work_xij, ui, v) else work_sumbk[1] = work_sumbk[1] + - disc.butcher_b[j] * CTModels.lagrange(ocp)(tij, work_xij, ui, v) + disc.butcher_b[j] * f(tij, work_xij, ui, v) end end # update lagrange cost - obj_lagrange = obj_lagrange + hi * work_sumbk[1] + value += hi * work_sumbk[1] end - return obj_lagrange + return value end """ diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index d2fe3438..6a202929 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -78,10 +78,9 @@ $(TYPEDSIGNATURES) Compute the running cost """ -function runningCost(docp::DOCP{Trapeze}, xu, v, time_grid) - ocp = ocp_model(docp) +function integral(docp::DOCP{Trapeze}, xu, v, time_grid, f) dims = docp.dims - obj_lagrange = 0.0 + value = 0.0 # sum_i=1..N h_i * (l_i + l_i+1) / 2 = (h_1 / 2) l_1 + sum_i=2..N (h_i-1+h_i)/2 * l_i + (h_N / 2) l_N+1 @@ -90,8 +89,8 @@ function runningCost(docp::DOCP{Trapeze}, xu, v, time_grid) ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) ui = get_OCP_control_at_time_step(xu, docp, i) - h = time_grid[i + 1] - time_grid[i] - obj_lagrange = obj_lagrange + h / 2.0 * CTModels.lagrange(ocp)(ti, xi, ui, v) + hi = time_grid[i + 1] - time_grid[i] + value += hi / 2.0 * f(ti, xi, ui, v) # loop over time steps for i in 2:docp.time.steps @@ -99,8 +98,8 @@ function runningCost(docp::DOCP{Trapeze}, xu, v, time_grid) ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) ui = get_OCP_control_at_time_step(xu, docp, i) - h2 = time_grid[i + 1] - time_grid[i - 1] - obj_lagrange = obj_lagrange + h2 / 2.0 * CTModels.lagrange(ocp)(ti, xi, ui, v) + hi2 = time_grid[i + 1] - time_grid[i - 1] + value += hi2 / 2.0 * f(ti, xi, ui, v) end # last term @@ -108,10 +107,10 @@ function runningCost(docp::DOCP{Trapeze}, xu, v, time_grid) ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) ui = get_OCP_control_at_time_step(xu, docp, i) - h = time_grid[i] - time_grid[i - 1] - obj_lagrange = obj_lagrange + h / 2.0 * CTModels.lagrange(ocp)(ti, xi, ui, v) + hi = time_grid[i] - time_grid[i - 1] + value += hi / 2.0 * f(ti, xi, ui, v) - return obj_lagrange + return value end """ From 995a80823580f128348a0a2a9bc18cda2fdb5055 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Mon, 23 Feb 2026 12:30:54 +0100 Subject: [PATCH 03/21] ci for develop too --- .github/workflows/CI.yml | 2 +- src/direct_shooting.jl | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 4547d7ef..93a2aa4a 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -3,7 +3,7 @@ name: CI on: push: branches: - - main + - main, develop tags: '*' pull_request: diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index 7cf8841c..f5dd1a90 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -11,7 +11,7 @@ Strategies.id(::Type{<:DirectShooting}) = :direct_shooting # default options __direct_shooting_grid_size()::Int = 100 __direct_shooting_scheme()::Symbol = :midpoint # later use variable step ode solver -__direct_shooting_scheme_steps() = 10 +__direct_shooting_control_steps() = 10 # ie number of controls per time step function Strategies.metadata(::Type{<:Collocation}) return Strategies.StrategyMetadata( @@ -22,11 +22,18 @@ function Strategies.metadata(::Type{<:Collocation}) description = "Number of time steps for the direct shooting grid", ), Options.OptionDefinition( + name = :control_steps, + type = Int, + default = __direct_shooting_control_steps(), + description = "Number of controls per time step for the direct shooting", + ), + Options.OptionDefinition( name = :scheme, type = Symbol, default = __direct_shooting_scheme(), description = "Time integration scheme (e.g., :midpoint, :trapeze)", ), + ) end From 01ecedac6db97e9030d53ae38fa6a296ef81cf74 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Tue, 24 Feb 2026 17:04:12 +0100 Subject: [PATCH 04/21] use DOCP substructs as arguments for scheme constructors; try to reuse same scheme structs for both collocation and direct shooting discretizers --- src/CTDirect.jl | 2 + src/collocation.jl | 7 +-- src/collocation_core.jl | 17 +++--- src/direct_shooting.jl | 7 ++- src/direct_shooting_core.jl | 114 ++++++++++++++++++++++++++++++++++++ src/ode/euler.jl | 23 +++----- src/ode/irk.jl | 56 ++++-------------- src/ode/midpoint.jl | 29 ++++----- src/ode/trapeze.jl | 19 +++--- 9 files changed, 170 insertions(+), 104 deletions(-) create mode 100644 src/direct_shooting_core.jl diff --git a/src/CTDirect.jl b/src/CTDirect.jl index c4371891..76af8a0a 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -69,5 +69,7 @@ include("ode/irk.jl") include("ode/midpoint.jl") include("ode/trapeze.jl") # ode/variable_step +#include("direct_shooting.jl") +#include("direct_shooting_core.jl") end diff --git a/src/collocation.jl b/src/collocation.jl index 35d3e547..ca6e83a9 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -48,10 +48,7 @@ end Strategies.options(c::Collocation) = c.options -# default options for modelers backend (now in CTSolvers ?) -#__adnlp_backend() = :optimized -#__exa_backend() = nothing - +# +++ move to collocation_core ? # ========================================================================================== # Build core DOCP structure with discretization information (ADNLP) @@ -64,7 +61,7 @@ function get_docp(discretizer::Collocation, ocp::AbstractModel) time_grid = Strategies.options(discretizer)[:time_grid] # initialize DOCP - docp = DOCP(ocp; grid_size=grid_size, time_grid=time_grid, scheme=scheme) + docp = DOCP(ocp, grid_size, time_grid, scheme) # set bounds in DOCP variables_bounds!(docp) diff --git a/src/collocation_core.jl b/src/collocation_core.jl index 221f4a64..66271d9b 100644 --- a/src/collocation_core.jl +++ b/src/collocation_core.jl @@ -146,6 +146,7 @@ DOCPtime(10, [0.0, 0.1, …, 1.0], [0.0, 0.1, …, 1.0]) """ struct DOCPtime steps::Int + control_steps::Int normalized_grid::Vector{Float64} fixed_grid::Vector{Float64} end @@ -172,7 +173,7 @@ julia> DOCPtime(ocp, 10, nothing) DOCPtime(10, [0.0, 0.1, …, 1.0], [0.0, 0.1, …, 1.0]) ``` """ -function DOCPtime(ocp::CTModels.Model, grid_size::Int, time_grid) +function DOCPtime(ocp::CTModels.Model, grid_size::Int, control_steps::Int, time_grid) # 1. build/recover normalized time grid if time_grid === nothing @@ -209,7 +210,7 @@ function DOCPtime(ocp::CTModels.Model, grid_size::Int, time_grid) NLP_fixed_time_grid = @. t0 + (NLP_normalized_time_grid * (tf - t0)) end - return DOCPtime(dim_NLP_steps, NLP_normalized_time_grid, NLP_fixed_time_grid) + return DOCPtime(dim_NLP_steps, control_steps, NLP_normalized_time_grid, NLP_fixed_time_grid) end """ @@ -289,12 +290,7 @@ mutable struct DOCP{ dim_NLP_constraints::Int # constructor - function DOCP( - ocp::CTModels.Model; - grid_size=__grid_size(), - time_grid=__time_grid(), - scheme=__scheme(), - ) + function DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) # boolean flags flags = DOCPFlags(ocp) @@ -303,10 +299,11 @@ mutable struct DOCP{ dims = DOCPdims(ocp) # time grid - time = DOCPtime(ocp, grid_size, time_grid) + control_steps = 1 + time = DOCPtime(ocp, grid_size, control_steps, time_grid) # discretization method - disc_args = [time.steps, dims.NLP_x, dims.NLP_u, dims.NLP_v, dims.path_cons, dims.boundary_cons] + disc_args = [dims, time] if scheme == :trapeze discretization, dim_NLP_variables, dim_NLP_constraints = diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index f5dd1a90..1168da9f 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -45,14 +45,13 @@ end Strategies.options(c::DirectShooting) = c.options - # ========================================================================================== # Build discretizer API (return sets of model/solution builders) # ========================================================================================== function (discretizer::DirectShooting)(ocp::AbstractModel) # common parts for builders - docp = get_docp(discretizer, ocp) ? + docp = get_docp(discretizer, ocp) exa_getter = nothing # will be set in build_exa_model # ========================================================================================== @@ -66,6 +65,8 @@ function (discretizer::DirectShooting)(ocp::AbstractModel) kwargs... )::ADNLPModels.ADNLPModel + docp = + # functions for objective and constraints f = x -> CTDirect.DirectShooting_objective(x, docp) c! = (c, x) -> CTDirect.DirectShooting_constraints!(c, x, docp) @@ -126,6 +127,8 @@ function (discretizer::DirectShooting)(ocp::AbstractModel) initial_guess::CTModels.AbstractInitialGuess; backend )::ExaModels.ExaModel where {BaseType<:AbstractFloat} + + # try to call constructor here with constraints component wise ? end # Solution builder for ExaModels diff --git a/src/direct_shooting_core.jl b/src/direct_shooting_core.jl new file mode 100644 index 00000000..bc35d3f7 --- /dev/null +++ b/src/direct_shooting_core.jl @@ -0,0 +1,114 @@ +# Direct Shooting discretizer: core functions + +# ========================================================================================== +# Build core DOCP structure with discretization information (ADNLP) +# ========================================================================================== +function get_docp(discretizer::DirectShooting, ocp::AbstractModel) + + # recover discretization scheme and options + scheme = Strategies.options(discretizer)[:scheme] + grid_size = Strategies.options(discretizer)[:grid_size] + control_steps = Strategies.options(discretizer)[:control_steps] + + # initialize DOCP + docp = DOCP_DirectShooting(ocp, grid_size, control_steps, scheme) + + # set bounds in DOCP + variables_bounds!(docp) + constraints_bounds!(docp) + + return docp +end + +# could we reuse the same DOCP as collocation +# pass a keyword for the discretizer, as for scheme ? +mutable struct DOCP_DirectShooting{ + D<:CTDirect.Discretization, O<:CTModels.Model + } + + # discretization scheme + discretization::D + + # OCP (for OCP functions calls and some getters eg free times) + # parametric instead of just qualifying reduces allocations (but not time). Specialization ? + ocp::O + + # boolean flags + flags::DOCPFlags + + # dimensions + dims::DOCPdims + + # time grid + time::DOCPtime + + # lower and upper bounds for variables and constraints + bounds::DOCPbounds + + # NLP variables and constraints + dim_NLP_variables::Int + dim_NLP_constraints::Int + + # constructor + function DOCP_DirectShooting(ocp::CTModels.Model, grid_size, control_steps, scheme) + + # boolean flags + flags = DOCPFlags(ocp) + + # dimensions + dims = DOCPdims(ocp) + + # time grid + time = DOCPtime(ocp, grid_size, nothing) + + # discretization method + disc_args = [:direct_shooting, dims, time.steps, control_steps] + + if scheme == :trapeze + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Trapeze(disc_args...) + + elseif scheme == :midpoint + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Midpoint(disc_args...) + + elseif scheme == :euler || scheme == :euler_explicit || scheme == :euler_forward + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Euler(disc_args...) + elseif scheme == :euler_implicit || scheme == :euler_backward + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Euler(disc_args...; explicit=false) + + elseif scheme == :gauss_legendre_2 + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Gauss_Legendre_2(disc_args...) + + elseif scheme == :gauss_legendre_3 + discretization, dim_NLP_variables, dim_NLP_constraints = + CTDirect.Gauss_Legendre_3(disc_args...) + + else + error( + "Unknown discretization method: ", + scheme, + "\nValid options are scheme={:trapeze, :midpoint, :euler | :euler_explicit | :euler_forward, :euler_implicit | :euler_backward, :gauss_legendre_2, :gauss_legendre_3}\n", + typeof(scheme), + ) + end + + # lower and upper bounds for variables and constraints + bounds = DOCPbounds( + -Inf * ones(dim_NLP_variables), + Inf * ones(dim_NLP_variables), + zeros(dim_NLP_constraints), + zeros(dim_NLP_constraints), + ) + + # call constructor with const fields + docp = new{typeof(discretization),typeof(ocp)}( + discretization, ocp, flags, dims, time, bounds, dim_NLP_variables, dim_NLP_constraints, + ) + + return docp + end +end diff --git a/src/ode/euler.jl b/src/ode/euler.jl index 046232bd..01f10876 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -16,29 +16,20 @@ struct Euler <: Discretization _explicit::Bool # constructor - function Euler( - dim_NLP_steps, - dim_NLP_x, - dim_NLP_u, - dim_NLP_v, - dim_path_cons, - dim_boundary_cons; - explicit=true, - ) + function Euler(dims::DOCPdims, time::DOCPtime; explicit=true) # aux variables - step_variables_block = dim_NLP_x + dim_NLP_u - state_stage_eqs_block = dim_NLP_x - step_pathcons_block = dim_path_cons + step_variables_block = dims.NLP_x + dims.NLP_u + state_stage_eqs_block = dims.NLP_x + step_pathcons_block = dims.path_cons # NLP variables size ([state, control]_1..N, final state, variable) - dim_NLP_variables = dim_NLP_steps * step_variables_block + dim_NLP_x + dim_NLP_v + dim_NLP_variables = time.steps * step_variables_block + dims.NLP_x + dims.NLP_v # NLP constraints size ([dynamics, path]_1..N, final path, boundary, variable) dim_NLP_constraints = - dim_NLP_steps * (state_stage_eqs_block + step_pathcons_block) + - step_pathcons_block + - dim_boundary_cons + time.steps * (state_stage_eqs_block + step_pathcons_block) + + step_pathcons_block + dims.boundary_cons if explicit info = "Euler (explicit), 1st order" diff --git a/src/ode/irk.jl b/src/ode/irk.jl index 1ccd25dc..74b8b593 100644 --- a/src/ode/irk.jl +++ b/src/ode/irk.jl @@ -28,20 +28,11 @@ struct Gauss_Legendre_1 <: GenericIRK _step_pathcons_block::Int _final_control::Bool - function Gauss_Legendre_1( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons - ) + function Gauss_Legendre_1(dims::DOCPdims, time::DOCPtime) stage = 1 step_variables_block, state_stage_eqs_block, step_pathcons_block, dim_NLP_variables, dim_NLP_constraints = IRK_dims( - dim_NLP_steps, - dim_NLP_x, - dim_NLP_u, - dim_NLP_v, - dim_path_cons, - dim_boundary_cons, - stage, - ) + dims, time, stage) disc = new( "[test only] Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic, A-stable", @@ -75,20 +66,11 @@ struct Gauss_Legendre_2 <: GenericIRK _step_pathcons_block::Int _final_control::Bool - function Gauss_Legendre_2( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons - ) + function Gauss_Legendre_2(dims::DOCPdims, time::DOCPtime) stage = 2 step_variables_block, state_stage_eqs_block, step_pathcons_block, dim_NLP_variables, dim_NLP_constraints = IRK_dims( - dim_NLP_steps, - dim_NLP_x, - dim_NLP_u, - dim_NLP_v, - dim_path_cons, - dim_boundary_cons, - stage, - ) + dims, time, stage) disc = new( "Implicit Gauss-Legendre collocation for s=2, 4th order, symplectic, A-stable", @@ -122,20 +104,11 @@ struct Gauss_Legendre_3 <: GenericIRK _step_pathcons_block::Int _final_control::Bool - function Gauss_Legendre_3( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons - ) + function Gauss_Legendre_3(dims::DOCPdims, time::DOCPtime) stage = 3 step_variables_block, state_stage_eqs_block, step_pathcons_block, dim_NLP_variables, dim_NLP_constraints = IRK_dims( - dim_NLP_steps, - dim_NLP_x, - dim_NLP_u, - dim_NLP_v, - dim_path_cons, - dim_boundary_cons, - stage, - ) + dims, time, stage) disc = new( "Implicit Gauss-Legendre collocation for s=3, 6th order, symplectic, A-stable", @@ -162,27 +135,24 @@ $(TYPEDSIGNATURES) Return the dimension of the NLP variables and constraints for a generic IRK discretizion, with the control taken constant per step (ie not distinct controls at time stages) """ -function IRK_dims( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons, stage -) +function IRK_dims(dims::DOCPdims, time::DOCPtime, stage::Int) # size of variables block for one step: x, u, k - step_variables_block = dim_NLP_x + dim_NLP_u + dim_NLP_x * stage + step_variables_block = dims.NLP_x + dims.NLP_u + dims.NLP_x * stage # size of state + stage equations for one step - state_stage_eqs_block = dim_NLP_x * (1 + stage) + state_stage_eqs_block = dims.NLP_x * (1 + stage) # size of path constraints block for one step: u, x, xu - step_pathcons_block = dim_path_cons + step_pathcons_block = dims.path_cons # NLP variables size ([state, control, stage]_1..N, final state and control, variable) - dim_NLP_variables = dim_NLP_steps * step_variables_block + dim_NLP_x + dim_NLP_v + dim_NLP_variables = time.steps * step_variables_block + dims.NLP_x + dims.NLP_v # NLP constraints size ([dynamics, stage, path]_1..N, final path, boundary, variable) dim_NLP_constraints = - dim_NLP_steps * (state_stage_eqs_block + step_pathcons_block) + - step_pathcons_block + - dim_boundary_cons + time.steps * (state_stage_eqs_block + step_pathcons_block) + + step_pathcons_block + dims.boundary_cons return step_variables_block, state_stage_eqs_block, step_pathcons_block, dim_NLP_variables, diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 14bbd773..336adbb7 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -14,23 +14,18 @@ struct Midpoint <: Discretization _final_control::Bool # constructor - function Midpoint( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons - ) - - # aux variables - step_variables_block = dim_NLP_x + dim_NLP_u - state_stage_eqs_block = dim_NLP_x - step_pathcons_block = dim_path_cons - - # NLP variables size ([state, control]_1..N, final state, variable) - dim_NLP_variables = dim_NLP_steps * step_variables_block + dim_NLP_x + dim_NLP_v - - # NLP constraints size ([dynamics, path]_1..N, final path, boundary, variable) - dim_NLP_constraints = - dim_NLP_steps * (state_stage_eqs_block + step_pathcons_block) + - step_pathcons_block + - dim_boundary_cons + function Midpoint(dims::DOCPdims, time::DOCPtime) + + step_variables_block = dims.NLP_x + time.control_steps * dims.NLP_u + state_stage_eqs_block = dims.NLP_x + step_pathcons_block = dims.path_cons + + # NLP variables size ([state, controls]_1..N, final state, variable) + dim_NLP_variables = time.steps * step_variables_block + dims.NLP_x + dims.NLP_v + + # NLP constraints size ([state eq, path]_1..N, final_path, boundary) + dim_NLP_constraints = time.steps * (dims.NLP_x + step_pathcons_block) + + step_pathcons_block + dims.boundary_cons disc = new( "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic", diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index 6a202929..0b59fdde 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -11,27 +11,24 @@ struct Trapeze <: Discretization _final_control::Bool # constructor - function Trapeze( - dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons - ) + function Trapeze(dims::DOCPdims, time::DOCPtime) # Trapeze is better with final control (used in final dynamics) final_control = true #false about 10% slower # aux variables - step_variables_block = dim_NLP_x + dim_NLP_u - state_stage_eqs_block = dim_NLP_x - step_pathcons_block = dim_path_cons + step_variables_block = dims.NLP_x + dims.NLP_u + state_stage_eqs_block = dims.NLP_x + step_pathcons_block = dims.path_cons # NLP variables size ([state, control]_1..N+1, variable) - dim_NLP_variables = dim_NLP_steps * step_variables_block + dim_NLP_x + dim_NLP_v - final_control && (dim_NLP_variables += dim_NLP_u) + dim_NLP_variables = time.steps * step_variables_block + dims.NLP_x + dims.NLP_v + final_control && (dim_NLP_variables += dims.NLP_u) # NLP constraints size ([dynamics, stage, path]_1..N, final path, boundary, variable) dim_NLP_constraints = - dim_NLP_steps * (state_stage_eqs_block + step_pathcons_block) + - step_pathcons_block + - dim_boundary_cons + time.steps * (state_stage_eqs_block + step_pathcons_block) + + step_pathcons_block + dims.boundary_cons disc = new( "Implicit Trapeze aka Crank-Nicolson, 2nd order, A-stable", From 0a670e2364718cc7b0aa5439381415f4f97f8ea2 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 09:59:04 +0100 Subject: [PATCH 05/21] local tests ok --- src/CTDirect.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/CTDirect.jl b/src/CTDirect.jl index 76af8a0a..751aa323 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -55,7 +55,7 @@ julia> struct MyDiscretization <: Discretization end MyDiscretization ``` """ -abstract type Discretization end +abstract type Discretization end #+++ rename as Scheme ? # includes From 76b6deb641c66b75d5d48390ee30d76f52f1a9b0 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 10:20:53 +0100 Subject: [PATCH 06/21] some renaming; split path constraints from state equation --- src/CTDirect.jl | 22 +++-- src/{collocation_core.jl => DOCP_core.jl} | 6 +- src/collocation.jl | 6 +- src/collocation_functions.jl | 34 +++++-- src/direct_shooting.jl | 20 ++++ src/direct_shooting_core.jl | 112 ---------------------- src/ode/common.jl | 8 +- src/ode/euler.jl | 2 +- src/ode/irk.jl | 2 +- src/ode/midpoint.jl | 12 +-- src/ode/trapeze.jl | 2 +- 11 files changed, 75 insertions(+), 151 deletions(-) rename src/{collocation_core.jl => DOCP_core.jl} (99%) diff --git a/src/CTDirect.jl b/src/CTDirect.jl index 751aa323..a8da2ce3 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -36,38 +36,40 @@ function discretize( end # --------------------------------------------------------------------------- -# Discretization schemes: see disc/ +# Discretization schemes: see ode/ # --------------------------------------------------------------------------- """ $(TYPEDEF) -Abstract type representing a discretization strategy for an optimal +Abstract type representing a discretization scheme strategy for an optimal control problem. -Concrete subtypes of `Discretization` define specific schemes for +Concrete subtypes of `Scheme` define specific schemes for transforming a continuous-time problem into a discrete-time representation suitable for numerical solution. # Example ```julia-repl -julia> struct MyDiscretization <: Discretization end -MyDiscretization +julia> struct MyScheme <: Scheme end +MyScheme ``` """ -abstract type Discretization end #+++ rename as Scheme ? +abstract type Scheme end # includes -include("collocation.jl") -include("collocation_core.jl") -include("collocation_variables.jl") -include("collocation_functions.jl") +include("DOCP_core.jl") include("ode/common.jl") include("ode/euler.jl") include("ode/irk.jl") include("ode/midpoint.jl") include("ode/trapeze.jl") + +include("collocation.jl") +include("collocation_variables.jl") +include("collocation_functions.jl") + # ode/variable_step #include("direct_shooting.jl") #include("direct_shooting_core.jl") diff --git a/src/collocation_core.jl b/src/DOCP_core.jl similarity index 99% rename from src/collocation_core.jl rename to src/DOCP_core.jl index 66271d9b..28c87271 100644 --- a/src/collocation_core.jl +++ b/src/DOCP_core.jl @@ -262,8 +262,8 @@ julia> DOCP(ocp, nlp_model_backend) DOCP{...}(...) ``` """ -mutable struct DOCP{ - D<:CTDirect.Discretization, O<:CTModels.Model +mutable struct __DOCP{ + D<:CTDirect.Scheme, O<:CTModels.Model } # discretization scheme @@ -290,7 +290,7 @@ mutable struct DOCP{ dim_NLP_constraints::Int # constructor - function DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) + function __DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) # boolean flags flags = DOCPFlags(ocp) diff --git a/src/collocation.jl b/src/collocation.jl index ca6e83a9..ed36bcfb 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -61,11 +61,11 @@ function get_docp(discretizer::Collocation, ocp::AbstractModel) time_grid = Strategies.options(discretizer)[:time_grid] # initialize DOCP - docp = DOCP(ocp, grid_size, time_grid, scheme) + docp = __DOCP(ocp, grid_size, time_grid, scheme) # set bounds in DOCP - variables_bounds!(docp) - constraints_bounds!(docp) + __variables_bounds!(docp) + __constraints_bounds!(docp) return docp end diff --git a/src/collocation_functions.jl b/src/collocation_functions.jl index 5891aff6..18d90a40 100644 --- a/src/collocation_functions.jl +++ b/src/collocation_functions.jl @@ -20,7 +20,7 @@ julia> DOCP_objective(xu, docp) 12.34 ``` """ -function DOCP_objective(xu, docp::DOCP) +function __objective(xu, docp::DOCP) # initialization if docp.flags.freet0 || docp.flags.freetf @@ -77,7 +77,7 @@ julia> DOCP_constraints!(zeros(docp.dim_NLP_constraints), xu, docp) [0.0, 0.1, …] ``` """ -function DOCP_constraints!(c, xu, docp::DOCP) +function __constraints!(c, xu, docp::DOCP) # initialization if docp.flags.freet0 || docp.flags.freetf @@ -89,10 +89,17 @@ function DOCP_constraints!(c, xu, docp::DOCP) work = setWorkArray(docp, xu, time_grid, v) # main loop on time steps - for i in 1:(docp.time.steps + 1) - setStepConstraints!(docp, c, xu, v, time_grid, i, work) + for i in 1:docp.time.steps + # state equation (includes stage equation depending on scheme) + stepStateConstraints!(docp, c, xu, v, time_grid, i, work) + + #path constraints + stepPathConstraints!(docp, c, xu, v, time_grid, i, work) end - + # path constraints at final time + stepPathConstraints!(docp, c, xu, v, time_grid, docp.time.steps+1, work) + + # boundary constraints if docp.dims.boundary_cons > 0 offset = docp.dim_NLP_constraints - docp.dims.boundary_cons @@ -109,6 +116,21 @@ function DOCP_constraints!(c, xu, docp::DOCP) end +""" +$(TYPEDSIGNATURES) +Set path constraints at given time step +""" +function stepPathConstraints!(docp, c, xu, v, time_grid, i, work) + if docp.dims.path_cons > 0 + ui = get_OCP_control_at_time_step(xu, docp, i) + CTModels.path_constraints_nl(ocp)[2]( + (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v + ) + end + return +end + + """ $(TYPEDSIGNATURES) @@ -129,7 +151,7 @@ julia> constraints_bounds!(docp) ([-1.0, …], [1.0, …]) ``` """ -function constraints_bounds!(docp::DOCP) +function __constraints_bounds!(docp::DOCP) lb = docp.bounds.con_l ub = docp.bounds.con_u disc = disc_model(docp) diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index 1168da9f..7f5ba9a8 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -45,6 +45,26 @@ end Strategies.options(c::DirectShooting) = c.options +# ========================================================================================== +# Build core DOCP structure with discretization information (ADNLP) +# ========================================================================================== +function get_docp(discretizer::DirectShooting, ocp::AbstractModel) + + # recover discretization scheme and options + scheme = Strategies.options(discretizer)[:scheme] + grid_size = Strategies.options(discretizer)[:grid_size] + control_steps = Strategies.options(discretizer)[:control_steps] + + # initialize DOCP + docp = __DOCP(ocp, grid_size, control_steps, scheme) + + # set bounds in DOCP + variables_bounds!(docp) + constraints_bounds!(docp) + + return docp +end + # ========================================================================================== # Build discretizer API (return sets of model/solution builders) # ========================================================================================== diff --git a/src/direct_shooting_core.jl b/src/direct_shooting_core.jl index bc35d3f7..74b28cc3 100644 --- a/src/direct_shooting_core.jl +++ b/src/direct_shooting_core.jl @@ -1,114 +1,2 @@ # Direct Shooting discretizer: core functions -# ========================================================================================== -# Build core DOCP structure with discretization information (ADNLP) -# ========================================================================================== -function get_docp(discretizer::DirectShooting, ocp::AbstractModel) - - # recover discretization scheme and options - scheme = Strategies.options(discretizer)[:scheme] - grid_size = Strategies.options(discretizer)[:grid_size] - control_steps = Strategies.options(discretizer)[:control_steps] - - # initialize DOCP - docp = DOCP_DirectShooting(ocp, grid_size, control_steps, scheme) - - # set bounds in DOCP - variables_bounds!(docp) - constraints_bounds!(docp) - - return docp -end - -# could we reuse the same DOCP as collocation -# pass a keyword for the discretizer, as for scheme ? -mutable struct DOCP_DirectShooting{ - D<:CTDirect.Discretization, O<:CTModels.Model - } - - # discretization scheme - discretization::D - - # OCP (for OCP functions calls and some getters eg free times) - # parametric instead of just qualifying reduces allocations (but not time). Specialization ? - ocp::O - - # boolean flags - flags::DOCPFlags - - # dimensions - dims::DOCPdims - - # time grid - time::DOCPtime - - # lower and upper bounds for variables and constraints - bounds::DOCPbounds - - # NLP variables and constraints - dim_NLP_variables::Int - dim_NLP_constraints::Int - - # constructor - function DOCP_DirectShooting(ocp::CTModels.Model, grid_size, control_steps, scheme) - - # boolean flags - flags = DOCPFlags(ocp) - - # dimensions - dims = DOCPdims(ocp) - - # time grid - time = DOCPtime(ocp, grid_size, nothing) - - # discretization method - disc_args = [:direct_shooting, dims, time.steps, control_steps] - - if scheme == :trapeze - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Trapeze(disc_args...) - - elseif scheme == :midpoint - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Midpoint(disc_args...) - - elseif scheme == :euler || scheme == :euler_explicit || scheme == :euler_forward - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Euler(disc_args...) - elseif scheme == :euler_implicit || scheme == :euler_backward - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Euler(disc_args...; explicit=false) - - elseif scheme == :gauss_legendre_2 - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Gauss_Legendre_2(disc_args...) - - elseif scheme == :gauss_legendre_3 - discretization, dim_NLP_variables, dim_NLP_constraints = - CTDirect.Gauss_Legendre_3(disc_args...) - - else - error( - "Unknown discretization method: ", - scheme, - "\nValid options are scheme={:trapeze, :midpoint, :euler | :euler_explicit | :euler_forward, :euler_implicit | :euler_backward, :gauss_legendre_2, :gauss_legendre_3}\n", - typeof(scheme), - ) - end - - # lower and upper bounds for variables and constraints - bounds = DOCPbounds( - -Inf * ones(dim_NLP_variables), - Inf * ones(dim_NLP_variables), - zeros(dim_NLP_constraints), - zeros(dim_NLP_constraints), - ) - - # call constructor with const fields - docp = new{typeof(discretization),typeof(ocp)}( - discretization, ocp, flags, dims, time, bounds, dim_NLP_variables, dim_NLP_constraints, - ) - - return docp - end -end diff --git a/src/ode/common.jl b/src/ode/common.jl index 1ea4c914..0ce88668 100644 --- a/src/ode/common.jl +++ b/src/ode/common.jl @@ -203,7 +203,7 @@ $(TYPEDSIGNATURES) Set work array for all dynamics evaluations """ -function setWorkArray(docp::DOCP{<: Discretization}, xu, time_grid, v) +function setWorkArray(docp::DOCP{<: Scheme}, xu, time_grid, v) return nothing end @@ -213,7 +213,7 @@ $(TYPEDSIGNATURES) Compute the running cost (must be implemented for each discretization scheme) """ -function runningCost(docp::DOCP{D}, xu, v, time_grid) where {(D<:Discretization)} +function runningCost(docp::DOCP{D}, xu, v, time_grid) where {(D<:Scheme)} return integral(docp, xu, v, time_grid, CTModels.lagrange(ocp_model(docp))) end @@ -223,7 +223,7 @@ $(TYPEDSIGNATURES) Build sparsity pattern for Jacobian of constraints (to be implemented for each discretization scheme) """ -function DOCP_Jacobian_pattern(docp::DOCP{D}) where {(D<:Discretization)} +function DOCP_Jacobian_pattern(docp::DOCP{D}) where {(D<:Scheme)} error( "DOCP_Jacobian_pattern not implemented for discretization ", D, @@ -237,7 +237,7 @@ $(TYPEDSIGNATURES) Build sparsity pattern for Hessian of Lagrangian (to be implemented for each discretization scheme) """ -function DOCP_Hessian_pattern(docp::DOCP{D}) where {(D<:Discretization)} +function DOCP_Hessian_pattern(docp::DOCP{D}) where {(D<:Scheme)} error( "DOCP_Hessian_pattern not implemented for discretization ", D, diff --git a/src/ode/euler.jl b/src/ode/euler.jl index 01f10876..7ae86eee 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -7,7 +7,7 @@ with the convention Note that both the explicit and implicit versions therefore use the same variables layout. =# -struct Euler <: Discretization +struct Euler <: Scheme info::String _step_variables_block::Int _state_stage_eqs_block::Int diff --git a/src/ode/irk.jl b/src/ode/irk.jl index 74b8b593..132b1db7 100644 --- a/src/ode/irk.jl +++ b/src/ode/irk.jl @@ -9,7 +9,7 @@ with s the stage number and U piecewise constant equal to U_i in [t_i, t_i+1] Path constraints are all evaluated at time steps, including final time. =# -abstract type GenericIRK <: Discretization end +abstract type GenericIRK <: Scheme end """ $(TYPEDSIGNATURES) diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 336adbb7..8f0cee9f 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -6,7 +6,7 @@ NB. This version is much faster than the one using stage variables Version without work array gives identical performance but is less readable =# -struct Midpoint <: Discretization +struct Midpoint <: Scheme info::String _step_variables_block::Int _state_stage_eqs_block::Int @@ -98,7 +98,7 @@ $(TYPEDSIGNATURES) Set the constraints corresponding to the state equation Convention: 1 <= i <= dim_NLP_steps+1 """ -function setStepConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work) +function stepStateConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work) ocp = ocp_model(docp) disc = disc_model(docp) @@ -122,14 +122,6 @@ function setStepConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work) xip1 - (xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)]) offset += docp.dims.NLP_x end - - # 2. path constraints - if docp.dims.path_cons > 0 - ui = get_OCP_control_at_time_step(xu, docp, i) - CTModels.path_constraints_nl(ocp)[2]( - (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v - ) - end end """ diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index 0b59fdde..1ccfeadf 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -3,7 +3,7 @@ Internal layout for NLP variables: [X_1,U_1, .., X_N+1,U_N+1, V] =# -struct Trapeze <: Discretization +struct Trapeze <: Scheme info::String _step_variables_block::Int _state_stage_eqs_block::Int From 4578a7b09b2256121d902936b20344e68ca8730e Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 12:06:20 +0100 Subject: [PATCH 07/21] DOCP_data --- src/CTDirect.jl | 2 +- src/{DOCP_core.jl => DOCP_data.jl} | 6 +++--- src/collocation.jl | 2 +- src/direct_shooting.jl | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) rename src/{DOCP_core.jl => DOCP_data.jl} (99%) diff --git a/src/CTDirect.jl b/src/CTDirect.jl index a8da2ce3..a4ed9eaa 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -59,7 +59,7 @@ abstract type Scheme end # includes -include("DOCP_core.jl") +include("DOCP_data.jl") include("ode/common.jl") include("ode/euler.jl") include("ode/irk.jl") diff --git a/src/DOCP_core.jl b/src/DOCP_data.jl similarity index 99% rename from src/DOCP_core.jl rename to src/DOCP_data.jl index 28c87271..05b9369c 100644 --- a/src/DOCP_core.jl +++ b/src/DOCP_data.jl @@ -1,4 +1,4 @@ -# Discretized Optimal Control Problem DOCP +# Data for iscretized Optimal Control Problem DOCP """ $(TYPEDEF) @@ -262,7 +262,7 @@ julia> DOCP(ocp, nlp_model_backend) DOCP{...}(...) ``` """ -mutable struct __DOCP{ +mutable struct DOCP{ D<:CTDirect.Scheme, O<:CTModels.Model } @@ -290,7 +290,7 @@ mutable struct __DOCP{ dim_NLP_constraints::Int # constructor - function __DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) + function DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) # boolean flags flags = DOCPFlags(ocp) diff --git a/src/collocation.jl b/src/collocation.jl index ed36bcfb..d30a00e6 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -61,7 +61,7 @@ function get_docp(discretizer::Collocation, ocp::AbstractModel) time_grid = Strategies.options(discretizer)[:time_grid] # initialize DOCP - docp = __DOCP(ocp, grid_size, time_grid, scheme) + docp = DOCP(ocp, grid_size, time_grid, scheme) # set bounds in DOCP __variables_bounds!(docp) diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index 7f5ba9a8..655cda97 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -56,7 +56,7 @@ function get_docp(discretizer::DirectShooting, ocp::AbstractModel) control_steps = Strategies.options(discretizer)[:control_steps] # initialize DOCP - docp = __DOCP(ocp, grid_size, control_steps, scheme) + docp = DOCP(ocp, grid_size, control_steps, scheme) # set bounds in DOCP variables_bounds!(docp) From d528ae40a0afb4b5ec5f915bfabbeccc42d56141 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 12:11:14 +0100 Subject: [PATCH 08/21] local functions renaming --- src/collocation.jl | 6 +++--- src/collocation_variables.jl | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/collocation.jl b/src/collocation.jl index d30a00e6..fbf5e18f 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -83,7 +83,7 @@ function get_docp_initial_guess(modeler::Symbol, docp, functional_init = CTModels.build_initial_guess(ocp, initial_guess) # build discretized initial guess - x0 = DOCP_initial_guess(docp, functional_init) + x0 = __initial_guess(docp, functional_init) if modeler == :adnlp return x0 @@ -132,8 +132,8 @@ function (discretizer::Collocation)(ocp::AbstractModel) )::ADNLPModels.ADNLPModel # functions for objective and constraints - f = x -> CTDirect.DOCP_objective(x, docp) - c! = (c, x) -> CTDirect.DOCP_constraints!(c, x, docp) + f = x -> CTDirect.__objective(x, docp) + c! = (c, x) -> CTDirect.__constraints!(c, x, docp) # build initial guess init = get_docp_initial_guess(:adnlp, docp, initial_guess) diff --git a/src/collocation_variables.jl b/src/collocation_variables.jl index 9417dcbe..6d37c658 100644 --- a/src/collocation_variables.jl +++ b/src/collocation_variables.jl @@ -21,7 +21,7 @@ julia> DOCP_initial_guess(docp) [0.1, 0.1, …] ``` """ -function DOCP_initial_guess(docp::DOCP, init::CTModels.InitialGuess) +function __initial_guess(docp::DOCP, init::CTModels.InitialGuess) # default initialization (including internal variables such as k_i for RK schemes) # NB. passing nothing to the setters below will leave this default values @@ -63,7 +63,7 @@ julia> variables_bounds!(docp) ([-Inf, …], [Inf, …]) ``` """ -function variables_bounds!(docp::DOCP) +function __variables_bounds!(docp::DOCP) N = docp.time.steps var_l = docp.bounds.var_l var_u = docp.bounds.var_u From cf8f803755102299a83430b8e54603ae7d200988 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 12:29:43 +0100 Subject: [PATCH 09/21] objective for direct shooting with midpoint --- src/collocation_functions.jl | 1 - src/ode/common.jl | 9 +++++++- src/ode/euler.jl | 1 - src/ode/midpoint.jl | 44 ++++++++++++++++++++++++++---------- src/ode/trapeze.jl | 1 - 5 files changed, 40 insertions(+), 16 deletions(-) diff --git a/src/collocation_functions.jl b/src/collocation_functions.jl index 18d90a40..18d7636c 100644 --- a/src/collocation_functions.jl +++ b/src/collocation_functions.jl @@ -99,7 +99,6 @@ function __constraints!(c, xu, docp::DOCP) # path constraints at final time stepPathConstraints!(docp, c, xu, v, time_grid, docp.time.steps+1, work) - # boundary constraints if docp.dims.boundary_cons > 0 offset = docp.dim_NLP_constraints - docp.dims.boundary_cons diff --git a/src/ode/common.jl b/src/ode/common.jl index 0ce88668..a87f220f 100644 --- a/src/ode/common.jl +++ b/src/ode/common.jl @@ -123,13 +123,20 @@ Retrieve control variables at given time step from the NLP variables. Convention: 1 <= i <= dim_NLP_steps(+1), with convention u(tf) = U_N Vector output """ -function get_OCP_control_at_time_step(xu, docp::DOCP, i) +function get_OCP_control_at_time_step(xu, docp::DOCP, i; j=1) disc = disc_model(docp) + # final time case, pick U_N unless U_N+1 is present if !disc._final_control && i == docp.time.steps + 1 i = docp.time.steps end + + # skip previous time steps offset = (i-1) * disc._step_variables_block + docp.dims.NLP_x + + # skip previous controls for this time step + offset += (j-1) * docp.dims.NLP_u + return @view xu[(offset + 1):(offset + docp.dims.NLP_u)] end diff --git a/src/ode/euler.jl b/src/ode/euler.jl index 7ae86eee..0a13348c 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -116,7 +116,6 @@ function integral(docp::DOCP{Euler}, xu, v, time_grid, f) # loop over time steps for i in 1:docp.time.steps - offset = (i-1) * dims.NLP_x # get variables at t_i or t_i+1 if disc._explicit index = i diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 8f0cee9f..c2d35f5c 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -75,18 +75,38 @@ Compute the running cost function integral(docp::DOCP{Midpoint}, xu, v, time_grid, f) value = 0.0 - # loop over time steps - for i in 1:docp.time.steps - offset = (i-1) * docp.dims.NLP_x - hi = time_grid[i + 1] - time_grid[i] - ts = 0.5 * (time_grid[i] + time_grid[i + 1]) - xs = - 0.5 * ( - get_OCP_state_at_time_step(xu, docp, i) + - get_OCP_state_at_time_step(xu, docp, i+1) - ) - ui = get_OCP_control_at_time_step(xu, docp, i) - value += hi * f(ts, xs, ui, v) + # +++ add inner loop over control steps ? + # nb if more than one, need to store value + # split both cases control_steps = 1 or > 1 ? + + if docp.time.control_steps == 1 + # loop over time steps + for i in 1:docp.time.steps + hi = time_grid[i + 1] - time_grid[i] + ts = 0.5 * (time_grid[i] + time_grid[i + 1]) + xs = + 0.5 * ( + get_OCP_state_at_time_step(xu, docp, i) + + get_OCP_state_at_time_step(xu, docp, i+1) + ) + ui = get_OCP_control_at_time_step(xu, docp, i) + value += hi * f(ts, xs, ui, v) + end + else + # loop over time steps + for i in 1:docp.time.steps + hi = (time_grid[i + 1] - time_grid[i]) / docp.time.control_steps + xs = 0.5 * ( + get_OCP_state_at_time_step(xu, docp, i) + + get_OCP_state_at_time_step(xu, docp, i+1) + ) + # loop over control steps + for j in 1:docp.time.control_steps + tij = time_grid[i] + (j - 0.5) * hi + uij = get_OCP_control_at_time_step(xu, docp, i; j=j) + value += hi * f(tij, xs, uij, v) + end + end end return value diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index 1ccfeadf..07c5ac90 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -91,7 +91,6 @@ function integral(docp::DOCP{Trapeze}, xu, v, time_grid, f) # loop over time steps for i in 2:docp.time.steps - offset = (i-1) * dims.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) From 6ce92de81f0147dda80483ff153b601a9e17565d Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 12:48:03 +0100 Subject: [PATCH 10/21] state eq for direct shooting with midpoint --- src/ode/midpoint.jl | 53 ++++++++++++++++++++++----------------------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index c2d35f5c..65f3bc5f 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -45,23 +45,27 @@ $(TYPEDSIGNATURES) Set work array for all dynamics cost evaluations """ function setWorkArray(docp::DOCP{Midpoint}, xu, time_grid, v) - work = similar(xu, docp.dims.NLP_x * docp.time.steps) + work = similar(xu, docp.dims.NLP_x * docp.time.steps * docp.time.control_steps) ocp = ocp_model(docp) # loop over time steps + offset = 0 for i in 1:docp.time.steps - offset = (i-1) * docp.dims.NLP_x ts = 0.5 * (time_grid[i] + time_grid[i + 1]) xs = 0.5 * ( get_OCP_state_at_time_step(xu, docp, i) + get_OCP_state_at_time_step(xu, docp, i+1) ) - ui = get_OCP_control_at_time_step(xu, docp, i) - # OCP dynamics - CTModels.dynamics(ocp)( - (@view work[(offset + 1):(offset + docp.dims.NLP_x)]), ts, xs, ui, v - ) + # +++ loop over control steps + for j in 1:docp.time.control_steps + uij = get_OCP_control_at_time_step(xu, docp, i; j=j) + # OCP dynamics + CTModels.dynamics(ocp)( + (@view work[(offset + 1):(offset + docp.dims.NLP_x)]), ts, xs, uij, v + ) + offset += docp.dims.NLP_x + end end return work @@ -75,9 +79,8 @@ Compute the running cost function integral(docp::DOCP{Midpoint}, xu, v, time_grid, f) value = 0.0 - # +++ add inner loop over control steps ? - # nb if more than one, need to store value - # split both cases control_steps = 1 or > 1 ? + # Collocation: 1 control per step + # Direct shooting: >=1 controls per step if docp.time.control_steps == 1 # loop over time steps @@ -122,26 +125,22 @@ function stepStateConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, wor ocp = ocp_model(docp) disc = disc_model(docp) - # offset for previous steps - offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) - - # 0. variables + # compute state variables at next step: midpoint rule ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) - - # 1. state equation - if i <= docp.time.steps - # more variables - 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.dims.NLP_x - - # state equation: midpoint rule - @views @. c[(offset + 1):(offset + docp.dims.NLP_x)] = - xip1 - (xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)]) - offset += docp.dims.NLP_x + tip1 = time_grid[i + 1] + xip1 = get_OCP_state_at_time_step(xu, docp, i+1) + hi = (tip1 - ti) / docp.time.control_steps + offset_dyn_i = (i-1) * docp.dims.NLP_x * docp.time.control_steps + for j in 1:docp.time.control_steps + x_next = xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] + offset_dyn_i += docp.dims.NLP_x end + + # set state equation as constraints (equal to 0) + offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + @views @. c[(offset + 1):(offset + docp.dims.NLP_x)] = xip1 - x_next + end """ From 63fd14be6b8c739007d1f139f37bab37347660a3 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 12:59:30 +0100 Subject: [PATCH 11/21] fix --- src/ode/midpoint.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 65f3bc5f..9b5fb548 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -132,8 +132,9 @@ function stepStateConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, wor xip1 = get_OCP_state_at_time_step(xu, docp, i+1) hi = (tip1 - ti) / docp.time.control_steps offset_dyn_i = (i-1) * docp.dims.NLP_x * docp.time.control_steps + x_next = xi for j in 1:docp.time.control_steps - x_next = xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] + x_next += hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] offset_dyn_i += docp.dims.NLP_x end From afb475551b5447983e0a81808074e2515e9c3812 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 13:04:26 +0100 Subject: [PATCH 12/21] variables bounds and initial guess for direct shooting --- src/collocation_variables.jl | 99 ++++++++++++++++++------------------ 1 file changed, 50 insertions(+), 49 deletions(-) diff --git a/src/collocation_variables.jl b/src/collocation_variables.jl index 6d37c658..267cc804 100644 --- a/src/collocation_variables.jl +++ b/src/collocation_variables.jl @@ -1,48 +1,3 @@ - - -""" -$(TYPEDSIGNATURES) - -Build an initial guess vector for the discretized OCP. - -# Arguments - -- `docp::DOCP`: The discretized OCP. -- `init::CTModels.Init`: Initialization settings (default: `CTModels.Init()`). - -# Returns - -- `NLP_X::Vector{Float64}`: Initial guess vector. - -# Example - -```julia-repl -julia> DOCP_initial_guess(docp) -[0.1, 0.1, …] -``` -""" -function __initial_guess(docp::DOCP, init::CTModels.InitialGuess) - - # default initialization (including internal variables such as k_i for RK schemes) - # NB. passing nothing to the setters below will leave this default values - NLP_X = 0.1 * ones(docp.dim_NLP_variables) - - # set variables if provided (needed first in case of free times !) - set_optim_variable!(NLP_X, init.variable, docp) - - # set state / control variables if provided - # NB. setter handles the final control case - time_grid = get_time_grid(NLP_X, docp) - for i in 1:(docp.time.steps + 1) - ti = time_grid[i] - set_state_at_time_step!(NLP_X, init.state(ti), docp, i) - set_control_at_time_step!(NLP_X, init.control(ti), docp, i) - end - - return NLP_X -end - - """ $(TYPEDSIGNATURES) @@ -64,7 +19,7 @@ julia> variables_bounds!(docp) ``` """ function __variables_bounds!(docp::DOCP) - N = docp.time.steps + var_l = docp.bounds.var_l var_u = docp.bounds.var_u ocp = ocp_model(docp) @@ -82,11 +37,13 @@ function __variables_bounds!(docp::DOCP) ) # set state / control box along time steps - for i in 1:(N + 1) + for i in 1:(docp.time.steps + 1) set_state_at_time_step!(var_l, x_lb, docp, i) set_state_at_time_step!(var_u, x_ub, docp, i) - set_control_at_time_step!(var_l, u_lb, docp, i) - set_control_at_time_step!(var_u, u_ub, docp, i) + for j in 1:docp.time.control_steps + set_control_at_time_step!(var_l, u_lb, docp, i; j=j) + set_control_at_time_step!(var_u, u_ub, docp, i; j=j) + end end # variable box @@ -138,3 +95,47 @@ function build_bounds_block(dim_var, dim_box, box_triplet) return x_lb, x_ub end + +""" +$(TYPEDSIGNATURES) + +Build an initial guess vector for the discretized OCP. + +# Arguments + +- `docp::DOCP`: The discretized OCP. +- `init::CTModels.Init`: Initialization settings (default: `CTModels.Init()`). + +# Returns + +- `NLP_X::Vector{Float64}`: Initial guess vector. + +# Example + +```julia-repl +julia> DOCP_initial_guess(docp) +[0.1, 0.1, …] +``` +""" +function __initial_guess(docp::DOCP, init::CTModels.InitialGuess) + + # default initialization (including internal variables such as k_i for RK schemes) + # NB. passing nothing to the setters below will leave this default values + NLP_X = 0.1 * ones(docp.dim_NLP_variables) + + # set variables if provided (needed first in case of free times !) + set_optim_variable!(NLP_X, init.variable, docp) + + # set state / control variables if provided + # NB. setter handles the final control case + time_grid = get_time_grid(NLP_X, docp) + for i in 1:(docp.time.steps + 1) + ti = time_grid[i] + set_state_at_time_step!(NLP_X, init.state(ti), docp, i) + for j in 1:docp.time.control_steps + set_control_at_time_step!(NLP_X, init.control(ti), docp, i; j=j) + end + end + + return NLP_X +end \ No newline at end of file From 1a271c21bcb97488ac56071fe8adea865633cffb Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 13:13:42 +0100 Subject: [PATCH 13/21] rename files; fix for control setter --- src/CTDirect.jl | 10 ++++------ src/{collocation_functions.jl => DOCP_functions.jl} | 0 src/{collocation_variables.jl => DOCP_variables.jl} | 0 src/direct_shooting.jl | 2 +- src/direct_shooting_core.jl | 2 -- src/ode/common.jl | 6 +++++- 6 files changed, 10 insertions(+), 10 deletions(-) rename src/{collocation_functions.jl => DOCP_functions.jl} (100%) rename src/{collocation_variables.jl => DOCP_variables.jl} (100%) delete mode 100644 src/direct_shooting_core.jl diff --git a/src/CTDirect.jl b/src/CTDirect.jl index a4ed9eaa..eae8d457 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -60,6 +60,9 @@ abstract type Scheme end # includes include("DOCP_data.jl") +include("DOCP_variables.jl") +include("DOCP_functions.jl") + include("ode/common.jl") include("ode/euler.jl") include("ode/irk.jl") @@ -67,11 +70,6 @@ include("ode/midpoint.jl") include("ode/trapeze.jl") include("collocation.jl") -include("collocation_variables.jl") -include("collocation_functions.jl") - -# ode/variable_step -#include("direct_shooting.jl") -#include("direct_shooting_core.jl") +include("direct_shooting.jl") end diff --git a/src/collocation_functions.jl b/src/DOCP_functions.jl similarity index 100% rename from src/collocation_functions.jl rename to src/DOCP_functions.jl diff --git a/src/collocation_variables.jl b/src/DOCP_variables.jl similarity index 100% rename from src/collocation_variables.jl rename to src/DOCP_variables.jl diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index 655cda97..3c482676 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -13,7 +13,7 @@ __direct_shooting_grid_size()::Int = 100 __direct_shooting_scheme()::Symbol = :midpoint # later use variable step ode solver __direct_shooting_control_steps() = 10 # ie number of controls per time step -function Strategies.metadata(::Type{<:Collocation}) +function Strategies.metadata(::Type{<:DirectShooting}) return Strategies.StrategyMetadata( Options.OptionDefinition( name = :grid_size, diff --git a/src/direct_shooting_core.jl b/src/direct_shooting_core.jl deleted file mode 100644 index 74b28cc3..00000000 --- a/src/direct_shooting_core.jl +++ /dev/null @@ -1,2 +0,0 @@ -# Direct Shooting discretizer: core functions - diff --git a/src/ode/common.jl b/src/ode/common.jl index a87f220f..ceeaa8d4 100644 --- a/src/ode/common.jl +++ b/src/ode/common.jl @@ -194,11 +194,15 @@ $(TYPEDSIGNATURES) Set initial guess for control variables at given time step Convention: 1 <= i <= dim_NLP_steps(+1) """ -function set_control_at_time_step!(xu, u_init, docp::DOCP, i) +function set_control_at_time_step!(xu, u_init, docp::DOCP, i; j=1) if !isnothing(u_init) disc = disc_model(docp) + # NB ignore control at final time if not handled by scheme if i <= docp.time.steps || (disc._final_control && i <= docp.time.steps + 1) + # skip: previous stepsand state for this step offset = (i-1) * disc._step_variables_block + docp.dims.NLP_x + # skip previous controls for this stpe + offset += (j-1) * docp.dims.NLP_u xu[(offset + 1):(offset + docp.dims.NLP_u)] .= u_init end end From 0084836b53e8205d11c0833d487f7a6a270b28b3 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 14:58:22 +0100 Subject: [PATCH 14/21] other schemes update --- src/ode/euler.jl | 37 +++++---------- src/ode/irk.jl | 110 ++++++++++++++++++++------------------------ src/ode/midpoint.jl | 5 +- src/ode/trapeze.jl | 47 ++++++------------- 4 files changed, 79 insertions(+), 120 deletions(-) diff --git a/src/ode/euler.jl b/src/ode/euler.jl index 0a13348c..ee0407e3 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -19,7 +19,7 @@ struct Euler <: Scheme function Euler(dims::DOCPdims, time::DOCPtime; explicit=true) # aux variables - step_variables_block = dims.NLP_x + dims.NLP_u + step_variables_block = dims.NLP_x + dims.NLP_u * time.control_steps state_stage_eqs_block = dims.NLP_x step_pathcons_block = dims.path_cons @@ -56,7 +56,7 @@ Retrieve control variables at given time step from the NLP variables. Convention: see above for explicit / implicit versions Vector output """ -function get_OCP_control_at_time_step(xu, docp::DOCP{Euler}, i) +function get_OCP_control_at_time_step(xu, docp::DOCP{Euler}, i; j=1) disc = disc_model(docp) dims = docp.dims if disc._explicit @@ -143,34 +143,19 @@ function setStepConstraints!(docp::DOCP{Euler}, c, xu, v, time_grid, i, work) disc = disc_model(docp) dims = docp.dims - # offset for previous steps - offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) - - # 0. variables + # compute state variables at next step: euler rule ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) + tip1 = time_grid[i + 1] + xip1 = get_OCP_state_at_time_step(xu, docp, i+1) + hi = tip1 - ti + offset_dyn_i = (i-1)*dims.NLP_x - # 1. state equation - if i <= docp.time.steps - # more variables - tip1 = time_grid[i + 1] - xip1 = get_OCP_state_at_time_step(xu, docp, i+1) - hi = tip1 - ti - offset_dyn_i = (i-1)*dims.NLP_x - - # state equation: euler rule - @views @. c[(offset + 1):(offset + dims.NLP_x)] = - xip1 - (xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + dims.NLP_x)]) - offset += dims.NLP_x - end + # set state equation as constraints (equal to 0) + offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + @views @. c[(offset + 1):(offset + dims.NLP_x)] = + xip1 - (xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + dims.NLP_x)]) - # 2. path constraints - if disc._step_pathcons_block > 0 - ui = get_OCP_control_at_time_step(xu, docp, i) - CTModels.path_constraints_nl(ocp)[2]( - (@view c[(offset + 1):(offset + dims.path_cons)]), ti, xi, ui, v - ) - end end """ diff --git a/src/ode/irk.jl b/src/ode/irk.jl index 132b1db7..0465f3f6 100644 --- a/src/ode/irk.jl +++ b/src/ode/irk.jl @@ -138,7 +138,7 @@ Return the dimension of the NLP variables and constraints for a generic IRK disc function IRK_dims(dims::DOCPdims, time::DOCPtime, stage::Int) # size of variables block for one step: x, u, k - step_variables_block = dims.NLP_x + dims.NLP_u + dims.NLP_x * stage + step_variables_block = dims.NLP_x + dims.NLP_u * time.control_steps + dims.NLP_x * stage # size of state + stage equations for one step state_stage_eqs_block = dims.NLP_x * (1 + stage) @@ -220,7 +220,7 @@ function integral(docp::DOCP{<: GenericIRK}, xu, v, time_grid, f) end end - # update lagrange cost + # update value value += hi * work_sumbk[1] end @@ -251,70 +251,60 @@ function setStepConstraints!(docp::DOCP{<: GenericIRK}, c, xu, v, time_grid, i, ui = get_OCP_control_at_time_step(xu, docp, i) # 1. state and stage equations - if i <= docp.time.steps - # more variables - tip1 = time_grid[i + 1] - xip1 = get_OCP_state_at_time_step(xu, docp, i+1) - hi = tip1 - ti - offset_stage_eqs = dims.NLP_x - # loop over stages - for j in 1:disc.stage - - # time at stage: t_i^j = t_i + c[j] h_i - cj = disc.butcher_c[j] - tij = ti + cj * hi - - # stage variables - kij = get_stagevars_at_time_step(xu, docp, i, j) - - # update sum b_j k_i^j (w/ lagrange term) for state equation after loop - # split to avoid dual tag ordering error in AD - if j == 1 - @views @. work_sumbk[1:dims.NLP_x] = disc.butcher_b[j] * kij[1:dims.NLP_x] - else - @views @. work_sumbk[1:dims.NLP_x] = - work_sumbk[1:dims.NLP_x] + disc.butcher_b[j] * kij[1:dims.NLP_x] - end - - # state at stage: x_i^j = x_i + h_i sum a_jl k_i^l - # +++ still some allocations here - @views @. work_xij[1:dims.NLP_x] = xi - for l in 1:disc.stage - kil = get_stagevars_at_time_step(xu, docp, i, l) - @views @. work_xij[1:dims.NLP_x] = - work_xij[1:dims.NLP_x] + hi * disc.butcher_a[j, l] * kil[1:dims.NLP_x] - end - - # stage equations k_i^j = f(t_i^j, x_i^j, u_i, v) as c[] = k - f - # NB. we skip the state equation here, which will be set below - CTModels.dynamics(ocp)( - (@view c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)]), - tij, - work_xij, - ui, - v, - ) - @views @. c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)] = - kij - - c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)] - offset_stage_eqs += dims.NLP_x + # more variables + tip1 = time_grid[i + 1] + xip1 = get_OCP_state_at_time_step(xu, docp, i+1) + hi = tip1 - ti + offset_stage_eqs = dims.NLP_x + + # loop over stages + for j in 1:disc.stage + + # time at stage: t_i^j = t_i + c[j] h_i + cj = disc.butcher_c[j] + tij = ti + cj * hi + + # stage variables + kij = get_stagevars_at_time_step(xu, docp, i, j) + + # update sum b_j k_i^j (w/ lagrange term) for state equation after loop + # split to avoid dual tag ordering error in AD + if j == 1 + @views @. work_sumbk[1:dims.NLP_x] = disc.butcher_b[j] * kij[1:dims.NLP_x] + else + @views @. work_sumbk[1:dims.NLP_x] = + work_sumbk[1:dims.NLP_x] + disc.butcher_b[j] * kij[1:dims.NLP_x] end - # state equation x_i+1 = x_i + h_i sum b_j k_i^j - @views @. c[(offset + 1):(offset + dims.NLP_x)] = - xip1 - (xi + hi * work_sumbk[1:dims.NLP_x]) - - # update offset for stage and state equations - offset += dims.NLP_x * (1 + disc.stage) - end + # state at stage: x_i^j = x_i + h_i sum a_jl k_i^l + # +++ still some allocations here + @views @. work_xij[1:dims.NLP_x] = xi + for l in 1:disc.stage + kil = get_stagevars_at_time_step(xu, docp, i, l) + @views @. work_xij[1:dims.NLP_x] = + work_xij[1:dims.NLP_x] + hi * disc.butcher_a[j, l] * kil[1:dims.NLP_x] + end - #2. path constraints - if dims.path_cons > 0 - CTModels.path_constraints_nl(ocp)[2]( - (@view c[(offset + 1):(offset + dims.path_cons)]), ti, xi, ui, v + # stage equations k_i^j = f(t_i^j, x_i^j, u_i, v) as c[] = k - f + # NB. we skip the state equation here, which will be set below + CTModels.dynamics(ocp)( + (@view c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)]), + tij, + work_xij, + ui, + v, ) + @views @. c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)] = + kij - + c[(offset + offset_stage_eqs + 1):(offset + offset_stage_eqs + dims.NLP_x)] + offset_stage_eqs += dims.NLP_x end + + # state equation x_i+1 = x_i + h_i sum b_j k_i^j + @views @. c[(offset + 1):(offset + dims.NLP_x)] = + xip1 - (xi + hi * work_sumbk[1:dims.NLP_x]) + end """ diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 9b5fb548..55adcc72 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -16,7 +16,7 @@ struct Midpoint <: Scheme # constructor function Midpoint(dims::DOCPdims, time::DOCPtime) - step_variables_block = dims.NLP_x + time.control_steps * dims.NLP_u + step_variables_block = dims.NLP_x + dims.NLP_u * time.control_steps state_stage_eqs_block = dims.NLP_x step_pathcons_block = dims.path_cons @@ -57,7 +57,7 @@ function setWorkArray(docp::DOCP{Midpoint}, xu, time_grid, v) get_OCP_state_at_time_step(xu, docp, i) + get_OCP_state_at_time_step(xu, docp, i+1) ) - # +++ loop over control steps + # loop over control steps for j in 1:docp.time.control_steps uij = get_OCP_control_at_time_step(xu, docp, i; j=j) # OCP dynamics @@ -132,6 +132,7 @@ function stepStateConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, wor xip1 = get_OCP_state_at_time_step(xu, docp, i+1) hi = (tip1 - ti) / docp.time.control_steps offset_dyn_i = (i-1) * docp.dims.NLP_x * docp.time.control_steps + # +++ allocations here ? x_next = xi for j in 1:docp.time.control_steps x_next += hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index 07c5ac90..c6ea4af2 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -17,7 +17,7 @@ struct Trapeze <: Scheme final_control = true #false about 10% slower # aux variables - step_variables_block = dims.NLP_x + dims.NLP_u + step_variables_block = dims.NLP_x + dims.NLP_u * time.control_steps state_stage_eqs_block = dims.NLP_x step_pathcons_block = dims.path_cons @@ -120,42 +120,25 @@ function setStepConstraints!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work) disc = disc_model(docp) dims = docp.dims - # offset for previous steps - offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) - - # 0. variables + # compute state variables at next step: trapeze rule ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) - - # 1. state equation - if i <= docp.time.steps - # 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)*dims.NLP_x - offset_dyn_ip1 = i*dims.NLP_x - - # trapeze rule (no allocations ^^) - @views @. c[(offset + 1):(offset + dims.NLP_x)] = - xip1 - ( - xi + - half_hi * ( - work[(offset_dyn_i + 1):(offset_dyn_i + dims.NLP_x)] + - work[(offset_dyn_ip1 + 1):(offset_dyn_ip1 + dims.NLP_x)] - ) + 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)*dims.NLP_x + offset_dyn_ip1 = i*dims.NLP_x + # +++ allocations here ? + x_next = xi + x_next += half_hi * ( + work[(offset_dyn_i + 1):(offset_dyn_i + dims.NLP_x)] + + work[(offset_dyn_ip1 + 1):(offset_dyn_ip1 + dims.NLP_x)] ) - offset += dims.NLP_x - end + # set state equation as constraints (equal to 0) + offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + @views @. c[(offset + 1):(offset + dims.NLP_x)] = xip1 - x_next - # 2. path constraints - if dims.path_cons > 0 - ui = get_OCP_control_at_time_step(xu, docp, i) - CTModels.path_constraints_nl(ocp)[2]( - (@view c[(offset + 1):(offset + dims.path_cons)]), ti, xi, ui, v - ) - end end """ From 95849813d36c0d17afa6e4663c4475d430680efd Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 17:03:09 +0100 Subject: [PATCH 15/21] bugfix --- src/ode/euler.jl | 2 +- src/ode/irk.jl | 2 +- src/ode/trapeze.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ode/euler.jl b/src/ode/euler.jl index ee0407e3..f186aa65 100644 --- a/src/ode/euler.jl +++ b/src/ode/euler.jl @@ -138,7 +138,7 @@ $(TYPEDSIGNATURES) Set the constraints corresponding to the state equation Convention: 1 <= i <= dim_NLP_steps+1 """ -function setStepConstraints!(docp::DOCP{Euler}, c, xu, v, time_grid, i, work) +function stepStateConstraints!(docp::DOCP{Euler}, c, xu, v, time_grid, i, work) ocp = ocp_model(docp) disc = disc_model(docp) dims = docp.dims diff --git a/src/ode/irk.jl b/src/ode/irk.jl index 0465f3f6..28c666eb 100644 --- a/src/ode/irk.jl +++ b/src/ode/irk.jl @@ -233,7 +233,7 @@ $(TYPEDSIGNATURES) Set the constraints corresponding to the state equation Convention: 1 <= i <= dim_NLP_steps (+1) """ -function setStepConstraints!(docp::DOCP{<: GenericIRK}, c, xu, v, time_grid, i, work) +function stepStateConstraints!(docp::DOCP{<: GenericIRK}, c, xu, v, time_grid, i, work) ocp = ocp_model(docp) disc = disc_model(docp) dims = docp.dims diff --git a/src/ode/trapeze.jl b/src/ode/trapeze.jl index c6ea4af2..b5cb963e 100644 --- a/src/ode/trapeze.jl +++ b/src/ode/trapeze.jl @@ -115,7 +115,7 @@ $(TYPEDSIGNATURES) Set the constraints corresponding to the state equation Convention: 1 <= i <= dim_NLP_steps+1 """ -function setStepConstraints!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work) +function stepStateConstraints!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work) ocp = ocp_model(docp) disc = disc_model(docp) dims = docp.dims From c5b232514d314e2017fff3b0989e20ba511c6728 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 17:58:06 +0100 Subject: [PATCH 16/21] fix for path constraints --- src/DOCP_functions.jl | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/DOCP_functions.jl b/src/DOCP_functions.jl index 18d7636c..28f38957 100644 --- a/src/DOCP_functions.jl +++ b/src/DOCP_functions.jl @@ -94,10 +94,10 @@ function __constraints!(c, xu, docp::DOCP) stepStateConstraints!(docp, c, xu, v, time_grid, i, work) #path constraints - stepPathConstraints!(docp, c, xu, v, time_grid, i, work) + (docp.dims.path_cons > 0) && stepPathConstraints!(docp, c, xu, v, time_grid, i) end # path constraints at final time - stepPathConstraints!(docp, c, xu, v, time_grid, docp.time.steps+1, work) + (docp.dims.path_cons > 0) && stepPathConstraints!(docp, c, xu, v, time_grid, docp.time.steps+1) # boundary constraints if docp.dims.boundary_cons > 0 @@ -119,13 +119,19 @@ end $(TYPEDSIGNATURES) Set path constraints at given time step """ -function stepPathConstraints!(docp, c, xu, v, time_grid, i, work) - if docp.dims.path_cons > 0 - ui = get_OCP_control_at_time_step(xu, docp, i) - CTModels.path_constraints_nl(ocp)[2]( - (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v - ) - end +function stepPathConstraints!(docp, c, xu, v, time_grid, i) + + ocp = ocp_model(docp) + disc = disc_model(docp) + ti = time_grid[i] + xi = get_OCP_state_at_time_step(xu, docp, i) + ui = get_OCP_control_at_time_step(xu, docp, i) + offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + + disc._state_stage_eqs_block + + CTModels.path_constraints_nl(ocp)[2]( + (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v + ) return end From 0098f7d139518725b86c61691932a017aec5bc37 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 18:26:55 +0100 Subject: [PATCH 17/21] another fix for path constraints --- src/DOCP_functions.jl | 10 +++++++--- test/manual_test.jl | 5 +++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/DOCP_functions.jl b/src/DOCP_functions.jl index 28f38957..4f0a155f 100644 --- a/src/DOCP_functions.jl +++ b/src/DOCP_functions.jl @@ -123,12 +123,16 @@ function stepPathConstraints!(docp, c, xu, v, time_grid, i) ocp = ocp_model(docp) disc = disc_model(docp) + + # skip previous steps + offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + # skip state equation except at final time + (i <= docp.time.steps) && (offset += disc._state_stage_eqs_block) + + # set constraint ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) ui = get_OCP_control_at_time_step(xu, docp, i) - offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) - + disc._state_stage_eqs_block - CTModels.path_constraints_nl(ocp)[2]( (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v ) diff --git a/test/manual_test.jl b/test/manual_test.jl index 90afac5e..bd6d3c43 100644 --- a/test/manual_test.jl +++ b/test/manual_test.jl @@ -6,7 +6,8 @@ include("test_common.jl") include("./problems/beam.jl") include("./problems/goddard.jl") include("./problems/double_integrator.jl") +include("./problems/truck_trailer.jl") -sol = solve_problem(double_integrator_minenergy(); display=true) -sol1 = solve_problem(double_integrator_minenergy2(); display=true, modeler=:exa, solver=:madnlp) +sol = solve_problem(truck_trailer(); display=true) +#sol1 = solve_problem(double_integrator_minenergy2(); display=true, modeler=:exa, solver=:madnlp) From d3ea7a230dbc36fb9b1da5fac9533a1d1f945052 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Wed, 25 Feb 2026 19:09:41 +0100 Subject: [PATCH 18/21] madnlpgpu 0.8 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 2eaeb058..af14e89e 100644 --- a/Project.toml +++ b/Project.toml @@ -29,7 +29,7 @@ CUDA = "5" CommonSolve = "0.2" DocStringExtensions = "0.9" ExaModels = "0.9" -MadNLPGPU = "0.7" +MadNLPGPU = "0.8" MadNLPMumps = "0.5" NLPModels = "0.21" NLPModelsIpopt = "0.11" From 4c2c192b9b0175a83ed5b8a914c5a306bfd8d300 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Thu, 26 Feb 2026 10:25:03 +0100 Subject: [PATCH 19/21] put back compat for madnlpgpu 0.7 in addition to new 0.8 (dependency problem) --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index af14e89e..f98eb636 100644 --- a/Project.toml +++ b/Project.toml @@ -29,7 +29,7 @@ CUDA = "5" CommonSolve = "0.2" DocStringExtensions = "0.9" ExaModels = "0.9" -MadNLPGPU = "0.8" +MadNLPGPU = "0.7, 0.8" MadNLPMumps = "0.5" NLPModels = "0.21" NLPModelsIpopt = "0.11" From a35482b85a2ebd8c6efba47c279e7fcb782e5bbd Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Thu, 26 Feb 2026 11:30:53 +0100 Subject: [PATCH 20/21] direct shooting (midpoint) on truck trailer problem (7state, 2control, free tf, bolza, path constraints): first tests ok with 1 control step ie similar to collocation, and with 50 steps / 20 controls per step as well --- src/CTDirect.jl | 20 +++++--------------- src/DOCP_data.jl | 3 +-- src/collocation.jl | 5 +++-- src/direct_shooting.jl | 36 ++++++++++++------------------------ test/manual_test.jl | 9 ++++++--- test/test_common.jl | 11 +++++++++-- 6 files changed, 36 insertions(+), 48 deletions(-) diff --git a/src/CTDirect.jl b/src/CTDirect.jl index eae8d457..a80cc8ff 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -10,28 +10,18 @@ using SparseArrays using SolverCore using NLPModels -# ---------------------------------------------------------------------- -# TYPES -const AbstractModel = CTModels.AbstractModel - # --------------------------------------------------------------------------- # Abstract discretizer type # --------------------------------------------------------------------------- +const AbstractModel = CTModels.AbstractModel abstract type AbstractDiscretizer <: Strategies.AbstractStrategy end -function discretize( - ocp::AbstractModel, - discretizer::AbstractDiscretizer -) - return discretizer(ocp) -end - __discretizer()::AbstractDiscretizer = Collocation() -function discretize( - ocp::AbstractModel; - discretizer::AbstractDiscretizer=__discretizer(), -) +function discretize(ocp::AbstractModel, discretizer::AbstractDiscretizer) + return discretizer(ocp) +end +function discretize(ocp::AbstractModel; discretizer::AbstractDiscretizer=__discretizer()) return discretize(ocp, discretizer) end diff --git a/src/DOCP_data.jl b/src/DOCP_data.jl index 05b9369c..4c9a207e 100644 --- a/src/DOCP_data.jl +++ b/src/DOCP_data.jl @@ -290,7 +290,7 @@ mutable struct DOCP{ dim_NLP_constraints::Int # constructor - function DOCP(ocp::CTModels.Model, grid_size, time_grid, scheme) + function DOCP(ocp::CTModels.Model, grid_size::Int, control_steps::Int, scheme::Symbol, time_grid) # boolean flags flags = DOCPFlags(ocp) @@ -299,7 +299,6 @@ mutable struct DOCP{ dims = DOCPdims(ocp) # time grid - control_steps = 1 time = DOCPtime(ocp, grid_size, control_steps, time_grid) # discretization method diff --git a/src/collocation.jl b/src/collocation.jl index fbf5e18f..30a0479c 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -48,7 +48,7 @@ end Strategies.options(c::Collocation) = c.options -# +++ move to collocation_core ? +# +++ todo if possible: unify get_docp for Collocation / directshooting and move to DOCP_data.jl ? # ========================================================================================== # Build core DOCP structure with discretization information (ADNLP) @@ -61,7 +61,8 @@ function get_docp(discretizer::Collocation, ocp::AbstractModel) time_grid = Strategies.options(discretizer)[:time_grid] # initialize DOCP - docp = DOCP(ocp, grid_size, time_grid, scheme) + control_steps = 1 + docp = DOCP(ocp, grid_size, control_steps, scheme, time_grid) # set bounds in DOCP __variables_bounds!(docp) diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index 3c482676..f3bbef45 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -9,9 +9,9 @@ end Strategies.id(::Type{<:DirectShooting}) = :direct_shooting # default options -__direct_shooting_grid_size()::Int = 100 +__direct_shooting_grid_size()::Int = 250 +__direct_shooting_control_steps()::Int = 1 # ie number of controls per time step __direct_shooting_scheme()::Symbol = :midpoint # later use variable step ode solver -__direct_shooting_control_steps() = 10 # ie number of controls per time step function Strategies.metadata(::Type{<:DirectShooting}) return Strategies.StrategyMetadata( @@ -45,6 +45,8 @@ end Strategies.options(c::DirectShooting) = c.options +# +++ todo if possible: unify get_docp for Collocation / directshooting and move to DOCP_data.jl ? + # ========================================================================================== # Build core DOCP structure with discretization information (ADNLP) # ========================================================================================== @@ -56,11 +58,12 @@ function get_docp(discretizer::DirectShooting, ocp::AbstractModel) control_steps = Strategies.options(discretizer)[:control_steps] # initialize DOCP - docp = DOCP(ocp, grid_size, control_steps, scheme) + time_grid = nothing + docp = DOCP(ocp, grid_size, control_steps, scheme, time_grid) # set bounds in DOCP - variables_bounds!(docp) - constraints_bounds!(docp) + __variables_bounds!(docp) + __constraints_bounds!(docp) return docp end @@ -85,26 +88,12 @@ function (discretizer::DirectShooting)(ocp::AbstractModel) kwargs... )::ADNLPModels.ADNLPModel - docp = - # functions for objective and constraints - f = x -> CTDirect.DirectShooting_objective(x, docp) - c! = (c, x) -> CTDirect.DirectShooting_constraints!(c, x, docp) + f = x -> CTDirect.__objective(x, docp) + c! = (c, x) -> CTDirect.__constraints!(c, x, docp) # build initial guess - init = get_docp_initial_guess(:adnlp, docp, initial_guess) - - # unused backends (option excluded_backend = [:jprod_backend, :jtprod_backend, :hprod_backend, :ghjvprod_backend] does not seem to work) - unused_backends = ( - hprod_backend=ADNLPModels.EmptyADbackend, - jtprod_backend=ADNLPModels.EmptyADbackend, - jprod_backend=ADNLPModels.EmptyADbackend, - ghjvprod_backend=ADNLPModels.EmptyADbackend, - ) - - - # use backend preset - backend_options = (backend=backend,) + init = ones(docp.dim_NLP_variables) # build NLP nlp = ADNLPModel!( @@ -116,8 +105,7 @@ function (discretizer::DirectShooting)(ocp::AbstractModel) docp.bounds.con_l, docp.bounds.con_u; minimize=(!docp.flags.max), - backend_options..., - unused_backends..., + backend=:optimized, kwargs..., ) diff --git a/test/manual_test.jl b/test/manual_test.jl index bd6d3c43..bb12fdf2 100644 --- a/test/manual_test.jl +++ b/test/manual_test.jl @@ -8,6 +8,9 @@ include("./problems/goddard.jl") include("./problems/double_integrator.jl") include("./problems/truck_trailer.jl") -sol = solve_problem(truck_trailer(); display=true) -#sol1 = solve_problem(double_integrator_minenergy2(); display=true, modeler=:exa, solver=:madnlp) - +sol = solve_problem(truck_trailer(); display=false) +println("J ", objective(sol), " tf ", variable(sol)) +sol1 = solve_problem(truck_trailer(); discretizer=:direct_shooting, display=false) +println("J ", objective(sol1), " tf ", variable(sol1)) +sol2 = solve_problem(truck_trailer(); discretizer=:direct_shooting, grid_size=50, control_steps=10, display=true) +println("J ", objective(sol1), " tf ", variable(sol1)) \ No newline at end of file diff --git a/test/test_common.jl b/test/test_common.jl index 9781d318..0b1f91d6 100644 --- a/test/test_common.jl +++ b/test/test_common.jl @@ -25,6 +25,7 @@ using SplitApplyCombine # for flatten in some tests function solve_problem(prob; max_iter=1000, tol=1e-6, + discretizer=:collocation, modeler=:adnlp, #+++exa solver=:ipopt, #+++madnlp display=false, @@ -37,8 +38,14 @@ function solve_problem(prob; kwargs...) # discretized problem (model and solution builders) - discretizer = CTDirect.Collocation(; kwargs...) # kwargs here - docp = CTDirect.discretize(prob.ocp, discretizer) + if discretizer == :collocation + my_discretizer = CTDirect.Collocation(; kwargs...) + elseif discretizer == :direct_shooting + my_discretizer = CTDirect.DirectShooting(; kwargs...) + else + error("Unknown discretizer: ", discretizer) + end + docp = CTDirect.discretize(prob.ocp, my_discretizer) # initial guess for model builders if isnothing(init) From 35db9ff0480f21c50dbe3c81496874010b2dc756 Mon Sep 17 00:00:00 2001 From: Pierre Martinon Date: Fri, 27 Feb 2026 12:18:55 +0100 Subject: [PATCH 21/21] split collocation / direct shooting case for midpoint --- src/collocation.jl | 64 +++++++++++----------------------- src/direct_shooting.jl | 7 ++-- src/ode/midpoint.jl | 32 +++++++++++------ test/benchmark.jl | 4 ++- test/manual_test.jl | 10 +++--- test/problems/truck_trailer.jl | 1 + 6 files changed, 55 insertions(+), 63 deletions(-) diff --git a/src/collocation.jl b/src/collocation.jl index 30a0479c..6fff3362 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -71,46 +71,6 @@ function get_docp(discretizer::Collocation, ocp::AbstractModel) return docp end -# ========================================================================================== -# Build initial guess for discretized problem -# ========================================================================================== -function get_docp_initial_guess(modeler::Symbol, docp, - initial_guess::Union{CTModels.AbstractInitialGuess,Nothing}, - ) - - ocp = ocp_model(docp) - - # build functional initial guess - functional_init = CTModels.build_initial_guess(ocp, initial_guess) - - # build discretized initial guess - x0 = __initial_guess(docp, functional_init) - - if modeler == :adnlp - return x0 - elseif modeler == :exa - # reshape initial guess for ExaModel variables layout - # - do not broadcast, apparently fails on GPU arrays - # - unused final control in examodel / euler, hence the different x0 sizes - n = CTModels.state_dimension(ocp) - m = CTModels.control_dimension(ocp) - q = CTModels.variable_dimension(ocp) - N = docp.time.steps - # N + 1 states, N controls - state = hcat([x0[(1 + i * (n + m)):(1 + i * (n + m) + n - 1)] - for i in 0:N]...) - control = hcat([x0[(n + 1 + i * (n + m)):(n + 1 + i * (n + m) + m - 1)] - for i in 0:(N - 1)]...,) - # see with JB: pass indeed to grid_size only for euler(_b), trapeze and midpoint - control = [control control[:, end]] - variable = x0[(end - q + 1):end] - - return (variable, state, control) - else - error("unknown modeler in get_docp_initial_guess: ", modeler) - end - end - # ========================================================================================== # Build discretizer API (return sets of model/solution builders) @@ -137,7 +97,8 @@ function (discretizer::Collocation)(ocp::AbstractModel) c! = (c, x) -> CTDirect.__constraints!(c, x, docp) # build initial guess - init = get_docp_initial_guess(:adnlp, docp, initial_guess) + functional_init = CTModels.build_initial_guess(ocp, initial_guess) + x0 = __initial_guess(docp, functional_init) # unused backends (option excluded_backend = [:jprod_backend, :jtprod_backend, :hprod_backend, :ghjvprod_backend] does not seem to work) unused_backends = ( @@ -174,7 +135,7 @@ function (discretizer::Collocation)(ocp::AbstractModel) # build NLP nlp = ADNLPModel!( f, - init, + x0, docp.bounds.var_l, docp.bounds.var_u, c!, @@ -218,7 +179,24 @@ function (discretizer::Collocation)(ocp::AbstractModel) grid_size = Strategies.options(discretizer)[:grid_size] # build initial guess - init = get_docp_initial_guess(:exa, docp, initial_guess) + functional_init = CTModels.build_initial_guess(ocp, initial_guess) + x0 = __initial_guess(docp, functional_init) + # reshape initial guess for ExaModel variables layout + # - do not broadcast, apparently fails on GPU arrays + # - unused final control in examodel / euler, hence the different x0 sizes + n = CTModels.state_dimension(ocp) + m = CTModels.control_dimension(ocp) + q = CTModels.variable_dimension(ocp) + N = docp.time.steps + # N + 1 states, N controls + state = hcat([x0[(1 + i * (n + m)):(1 + i * (n + m) + n - 1)] + for i in 0:N]...) + control = hcat([x0[(n + 1 + i * (n + m)):(n + 1 + i * (n + m) + m - 1)] + for i in 0:(N - 1)]...,) + # see with JB: pass indeed to grid_size only for euler(_b), trapeze and midpoint + control = [control control[:, end]] + variable = x0[(end - q + 1):end] + init = (variable, state, control) # build Exa model and getters # see with JB. later try to call Exa constructor here if possible, reusing existing functions... diff --git a/src/direct_shooting.jl b/src/direct_shooting.jl index f3bbef45..670060be 100644 --- a/src/direct_shooting.jl +++ b/src/direct_shooting.jl @@ -93,12 +93,13 @@ function (discretizer::DirectShooting)(ocp::AbstractModel) c! = (c, x) -> CTDirect.__constraints!(c, x, docp) # build initial guess - init = ones(docp.dim_NLP_variables) - + functional_init = CTModels.build_initial_guess(ocp, initial_guess) + x0 = __initial_guess(docp, functional_init) + # build NLP nlp = ADNLPModel!( f, - init, + x0, docp.bounds.var_l, docp.bounds.var_u, c!, diff --git a/src/ode/midpoint.jl b/src/ode/midpoint.jl index 55adcc72..dba8659c 100644 --- a/src/ode/midpoint.jl +++ b/src/ode/midpoint.jl @@ -125,24 +125,34 @@ function stepStateConstraints!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, wor ocp = ocp_model(docp) disc = disc_model(docp) - # compute state variables at next step: midpoint rule + # set state equation as constraints (equal to 0) ti = time_grid[i] xi = get_OCP_state_at_time_step(xu, docp, i) tip1 = time_grid[i + 1] xip1 = get_OCP_state_at_time_step(xu, docp, i+1) hi = (tip1 - ti) / docp.time.control_steps offset_dyn_i = (i-1) * docp.dims.NLP_x * docp.time.control_steps - # +++ allocations here ? - x_next = xi - for j in 1:docp.time.control_steps - x_next += hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] - offset_dyn_i += docp.dims.NLP_x - end - - # set state equation as constraints (equal to 0) - offset = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) - @views @. c[(offset + 1):(offset + docp.dims.NLP_x)] = xip1 - x_next + offset_c = (i-1)*(disc._state_stage_eqs_block + disc._step_pathcons_block) + if docp.time.control_steps == 1 + # use c directly to avoid allocations + @views @. c[(offset_c + 1):(offset_c + docp.dims.NLP_x)] = xip1 - (xi + hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)]) + return + # NB. benchmark is close with version below for 1 step + # SUCCESS 30/30 0.61(107) 1.44(117) 3.76(121) + # vs SUCCESS 30/30 0.67(107) 1.50(117) 4.00(121) + # +++recheck vs on the fly dynamics instead of work array + else + # NB. the dynamics in work use only the x values at time steps + # +++update: recompute dynamics here instead of using work array ? + x_next = xi + for j in 1:docp.time.control_steps + x_next += hi * work[(offset_dyn_i + 1):(offset_dyn_i + docp.dims.NLP_x)] + offset_dyn_i += docp.dims.NLP_x + end + @views @. c[(offset_c + 1):(offset_c + docp.dims.NLP_x)] = xip1 - x_next + return + end end """ diff --git a/test/benchmark.jl b/test/benchmark.jl index e25e3910..afcf95aa 100644 --- a/test/benchmark.jl +++ b/test/benchmark.jl @@ -142,7 +142,7 @@ end # verbose <= 1: no output # verbose > 1: print summary (iter, obj, time) # verbose > 2: print NLP iterations also -function bench_problem(problem; verbose=1, solver, grid_size, timer=false, kwargs...) +function bench_problem(problem; verbose=1, solver, grid_size, timer=true, kwargs...) if verbose > 2 display = true else @@ -206,6 +206,7 @@ function bench(; solver=:ipopt, return_sols=false, save_sols=false, + timer=false, kwargs..., ) @@ -239,6 +240,7 @@ function bench(; grid_size=grid_size, verbose=verbose-1, solver=solver, + timer=timer, kwargs..., ) t_bench[i, j] = time diff --git a/test/manual_test.jl b/test/manual_test.jl index bb12fdf2..7e5dca39 100644 --- a/test/manual_test.jl +++ b/test/manual_test.jl @@ -8,9 +8,9 @@ include("./problems/goddard.jl") include("./problems/double_integrator.jl") include("./problems/truck_trailer.jl") -sol = solve_problem(truck_trailer(); display=false) -println("J ", objective(sol), " tf ", variable(sol)) -sol1 = solve_problem(truck_trailer(); discretizer=:direct_shooting, display=false) -println("J ", objective(sol1), " tf ", variable(sol1)) +#sol = solve_problem(truck_trailer(); display=true) +#println("J ", objective(sol), " tf ", variable(sol), " iter ", iterations(sol)) +sol1 = solve_problem(truck_trailer(); discretizer=:direct_shooting, display=true) +println("J ", objective(sol1), " tf ", variable(sol1), " iter ", iterations(sol1)) sol2 = solve_problem(truck_trailer(); discretizer=:direct_shooting, grid_size=50, control_steps=10, display=true) -println("J ", objective(sol1), " tf ", variable(sol1)) \ No newline at end of file +println("J ", objective(sol2), " tf ", variable(sol2)) \ No newline at end of file diff --git a/test/problems/truck_trailer.jl b/test/problems/truck_trailer.jl index c0841e65..8a38d740 100644 --- a/test/problems/truck_trailer.jl +++ b/test/problems/truck_trailer.jl @@ -111,6 +111,7 @@ function truck_trailer(; data=[0.4 0.1 0.2; 1.1 0.2 0.2; 0.8 0.1 0.2]) # initial guess tf_init = 10 + #init = (state=[0.1,0.2,0.3,0.4,0.5,0.6,0.7], control=[0.1,0.2], variable=[tf_init],) init = (variable=[tf_init],) return ((ocp=ocp, obj=59.28, name="truck_trailer", init=init))