todo: check GL2 with stage controls

src/irk.jl

@@ -18,14 +18,14 @@ $(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_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons, stage, control_disc)
+function IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons, stage; control_disc=:step)
 
     if control_disc == :step
         # NLP variables size (state, control, variable, stage)
         dim_NLP_variables = (dim_NLP_steps + 1) * dim_NLP_x + dim_NLP_steps * dim_NLP_u + dim_NLP_v + dim_NLP_steps * dim_NLP_x * stage
 
         # Path constraints (control, state, mixed) 
-        dim_path_cons = dim_control_cons + dim_state_cons + dim_mixed_cons
+        dim_path_cons = dim_u_cons + dim_x_cons + dim_xu_cons
 
         # size of variables block for one step
         step_block = dim_NLP_x + dim_NLP_u + dim_NLP_x * stage
@@ -35,7 +35,7 @@ function IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_co
         dim_NLP_variables = (dim_NLP_steps + 1) * dim_NLP_x + dim_NLP_steps * dim_NLP_u * stage + dim_NLP_v + dim_NLP_steps * dim_NLP_x * stage
 
         # Path constraints (control, state, mixed) 
-        dim_path_cons = dim_control_cons * stage + dim_state_cons + dim_mixed_cons * stage
+        dim_path_cons = dim_u_cons * stage + dim_x_cons + dim_xu_cons * stage
 
         # size of variables block for one step
         step_block = dim_NLP_x + dim_NLP_u + dim_NLP_x * stage
@@ -55,24 +55,24 @@ function IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_co
 $(TYPEDSIGNATURES)
 
 Implicit Midpoint discretization, formulated as a generic IRK
