type stable usage of hessian!

Author: tmigot

src/ARCTR.jl

@@ -3,7 +3,7 @@ module ARCTR
 # stdlib
 using LinearAlgebra, SparseArrays
 # JSO
-using Krylov, NLPModels, SparseMatricesCOO, SolverCore
+using Krylov, LinearOperators, NLPModels, SparseMatricesCOO, SolverCore
 # Stopping
 using Stopping, StoppingInterface
 

src/main.jl

@@ -86,7 +86,7 @@ function preprocess(
     α,
 )
     max_hprod = stp.meta.max_cntrs[:neval_hprod]
-    Hx = stp.current_state.Hx
+    Hx = workspace.Hstruct.H
     PData = preprocess(PData, Hx, ∇f, norm_∇f, neval_hprod(stp.pb), max_hprod, α)
     return PData
 end
@@ -101,7 +101,7 @@ function compute_direction(
     solve_model,
 )
     max_hprod = stp.meta.max_cntrs[:neval_hprod]
-    Hx = stp.current_state.Hx
+    Hx = workspace.Hstruct.H
     return solve_model(PData, Hx, ∇f, norm_∇f, neval_hprod(stp.pb), max_hprod, α)
 end
 
@@ -162,11 +162,11 @@ end
 
 Update `Δq = -(∇f + 0.5 * (Hx * d)) ⋅ d` in-place.
 """
-function compute_Δq(workspace, Hx, d, ∇f) # -(∇f + 0.5 * (nlp_at_x.Hx * d)) ⋅ d
-    mul!(workspace.Hd, Hx, d)
-    workspace.Hd .*= 0.5
-    workspace.Hd .+= ∇f
-    return -dot(workspace.Hd, d)
+function compute_Δq(workspace, Hx, d, ∇f)
+  mul!(workspace.Hd, Hx, d)
+  workspace.Hd .*= 0.5
+  workspace.Hd .+= ∇f
+  return -dot(workspace.Hd, d)
 end
 
 function TRARC(
@@ -180,8 +180,8 @@ function TRARC(
     kwargs...,
 ) where {Pb,M,SRC,MStp,LoS,S,T,Hess,ParamData}
     nlp, nlp_at_x = nlp_stop.pb, nlp_stop.current_state
-    xt, xtnext, d, ∇f, ∇fnext =
-        workspace.xt, workspace.xtnext, workspace.d, workspace.∇f, workspace.∇fnext
+    xt, xtnext, ∇f, ∇fnext = workspace.xt, workspace.xtnext, workspace.∇f, workspace.∇fnext
+    d, Hx = workspace.d, workspace.Hstruct.H
 
     α = TR.α₀
     max_unsuccinarow = TR.max_unsuccinarow
@@ -197,7 +197,8 @@ function TRARC(
 
     Stopping._smart_update!(nlp_at_x, x = xt, fx = ft, gx = ∇f)
     OK = start!(nlp_stop)
-    !OK && Stopping._smart_update!(nlp_at_x, Hx = hessian!(workspace, nlp, xt))
+    Hx = hessian!(workspace, nlp, xt)
+    !OK && Stopping._smart_update!(nlp_at_x, Hx = Hx)
 
     iter = 0 # counter different than stop count
     succ, unsucc, verysucc, unsuccinarow = 0, 0, 0, 0
@@ -222,7 +223,7 @@ function TRARC(
                 compute_direction(nlp_stop, PData, workspace, ∇f, norm_∇f, α, solve_model)
 
             slope = ∇f ⋅ d
-            Δq = compute_Δq(workspace, nlp_at_x.Hx, d, ∇f) # -(∇f + 0.5 * (nlp_at_x.Hx * d)) ⋅ d
+            Δq = compute_Δq(workspace, Hx, d, ∇f)
 
             xtnext .= xt .+ d
             ftnext = obj(nlp, xtnext, workspace)
@@ -238,7 +239,8 @@ function TRARC(
                 unsucc += 1
                 unsuccinarow += 1
                 η = (1 - acceptance_threshold) / 10 # ∈ (acceptance_threshold, 1)
-                qksk = ft + slope + ((nlp_at_x.Hx * d) ⋅ d) / 2
+                mul!(workspace.Hd, Hx, d)
+                qksk = ft + slope + (workspace.Hd ⋅ d) / 2
                 αbad = (1 - η) * slope / ((1 - η) * (ft + slope) + η * qksk - ftnext)
                 α = min(decrease(PData, α, TR), max(TR.large_decrease_factor, αbad) * α)
             elseif r < acceptance_threshold # unsucessful
@@ -276,7 +278,8 @@ function TRARC(
         nlp_stop.meta.nb_of_stop = iter
         Stopping._smart_update!(nlp_at_x, x = xt, fx = ft, gx = ∇f)
         OK = stop!(nlp_stop)
-        success && Stopping._smart_update!(nlp_at_x, Hx = hessian!(workspace, nlp, xt))
+        Hx = hessian!(workspace, nlp, xt)
+        success && Stopping._smart_update!(nlp_at_x, Hx = Hx)
     end # while !OK
 
     return nlp_stop

src/utils/hessian_rep.jl

@@ -2,24 +2,28 @@
     HessDense(::AbstractNLPModel{T,S}, n)
 Return a structure used for the evaluation of dense Hessian matrix.
 """
