Merge pull request #25 from RoboticExplorationLab/bilinear

More fixes and tweaks

Project.toml

@@ -1,7 +1,7 @@
 name = "RobotDynamics"
 uuid = "38ceca67-d8d3-44e8-9852-78a5596522e1"
 authors = ["Brian Jackson <bjack205@gmail.com>"]
-version = "0.4.5"
+version = "0.4.6"
 
 [deps]
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"

src/discretized_dynamics.jl

@@ -188,6 +188,7 @@ for method in (:state_dim, :control_dim, :output_dim, :errstate_dim, :statevecto
     @eval $method(model::DiscretizedDynamics) = $method(model.continuous_dynamics)
 end
 default_diffmethod(::DiscretizedDynamics) = ForwardAD()
+default_signature(model::DiscretizedDynamics) = default_signature(model.continuous_dynamics)
 
 # Shallow copy
 function Base.copy(model::DiscretizedDynamics)
@@ -229,8 +230,11 @@ const ImplicitDynamicsModel{L,Q} = DiscretizedDynamics{L,Q} where {L,Q<:Implicit
 
 discrete_dynamics(model::ImplicitDynamicsModel, z::AbstractKnotPoint) =
     integrate(integration(model), model, z)
-discrete_dynamics!(model::ImplicitDynamicsModel, xn, z::AbstractKnotPoint) =
-    integrate!(integration(model), model, xn, z)
+# discrete_dynamics!(model::ImplicitDynamicsModel, xn, z::AbstractKnotPoint) =
+#     integrate!(integration(model), model, xn, z)
+
+discrete_dynamics!(model::ImplicitDynamicsModel, xn, x, u, t, dt) =
+    integrate!(integration(model), model, xn, x, u, t, dt)
 
 function jacobian!(sig::FunctionSignature, ::ImplicitFunctionTheorem{D}, model::ImplicitDynamicsModel, J, xn, z) where D
     jacobian!(integration(model), sig, D(), model, J, xn, z)

src/integration.jl

@@ -431,9 +431,11 @@ function integrate(integrator::Implicit, model::ImplicitDynamicsModel,
     y1 = cache.y1
     z1 = z
     z2 = StaticKnotPoint(z)
-    xn = state(z2)
     ix = SVector{Nx}(1:Nx) 
 
+    # Use the current state as the current guess
+    xn = state(z1)
+
     diff = default_diffmethod(model.continuous_dynamics)
     for iter = 1:newton_iters
         # Set the guess for the next state
@@ -462,22 +464,35 @@ end
 
 function integrate!(integrator::Implicit, model::ImplicitDynamicsModel, xn, 
                     z::AbstractKnotPoint)
+    integrate!(integrator, model, xn, state(z), control(z), z.t, z.dt)
+end
+function integrate!(integrator::Implicit, model::ImplicitDynamicsModel, xn, 
+                    x, u, t, h)
     cache = getnewtoncache(integrator) 
     newton_iters = cache.newton_iters
     tol = cache.newton_tol
 
-    z1 = z
-    n,m = dims(z)
+    z1 = cache.z1
+    z2 = cache.z2
+    # copyto!(z1, z)
+    setstate!(z1, x)
+    setcontrol!(z1, u)
+    z1.t = t
+    z1.dt = h
+    copyto!(z2, z1)
+
+    n,m = dims(model)
     J2 = cache.J2
     J1 = cache.J1
     r = cache.y2
     dx = cache.y1
-    z2 = cache.z2
     ipiv = cache.ipiv
     A = cache.A
     Aview = @view J2[:, 1:n]
 
-    copyto!(z2, z1)
+    # Use the current state as the guess
+    copyto!(xn, state(z1))
+
     diff = default_diffmethod(model.continuous_dynamics)
     for iter = 1:newton_iters
         # Set the guess for the next state
@@ -502,7 +517,7 @@ function integrate!(integrator::Implicit, model::ImplicitDynamicsModel, xn,
         # Apply the step
         xn .-= dx
     end
-    setdata!(cache.z2, getdata(z))  # copy input to cache for Jacobian check
+    setdata!(cache.z2, getdata(z1))  # copy input to cache for Jacobian check
 end
 
 # Use Implicit Function Theorem to calculate the dynamics Jacobians
@@ -569,6 +584,7 @@ mutable struct ImplicitNewtonCache
     y2::Vector{Float64}
     y1::Vector{Float64}
     z2::StaticKnotPoint{Any,Any,Vector{Float64},Float64}
+    z1::KnotPoint{Any,Any,Vector{Float64},Float64}
     ipiv::Vector{BlasInt}
     A::Matrix{Float64}
     F::LinearAlgebra.LU{Float64, Matrix{Float64}} 
@@ -582,12 +598,13 @@ function ImplicitNewtonCache(n::Integer, m::Integer)
     y1 = zeros(n)
     v = zeros(n+m)
     z2 = StaticKnotPoint{Any,Any}(n, m, v, 0.0, NaN)
+    z1 = KnotPoint{Any,Any}(n, m, copy(v), 0.0, NaN)
     ipiv = zeros(BlasInt, n)
     A = zeros(n,n) 
     F = lu!(A, check=false)
     iters = 10    # Default number of Newton iterations
     tol = 1e-12   # Default Newton tolerance
-    ImplicitNewtonCache(J2, J1, y2, y1, z2, ipiv, A, F, iters, tol)
+    ImplicitNewtonCache(J2, J1, y2, y1, z2, z1, ipiv, A, F, iters, tol)
 end
 
 """

src/knotpoint.jl

@@ -195,8 +195,10 @@ for (name,mutable) in [(:KnotPoint, true), (:StaticKnotPoint, false)]
         @inline getparams(z::$name) = (t=z.t, dt=z.dt)
         @inline getdata(z::$name) = z.z
         function Base.copyto!(dest::$name, src::$name)
-            dest.t = src.t
-            dest.dt = src.dt
+            if ismutabletype($name)
+                dest.t = src.t
+                dest.dt = src.dt
+            end
             copyto!(dest.z, src.z)
             dest
         end

src/trajectories.jl

@@ -185,7 +185,7 @@ function controls(Z::SampledTrajectory)
     return [control(Z[k]) for k in eachcontrol(Z) ]
 end
 controls(Z::SampledTrajectory, inds::AbstractVector{<:Integer}) = [controls(Z, i) for i in inds]
-controls(Z::SampledTrajectory, ind::Integer) = [control(z)[ind] for z in Z] 
+controls(Z::SampledTrajectory, ind::Integer) = [control(Z[k])[ind] for k in eachcontrol(Z)] 
 
 """
     gettimes(Z)

test/implicit_dynamics_test.jl

@@ -72,6 +72,7 @@ x2 = newton_solve(RD.getdata(z1))
 x2_inplace = copy(x1)
 RD.discrete_dynamics!(dmodel, x2_inplace, z1)
 @test norm(midpoint_residual(x2_inplace, z1)) < √eps() 
+x2_inplace .= 0
 RD.evaluate!(dmodel, x2_inplace, z1)
 @test norm(midpoint_residual(x2_inplace, z1)) < √eps()