JSO Compliance : LM

src/LMModel.jl

@@ -0,0 +1,57 @@
+export LMModel
+
+@doc raw"""
+    LMModel(j_prod!, jt_prod, F, v, σ, xk)
+
+Given the unconstrained optimization problem:
+```math
+\min \tfrac{1}{2} \| F(x) \|^2,
+```
+this model represents the smooth LM subproblem:
+```math
+\min_s \ \tfrac{1}{2} \| F(x) + J(x)s \|^2 + \tfrac{1}{2} σ \|s\|^2
+```
+where `J` is the Jacobian of `F` at `xk`, represented via matrix-free operations.
+`j_prod!(xk, s, out)` computes `J(xk) * s`, and `jt_prod!(xk, r, out)` computes `J(xk)' * r`.
+
+`σ > 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} <:
+               AbstractNLPModel{T, V}
+  j_prod!::J
+  jt_prod!::Jt
+  F::V
+  v::V
+  xk::V
+  σ::T
+  meta::NLPModelMeta{T, V}
+  counters::Counters
+end
+
+function LMModel(j_prod!::J, jt_prod!::Jt, F::V, σ::T, xk::V) where {T, V, J, Jt}
+  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())
+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
+  nlp.v .+= nlp.F
+  return ( dot(nlp.v, nlp.v) + nlp.σ * dot(x, x) ) / 2
+end
+
+function NLPModels.grad!(nlp::LMModel, x::AbstractVector{T}, g::AbstractVector{T}) where{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
+  nlp.v .+= nlp.F
+  nlp.jt_prod!(nlp.xk, nlp.v, g) # g = J^T(xk) v
+  @. g += nlp.σ .* x
+  return g
+end