Merge pull request #279 from JuliaControl/nonzero_operator_hessian

debug: only fills non-zero in `∇²f` of `@operator`

Author: franckgaga

Project.toml

@@ -1,6 +1,6 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
-version = "1.13.0"
+version = "1.13.1"
 authors = ["Francis Gagnon"]
 
 [deps]

src/controller/nonlinmpc.jl

@@ -651,6 +651,7 @@ function get_nonlinobj_op(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where J
     if !isnothing(hess)
         ∇²J_prep = prepare_hessian(J!, hess, Z̃_J, J_cache...; strict)
         ∇²J = init_diffmat(JNT, hess, ∇²J_prep, nZ̃, nZ̃)
+        ∇²J_structure = lowertriangle_indices(init_diffstructure(∇²J))
     end
     update_objective! = if !isnothing(hess)
         function (J, ∇J, ∇²J, Z̃_J, Z̃_arg)
@@ -715,7 +716,7 @@ function get_nonlinobj_op(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where J
     else                        # multivariate syntax (see JuMP.@operator doc):
         function (∇²J_arg::AbstractMatrix{T}, Z̃_arg::Vararg{T, N}) where {N, T<:Real}
             update_objective!(J, ∇J, ∇²J, Z̃_J, Z̃_arg)
-            return fill_lowertriangle!(∇²J_arg, ∇²J)
+            return fill_diffstructure!(∇²J_arg, ∇²J, ∇²J_structure)
         end
     end
     if !isnothing(hess)
@@ -807,13 +808,13 @@ function get_nonlincon_oracle(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JN
     end
     function ∇gi_func!(∇gi_arg, Z̃_arg)
         update_con!(gi, ∇gi, Z̃_∇gi, Z̃_arg) 
-        return diffmat2vec!(∇gi_arg, ∇gi, ∇gi_structure)
+        return fill_diffstructure!(∇gi_arg, ∇gi, ∇gi_structure)
     end
     function ∇²gi_func!(∇²ℓ_arg, Z̃_arg, λ_arg)
         Z̃_∇gi  .= Z̃_arg
         λi     .= λ_arg
         hessian!(ℓ_gi, ∇²ℓ_gi, ∇²gi_prep, hess, Z̃_∇gi, Constant(λi), ∇²gi_cache...)
-        return diffmat2vec!(∇²ℓ_arg, ∇²ℓ_gi, ∇²gi_structure)
+        return fill_diffstructure!(∇²ℓ_arg, ∇²ℓ_gi, ∇²gi_structure)
     end
     gi_min = fill(-myInf, ngi)
     gi_max = zeros(JNT,   ngi)
@@ -870,13 +871,13 @@ function get_nonlincon_oracle(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JN
     end
     function ∇geq_func!(∇geq_arg, Z̃_arg)
         update_con_eq!(geq, ∇geq, Z̃_∇geq, Z̃_arg)
-        return diffmat2vec!(∇geq_arg, ∇geq, ∇geq_structure)
+        return fill_diffstructure!(∇geq_arg, ∇geq, ∇geq_structure)
     end
     function ∇²geq_func!(∇²ℓ_arg, Z̃_arg, λ_arg)
         Z̃_∇geq .= Z̃_arg
         λeq    .= λ_arg
         hessian!(ℓ_geq, ∇²ℓ_geq, ∇²geq_prep, hess, Z̃_∇geq, Constant(λeq), ∇²geq_cache...)
-        return diffmat2vec!(∇²ℓ_arg, ∇²ℓ_geq, ∇²geq_structure)
+        return fill_diffstructure!(∇²ℓ_arg, ∇²ℓ_geq, ∇²geq_structure)
     end
     geq_min = geq_max = zeros(JNT, neq)
     geq_oracle = MOI.VectorNonlinearOracle(;

src/estimator/mhe/construct.jl

@@ -1395,6 +1395,7 @@ function get_nonlinobj_op(
         ∇²J_prep = prepare_hessian(J!, hess, Z̃_J, J_cache...; strict)
         estim.Nk[] = 0
         ∇²J = init_diffmat(JNT, hess, ∇²J_prep, nZ̃, nZ̃)
+        ∇²J_structure = lowertriangle_indices(init_diffstructure(∇²J))
     end
     update_objective! = if !isnothing(hess)
         function (J, ∇J, ∇²J, Z̃_∇J, Z̃_arg)
@@ -1439,7 +1440,7 @@ function get_nonlinobj_op(
     end
     function ∇²J_func!(∇²J_arg::AbstractMatrix{T}, Z̃_arg::Vararg{T, N}) where {N, T<:Real}
         update_objective!(J, ∇J, ∇²J, Z̃_J, Z̃_arg)
-        return fill_lowertriangle!(∇²J_arg, ∇²J)
+        return fill_diffstructure!(∇²J_arg, ∇²J, ∇²J_structure)
     end
     if !isnothing(hess)
         @operator(optim, J_op, nZ̃, J_func, ∇J_func!, ∇²J_func!)
@@ -1525,13 +1526,13 @@ function get_nonlincon_oracle(
     end
     function ∇gi_func!(∇gi_arg, Z̃_arg)
         update_con!(gi, ∇gi, Z̃_∇gi, Z̃_arg)
-        return diffmat2vec!(∇gi_arg, ∇gi, ∇gi_structure)
+        return fill_diffstructure!(∇gi_arg, ∇gi, ∇gi_structure)
     end
     function ∇²gi_func!(∇²ℓ_arg, Z̃_arg, λ_arg)
         Z̃_∇gi  .= Z̃_arg
         λi     .= λ_arg
         hessian!(ℓ_gi, ∇²ℓ_gi, ∇²gi_prep, hess, Z̃_∇gi, Constant(λi), ∇²gi_cache...)
-        return diffmat2vec!(∇²ℓ_arg, ∇²ℓ_gi, ∇²gi_structure)
+        return fill_diffstructure!(∇²ℓ_arg, ∇²ℓ_gi, ∇²gi_structure)
     end
     gi_min = fill(-myInf, ngi)
     gi_max = zeros(JNT,   ngi)

src/general.jl

@@ -76,21 +76,26 @@ function lowertriangle_indices(i_vec::Vector{Tuple{Int, Int}})
     return [(i,j) for (i,j) in i_vec if i ≥ j]
 end
 
-"Fill the lower triangular part of A in-place with the corresponding part in B."
-function fill_lowertriangle!(A::AbstractMatrix, B::AbstractMatrix)
-    for j in axes(A, 2), i in axes(A, 1)
-        (i ≥ j) && (A[i, j] = B[i, j])
+"Store the diff. matrix `A` in the vector `v` with list of nonzero indices `i_vec`." 
+function fill_diffstructure!(
+    v::AbstractVector, A::AbstractMatrix, i_vec::Vector{Tuple{Int, Int}}
+)
+    for i in eachindex(i_vec)
+        i_A, j_A = i_vec[i]
+        v[i] = A[i_A, j_A]
     end
-    return A
+    return v
 end
 
-"Store the diff. matrix `A` in the vector `v` with list of nonzero indices `i_vec`" 
-function diffmat2vec!(v::AbstractVector, A::AbstractMatrix, i_vec::Vector{Tuple{Int, Int}})
-    for i in eachindex(v)
+"Store the diff. matrix `A` in the matrix `T` with list of nonzero indices `i_vec`." 
+function fill_diffstructure!(
+    T::AbstractMatrix, A::AbstractMatrix, i_vec::Vector{Tuple{Int, Int}}
+)
+    for i in eachindex(i_vec)
         i_A, j_A = i_vec[i]
-        v[i] = A[i_A, j_A]
+        T[i_A, j_A] = A[i_A, j_A]
     end
-    return v
+    return T
 end
 
 backend_str(backend::AbstractADType) = string(nameof(typeof(backend)))