update for parametric NT, test global obj and constraints

Project.toml

@@ -1,7 +1,7 @@
 name = "DirectTrajOpt"
 uuid = "c823fa1f-8872-4af5-b810-2b9b72bbbf56"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.2.7"
+version = "0.3.0"
 
 [deps]
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
@@ -31,7 +31,7 @@ JLD2 = "0.5"
 Libdl = "1.10, 1.11"
 LinearAlgebra = "1.10, 1.11"
 MathOptInterface = "1.40"
-NamedTrajectories = "0.3"
+NamedTrajectories = "0.4"
 OrdinaryDiffEq = "6.97"
 Reexport = "1.2"
 SciMLBase = "2.86"

src/constraints/linear_constraints.jl

@@ -64,8 +64,6 @@ function EqualityConstraint(
     return EqualityConstraint(name, ts, fill(val, traj.dims[name]), traj; label=label)
 end
 
-
-
 """
     GlobalEqualityConstraint(
         name::Symbol,
@@ -82,20 +80,36 @@ function GlobalEqualityConstraint(
     traj::NamedTrajectory;
     label="equality constraint on global variable $name"
 )
-    @assert name ∈ keys(traj.global_data)
+    @assert name ∈ traj.global_names
     @assert length(val) == traj.global_dims[name]
 
-    indices = traj.global_components[name]
+    indices = traj.dim * traj.T .+ traj.global_components[name]
     return EqualityConstraint(indices, val, label)
 end
 
+struct AllEqualConstraint <: AbstractLinearConstraint
+    indices::Vector{Int}
+    bar_index::Int
+    label::String
+end
 
+function TimeStepsAllEqualConstraint(
+    traj::NamedTrajectory;
+    label="timesteps all equal constraint"
+)
+    @assert traj.timestep isa Symbol
+    indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T-1]
+    bar_index = index(traj.T, traj.components[traj.timestep][1], traj.dim)
+    return AllEqualConstraint(indices, bar_index, label)
+end
 
 
 ### 
 ### BoundsConstraint
 ###
 
+# TODO: Refactor using `get_bounds_from_dims` from NamedTrajectories.jl
+
 struct BoundsConstraint <: AbstractLinearConstraint
     indices::Vector{Int}
     bounds::Vector{Tuple{Float64, Float64}}
@@ -158,11 +172,11 @@ function GlobalBoundsConstraint(
     traj::NamedTrajectory;
     label="bounds constraint on global variable $name"
 )
-    @assert name ∈ keys(traj.global_data)
+    @assert name ∈ traj.global_names
     @assert length(bounds[1]) == length(bounds[2]) == traj.global_dims[name]
     @assert all(bounds[1] .<= bounds[2])
 
-    indices = traj.global_components[name]
+    indices = traj.dim * traj.T .+ traj.global_components[name]
 
     bounds = collect(zip(bounds...))
 
@@ -175,7 +189,7 @@ function GlobalBoundsConstraint(
     traj::NamedTrajectory;
     label="bounds constraint on global variable $name"
 )
-    @assert name ∈ keys(traj.global_data)
+    @assert name ∈ traj.global_names
     @assert length(bound) == traj.global_dims[name]
     @assert all(bound .>= 0) "bound must be non-negative when only one bound is provided"
 
@@ -190,7 +204,7 @@ function GlobalBoundsConstraint(
     traj::NamedTrajectory;
     label="bounds constraint on global variable $name"
 )
-    @assert name ∈ keys(traj.global_data)
+    @assert name ∈ traj.global_names
     @assert bound >= 0 "bound must be non-negative when only one bound is provided"
 
     bounds = (-fill(bound, traj.global_dims[name]), fill(bound, traj.global_dims[name]))
@@ -199,22 +213,9 @@ function GlobalBoundsConstraint(
 end
 
 
-
-struct AllEqualConstraint <: AbstractLinearConstraint
-    indices::Vector{Int}
-    bar_index::Int
-    label::String
-end
-
-function TimeStepsAllEqualConstraint(
-    traj::NamedTrajectory;
-    label="timesteps all equal constraint"
-)
-    @assert traj.timestep isa Symbol
-    indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T-1]
-    bar_index = index(traj.T, traj.components[traj.timestep][1], traj.dim)
-    return AllEqualConstraint(indices, bar_index, label)
-end
+### 
+### L1SlackConstraint
+###
 
 # TODO: Doesn't work with parametric trajectory
 # struct L1SlackConstraint <: AbstractLinearConstraint

src/constraints/nonlinear_constraints.jl

@@ -26,7 +26,7 @@ struct NonlinearKnotPointConstraint{F1, F2, F3} <: AbstractNonlinearConstraint
     Create a NonlinearKnotPointConstraint object that represents a nonlinear constraint on a trajectory.
 
     # Arguments
-    - `g::Function`: Function that defines the constraint. If `equality=false`, the constraint is `g(x) ≤ 0`.
+    - `g::Function`: Function over knot point variable(s) that defines the constraint, `g(x)`.
     - `name::Symbol`: Name of the variable to be constrained.
     - `traj::NamedTrajectory`: The trajectory on which the constraint is defined.
 
@@ -52,8 +52,9 @@ struct NonlinearKnotPointConstraint{F1, F2, F3} <: AbstractNonlinearConstraint
         x_slices = [slice(t, x_comps, traj.dim) for t in times]
 
         # inspect view of knot point data
-        @assert g(traj.datavec[x_slices[1]], params[1]) isa AbstractVector{Float64}
-        g_dim = length(g(traj.datavec[x_slices[1]], params[1]))
+        Z⃗ = vec(traj)
+        @assert g(Z⃗[x_slices[1]], params[1]) isa AbstractVector{Float64}
+        g_dim = length(g(Z⃗[x_slices[1]], params[1]))
 
         @views function g!(δ::AbstractVector, Z⃗::AbstractVector)
             for (i, x_slice) ∈ enumerate(x_slices)
@@ -172,7 +173,7 @@ end
 
 # ============================================================================= #
 
-@testitem "testing NonlinearConstraint" begin
+@testitem "testing NonlinearKnotPointConstraint" begin
 
     using TrajectoryIndexingUtils
 
@@ -182,31 +183,33 @@ end
 
     g(a) = [norm(a) - 1.0]
 
+    g_dim = 1
     times = 1:traj.T
 
     NLC = NonlinearKnotPointConstraint(g, :u, traj; times=times, equality=false)
+    U_SLICE(k) = slice(k, traj.components[:u], traj.dim)
 
-    ĝ(Z⃗) = vcat([g(Z⃗[slice(k, traj.components[:u], traj.dim)]) for k ∈ times]...)
+    ĝ(Z⃗) = vcat([g(Z⃗[U_SLICE(k)]) for k ∈ times]...)
 
-    δ = zeros(length(times))
+    δ = zeros(g_dim * traj.T)
 
-    NLC.g!(δ, traj.datavec)
+    NLC.g!(δ, vec(traj))
 
-    @test δ ≈ ĝ(traj.datavec)
+    @test δ ≈ ĝ(vec(traj))
 
-    NLC.∂g!(NLC.∂gs, traj.datavec)
+    NLC.∂g!(NLC.∂gs, vec(traj))
 
     ∂g_full = Constraints.get_full_jacobian(NLC, traj)
 
-    ∂g_autodiff = ForwardDiff.jacobian(ĝ, traj.datavec)
+    ∂g_autodiff = ForwardDiff.jacobian(ĝ, vec(traj))
 
     @test ∂g_full ≈ ∂g_autodiff
 
-    μ = randn(length(times))
+    μ = randn(g_dim * traj.T)
 
-    NLC.μ∂²g!(NLC.μ∂²gs, traj.datavec, μ)
+    NLC.μ∂²g!(NLC.μ∂²gs, vec(traj), μ)
 
-    hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), traj.datavec)
+    hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), vec(traj))
 
     μ∂²g_full = Constraints.get_full_hessian(NLC, traj)