added: number of colors in getinfo for MPC

Author: franckgaga

src/ModelPredictiveControl.jl

@@ -9,6 +9,7 @@ using RecipesBase
 
 using DifferentiationInterface: ADTypes.AbstractADType, AutoForwardDiff
 using DifferentiationInterface: AutoSparse, SecondOrder
+using DifferentiationInterface: Prep, SparseJacobianPrep, SparseHessianPrep
 using DifferentiationInterface: gradient, jacobian, hessian
 using DifferentiationInterface: value_and_gradient, value_and_jacobian
 using DifferentiationInterface: value_gradient_and_hessian
@@ -20,6 +21,7 @@ using SparseConnectivityTracer: TracerSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
 using SparseMatrixColorings: NaturalOrder, LargestFirst, SmallestLast
 using SparseMatrixColorings: IncidenceDegree, DynamicLargestFirst, RandomOrder
+using SparseMatrixColorings: ncolors
 
 import ProgressLogging
 

src/controller/execute.jl

@@ -117,12 +117,17 @@ Lastly, the following fields are also available for [`NonLinMPC`](@ref) only:
 - `:gc`: custom nonlinear constraints values at the optimum, ``\mathbf{g_c}``
 - `:∇J` or *`:nablaJ`* : optimal gradient of the objective function, ``\mathbf{\nabla} J``
 - `:∇²J` or *`:nabla2J`* : optimal Hessian of the objective function, ``\mathbf{\nabla^2}J``
+- `:∇²J_ncolors` or *`:nabla2J_ncolors`* : number of colors in `:∇²J` sparsity pattern
 - `:g` : optimal nonlinear inequality constraint values, ``\mathbf{g}``
 - `:∇g` or *`:nablag`* : optimal Jacobian of the inequality constraint, ``\mathbf{\nabla g}``
+- `:∇g_ncolors` or *`:nablag_ncolors`* : number of colors in `:∇g` sparsity pattern
 - `:∇²ℓg` or *`:nabla2lg`* : optimal Hessian of the inequality Lagrangian, ``\mathbf{\nabla^2}\ell_{\mathbf{g}}``
+- `:∇²ℓg_ncolors` or *`:nabla2lg_ncolors`* : number of colors in `:∇²ℓg` sparsity pattern
 - `:geq` : optimal nonlinear equality constraint values, ``\mathbf{g_{eq}}``
 - `:∇geq` or *`:nablageq`* : optimal Jacobian of the equality constraint, ``\mathbf{\nabla g_{eq}}``
+- `:∇geq_ncolors` or *`:nablageq_ncolors`* : number of colors in `:∇geq` sparsity pattern
 - `:∇²ℓgeq` or *`:nabla2lgeq`* : optimal Hessian of the equality Lagrangian, ``\mathbf{\nabla^2}\ell_{\mathbf{g_{eq}}}``
+- `:∇²ℓgeq_ncolors` or *`:nabla2lgeq_ncolors`* : number of colors in `:∇²ℓgeq` sparsity pattern
 
 Note that the inequality constraint vectors and matrices only include the non-`Inf` values.
 

src/controller/nonlinmpc.jl

@@ -590,9 +590,13 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
         return obj_nonlinprog!(Ŷ0, U0, mpc, Ue, Ŷe, ΔŨ)
     end
     if !isnothing(hess)
-        _, ∇J_opt, ∇²J_opt = value_gradient_and_hessian(J!, hess, mpc.Z̃, J_cache...)
+        prep_∇²J = prepare_hessian(J!, hess, mpc.Z̃, J_cache...)
+        _, ∇J_opt, ∇²J_opt = value_gradient_and_hessian(J!, prep_∇²J, hess, mpc.Z̃, J_cache...)
+        ∇²J_ncolors = get_ncolors(prep_∇²J)
     else
