Add support for ArrayDiff

Project.toml

@@ -11,7 +11,14 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[weakdeps]
+ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50"
+
+[extensions]
+NLPModelsJuMPArrayDiffExt = "ArrayDiff"
+
 [compat]
+ArrayDiff = "0.1"
 JuMP = "1.25"
 LinearAlgebra = "1.10"
 MathOptInterface = "1.46"

ext/NLPModelsJuMPArrayDiffExt.jl

@@ -0,0 +1,108 @@
+module NLPModelsJuMPArrayDiffExt
+
+import NLPModelsJuMP
+import ArrayDiff
+import MathOptInterface as MOI
+import NLPModels
+import LinearAlgebra
+
+NLPModelsJuMP._nonlinear_model(ad::ArrayDiff.Mode) = ArrayDiff.model(ad)
+
+# Detect `(...).^2` (broadcast `:^` with exponent 2) and return the residual `...`.
+function NLPModelsJuMP._detect_squared_residual(inner::ArrayDiff.ArrayNonlinearFunction)
+    if inner.head !== :^ || !inner.broadcasted
+        return nothing
+    end
+    if length(inner.args) != 2
+        return nothing
+    end
+    exponent = inner.args[2]
+    if !(exponent isa Number) || exponent != 2
+        return nothing
+    end
+    return inner.args[1]
+end
+
+mutable struct ArrayDiffNLSModel{T, V <: AbstractVector{T}, R} <: NLPModels.AbstractNLSModel{T, V}
+    meta::NLPModels.NLPModelMeta{T, V}
+    nls_meta::NLPModels.NLSMeta{T, V}
+    counters::NLPModels.NLSCounters
+    evaluator::ArrayDiff.Evaluator{T, R}
+end
+
+function NLPModelsJuMP._build_nls_from_residual(
+    moimodel::MOI.ModelLike,
+    residual::ArrayDiff.ArrayNonlinearFunction,
+    ad::ArrayDiff.Mode{S},
+) where {S <: AbstractVector{<:Real}}
+    T = eltype(S)
+    V = S
+    _, nvar, lvar, uvar, x0 = NLPModelsJuMP.parser_variables(moimodel)
+    lvar = convert(V, lvar)
+    uvar = convert(V, uvar)
+    x0 = convert(V, x0)
+    model = ArrayDiff.model(ad)
+    ArrayDiff.set_residual!(model, residual)
+    vars = MOI.get(moimodel, MOI.ListOfVariableIndices())
+    evaluator = MOI.Nonlinear.Evaluator(model, ad, vars)
+    MOI.initialize(evaluator, [:Grad, :Jac, :JacVec])
+    nresid = ArrayDiff.residual_dimension(evaluator)
+    meta = NLPModels.NLPModelMeta{T, V}(
+        nvar;
+        x0 = x0,
+        lvar = lvar,
+        uvar = uvar,
+        minimize = MOI.get(moimodel, MOI.ObjectiveSense()) == MOI.MIN_SENSE,
+        islp = false,
+        name = "ArrayDiffNLS",
+        hprod_available = false,
+        hess_available = false,
+    )
+    nls_meta = NLPModels.NLSMeta{T, V}(
+        nresid,
+        nvar;
+        x0 = x0,
+        nnzj = nresid * nvar,
+        nnzh = 0,
+        jac_residual_available = false,
+        hess_residual_available = false,
+        jprod_residual_available = true,
+        jtprod_residual_available = true,
+        hprod_residual_available = false,
+    )
+    return ArrayDiffNLSModel(meta, nls_meta, NLPModels.NLSCounters(), evaluator)
+end
+
+function NLPModels.residual!(
+    nls::ArrayDiffNLSModel,
+    x::AbstractVector,
+    Fx::AbstractVector,
+)
+    NLPModels.increment!(nls, :neval_residual)
+    ArrayDiff.eval_residual!(nls.evaluator, Fx, x)
+    return Fx
+end
+
+function NLPModels.jprod_residual!(
+    nls::ArrayDiffNLSModel,
+    x::AbstractVector,
+    v::AbstractVector,
+    Jv::AbstractVector,
+)
+    NLPModels.increment!(nls, :neval_jprod_residual)
+    ArrayDiff.eval_residual_jprod!(nls.evaluator, Jv, x, v)
+    return Jv
+end
+
+function NLPModels.jtprod_residual!(
+    nls::ArrayDiffNLSModel,
+    x::AbstractVector,
+    v::AbstractVector,
+    Jtv::AbstractVector,
+)
+    NLPModels.increment!(nls, :neval_jtprod_residual)
+    ArrayDiff.eval_residual_jtprod!(nls.evaluator, Jtv, x, v)
+    return Jtv
+end
+
+end

