Add problems tests

Author: bjack205

src/problem.jl

@@ -62,19 +62,20 @@ function Problem(model::L, obj::O, xf::AbstractVector, tf;
         x0=zero(xf), N::Int=length(obj),
         X0=[x0*NaN for k = 1:N],
         U0=[@SVector zeros(size(model)[2]) for k = 1:N-1],
-        dt=fill((tf-t0)/(N-1),N),
+        dt=fill((tf-t0)/(N-1),N-1),
         integration=DEFAULT_Q) where {L,O}
     n,m = size(model)
     if dt isa Real
         dt = fill(dt,N)
     end
+	@assert sum(dt[1:N-1]) ≈ tf "Time steps are inconsistent with final time"
     if X0 isa AbstractMatrix
         X0 = [X0[:,k] for k = 1:size(X0,2)]
     end
     if U0 isa AbstractMatrix
         U0 = [U0[:,k] for k = 1:size(U0,2)]
     end
-    t = range(t0, tf, length=N)
+    t = pushfirst!(cumsum(dt), 0)
     Z = Traj(X0,U0,dt,t)
 
     Problem{integration}(model, obj, constraints, SVector{n}(x0), SVector{n}(xf),
@@ -112,6 +113,13 @@ Get the state trajectory
 "
 states(prob::Problem) = states(prob.Z)
 
+"""
+	get_times(::Problem)
+
+Get the times for all the knot points in the problem.
+"""
+@inline get_times(prob::Problem) = get_times(get_trajectory(prob))
+
 "```julia
 initial_trajectory!(::Problem, Z)
 initial_trajectory!(::AbstractSolver, Z)

src/trajectories.jl

@@ -61,7 +61,7 @@ end
 function Traj(n::Int, m::Int, dt::AbstractFloat, N::Int; equal=false)
     x = NaN*@SVector ones(n)
     u = @SVector zeros(m)
-    Traj(x,u,dt,N,equal)
+    Traj(x,u,dt,N; equal=equal)
 end
 
 function Traj(x::SVector, u::SVector, dt::AbstractFloat, N::Int; equal=false)

test/problems_tests.jl

@@ -0,0 +1,121 @@
+# Model and discretization
+model = Cartpole()
+n,m = size(model)
+N = 11
+tf = 5.
+dt = tf/(N-1)
+
+# Initial and Final conditions
+x0 = @SVector zeros(n)
+xf = @SVector [0, pi, 0, 0]
+
+# Objective
+Q = 1.0e-2*Diagonal(@SVector ones(n))
+Qf = 100.0*Diagonal(@SVector ones(n))
+R = 1.0e-1*Diagonal(@SVector ones(m))
+obj = LQRObjective(Q,R,Qf,xf,N)
+
+# Constraints
+u_bnd = 3.0
+conSet = ConstraintList(n,m,N)
+bnd = BoundConstraint(n,m, u_min=-u_bnd, u_max=u_bnd)
+goal = GoalConstraint(xf)
+add_constraint!(conSet, bnd, 1:N-1)
+add_constraint!(conSet, goal, N:N)
+
+# Initial conditions
+X0 = [@SVector fill(0.0,n) for k = 1:N]
+u0 = @SVector fill(0.01,m)
+U0 = [u0 for k = 1:N-1]
+Z = Traj(X0,U0, fill(dt, N))
+
+# Inner constructor
+prob = Problem{RK3}(model, obj, conSet, x0, xf, Z, N, 0.0, tf)
+@test TO.integration(prob) == RK3
+@test prob.x0 == x0
+@test prob.xf == xf
+@test prob.constraints === conSet
+add_constraint!(conSet, goal, N-1)
+@test length(TO.get_constraints(prob)) == length(conSet)
+@test TO.num_constraints(TO.get_constraints(prob)) === conSet.p
+@test prob.obj === obj
+@test prob.tf ≈ tf
+@test prob.N == N
+@test states(prob) ≈ X0
+@test controls(prob) ≈ U0
+
+# Alternate constructor
+prob = Problem(model, obj, xf, tf, x0=x0, constraints=conSet, X0=X0, U0=U0)
+@test TO.integration(prob) == RK3
+@test prob.x0 == x0
+@test prob.xf == xf
+@test prob.constraints === conSet
+add_constraint!(conSet, goal, N-1)
+@test length(TO.get_constraints(prob)) == length(conSet)
+@test TO.num_constraints(TO.get_constraints(prob)) === conSet.p
+@test prob.obj === obj
+@test prob.tf ≈ tf
+@test prob.N == N
+@test states(prob) ≈ X0
+@test controls(prob) ≈ U0
+
+# Change integration
+prob = Problem(model, obj, xf, tf, x0=x0, constraints=conSet, integration=RK2)
+@test TO.integration(prob) == RK2
+
+# Test defaults
+prob = Problem(model, obj, xf, tf)
+@test prob.x0 == zero(x0)
+@test prob.N == N
+@test all(all.(isnan, states(prob)))
+@test controls(prob) == [zeros(m) for k = 1:N-1]
+@test isempty(prob.constraints)
+@test TO.integration(prob) == RK3
+@test TO.get_times(prob) ≈ range(0, tf; step=dt)
+
+# Set initial trajectories
+initial_states!(prob, 2 .* X0)
+@test states(prob) ≈ 2 .* X0
+initial_controls!(prob, 2 .* U0)
+@test controls(prob) ≈ 2 .* U0
+
+initial_trajectory!(prob, Z)
+@test controls(prob) ≈ U0
+@test states(prob) ≈ X0
+
+# Use 2D matrices
+X0_mat = hcat(X0...)
+U0_mat = hcat(U0...)
+initial_states!(prob, 3*X0)
+initial_controls!(prob, 4*U0)
+@test states(prob) ≈ 3 .* X0
+@test controls(prob) ≈ 4 .* U0
+
+prob = Problem(model, obj, xf, tf, X0=X0_mat, U0=U0_mat)
+@test states(prob) ≈ X0
+@test controls(prob) ≈ U0
+
+# Use variable time steps
+dts = rand(N-1)
+dts = dts / sum(dts) * tf
+prob = Problem(model, obj, xf, tf, dt=dts)
+times = TO.get_times(prob)
+@test times[end] ≈ tf
+@test times[1] ≈ 0
+@test times[2] ≈ dts[1]
+@test diff(times) ≈ dts
+
+# Test initial and final conditions
+prob = Problem(model, obj, Vector(xf), tf, x0=Vector(x0))
+@test prob.x0 ≈ x0
+@test prob.xf ≈ xf
+prob = Problem(model, obj, MVector(xf), tf, x0=MVector(x0))
+@test prob.x0 ≈ x0
+@test prob.xf ≈ xf
+@test prob.x0 isa MVector
+@test prob.xf isa MVector
+
+x0_ = rand(n)
+TO.set_initial_state!(prob, x0_)
+@test prob.x0 ≈ x0_
+@test_throws DimensionMismatch TO.set_initial_state!(prob, rand(2n))

test/runtests.jl

@@ -22,6 +22,15 @@ end
     include("constraint_sets.jl")
 end
 
+@testset "Problems" begin
+    include("problems_tests.jl")
+end
+
 @testset "Utils" begin
     include("trajectories.jl")
 end
+
+@testset "NLP" begin
+    include("nlp_tests.jl")
+    include("moi_test.jl")
+end