Merge pull request #98 from control-toolbox/96-general-bug-constraint

96 general bug constraint

Author: ocots

Project.toml

@@ -1,7 +1,7 @@
 name = "CTModels"
 uuid = "34c4fa32-2049-4079-8329-de33c2a22e2d"
 authors = ["Olivier Cots <olivier.cots@toulouse-inp.fr>"]
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 CTBase = "54762871-cc72-4466-b8e8-f6c8b58076cd"

src/constraints.jl

@@ -337,3 +337,92 @@ Return the dimension of variable box constraints.
 function dim_variable_constraints_box(model::ConstraintsModel)::Dimension
     return length(variable_constraints_box(model)[1])
 end
+
+# ------------------------------------------------------------------------------ #
+"""
+$(TYPEDSIGNATURES)
+
+Get a labelled constraint from the model.
+"""
+function constraint(model::Model, label::Symbol)::Function # not type stable
+
+    # check if the label is in the path constraints
+    cp = path_constraints_nl(model)
+    labels = cp[4] # vector of labels
+    if label in labels
+        # get all the indices of the label
+        indices = findall(x -> x == label, labels)
+        fc! = (r, t, x, u, v) -> begin
+            r_ = zeros(length(cp[1]))
+            cp[2](r_, t, x, u, v)
+            r .= r_[indices]
+        end
+        return to_out_of_place(fc!, length(indices))
+    end
+
+    # check if the label is in the boundary constraints
+    cp = boundary_constraints_nl(model)
+    labels = cp[4] # vector of labels
+    if label in labels
+        # get all the indices of the label
+        indices = findall(x -> x == label, labels)
+        fc! = (r, x0, xf, v) -> begin
+            r_ = zeros(length(cp[1]))
+            cp[2](r_, x0, xf, v)
+            r .= r_[indices]
+        end
+        return to_out_of_place(fc!, length(indices))
+    end
+
+    # check if the label is in the state constraints
+    cp = state_constraints_box(model)
+    labels = cp[4] # vector of labels
+    if label in labels
+        # get all the indices of the label
+        indices = Int[]
+        for i in eachindex(labels)
+            if labels[i] == label
+                push!(indices, cp[2][i])
+            end
+        end
+        fc = (t, x, u, v) -> begin
+            length(indices) == 1 ? x[indices[1]] : x[indices]
+        end
+        return fc
+    end
+
+    # check if the label is in the control constraints
+    cp = control_constraints_box(model)
+    labels = cp[4] # vector of labels
+    if label in labels
+        # get all the indices of the label
+        indices = Int[]
+        for i in eachindex(labels)
+            if labels[i] == label
+                push!(indices, cp[2][i])
+            end
+        end
+        fc = (t, x, u, v) -> begin
+            length(indices) == 1 ? u[indices[1]] : u[indices]
+        end
+        return fc
+    end
+
+    # check if the label is in the variable constraints
+    cp = variable_constraints_box(model)
+    labels = cp[4] # vector of labels
+    if label in labels
+        # get all the indices of the label
+        indices = Int[]
+        for i in eachindex(labels)
+            if labels[i] == label
+                push!(indices, cp[2][i])
+            end
+        end
+        fc = (t, x, u, v) -> begin
+            length(indices) == 1 ? v[indices[1]] : v[indices]
+        end
+        return fc
+    end
+
+end

src/model.jl

