Direct shooting

.github/workflows/CI.yml

@@ -3,7 +3,7 @@ name: CI
 on:
   push:
     branches:
-      - main
+      - main, develop
     tags: '*'
   pull_request:
 

Project.toml

@@ -29,7 +29,7 @@ CUDA = "5"
 CommonSolve = "0.2"
 DocStringExtensions = "0.9"
 ExaModels = "0.9"
-MadNLPGPU = "0.7"
+MadNLPGPU = "0.7, 0.8"
 MadNLPMumps = "0.5"
 NLPModels = "0.21"
 NLPModelsIpopt = "0.11"

src/CTDirect.jl

@@ -10,63 +10,56 @@ 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
 
 # ---------------------------------------------------------------------------
-# 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
+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")
+include("ode/midpoint.jl")
+include("ode/trapeze.jl")
+
 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("direct_shooting.jl")
 
 end

src/collocation_core.jl → src/DOCP_data.jl

@@ -1,4 +1,4 @@
-# Discretized Optimal Control Problem DOCP
+# Data for iscretized Optimal Control Problem DOCP
 
 """
 $(TYPEDEF)
@@ -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
 
 """
@@ -262,7 +263,7 @@ DOCP{...}(...)
 ```
 """
 mutable struct DOCP{
-    D<:CTDirect.Discretization, O<:CTModels.Model
+    D<:CTDirect.Scheme, O<:CTModels.Model
     }
 
     # discretization scheme
@@ -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::Int, control_steps::Int, scheme::Symbol, time_grid)
 
         # boolean flags
         flags = DOCPFlags(ocp)
@@ -303,10 +299,10 @@ mutable struct DOCP{
         dims = DOCPdims(ocp)
 
         # time grid
-        time = DOCPtime(ocp, grid_size, time_grid)
+        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 = 

src/collocation_functions.jl → src/DOCP_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,16 @@ 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
+        (docp.dims.path_cons > 0) && stepPathConstraints!(docp, c, xu, v, time_grid, i)
     end
-
+    # path constraints at final time
+    (docp.dims.path_cons > 0) && stepPathConstraints!(docp, c, xu, v, time_grid, docp.time.steps+1)
+
     # boundary constraints
     if docp.dims.boundary_cons > 0
         offset = docp.dim_NLP_constraints - docp.dims.boundary_cons
@@ -109,6 +115,31 @@ 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)
+
+    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)
+    CTModels.path_constraints_nl(ocp)[2](
+    (@view c[(offset + 1):(offset + docp.dims.path_cons)]), ti, xi, ui, v
+    )
+    return
+end
+
+
 """
 $(TYPEDSIGNATURES)
 
@@ -129,7 +160,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)

src/collocation_variables.jl → src/DOCP_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 DOCP_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)
 
@@ -63,8 +18,8 @@ julia> variables_bounds!(docp)
 ([-Inf, …], [Inf, …])
 ```
 """
-function variables_bounds!(docp::DOCP)
-    N = docp.time.steps
+function __variables_bounds!(docp::DOCP)
+
     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