diff --git a/Project.toml b/Project.toml index c8defac2..00299a11 100644 --- a/Project.toml +++ b/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" diff --git a/ext/NLPModelsJuMPArrayDiffExt.jl b/ext/NLPModelsJuMPArrayDiffExt.jl new file mode 100644 index 00000000..6e75443e --- /dev/null +++ b/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 diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 9ec6d2ce..ec51ef42 100644 --- a/src/MOI_wrapper.jl +++ b/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" ) diff --git a/src/moi_nlp_model.jl b/src/moi_nlp_model.jl index a4975aa3..b952a013 100644 --- a/src/moi_nlp_model.jl +++ b/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() diff --git a/src/utils.jl b/src/utils.jl index 9013716c..bf4232fd 100644 --- a/src/utils.jl +++ b/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,6 +592,11 @@ 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 @@ -557,11 +604,25 @@ function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model} 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) @@ -569,8 +630,7 @@ function _nlp_block(model::MOI.ModelLike) 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