-struct HessDense
+struct HessDense{T}
+    H::Matrix{T}
     function HessDense(::AbstractNLPModel{T,S}, n) where {T,S}
-        return new()
+        H = Matrix{Float64}(undef, n, n)
+        return new{T}(H)
     end
 end
 
 """
     HessSparse(::AbstractNLPModel{T,S}, n)
 Return a structure used for the evaluation of sparse Hessian matrix.
 """
-struct HessSparse{S,Vi}
+struct HessSparse{T,S,Vi,It<:Integer}
     rows::Vi
     cols::Vi
     vals::S
+    H::Symmetric{T, SparseMatrixCSC{T, It}}
     function HessSparse(nlp::AbstractNLPModel{T,S}, n) where {T,S}
         rows, cols = hess_structure(nlp)
         vals = S(undef, nlp.meta.nnzh)
-        return new{S,typeof(rows)}(rows, cols, vals)
+        H = Symmetric(spzeros(T, n, n), :L)
+        return new{T,S,typeof(rows),eltype(rows)}(rows, cols, vals, H)
     end
 end
 
@@ -42,22 +46,26 @@ end
     HessOp(::AbstractNLPModel{T,S}, n)
 Return a structure used for the evaluation of the Hessian matrix as an operator.
 """
-struct HessOp{S}
+mutable struct HessOp{S}
     Hv::S
+    H
     function HessOp(::AbstractNLPModel{T,S}, n) where {T,S}
-        return new{S}(S(undef, n))
+        H = LinearOperator{T}(n, n, true, true, v -> v, v -> v, v -> v)
+        return new{S}(S(undef, n), H)
     end
 end
 
 """
     HessGaussNewtonOp(::AbstractNLSModel{T,S}, n)
 Return a structure used for the evaluation of the Hessian matrix as an operator.
 """
-struct HessGaussNewtonOp{S}
+mutable struct HessGaussNewtonOp{S}
     Jv::S
     Jtv::S
+    H
     function HessGaussNewtonOp(nls::AbstractNLSModel{T,S}, n) where {T,S}
-        return new{S}(S(undef, nls.nls_meta.nequ), S(undef, n))
+        Jx = LinearOperator{T}(nls.nls_meta.nequ, n, false, false, v -> v, v -> v, v -> v)
+        return new{S}(S(undef, nls.nls_meta.nequ), S(undef, n), Jx' * Jx)
     end
 end
 
@@ -81,23 +89,26 @@ Return the Hessian matrix of `nlp` at `x` in-place with memory update of `worksp
 function hessian! end
 
 function hessian!(workspace::HessDense, nlp, x)
-    H = Matrix(hess(nlp, x))
-    return H
+    workspace.H .= Matrix(hess(nlp, x))
+    return workspace.H
 end
 
 function hessian!(workspace::HessOp, nlp, x)
-    return hess_op!(nlp, x, workspace.Hv)
+    workspace.H = hess_op!(nlp, x, workspace.Hv)
+    return workspace.H
 end
 
 function hessian!(workspace::HessGaussNewtonOp, nlp, x)
     Jx = jac_op_residual!(nlp, x, workspace.Jv, workspace.Jtv)
-    return Jx' * Jx
+    workspace.H = Jx' * Jx
+    return workspace.H
 end
 
 function hessian!(workspace::HessSparse, nlp, x)
     hess_coord!(nlp, x, workspace.vals)
     n = nlp.meta.nvar
-    return Symmetric(sparse(workspace.rows, workspace.cols, workspace.vals, n, n), :L)
+    workspace.H .= Symmetric(sparse(workspace.rows, workspace.cols, workspace.vals, n, n), :L)
+    return workspace.H
 end
 
 function hessian!(workspace::HessSparseCOO, nlp, x)