Add support for ArrayDiff

.github/workflows/ci.yml

@@ -54,6 +54,13 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
+      - name: MOI
+        shell: julia --project=@. {0}
+        run: |
+          using Pkg
+          Pkg.add([
+              PackageSpec(name="MathOptInterface", rev="bl/linearity"),
+          ])
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1

Project.toml

@@ -11,7 +11,15 @@ 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"
+ExaModels = "0.11"
 JuMP = "1.25"
 LinearAlgebra = "1.10"
 MathOptInterface = "1.46"
@@ -25,10 +33,11 @@ Test = "1.10"
 julia = "1.10"
 
 [extras]
+ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856"
 Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["LinearAlgebra", "NLPModelsTest", "Percival", "Test"]
+test = ["ExaModels", "LinearAlgebra", "NLPModelsTest", "Percival", "Test"]

ext/NLPModelsJuMPArrayDiffExt.jl

@@ -0,0 +1,110 @@
+module NLPModelsJuMPArrayDiffExt
+
+import NLPModelsJuMP
+import ArrayDiff
+import MathOptInterface as MOI
+import NLPModels
+import LinearAlgebra
+
+NLPModelsJuMP._nonlinear_model(ad::ArrayDiff.Mode) = ArrayDiff.model(ad)
+
+NLPModelsJuMP._supports_vector_objective(::ArrayDiff.Mode) = true
+
+# 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,11 +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}
+  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
 
@@ -48,6 +56,32 @@ end
 
 MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent
 
+function MOI.supports(
+  optimizer::Optimizer,
+  ::MOI.ObjectiveFunction{F},
+) where {F <: MOI.AbstractVectorFunction}
+  return _supports_vector_objective(optimizer.ad_backend)
+end
+
+###
+### 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
 ###
@@ -89,12 +123,32 @@ 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
+  if _route_all_to_evaluator(dest.ad_backend)
+    dest.nlp, index_map = evaluator_nlp_model(src; ad_backend = dest.ad_backend)
+  else
+    dest.nlp, index_map = nlp_model(src; ad_backend = dest.ad_backend)
+  end
   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/NLPModelsJuMP.jl

@@ -3,6 +3,7 @@ module NLPModelsJuMP
 include("utils.jl")
 include("moi_nlp_model.jl")
 include("moi_nls_model.jl")
+include("moi_evaluator_model.jl")
 include("MOI_wrapper.jl")
 
 end