+NB. does not use the simplification xs = 0.5 * (xi + xip1) as in midpoint.jl
 """
 struct Midpoint_IRK <: GenericIRK
 
     stage::Int
     butcher_a::Matrix{Float64}
     butcher_b::Vector{Float64}
     butcher_c::Vector{Float64}
-    control_disc::Symbol
     _step_block::Int
     info::String
 
-    function Midpoint_IRK(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons; control_disc = :step)
+    function Midpoint_IRK(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
 
         stage = 1
 
-        dim_NLP_variables, dim_NLP_constraints, dim_path_cons, _step_block = IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons, stage, control_disc)
+        dim_NLP_variables, dim_NLP_constraints, dim_path_cons, _step_block = IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons, stage)
 
-        disc = new(stage, hcat(0.5), [1], [0.5], control_disc, _step_block,
+        disc = new(stage, hcat(0.5), [1], [0.5], _step_block,
         "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
 
         return disc, dim_NLP_variables, dim_NLP_constraints, dim_path_cons
@@ -94,11 +94,11 @@ struct Gauss_Legendre_2 <: GenericIRK
     _step_block::Int
     info::String
 
-    function Gauss_Legendre_2(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons; control_disc=:step)
+    function Gauss_Legendre_2(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons; control_disc=:step)
 
         stage = 2
 
-        dim_NLP_variables, dim_NLP_constraints, dim_path_cons, _step_block = IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons, stage, control_disc)
+        dim_NLP_variables, dim_NLP_constraints, dim_path_cons, _step_block = IRK_dims(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons, stage; control_disc)
 
         disc = new(stage,
         [0.25 (0.25-sqrt(3) / 6); (0.25+sqrt(3) / 6) 0.25],
@@ -312,8 +312,8 @@ function setConstraintBlock!(docp::DOCP{ <: GenericIRK}, c, xu, v, time_grid, i,
     if docp.dim_x_cons > 0 
         docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, xi, v)
     end
-    if docp.dim_mixed_cons > 0 
-        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, xi, ui, v)
+    if docp.dim_xu_cons > 0 
+        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_xu_cons]), ti, xi, ui, v)
     end
 
 end

src/midpoint.jl

@@ -4,30 +4,28 @@ Internal layout for NLP variables:
 with the convention u([t_i,t_i+1[) = U_i and u(tf) = U_N-1
 =#
 
-# NB. could be defined as a generic IRK in future IRK.jl for comparison purposes
-#butcher_a::Matrix{Float64}
-#butcher_b::Vector{Float64}
-#butcher_c::Vector{Float64}
-#return new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
 struct Midpoint <: Discretization
 
     stage::Int
     info::String
 
     # constructor
-    function Midpoint(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons, dim_v_cons)
+    function Midpoint(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
 
         stage = 1
 
         disc = new(stage, "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
 
         # NLP variables size (state, control, variable, stage)
         dim_NLP_variables = (dim_NLP_steps + 1) * dim_NLP_x + dim_NLP_steps * dim_NLP_u + dim_NLP_v + dim_NLP_steps * dim_NLP_x * stage
+
+        # Path constraints (control, state, mixed) 
+        dim_path_cons = dim_u_cons + dim_x_cons + dim_xu_cons
 
         # NLP constraints size (dynamics, stage, path, boundary, variable)
         dim_NLP_constraints = dim_NLP_steps * (dim_NLP_x + (dim_NLP_x * stage) + dim_path_cons) + dim_path_cons + dim_boundary_cons + dim_v_cons
 
-        return disc, dim_NLP_variables, dim_NLP_constraints
+        return disc, dim_NLP_variables, dim_NLP_constraints, dim_path_cons
     end
 end
 
@@ -216,8 +214,8 @@ function setConstraintBlock!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work)
     if docp.dim_x_cons > 0 
         docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, xi, v)
     end
-    if docp.dim_mixed_cons > 0 
-        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, xi, ui, v)
+    if docp.dim_xu_cons > 0 
+        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_xu_cons]), ti, xi, ui, v)
     end
 
 end

src/problem.jl

@@ -1,7 +1,4 @@
 # Discretized Optimal Control Problem DOCP
-# Notes:
-# - for now the path constraints are checked on the time steps ie g(t_i, x_i, u_i). This requires a getter that provide a control value at each time step, which may not coincide with the actual control discretization (eg RK stages). In this case we will use the 'average' control using the same coefficients as the RK method. Later we can add an option for control discretization, step or stage. Taking a control constant per step would solve the question for the path constraints evaluation and reduces the number of variables. Compared to the more standard stage control discretization, the consistency of the costate would need to be checked, as well as the oscillations in the trajectory (which may actually be better). Further options may include CVP (control vector parametrization) on a coarser grid.
-# - for the other choice of enforcing path constraints at the time stages, the symmetric question of getting state values at time stages can be solved by reusing the states used in the evaluation of the stage dynamics. This second choice (as in Bocop2) has the drawback of a larger problem size, and does not check the constraints at the points of the actual trajectory computed (including tf).
 
 # generic discretization
 abstract type Discretization end
@@ -56,7 +53,7 @@ struct DOCP{T <: Discretization, X <: ScalVect, U <: ScalVect, V <: ScalVect}
     dim_x_cons::Int
     dim_u_cons::Int
     dim_v_cons::Int
-    dim_mixed_cons::Int
+    dim_xu_cons::Int
     dim_boundary_cons::Int
 
     ## NLP  
@@ -175,18 +172,20 @@ struct DOCP{T <: Discretization, X <: ScalVect, U <: ScalVect, V <: ScalVect}
         dim_x_cons = dim_state_constraints(ocp)
         dim_u_cons = dim_control_constraints(ocp)
         dim_v_cons = dim_variable_constraints(ocp)
-        dim_mixed_cons = dim_mixed_constraints(ocp)
+        dim_xu_cons = dim_mixed_constraints(ocp)
         dim_boundary_cons = dim_boundary_constraints(ocp)
 
         # parameter: discretization method
         if disc_method == :midpoint
-            discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Midpoint(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons)
+            discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Midpoint(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
         elseif disc_method == :trapeze
-            discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Trapeze(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons)
+            discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Trapeze(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
         elseif disc_method == :midpoint_irk
-                discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Midpoint_IRK(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons)
+                discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Midpoint_IRK(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
         elseif disc_method == :gauss_legendre_2
-                discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Gauss_Legendre_2(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_control_cons, dim_state_cons, dim_mixed_cons, dim_boundary_cons, dim_v_cons)                
+                discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Gauss_Legendre_2(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
+        elseif disc_method == :gauss_legendre_2_stage_control
+                discretization, dim_NLP_variables, dim_NLP_constraints, dim_path_cons = CTDirect.Gauss_Legendre_2(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_con, dim_v_cons, control_disc=:stage)                             
         else           
             error("Unknown discretization method: ", disc_method, "\nValid options are disc_method={:trapeze, :midpoint, :midpoint_irk, :gauss_legendre_2}\n", typeof(disc_method))
         end
@@ -229,7 +228,7 @@ struct DOCP{T <: Discretization, X <: ScalVect, U <: ScalVect, V <: ScalVect}
             dim_x_cons,
             dim_u_cons,
             dim_v_cons,
-            dim_mixed_cons,
+            dim_xu_cons,
             dim_boundary_cons,
             dim_NLP_x,
             dim_NLP_u,
@@ -413,10 +412,10 @@ function setPathBounds!(docp::DOCP, index::Int, lb, ub)
     end
 
     # mixed state / control constraints
-    if docp.dim_mixed_cons > 0
-        lb[index:(index + docp.dim_mixed_cons - 1)] = docp.mixed_constraints[1]
-        ub[index:(index + docp.dim_mixed_cons - 1)] = docp.mixed_constraints[3]
-        index = index + docp.dim_mixed_cons
+    if docp.dim_xu_cons > 0
+        lb[index:(index + docp.dim_xu_cons - 1)] = docp.mixed_constraints[1]
+        ub[index:(index + docp.dim_xu_cons - 1)] = docp.mixed_constraints[3]
+        index = index + docp.dim_xu_cons
     end
 
     return index

src/trapeze.jl

@@ -6,21 +6,24 @@ Internal layout for NLP variables:
 # NB. could be defined as a generic IRK
 struct Trapeze <: Discretization
 
-    stage::Int # 0 but value used for constraints bounds...
+    stage::Int # 0 but value used for some index computations in problem.jl
     info::String
 
     # constructor
-    function Trapeze(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_path_cons, dim_boundary_cons, dim_v_cons)
+    function Trapeze(dim_NLP_steps, dim_NLP_x, dim_NLP_u, dim_NLP_v, dim_u_cons, dim_x_cons, dim_xu_cons, dim_boundary_cons, dim_v_cons)
 
         disc = new(0, "Implicit Trapeze aka Crank-Nicolson, 2nd order, A-stable")
 
         # NLP variables size (state, control, variable, stage)
         dim_NLP_variables = (dim_NLP_steps + 1) * (dim_NLP_x + dim_NLP_u) + dim_NLP_v
 
+        # Path constraints (control, state, mixed) 
+        dim_path_cons = dim_u_cons + dim_x_cons + dim_xu_cons
+
         # NLP constraints size (dynamics, stage, path, boundary, variable)
         dim_NLP_constraints = dim_NLP_steps * (dim_NLP_x + dim_path_cons) + dim_path_cons + dim_boundary_cons + dim_v_cons
 
-        return disc, dim_NLP_variables, dim_NLP_constraints
+        return disc, dim_NLP_variables, dim_NLP_constraints, dim_path_cons
     end
 end
 
@@ -167,8 +170,8 @@ function setConstraintBlock!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work)
     if docp.dim_x_cons > 0 
         docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, xi, v)
     end
-    if docp.dim_mixed_cons > 0 
-        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, xi, ui, v)
+    if docp.dim_xu_cons > 0 
+        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_xu_cons]), ti, xi, ui, v)
     end
 
 end

test/suite/test_discretization.jl

@@ -72,7 +72,9 @@ end
     sol = direct_solve(prob.ocp, display = false, disc_method = :midpoint_irk)
     @test sol.objective ≈ prob.obj rtol = 1e-2
     sol = direct_solve(prob.ocp, display = false, disc_method = :gauss_legendre_2)
-    @test sol.objective ≈ prob.obj rtol = 1e-2   
+    @test sol.objective ≈ prob.obj rtol = 1e-2
+    sol = direct_solve(prob.ocp, display = false, disc_method = :gauss_legendre_2_stage_control)
+    @test sol.objective ≈ prob.obj rtol = 1e-2
 end
 
 #=@testset verbose = true showtiming = true ":double_integrator :trapeze :midpoint :gl2" begin
@@ -96,6 +98,8 @@ end=#
     sol = direct_solve(prob.ocp, display = false, disc_method = :midpoint_irk)
     @test sol.objective ≈ prob.obj rtol = 1e-2
     sol = direct_solve(prob.ocp, display = false, disc_method = :gauss_legendre_2)
-    @test sol.objective ≈ prob.obj rtol = 1e-2   
+    @test sol.objective ≈ prob.obj rtol = 1e-2
+    sol = direct_solve(prob.ocp, display = false, disc_method = :gauss_legendre_2_stage_control)
+    @test sol.objective ≈ prob.obj rtol = 1e-2  
 end