src/MOI_wrapper.jl

@@ -3,14 +3,19 @@ import SolverCore
 mutable struct Optimizer <: MOI.AbstractOptimizer
   options::Dict{String, Any}
   silent::Bool
+  ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation
   solver
-  nlp::Union{Nothing, MathOptNLPModel}
-  stats::Union{
-    Nothing,
-    SolverCore.GenericExecutionStats{Float64, Vector{Float64}, Vector{Float64}, Any},
-  }
+  nlp::Union{Nothing, AbstractNLPModel}
+  stats::Union{Nothing, SolverCore.GenericExecutionStats}
   function Optimizer()
-    return new(Dict{String, Any}(), false, nothing, nothing, nothing)
+    return new(
+      Dict{String, Any}(),
+      false,
+      MOI.Nonlinear.SparseReverseMode(),
+      nothing,
+      nothing,
+      nothing,
+    )
   end
 end
 
@@ -51,6 +56,25 @@ end
 
 MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent
 
+###
+### MOI.AutomaticDifferentiationBackend
+###
+
+MOI.supports(::Optimizer, ::MOI.AutomaticDifferentiationBackend) = true
+
+function MOI.get(optimizer::Optimizer, ::MOI.AutomaticDifferentiationBackend)
+  return optimizer.ad_backend
+end
+
+function MOI.set(
+  optimizer::Optimizer,
+  ::MOI.AutomaticDifferentiationBackend,
+  backend::MOI.Nonlinear.AbstractAutomaticDifferentiation,
+)
+  optimizer.ad_backend = backend
+  return
+end
+
 ###
 ### MOI.AbstractModelAttribute
 ###
@@ -60,6 +84,7 @@ function MOI.supports(
   ::Union{
     MOI.ObjectiveSense,
     MOI.ObjectiveFunction{<:Union{LinQuad, MOI.ScalarNonlinearFunction}},
+    MOI.ObjectiveFunction{<:MOI.AbstractVectorFunction},
     MOI.NLPBlock,
     MOI.UserDefinedFunction,
   },
@@ -92,12 +117,28 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
       "No solver specified, use for instance `using Percival; JuMP.set_attribute(model, \"solver\", PercivalSolver)`",
     )
   end
-  dest.nlp, index_map = nlp_model(src)
+  nls = _try_nls_model(src, dest.ad_backend)
+  if nls !== nothing
+    dest.nlp = nls
+    dest.solver = dest.options["solver"](dest.nlp)
+    return parser_variables(src)[1]
+  end
+  dest.nlp, index_map = nlp_model(src; ad_backend = dest.ad_backend)
   dest.solver = dest.options["solver"](dest.nlp)
   return index_map
 end
 
 function MOI.optimize!(model::Optimizer)
+  if model.nlp === nothing
+    # Direct mode: build NLPModel from the optimizer itself
+    if !haskey(model.options, "solver")
+      error(
+        "No solver specified, use for instance `using Percival; JuMP.set_attribute(model, \"solver\", PercivalSolver)`",
+      )
+    end
+    model.nlp, _ = nlp_model(model; ad_backend = model.ad_backend)
+    model.solver = model.options["solver"](model.nlp)
+  end
   options = Dict{Symbol, Any}(
     Symbol(key) => model.options[key] for key in keys(model.options) if key != "solver"
   )

src/moi_nlp_model.jl

@@ -2,7 +2,7 @@ export MathOptNLPModel
 
 mutable struct MathOptNLPModel <: AbstractNLPModel{Float64, Vector{Float64}}
   meta::NLPModelMeta{Float64, Vector{Float64}}
-  eval::MOI.Nonlinear.Evaluator
+  eval::MOI.AbstractNLPEvaluator
   lincon::LinearConstraints
   quadcon::QuadraticConstraints
   nlcon::NonLinearStructure
@@ -29,12 +29,17 @@ function MathOptNLPModel(moimodel::MOI.ModelLike; kws...)
   return nlp_model(moimodel; kws...)[1]
 end
 
-function nlp_model(moimodel::MOI.ModelLike; hessian::Bool = true, name::String = "Generic")
+function nlp_model(
+  moimodel::MOI.ModelLike;
+  hessian::Bool = true,
+  name::String = "Generic",
+  ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation = _get_ad_backend(moimodel),
+)
   index_map, nvar, lvar, uvar, x0 = parser_variables(moimodel)
   nlin, lincon, lin_lcon, lin_ucon, quadcon, quad_lcon, quad_ucon =
     parser_MOI(moimodel, index_map, nvar)
 
