implicit jacobian update

Author: MaxenceGollier

src/LMModel.jl

@@ -16,10 +16,9 @@ where `J` is the Jacobian of `F` at `xk`, represented via matrix-free operations
 
 `σ > 0` is a regularization parameter and `v` is a vector of the same size as `F(xk)` used for intermediary computations.
 """
-mutable struct LMModel{T <: Real, V <: AbstractVector{T}, J <: Function, Jt <: Function} <:
+mutable struct LMModel{T <: Real, V <: AbstractVector{T}, Jac <: Union{AbstractMatrix, AbstractLinearOperator}} <:
                AbstractNLPModel{T, V}
-  j_prod!::J
-  jt_prod!::Jt
+  J::Jac
   F::V
   v::V
   xk::V
@@ -28,19 +27,19 @@ mutable struct LMModel{T <: Real, V <: AbstractVector{T}, J <: Function, Jt <: F
   counters::Counters
 end
 
-function LMModel(j_prod!::J, jt_prod!::Jt, F::V, σ::T, xk::V) where {T, V, J, Jt}
+function LMModel(J::Jac, F::V, σ::T, xk::V) where {T, V, Jac}
   meta = NLPModelMeta(
     length(xk),
     x0 = xk, # Perhaps we should add lvar and uvar as well here.
   )
   v = similar(F)
-  return LMModel(j_prod!, jt_prod!, F, v, xk, σ, meta, Counters())
+  return LMModel(J, F, v, xk, σ, meta, Counters())
 end
 
 function NLPModels.obj(nlp::LMModel, x::AbstractVector{T}) where {T}
   @lencheck nlp.meta.nvar x
   increment!(nlp, :neval_obj)
-  nlp.j_prod!(nlp.xk, x, nlp.v) # v = J(xk)x
+  mul!(nlp.v, nlp.J, x)
   nlp.v .+= nlp.F
   return (dot(nlp.v, nlp.v) + nlp.σ * dot(x, x)) / 2
 end
@@ -49,9 +48,9 @@ function NLPModels.grad!(nlp::LMModel, x::AbstractVector{T}, g::AbstractVector{T
   @lencheck nlp.meta.nvar x
   @lencheck nlp.meta.nvar g
   increment!(nlp, :neval_grad)
-  nlp.j_prod!(nlp.xk, x, nlp.v) # v = J(xk)x + F
+  mul!(nlp.v, nlp.J, x)
   nlp.v .+= nlp.F
-  nlp.jt_prod!(nlp.xk, nlp.v, g) # g = J^T(xk) v
+  mul!(g, nlp.J', nlp.v)
   @. g += nlp.σ .* x
   return g
 end

src/LM_alg.jl

@@ -61,14 +61,8 @@ function LMSolver(reg_nls::AbstractRegularizedNLPModel{T, V}; subsolver = R2Solv
     has_bnds ? shifted(reg_nls.h, xk, l_bound_m_x, u_bound_m_x, reg_nls.selected) :
     shifted(reg_nls.h, xk)
 
-  jprod! = let nls = reg_nls.model
-    (x, v, Jv) -> jprod_residual!(nls, x, v, Jv)
-  end
-  jt_prod! = let nls = reg_nls.model
-    (x, v, Jtv) -> jtprod_residual!(nls, x, v, Jtv)
-  end
-
-  sub_nlp = LMModel(jprod!, jt_prod!, Fk, T(1), xk)
+  Jk = jac_op_residual(reg_nls.model, xk)
+  sub_nlp = LMModel(Jk, Fk, T(1), xk)
   subpb = RegularizedNLPModel(sub_nlp, ψ)
   substats = RegularizedExecutionStats(subpb)
   subsolver = subsolver(subpb)
@@ -283,7 +277,7 @@ function SolverCore.solve!(
   jtprod_residual!(nls, xk, Fk, ∇fk)
   fk = dot(Fk, Fk) / 2
 
-  σmax, found_σ = opnorm(jac_op_residual!(nls, xk, Jv, Jtv))
+  σmax, found_σ = opnorm(solver.subpb.model.J)
   found_σ || error("operator norm computation failed")
   ν = θ / (σmax^2 + σk) # ‖J'J + σₖ I‖ = ‖J‖² + σₖ
   sqrt_ξ1_νInv = one(T)
@@ -407,7 +401,7 @@ function SolverCore.solve!(
 
       # update opnorm if not linear least squares
       if nonlinear == true
-        σmax, found_σ = opnorm(jac_op_residual!(nls, xk, Jv, Jtv))
+        σmax, found_σ = opnorm(solver.subpb.model.J)
         found_σ || error("operator norm computation failed")
       end
     end