update for new NT

src/constraints/linear_constraints.jl

@@ -216,36 +216,37 @@ function TimeStepsAllEqualConstraint(
     return AllEqualConstraint(indices, bar_index, label)
 end
 
-struct L1SlackConstraint <: AbstractLinearConstraint
-    x_indices::Vector{Int}
-    s1_indices::Vector{Int}
-    s2_indices::Vector{Int}
-    label::String
-end
-
-function L1SlackConstraint(
-    name::Symbol,
-    traj::NamedTrajectory;
-    indices=1:traj.dims[name],
-    ts=(name ∈ keys(traj.initial) ? 2 : 1):(name ∈ keys(traj.final) ? traj.T-1 : traj.T),
-    label="L1 slack constraint on $name[$(indices)]"
-)
-    @assert all(i ∈ 1:traj.dims[name] for i ∈ indices)
-
-    s1_name = Symbol("s1_$name")
-    s2_name = Symbol("s2_$name")
-
-    add_component!(traj, s1_name, rand(length(indices), traj.T))
-    add_component!(traj, s2_name, rand(length(indices), traj.T))
-
-    x_indices = stack(slice(t, traj.components[name][indices], traj.dim) for t ∈ ts)
-    s1_indices = stack(slice(t, traj.components[s1_name], traj.dim) for t ∈ ts)
-    s2_indices = stack(slice(t, traj.components[s2_name], traj.dim) for t ∈ ts)
-
-    return L1SlackConstraint(
-        x_indices,
-        s1_indices,
-        s2_indices,
-        label
-    )
-end
+# TODO: Doesn't work with parametric trajectory
+# struct L1SlackConstraint <: AbstractLinearConstraint
+#     x_indices::Vector{Int}
+#     s1_indices::Vector{Int}
+#     s2_indices::Vector{Int}
+#     label::String
+# end
+
+# function L1SlackConstraint(
+#     name::Symbol,
+#     traj::NamedTrajectory;
+#     indices=1:traj.dims[name],
+#     ts=(name ∈ keys(traj.initial) ? 2 : 1):(name ∈ keys(traj.final) ? traj.T-1 : traj.T),
+#     label="L1 slack constraint on $name[$(indices)]"
+# )
+#     @assert all(i ∈ 1:traj.dims[name] for i ∈ indices)
+
+#     s1_name = Symbol("s1_$name")
+#     s2_name = Symbol("s2_$name")
+
+#     add_component!(traj, s1_name, rand(length(indices), traj.T))
+#     add_component!(traj, s2_name, rand(length(indices), traj.T))
+
+#     x_indices = stack(slice(t, traj.components[name][indices], traj.dim) for t ∈ ts)
+#     s1_indices = stack(slice(t, traj.components[s1_name], traj.dim) for t ∈ ts)
+#     s2_indices = stack(slice(t, traj.components[s2_name], traj.dim) for t ∈ ts)
+
+#     return L1SlackConstraint(
+#         x_indices,
+#         s1_indices,
+#         s2_indices,
+#         label
+#     )
+# end

src/constraints/nonlinear_global_constraints.jl

@@ -22,44 +22,44 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint
         equality::Bool=true,
     )
         global_comps = vcat([traj.global_components[name] for name in global_names]...)
-        local_comps = global_comps .- traj.dim * traj.T
+        offset_global_comps = global_comps .+ traj.dim * traj.T
 
-        g_eval = g(vec(traj)[global_comps])
+        g_eval = g(vec(traj)[offset_global_comps])
         @assert g_eval isa AbstractVector{Float64}
         g_dim = length(g_eval)
 
         @views function g!(δ::AbstractVector, Z⃗::AbstractVector)
-            δ[:] = g(Z⃗[global_comps])
+            δ[:] = g(Z⃗[offset_global_comps])
             return nothing
         end
 
         @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector)
             ForwardDiff.jacobian!(
                 ∂g, 
                 x -> g(x),
-                Z⃗[global_comps]
+                Z⃗[offset_global_comps]
             )
         end
 
         # global subspace
         jacobian_structure = spzeros(g_dim, traj.global_dim) 
