WIP: using Cache instead of DiffCache

Author: franckgaga

src/ModelPredictiveControl.jl

@@ -9,6 +9,7 @@ using ProgressLogging
 
 using DifferentiationInterface: ADTypes.AbstractADType, AutoForwardDiff, AutoSparse
 using DifferentiationInterface: gradient!, jacobian!, prepare_gradient, prepare_jacobian
+using DifferentiationInterface: Constant, Cache
 using SparseConnectivityTracer: TracerSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm
 

src/controller/nonlinmpc.jl

@@ -582,43 +582,14 @@ function get_optim_functions(
     jac_backend ::AbstractADType
 ) where JNT<:Real
     model, transcription = mpc.estim.model, mpc.transcription
-    nu, ny, nx̂, nϵ, Hp, Hc = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ, mpc.Hp, mpc.Hc
-    ng, nc, neq = length(mpc.con.i_g), mpc.con.nc, mpc.con.neq
-    nZ̃, nU, nŶ, nX̂ = length(mpc.Z̃), Hp*nu, Hp*ny, Hp*nx̂
-    nΔŨ, nUe, nŶe = nu*Hc + nϵ, nU + nu, nŶ + ny
-    Ncache = nZ̃ + 3 
-    myNaN = convert(JNT, NaN) # fill Z̃ with NaNs to force update_simulations! at 1st call:
-    # ---------------------- differentiation cache ---------------------------------------
-    Z̃_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(fill(myNaN, nZ̃), Ncache)
-    ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nΔŨ), Ncache)
-    x̂0end_cache::DiffCache{Vector{JNT}, Vector{JNT}}  = DiffCache(zeros(JNT, nx̂),  Ncache)
-    Ŷe_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶe), Ncache)
-    Ue_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nUe), Ncache)
-    Ŷ0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ),  Ncache)
-    U0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU),  Ncache)
-    Û0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU),  Ncache)
-    X̂0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nX̂),  Ncache)
-    gc_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nc),  Ncache)
-    g_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, ng),  Ncache)
-    geq_cache::DiffCache{Vector{JNT}, Vector{JNT}}    = DiffCache(zeros(JNT, neq), Ncache)
     # --------------------- update simulation function ------------------------------------
-    function update_simulations!(
-        Z̃arg::Union{NTuple{N, T}, AbstractVector{T}}, Z̃cache
-    ) where {N, T<:Real}
-        if isdifferent(Z̃cache, Z̃arg)
-            for i in eachindex(Z̃cache)
+    function update_simulations!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
+        if isdifferent(Z̃arg, Z̃)
+            for i in eachindex(Z̃)
                 # Z̃cache .= Z̃arg is type unstable with Z̃arg::NTuple{N, FowardDiff.Dual}
-                Z̃cache[i] = Z̃arg[i]
+                Z̃[i] = Z̃arg[i]
             end
-            Z̃ = Z̃cache
             ϵ = (nϵ ≠ 0) ? Z̃[end] : zero(T) # ϵ = 0 if nϵ == 0 (meaning no relaxation)
-            ΔŨ     = get_tmp(ΔŨ_cache, T)
-            x̂0end  = get_tmp(x̂0end_cache, T)
-            Ue, Ŷe = get_tmp(Ue_cache, T), get_tmp(Ŷe_cache, T)
-            U0, Ŷ0 = get_tmp(U0_cache, T), get_tmp(Ŷ0_cache, T)
-            X̂0, Û0 = get_tmp(X̂0_cache, T), get_tmp(Û0_cache, T) 
-            gc, g  = get_tmp(gc_cache, T), get_tmp(g_cache, T)
-            geq    = get_tmp(geq_cache, T)
             U0 = getU0!(U0, mpc, Z̃)
             ΔŨ = getΔŨ!(ΔŨ, mpc, transcription, Z̃)
             Ŷ0, x̂0end  = predict!(Ŷ0, x̂0end, X̂0, Û0, mpc, model, transcription, U0, Z̃)
