wip: exact hessian

Author: franckgaga

src/ModelPredictiveControl.jl

@@ -7,8 +7,9 @@ using Random: randn
 using RecipesBase
 
 using DifferentiationInterface: ADTypes.AbstractADType, AutoForwardDiff, AutoSparse
-using DifferentiationInterface: gradient!, jacobian!, prepare_gradient, prepare_jacobian
-using DifferentiationInterface: value_and_gradient!, value_and_jacobian!
+using DifferentiationInterface: gradient!, value_and_gradient!, prepare_gradient
+using DifferentiationInterface: jacobian!, value_and_jacobian!, prepare_jacobian 
+using DifferentiationInterface: hessian!, value_gradient_and_hessian!, prepare_hessian
 using DifferentiationInterface: Constant, Cache
 using SparseConnectivityTracer: TracerSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern

src/controller/nonlinmpc.jl

@@ -7,6 +7,7 @@ const DEFAULT_NONLINMPC_JACSPARSE = AutoSparse(
     sparsity_detector=TracerSparsityDetector(),
     coloring_algorithm=GreedyColoringAlgorithm(),
 )
+const DEFAULT_NONLINMPC_HESSIAN = DEFAULT_NONLINMPC_JACSPARSE
 
 struct NonLinMPC{
     NT<:Real,
@@ -15,7 +16,8 @@ struct NonLinMPC{
     TM<:TranscriptionMethod,
     JM<:JuMP.GenericModel,
     GB<:AbstractADType,
-    JB<:AbstractADType, 
+    JB<:AbstractADType,
+    HB<:Union{AbstractADType, Nothing}, 
     PT<:Any,
     JEfunc<:Function,
     GCfunc<:Function
@@ -28,6 +30,7 @@ struct NonLinMPC{
     con::ControllerConstraint{NT, GCfunc}
     gradient::GB
     jacobian::JB
+    hessian::HB
     oracle::Bool
     Z̃::Vector{NT}
     ŷ::Vector{NT}
@@ -117,9 +120,9 @@ struct NonLinMPC{
         nZ̃ = get_nZ(estim, transcription, Hp, Hc) + nϵ
         Z̃ = zeros(NT, nZ̃)
         buffer = PredictiveControllerBuffer(estim, transcription, Hp, Hc, nϵ)
-        mpc = new{NT, SE, CW, TM, JM, GB, JB, PT, JEfunc, GCfunc}(
+        mpc = new{NT, SE, CW, TM, JM, GB, JB, HB, PT, JEfunc, GCfunc}(
             estim, transcription, optim, con,
-            gradient, jacobian, oracle,
+            gradient, jacobian, hessian, oracle,
             Z̃, ŷ,
             Hp, Hc, nϵ, nb,
             weights,
@@ -619,7 +622,7 @@ function get_nonlinops(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where JNT<
     # ----------- common cache for all functions  ----------------------------------------
     model = mpc.estim.model
     transcription = mpc.transcription
-    grad, jac = mpc.gradient, mpc.jacobian
+    grad, jac, hess = mpc.gradient, mpc.jacobian, mpc.hessian
     nu, ny, nx̂, nϵ = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ
     nk = get_nk(model, transcription)
     Hp, Hc = mpc.Hp, mpc.Hc
@@ -645,14 +648,28 @@ function get_nonlinops(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where JNT<
         gi .= @views g[i_g]
         return nothing
     end
-    Z̃_∇gi = fill(myNaN, nZ̃)      # NaN to force update_predictions! at first call
+    function ℓ_gi(Z̃_λ_gi, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K0, X̂0, gc, geq, g, gi)
+        Z̃, λ = @views Z̃_λ_gi[begin:begin+nZ̃-1], Z̃_λ_gi[begin+nZ̃:end]
+        update_predictions!(ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K0, X̂0, gc, g, geq, mpc, Z̃)
+        gi .= @views g[i_g]
+        return dot(λ, gi)
+    end
+    Z̃_∇gi  = fill(myNaN, nZ̃)      # NaN to force update at first call
+    Z̃_λ_gi = fill(myNaN, nZ̃ + ngi) 
     ∇gi_context = (
         Cache(ΔŨ), Cache(x̂0end), Cache(Ue), Cache(Ŷe), Cache(U0), Cache(Ŷ0), 
         Cache(Û0), Cache(K0), Cache(X̂0), 
         Cache(gc), Cache(geq), Cache(g)
     )
-    ∇gi_prep = prepare_jacobian(gi!, gi, jac, Z̃_∇gi, ∇gi_context...; strict)
-    ∇gi      = init_diffmat(JNT, jac, ∇gi_prep, nZ̃, ngi)
+    ∇gi_prep  = prepare_jacobian(gi!, gi, jac, Z̃_∇gi, ∇gi_context...; strict)
+    ∇²gi_context = (
+        Cache(ΔŨ), Cache(x̂0end), Cache(Ue), Cache(Ŷe), Cache(U0), Cache(Ŷ0), 
+        Cache(Û0), Cache(K0), Cache(X̂0), 
+        Cache(gc), Cache(geq), Cache(g), Cache(gi)
+    )
+    ∇²gi_prep = prepare_hessian(ℓ_gi, hess, Z̃_λ_gi, ∇²gi_context...; strict)
+    ∇gi       = init_diffmat(JNT, jac, ∇gi_prep, nZ̃, ngi)
+    ∇²ℓ_gi    = init_diffmat(JNT, hess, ∇²gi_prep, nZ̃ + ngi, nZ̃ + ngi)
     function update_con!(gi, ∇gi, Z̃_∇gi, Z̃_arg)
         if isdifferent(Z̃_arg, Z̃_∇gi)
             Z̃_∇gi .= Z̃_arg
@@ -668,23 +685,33 @@ function get_nonlinops(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where JNT<
         update_con!(gi, ∇gi, Z̃_∇gi, Z̃_arg)
         return diffmat2vec!(∇gi_arg, ∇gi)
     end
+    function ∇²gi_func!(∇²ℓ_arg, Z̃_arg, λ_arg)
+        Z̃_λ_gi[1:begin:begin+nZ̃-1]  .= Z̃_arg
+        Z̃_λ_gi[[begin+nZ̃:end]]      .= λ_arg
+        hessian!(ℓ_gi, ∇²ℓ_gi, ∇²gi_prep, hess, Z̃_λ_gi, ∇²gi_context)
+        return diffmat2vec!(∇²ℓ_arg, ∇²ℓ_gi)
+    end
     gi_min = fill(-myInf, ngi)
     gi_max = zeros(JNT,   ngi)
-    ∇gi_structure = init_diffstructure(∇gi)
+    ∇gi_structure  = init_diffstructure(∇gi)
+    ∇²gi_structure = init_diffstructure(∇²ℓ_gi)
+    display(∇²ℓ_gi)
     g_oracle = MOI.VectorNonlinearOracle(;
         dimension = nZ̃,
         l = gi_min,
         u = gi_max,
         eval_f = gi_func!,
         jacobian_structure = ∇gi_structure,
-        eval_jacobian = ∇gi_func!
+        eval_jacobian = ∇gi_func!,
+        hessian_lagrangian_structure = ∇²gi_structure,
+        eval_hessian_lagrangian = ∇²gi_func!
     )
     # ------------- equality constraints : nonlinear oracle ------------------------------
     function geq!(geq, Z̃, ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K0, X̂0, gc, g) 
         update_predictions!(ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K0, X̂0, gc, g, geq, mpc, Z̃)
         return nothing
     end
-    Z̃_∇geq = fill(myNaN, nZ̃)    # NaN to force update_predictions! at first call
+    Z̃_∇geq = fill(myNaN, nZ̃)    # NaN to force update at first call
     ∇geq_context = (
         Cache(ΔŨ), Cache(x̂0end), Cache(Ue), Cache(Ŷe), Cache(U0), Cache(Ŷ0),
         Cache(Û0), Cache(K0),   Cache(X̂0),
@@ -722,7 +749,7 @@ function get_nonlinops(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where JNT<
         update_predictions!(ΔŨ, x̂0end, Ue, Ŷe, U0, Ŷ0, Û0, K0, X̂0, gc, g, geq, mpc, Z̃)
         return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)
     end
-    Z̃_∇J = fill(myNaN, nZ̃)      # NaN to force update_predictions! at first call
+    Z̃_∇J = fill(myNaN, nZ̃)      # NaN to force update at first call
     ∇J_context = (
         Cache(ΔŨ), Cache(x̂0end), Cache(Ue), Cache(Ŷe), Cache(U0), Cache(Ŷ0), 
         Cache(Û0), Cache(K0), Cache(X̂0),