-        jacobian_structure[:, local_comps] .= 1.0
+        jacobian_structure[:, global_comps] .= 1.0
 
         @views function μ∂²g!(
             μ∂²g::AbstractMatrix,   
             Z⃗::AbstractVector, 
             μ::AbstractVector
         )
             ForwardDiff.hessian!(
-                μ∂²g[local_comps, local_comps], 
+                μ∂²g[global_comps, global_comps], 
                 x -> μ'g(x), 
-                Z⃗[global_comps]
+                Z⃗[offset_global_comps]
             )
         end
 
         # global subspace
         hessian_structure = spzeros(traj.global_dim, traj.global_dim) 
-        hessian_structure[local_comps, local_comps] .= 1.0
+        hessian_structure[global_comps, global_comps] .= 1.0
 
         return new(
             g!,
@@ -136,14 +136,14 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint
 
         x_comps = vcat([traj.components[name] for name in names]...)
         global_comps = vcat([traj.global_components[name] for name in global_names]...)
-        local_comps = global_comps .- traj.dim * traj.T
+        offset_global_comps = global_comps .+ traj.dim * traj.T
 
-        # (Rebased) global data is appended to the knot point
-        xg_comps = vcat([x_comps, traj.dim .+ local_comps]...)
+        # append global data to the knot point
+        xg_comps = vcat([x_comps, global_comps .+ traj.dim]...)
         z_dim = traj.dim + traj.global_dim
 
-        # Each slice indexes into Z⃗
-        xg_slices = [vcat([slice(t, x_comps, traj.dim), global_comps]...) for t in times]
+        # append global data to the trajectory (each slice indexes into Z⃗)
+        xg_slices = [vcat([slice(t, x_comps, traj.dim), offset_global_comps]...) for t in times]
 
         Z⃗ = vec(traj)
         @assert g(Z⃗[xg_slices[1]], params[1]) isa AbstractVector{Float64}

src/integrators/time_dependent_bilinear_integrator.jl

@@ -151,8 +151,7 @@ end
 @testitem "testing zoh TimeDependentBilinearIntegrator" begin
     include("../../test/test_utils.jl")
 
-    G, traj = bilinear_dynamics_and_trajectory()
-    add_component!(traj, :t, get_times(traj))
+    G, traj = bilinear_dynamics_and_trajectory(time=true)
 
     # zero order hold
     B = TimeDependentBilinearIntegrator((a, t) -> G(a), traj, :x, :u, :t)

src/objectives/global_objectives.jl

@@ -29,7 +29,7 @@ function GlobalObjective(
     Q::Float64=1.0
 )
     Z_dim = traj.dim * traj.T + traj.global_dim
-    g_comps = vcat([traj.global_components[name] for name in global_names]...)
+    g_comps = vcat([traj.dim * traj.T .+ traj.global_components[name] for name in global_names]...)
 
     L(Z⃗::AbstractVector{<:Real}) = Q * ℓ(Z⃗[g_comps])
 
@@ -76,7 +76,7 @@ function GlobalKnotPointObjective(
 
     Z_dim = traj.dim * traj.T + traj.global_dim
     x_comps = vcat([traj.components[name] for name in names]...)
-    g_comps = vcat([traj.global_components[name] for name in global_names]...)
+    g_comps = vcat([traj.dim * traj.T .+ traj.global_components[name] for name in global_names]...)
 
     xg_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times]
 

test/test_utils.jl

@@ -125,6 +125,7 @@ function bilinear_dynamics_and_trajectory(;
     Δt = 0.1,
     u_bound = 0.1,
     ω = 0.1,
+    time=false
 )
     Gx = sparse(Float64[
         0  0 0 1;
@@ -182,5 +183,9 @@ function bilinear_dynamics_and_trajectory(;
         )
     )
 
+    if time
+        traj = add_component(traj, :t, get_times(traj))
+    end
+
     return G, traj
 end