@@ -629,78 +600,151 @@ function get_optim_functions(
         end
         return nothing
     end
-    # --------------------- objective functions -------------------------------------------
+    # ---------------------- JNT vectors cache --------------------------------------------
+    nu, ny, nx̂, nϵ, Hp, Hc = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ, mpc.Hp, mpc.Hc
+    ng, nc, neq = length(mpc.con.i_g), mpc.con.nc, mpc.con.neq
+    nZ̃, nU, nŶ, nX̂ = length(mpc.Z̃), Hp*nu, Hp*ny, Hp*nx̂
+    nΔŨ, nUe, nŶe = nu*Hc + nϵ, nU + nu, nŶ + ny  
+    myNaN = convert(JNT, NaN)
+    Z̃      = fill(myNaN, nZ̃)
+    ΔŨ     = zeros(JNT, nΔŨ)
+    x̂0end  = zeros(JNT, nx̂)
+    Ue, Ŷe = zeros(JNT, nUe), zeros(JNT, nŶe)
+    U0, Ŷ0 = zeros(JNT, nU),  zeros(JNT, nŶ)
+    Û0, X̂0 = zeros(JNT, nU),  zeros(JNT, nX̂)
+    gc, g  = zeros(JNT, nc),  zeros(JNT, ng)
+    geq    = zeros(JNT, neq)
+
+    # WIP: still receive INVALID_MODEL once and a while, needs to investigate.
+
+
     function Jfunc(Z̃arg::Vararg{T, N}) where {N, T<:Real}
-        update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
-        ΔŨ = get_tmp(ΔŨ_cache, T)
-        Ue, Ŷe = get_tmp(Ue_cache, T), get_tmp(Ŷe_cache, T)
-        U0, Ŷ0 = get_tmp(U0_cache, T), get_tmp(Ŷ0_cache, T)
+        update_simulations!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
         return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)::T
     end
-    function Jfunc_vec(Z̃arg::AbstractVector{T}) where T<:Real 
-        update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
-        ΔŨ = get_tmp(ΔŨ_cache, T)
-        Ue, Ŷe = get_tmp(Ue_cache, T), get_tmp(Ŷe_cache, T)
-        U0, Ŷ0 = get_tmp(U0_cache, T), get_tmp(Ŷ0_cache, T)
-        return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)::T
+    function Jfunc_vec!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
+        update_simulations!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
+        return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)
     end
+
+
+
+    gfuncs = Vector{Function}(undef, ng)
+    for i in eachindex(gfuncs)
+        func_i = function (Z̃arg::Vararg{T, N}) where {N, T<:Real}
+            update_simulations!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
+            return g[i]::T
+        end
+        gfuncs[i] = func_i
+    end
+    function gfunc_vec!(g, Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, geq)
+        update_simulations!(Z̃arg, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, X̂0, gc, g, geq)
+        return g
+    end
+
+
+
+
+
+
+
+
+    Z̃_∇J      = fill(myNaN, nZ̃)
+    ΔŨ_∇J     = zeros(JNT, nΔŨ)
+    x̂0end_∇J  = zeros(JNT, nx̂)
+    Ue_∇J, Ŷe_∇J = zeros(JNT, nUe), zeros(JNT, nŶe)
+    U0_∇J, Ŷ0_∇J = zeros(JNT, nU),  zeros(JNT, nŶ)
+    Û0_∇J, X̂0_∇J = zeros(JNT, nU),  zeros(JNT, nX̂)
+    gc_∇J, g_∇J  = zeros(JNT, nc),  zeros(JNT, ng)
+    geq_∇J    = zeros(JNT, neq)
+
+
+
+
+
+
+
     Z̃_∇J    = fill(myNaN, nZ̃) 
     ∇J      = Vector{JNT}(undef, nZ̃)       # gradient of objective J
-    ∇J_prep = prepare_gradient(Jfunc_vec, grad_backend, Z̃_∇J)  
+    ∇J_context = (
+        Cache(Z̃_∇J),  Cache(ΔŨ_∇J), Cache(x̂0end_∇J), 
+        Cache(Ue_∇J), Cache(Ŷe_∇J), 
+        Cache(U0_∇J), Cache(Ŷ0_∇J), 
+        Cache(Û0_∇J), Cache(X̂0_∇J), 
+        Cache(gc_∇J), Cache(g_∇J), Cache(geq_∇J)
+    )
+    ∇J_prep = prepare_gradient(Jfunc_vec!, grad_backend, Z̃_∇J, ∇J_context...)
     ∇Jfunc! = if nZ̃ == 1
-        function (Z̃arg::T) where T<:Real 
+        function (Z̃arg)
             Z̃_∇J .= Z̃arg
-            gradient!(Jfunc_vec, ∇J, ∇J_prep, grad_backend, Z̃_∇J)
+            gradient!(Jfunc_vec!, ∇J, ∇J_prep, grad_backend, Z̃_∇J, ∇J_context...)
             return ∇J[begin]    # univariate syntax, see JuMP.@operator doc
         end
     else
         function (∇J::AbstractVector{T}, Z̃arg::Vararg{T, N}) where {N, T<:Real}
             Z̃_∇J .= Z̃arg
-            gradient!(Jfunc_vec, ∇J, ∇J_prep, grad_backend, Z̃_∇J)
+            gradient!(Jfunc_vec!, ∇J, ∇J_prep, grad_backend, Z̃_∇J, ∇J_context...)
             return ∇J           # multivariate syntax, see JuMP.@operator doc
         end
     end