-  nlp_data = _nlp_block(moimodel)
+  nlp_data = _nlp_block(moimodel, ad_backend)
   nlcon = parser_NL(nlp_data, hessian = hessian)
   oracles = parser_oracles(moimodel)
   counters = Counters()

src/utils.jl

@@ -533,8 +533,50 @@ function _nlp_model(dest::MOI.Nonlinear.Model, src::MOI.ModelLike, F::Type{SNF},
   return has_nonlinear
 end
 
-function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model}
-  nlp_model = MOI.Nonlinear.Model()
+_nonlinear_model(::MOI.Nonlinear.AbstractAutomaticDifferentiation) = MOI.Nonlinear.Model()
+
+"""
+    _detect_squared_residual(inner)
+
+Hook for AD extensions: given the `inner` function under a `:sum` root
+(i.e., the `?` in `sum(?)`), return the residual whose square sum is being
+minimized — typically the first argument of a broadcast `:^` with exponent 2.
+
+Default returns `nothing` (no NLS routing). Extensions for AD backends that
+carry vector-function types (e.g. `ArrayDiff.Mode`) override this.
+"""
+_detect_squared_residual(::Any) = nothing
+
+"""
+    _build_nls_from_residual(moimodel, residual, ad_backend)
+
+Hook for AD extensions: build an `AbstractNLSModel` that evaluates the given
+`residual` (a vector function) using `ad_backend`. Default returns `nothing`,
+which makes the optimizer fall back to `MathOptNLPModel`.
+"""
+_build_nls_from_residual(::Any, ::Any, ::MOI.Nonlinear.AbstractAutomaticDifferentiation) = nothing
+
+function _try_nls_model(
+  moimodel::MOI.ModelLike,
+  ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation,
+)
+  F = MOI.get(moimodel, MOI.ObjectiveFunctionType())
+  if !(F <: SNF)
+    return nothing
+  end
+  obj = MOI.get(moimodel, MOI.ObjectiveFunction{F}())
+  if obj.head !== :sum || length(obj.args) != 1
+    return nothing
+  end
+  residual = _detect_squared_residual(obj.args[1])
+  if residual === nothing
+    return nothing
+  end
+  return _build_nls_from_residual(moimodel, residual, ad_backend)
+end
+
+function _nlp_model(model::MOI.ModelLike, ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation)
+  nlp_model = _nonlinear_model(ad_backend)
   has_nonlinear = false
   for attr in MOI.get(model, MOI.ListOfModelAttributesSet())
     if attr isa MOI.UserDefinedFunction
@@ -550,27 +592,45 @@ function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model}
   if F <: SNF
     MOI.Nonlinear.set_objective(nlp_model, MOI.get(model, MOI.ObjectiveFunction{F}()))
     has_nonlinear = true
+  elseif F <: MOI.AbstractVectorFunction
+    # ArrayNonlinearFunction or similar: return the function directly.
+    # The ad_backend from the model will build the evaluator.
+    func = MOI.get(model, MOI.ObjectiveFunction{F}())
+    return func
   end
   if !has_nonlinear
     return nothing
   end
   return nlp_model
 end
 
-function _nlp_block(model::MOI.ModelLike)
+function _get_ad_backend(model::MOI.ModelLike)
+  if MOI.supports(model, MOI.AutomaticDifferentiationBackend())
+    return MOI.get(model, MOI.AutomaticDifferentiationBackend())
+  end
+  return MOI.Nonlinear.SparseReverseMode()
+end
+
+function _nlp_block(
+  model::MOI.ModelLike,
+  ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation = _get_ad_backend(model),
+)
   # Old interface with `@NL...`
-  nlp_data = MOI.get(model, MOI.NLPBlock())
+  nlp_data = if MOI.NLPBlock() in MOI.get(model, MOI.ListOfModelAttributesSet())
+    MOI.get(model, MOI.NLPBlock())
+  else
+    nothing
+  end
   # New interface with `@constraint` and `@objective`
-  nlp_model = _nlp_model(model)
+  nlp_model = _nlp_model(model, ad_backend)
   vars = MOI.get(model, MOI.ListOfVariableIndices())
   if isnothing(nlp_data)
     if isnothing(nlp_model)
       evaluator =
         MOI.Nonlinear.Evaluator(MOI.Nonlinear.Model(), MOI.Nonlinear.SparseReverseMode(), vars)
       nlp_data = MOI.NLPBlockData(evaluator)
     else
-      backend = MOI.Nonlinear.SparseReverseMode()
-      evaluator = MOI.Nonlinear.Evaluator(nlp_model, backend, vars)
+      evaluator = MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars)
       nlp_data = MOI.NLPBlockData(evaluator)
     end
   else