-        ∇J_opt, ∇²J_opt = gradient(J!, mpc.gradient, mpc.Z̃, J_cache...), nothing
+        prep_∇J = prepare_gradient(J!, mpc.gradient, mpc.Z̃, J_cache...)
+        ∇J_opt = gradient(J!, prep_∇J, mpc.gradient, mpc.Z̃, J_cache...)
+        ∇²J_opt, ∇²J_ncolors = nothing, nothing
     end
     # --- inequality constraint derivatives ---
     ∇g_cache = (
@@ -605,7 +609,9 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
         gi .= @views g[i_g]
         return nothing
     end
-    g_opt, ∇g_opt = value_and_jacobian(gi!, gi, mpc.jacobian, mpc.Z̃, ∇g_cache...)
+    prep_∇g = prepare_jacobian(gi!, gi, mpc.jacobian, mpc.Z̃, ∇g_cache...)
+    g_opt, ∇g_opt = value_and_jacobian(gi!, gi, prep_∇g, mpc.jacobian, mpc.Z̃, ∇g_cache...)
+    ∇g_ncolors = get_ncolors(prep_∇g)
     if !isnothing(hess) && ngi > 0
         nonlincon = optim[:nonlinconstraint]
         λi = try
@@ -631,9 +637,11 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
             gi .= @views g[i_g]
             return dot(λi, gi)
         end
-        ∇²ℓg_opt = hessian(ℓ_gi, hess, mpc.Z̃, Constant(λi), ∇²g_cache...)
+        prep_∇²ℓg = prepare_hessian(ℓ_gi, hess, mpc.Z̃, Constant(λi), ∇²g_cache...)
+        ∇²ℓg_opt = hessian(ℓ_gi, prep_∇²ℓg, hess, mpc.Z̃, Constant(λi), ∇²g_cache...)
+        ∇²ℓg_ncolors = get_ncolors(prep_∇²ℓg)
     else
-        ∇²ℓg_opt = nothing
+        ∇²ℓg_opt, ∇²ℓg_ncolors = nothing, nothing
     end
     # --- equality constraint derivatives ---
     geq_cache = (
@@ -645,7 +653,9 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
         update_predictions!(ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K, X̂0, gc, g, geq, mpc, Z̃)
         return nothing
     end
-    geq_opt, ∇geq_opt = value_and_jacobian(geq!, geq, mpc.jacobian, mpc.Z̃, geq_cache...)
+    prep_∇geq = prepare_jacobian(geq!, geq, mpc.jacobian, mpc.Z̃, geq_cache...)
+    geq_opt, ∇geq_opt = value_and_jacobian(geq!, geq, prep_∇geq, mpc.jacobian, mpc.Z̃, geq_cache...)
+    ∇geq_ncolors = get_ncolors(prep_∇geq)
     if !isnothing(hess) && con.neq > 0
         nonlinconeq = optim[:nonlinconstrainteq]
         λeq = try
@@ -670,25 +680,37 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
             update_predictions!(ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K, X̂0, gc, g, geq, mpc, Z̃)
             return dot(λeq, geq)
         end
-        ∇²ℓgeq_opt = hessian(ℓ_geq, hess, mpc.Z̃, Constant(λeq), ∇²geq_cache...)
+        prep_∇²ℓgeq = prepare_hessian(ℓ_geq, hess, mpc.Z̃, Constant(λeq), ∇²geq_cache...)
+        ∇²ℓgeq_opt = hessian(ℓ_geq, prep_∇²ℓgeq, hess, mpc.Z̃, Constant(λeq), ∇²geq_cache...)
+        ∇²ℓgeq_ncolors = get_ncolors(prep_∇²ℓgeq)
     else
-        ∇²ℓgeq_opt = nothing
+        ∇²ℓgeq_opt, ∇²ℓgeq_ncolors = nothing, nothing
     end
     info[:∇J] = ∇J_opt
     info[:∇²J] = ∇²J_opt
+    info[:∇²J_ncolors] = ∇²J_ncolors
     info[:g] = g_opt
     info[:∇g] = ∇g_opt
+    info[:∇g_ncolors] = ∇g_ncolors
     info[:∇²ℓg] = ∇²ℓg_opt
+    info[:∇²ℓg_ncolors] = ∇²ℓg_ncolors
     info[:geq] = geq_opt
     info[:∇geq] = ∇geq_opt
+    info[:∇geq_ncolors] = ∇geq_ncolors
     info[:∇²ℓgeq] = ∇²ℓgeq_opt
+    info[:∇²ℓgeq_ncolors] = ∇²ℓgeq_ncolors
     # --- non-Unicode fields ---
     info[:nablaJ] = ∇J_opt
     info[:nabla2J] = ∇²J_opt
+    info[:nabla2J_ncolors] = ∇²J_ncolors
     info[:nablag] = ∇g_opt
+    info[:nablag_ncolors] = ∇g_ncolors
     info[:nabla2lg] = ∇²ℓg_opt
+    info[:nabla2lg_ncolors] = ∇²ℓg_ncolors
     info[:nablageq] = ∇geq_opt
+    info[:nablageq_ncolors] = ∇geq_ncolors
     info[:nabla2lgeq] = ∇²ℓgeq_opt
+    info[:nabla2lgeq_ncolors] = ∇²ℓgeq_ncolors
     return info
 end
 

src/general.jl

@@ -140,6 +140,10 @@ dense_backend(backend::AbstractADType) = backend
 dense_backend(backend::AutoSparse) = backend.dense_ad
 dense_backend(backend::SecondOrder) = backend.inner
 
+"Get the number of colors in preparation object `prep`, or `nothing` if not applicable."
+get_ncolors(::Prep) = nothing
+get_ncolors(prep::Union{SparseJacobianPrep, SparseHessianPrep}) = ncolors(prep)
+
 "Validate `hessian` keyword argument and return the differentiation `backend`."
 function validate_hessian(hessian, gradient, default)
     if hessian == true