-    # --------------------- inequality constraint functions -------------------------------
-    gfuncs = Vector{Function}(undef, ng)
-    for i in eachindex(gfuncs)
-        func_i = function (Z̃arg::Vararg{T, N}) where {N, T<:Real}
-            update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
-            g = get_tmp(g_cache, T)
-            return g[i]::T
-        end
-        gfuncs[i] = func_i
-    end
-    function gfunc_vec!(g, Z̃vec::AbstractVector{T}) where T<:Real
-        update_simulations!(Z̃vec, get_tmp(Z̃_cache, T))
-        g .= get_tmp(g_cache, T)
-        return g
-    end
+
+
+
+
+
+
+
+    Z̃_∇g      = fill(myNaN, nZ̃)
+    ΔŨ_∇g     = zeros(JNT, nΔŨ)
+    x̂0end_∇g  = zeros(JNT, nx̂)
+    Ue_∇g, Ŷe_∇g = zeros(JNT, nUe), zeros(JNT, nŶe)
+    U0_∇g, Ŷ0_∇g = zeros(JNT, nU),  zeros(JNT, nŶ)
+    Û0_∇g, X̂0_∇g = zeros(JNT, nU),  zeros(JNT, nX̂)
+    gc_∇g, g_∇g  = zeros(JNT, nc),  zeros(JNT, ng)
+    geq_∇g    = zeros(JNT, neq)
+
+
+
+
+
     Z̃_∇g     = fill(myNaN, nZ̃)
     g_vec    = Vector{JNT}(undef, ng)
     ∇g       = Matrix{JNT}(undef, ng, nZ̃)   # Jacobian of inequality constraints g
-    ∇g_prep  = prepare_jacobian(gfunc_vec!, g_vec, jac_backend, Z̃_∇g)
+    ∇g_context = (
+        Cache(Z̃_∇g),  Cache(ΔŨ_∇g), Cache(x̂0end_∇g), 
+        Cache(Ue_∇g), Cache(Ŷe_∇g), 
+        Cache(U0_∇g), Cache(Ŷ0_∇g), 
+        Cache(Û0_∇g), Cache(X̂0_∇g), 
+        Cache(gc_∇g), Cache(geq_∇g)
+    )
+    ∇g_prep  = prepare_jacobian(gfunc_vec!, g_vec, jac_backend, Z̃_∇g, ∇g_context...)
     ∇gfuncs! = Vector{Function}(undef, ng)
     for i in eachindex(∇gfuncs!)
         ∇gfuncs![i] = if nZ̃ == 1
             function (Z̃arg::T) where T<:Real
                 if isdifferent(Z̃arg, Z̃_∇g)
                     Z̃_∇g .= Z̃arg
-                    jacobian!(gfunc_vec!, g_vec, ∇g, ∇g_prep, jac_backend, Z̃_∇g)
+                    jacobian!(
+                        gfunc_vec!, g_vec, ∇g, ∇g_prep, jac_backend, Z̃_∇g, ∇g_context...
+                    )
                 end
                 return ∇g[i, begin]            # univariate syntax, see JuMP.@operator doc
             end
         else
             function (∇g_i, Z̃arg::Vararg{T, N}) where {N, T<:Real}
                 if isdifferent(Z̃arg, Z̃_∇g)
                     Z̃_∇g .= Z̃arg
-                    jacobian!(gfunc_vec!, g_vec, ∇g, ∇g_prep, jac_backend, Z̃_∇g)
+                    jacobian!(
+                        gfunc_vec!, g_vec, ∇g, ∇g_prep, jac_backend, Z̃_∇g, ∇g_context...
+                    )
                 end
                 return ∇g_i .= @views ∇g[i, :] # multivariate syntax, see JuMP.@operator doc
             end
         end
     end
     # --------------------- equality constraint functions ---------------------------------
     geqfuncs = Vector{Function}(undef, neq)
+    #=
     for i in eachindex(geqfuncs)
         func_i = function (Z̃arg::Vararg{T, N}) where {N, T<:Real}
             update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
@@ -718,7 +762,9 @@ function get_optim_functions(
     geq_vec    = Vector{JNT}(undef, neq)
     ∇geq       = Matrix{JNT}(undef, neq, nZ̃)   # Jacobian of equality constraints geq
     ∇geq_prep  = prepare_jacobian(geqfunc_vec!, geq_vec, jac_backend, Z̃_∇geq)
+    =#
     ∇geqfuncs! = Vector{Function}(undef, neq)
+    #=
     for i in eachindex(∇geqfuncs!)
         # only multivariate syntax, univariate is impossible since nonlinear equality
         # constraints imply MultipleShooting, thus input increment ΔU and state X̂0 in Z̃:
@@ -731,6 +777,7 @@ function get_optim_functions(
                 return ∇geq_i .= @views ∇geq[i, :]
             end
     end
+    =#
     return Jfunc, ∇Jfunc!, gfuncs, ∇gfuncs!, geqfuncs, ∇geqfuncs!
 end