@@ -741,15 +741,6 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Get a labelled constraint from the model.
-"""
-function constraint(ocp::Model, label::Symbol)::Tuple # not type stable
-    return constraints(ocp).dict[label]
-end
-
-"""
-$(TYPEDSIGNATURES)
-
 Get the nonlinear path constraints from the model.
 """
 function path_constraints_nl(

src/utils.jl

@@ -19,3 +19,19 @@ function matrix2vec(
 )::Vector{<:Vector{<:ctNumber}}
     return dim==1 ? [A[i, :] for i in 1:size(A, 1)] : [A[:, i] for i in 1:size(A, 2)]
 end
+
+"""
+$(TYPEDSIGNATURES)
+
+Transform an in-place function `f!` to an out-of-place function `f`.
+The function `f` will return a vector of the same type as `T` and size `n`.
+"""
+function to_out_of_place(f!, n; T=Float64)
+    function f(args...; kwargs...)
+        r = zeros(T, n)
+        f!(r, args...; kwargs...)
+        return n == 1 ? r[1] : r
+        #return r # everything is now a vector
+    end
+    return isnothing(f!) ? nothing : f
+end

test/solution_test.jld2

@@ -131,4 +131,5 @@ function test_constraints()
     # test with :variable constraint and range
     CTModels.constraint!(ocp_set, :variable; rg=1:1, lb=[1], ub=[1], label=:variable_rg)
     @test ocp_set.constraints[:variable_rg] == (:variable, 1:1, [1], [1])
+
 end

test/test_constraints.jl

@@ -57,23 +57,51 @@ function test_model()
     # set some constraints
     f_path(r, t, x, u, v) = r .= x .+ u .+ v .+ t
     f_boundary(r, x0, xf, v) = r .= x0 .+ v .* (xf .- x0)
-    f_variable(r, t, v) = r .= v .+ t
+
     CTModels.constraint!(pre_ocp, :path; f=f_path, lb=[0, 1], ub=[1, 2], label=:path)
-    CTModels.constraint!(
-        pre_ocp, :boundary; f=f_boundary, lb=[0, 1], ub=[1, 2], label=:boundary
-    )
-    CTModels.constraint!(pre_ocp, :state; rg=1:2, lb=[0, 1], ub=[1, 2], label=:state_rg)
-    CTModels.constraint!(pre_ocp, :control; rg=1:2, lb=[0, 1], ub=[1, 2], label=:control_rg)
-    CTModels.constraint!(
-        pre_ocp, :variable; rg=1:2, lb=[0, 1], ub=[1, 2], label=:variable_rg
-    )
+    CTModels.constraint!(pre_ocp, :boundary; f=f_boundary, lb=[0, 1], ub=[1, 2], label=:boundary)
+    CTModels.constraint!(pre_ocp, :state; rg=1:2, lb=[0, 1], ub=[1, 2], label=:state)
+    CTModels.constraint!(pre_ocp, :control; rg=1:2, lb=[0, 1], ub=[1, 2], label=:control)
+    CTModels.constraint!(pre_ocp, :variable; rg=1:2, lb=[0, 1], ub=[1, 2], label=:variable)
+
+    f_path_scalar(r, t, x, u, v) = r .= x[1] + u[1] + v[1] + t
+    f_boundary_scalar(r, x0, xf, v) = r .= x0[1] + v[1] * (xf[1] - x0[1])
+    CTModels.constraint!(pre_ocp, :path; f=f_path_scalar, lb=0, ub=1, label=:path_scalar)
+    CTModels.constraint!(pre_ocp, :boundary; f=f_boundary_scalar, lb=0, ub=1, label=:boundary_scalar)
+    CTModels.constraint!(pre_ocp, :state; rg=1, lb=0, ub=1, label=:state_scalar)
+    CTModels.constraint!(pre_ocp, :control; rg=1, lb=0, ub=1, label=:control_scalar)
+    CTModels.constraint!(pre_ocp, :variable; rg=1, lb=0, ub=1, label=:variable_scalar)
+    CTModels.constraint!(pre_ocp, :state; rg=2, lb=0, ub=1, label=:state_scalar_2)
+    CTModels.constraint!(pre_ocp, :control; rg=2, lb=0, ub=1, label=:control_scalar_2)
+    CTModels.constraint!(pre_ocp, :variable; rg=2, lb=0, ub=1, label=:variable_scalar_2)
 
     # build the model
     model = CTModels.build_model(pre_ocp)
 
     # check the type of the model
     @test model isa CTModels.Model
 
+    # check retrieved constraints
+    t = 1
+    x = [2, 3]
+    u = [4, 5]
+    v = [6, 7]
+    x0 = [1, 2]
+    xf = [3, 4]
+    @test CTModels.constraint(model, :path)(t, x, u, v) == x .+ u .+ v .+ t
+    @test CTModels.constraint(model, :boundary)(x0, xf, v) == x0 .+ v .* (xf .- x0)
+    @test CTModels.constraint(model, :state)(t, x, u, v) == x
+    @test CTModels.constraint(model, :control)(t, x, u, v) == u
+    @test CTModels.constraint(model, :variable)(t, x, u, v) == v
+    @test CTModels.constraint(model, :path_scalar)(t, x, u, v) == x[1] + u[1] + v[1] + t
+    @test CTModels.constraint(model, :boundary_scalar)(x0, xf, v) == x0[1] + v[1] * (xf[1] - x0[1])
+    @test CTModels.constraint(model, :state_scalar)(t, x, u, v) == x[1]
+    @test CTModels.constraint(model, :control_scalar)(t, x, u, v) == u[1]
+    @test CTModels.constraint(model, :variable_scalar)(t, x, u, v) == v[1]
+    @test CTModels.constraint(model, :state_scalar_2)(t, x, u, v) == x[2]
+    @test CTModels.constraint(model, :control_scalar_2)(t, x, u, v) == u[2]
+    @test CTModels.constraint(model, :variable_scalar_2)(t, x, u, v) == v[2]
+
     # print the premodel
     display(pre_ocp)
 
@@ -100,4 +128,5 @@ function test_model()
     CTModels.objective!(pre_ocp, :min; mayer=mayer, lagrange=lagrange)
     CTModels.definition!(pre_ocp, quote end)
     display(pre_ocp)
+
 end