diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4c1c53a46..a6953e1c5 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -28,6 +28,13 @@ jobs: with: version: ${{ matrix.julia-version }} - uses: julia-actions/cache@v3 + - name: MOI + shell: julia --project=@. {0} + run: | + using Pkg + Pkg.add([ + PackageSpec(name="MathOptInterface", rev="bl/linearity"), + ]) - uses: julia-actions/julia-buildpkg@latest - name: Run tests env: diff --git a/ext/ExaModelsMOI.jl b/ext/ExaModelsMOI.jl index 33da172b2..eecdf9740 100644 --- a/ext/ExaModelsMOI.jl +++ b/ext/ExaModelsMOI.jl @@ -1,770 +1,719 @@ module ExaModelsMOI -import ExaModels: ExaModels, NLPModels, SolverCore +import ExaModels +import MathOptInterface as MOI -import MathOptInterface -const MOI = MathOptInterface -const MOIU = MathOptInterface.Utilities -const MOIB = MathOptInterface.Bridges +function __init__() + setglobal!(ExaModels, :Optimizer, Optimizer) + setglobal!(ExaModels, :SIMDMode, SIMDMode) + return +end -const SUPPORTED_FUNC_TYPE{T} = Union{ - MOI.ScalarAffineFunction{T}, - MOI.ScalarQuadraticFunction{T}, - MOI.ScalarNonlinearFunction, -} -const SUPPORTED_FUNC_TYPE_WITH_VAR{T} = Union{SUPPORTED_FUNC_TYPE{T},MOI.VariableIndex} -const SUPPORTED_FUNC_SET_TYPE{T} = - Union{MOI.GreaterThan{T},MOI.LessThan{T},MOI.EqualTo{T},MOI.Interval{T}} -const SUPPORTED_VAR_SET_TYPE{T} = - Union{MOI.GreaterThan{T},MOI.LessThan{T},MOI.EqualTo{T},MOI.Parameter{T}} -const PARAMETER_INDEX_THRESHOLD = Int64(4_611_686_018_427_387_904) # div(typemax(Int64),2)+1 -""" - Abstract data structure for storing expression tree and data arrays -""" -abstract type AbstractBin end +const PARAMETER_INDEX_THRESHOLD = div(typemax(Int64), 2) + 1 """ - struct Bin{E,P,I} <: AbstractBin + struct Bin{E,P} + head::E + data::Vector{P} + end + +This struct represents `head(d) for d in data` + +`head` will be one of two things: + + 1) `DataIndexed() => e`: this means that the generator is a constraint + augmentation. It maps the row index to an expression. The `e` is a symbolic + function. This de-duplicates structural constraints, so two constraints with + the same symbolic form will get automatically added as elements to the + vector `data`. -This linked list with `inner` represents a sum of expressions `∑_{i in data} head(i)` -It is a linked list and not a `Vector` as each `head` may have a different type. -We append new ones at the beginning because `Bin` is non-mutable and -its fields are concretely typed. + 2) `e::ExaModels.AbstractNode`: this means that the generator is part of a + summation. This is used for the objective function. """ -struct Bin{E,P,I} <: AbstractBin +struct Bin{E,P} head::E - data::P - inner::I + data::Vector{P} end -struct BinNull <: AbstractBin end +""" + update_bin!(bin::Vector{Bin}, head, data) -function update_bin!(bin, e, p) - if _update_bin!(bin, e, p) # if update succeeded, return the original bin - return bin - else # if update has failed, return a new bin - return Bin(e, [p], bin) +This function loops thorugh the list of `bin` looking for a matching `head`. If +found, it updates the bin in place. Othersise, it appends a new bin. +""" +function update_bin!( + bins::Vector{Bin}, + head::Union{ + ExaModels.AbstractNode, + Pair{<:ExaModels.AbstractNode,<:ExaModels.AbstractNode}, + }, + data::Tuple, +) + for bin in bins + if _update_bin!(bin, head, data) + return + end end + push!(bins, Bin(head, [data])) + return bins end -function _update_bin!(bin::Bin{E,P,I}, e, p) where {E,P,I} - if e == bin.head && p isa eltype(bin.data) - push!(bin.data, p) + +# The types match for `head(data)` to be appended to this bin. We check if the +# head's match with `==`, otherwise we recurse to the next bin. +function _update_bin!(bin::Bin{E,P}, head::E, data::P) where {E,P} + if head == bin.head + push!(bin.data, data) return true - else - return _update_bin!(bin.inner, e, p) end -end -function _update_bin!(::BinNull, e, p) return false end -function check_supported(T, moim) - con_types = MOI.get(moim, MOI.ListOfConstraintTypesPresent()) - for (F, S) in con_types - !(F <: SUPPORTED_FUNC_TYPE_WITH_VAR) && error("Unsupported function type $F.") - if F <: MOI.VariableIndex - !(S <: SUPPORTED_VAR_SET_TYPE) && - error("Unsupported variable index constraint $F in $S") - else - !(S <: SUPPORTED_FUNC_SET_TYPE) && error("Unsupported set type $S") - end - end +# The head does not match. We can't update this bin in-place. +_update_bin!(::Bin, ::Any, ::Any) = false - obj_type = MOI.get(moim, MOI.ObjectiveFunctionType()) - !(obj_type <: SUPPORTED_FUNC_TYPE_WITH_VAR) && - error("Unsupported objective function type $obj_type.") +# A method for the objective function. First convert the MOI function `f` into +# an `ExaModels.AbstractNode`, then add that. +function update_bin!(bins::Vector{Bin}, f) + head, data = _exafy(f, (), nothing) + return update_bin!(bins, head, data) +end - obj_sense = MOI.get(moim, MOI.ObjectiveSense()) - !(obj_sense in (MOI.MIN_SENSE, MOI.MAX_SENSE)) && - error("Unsupported objective sense $obj_sense.") - return obj_sense === MOI.MIN_SENSE +# A method for adding to a constraint. First convert the MOI function `f` into +# an `ExaModels.AbstractNode`, then add that. +function update_bin!(bins::Vector{Bin}, (row, f)::Pair{Int,F}) where {F} + head, data = _exafy(f, (), nothing) + e = ExaModels.DataIndexed(ExaModels.DataSource(), length(data) + 1) + return update_bin!(bins, e => head, (data..., row)) end -function ExaModels.ExaModel( - moim::MOI.ModelLike; - backend = nothing, - prod = false, - T = ExaModels.default_T(backend), - ) +# This is a type that lets us dispatch on the difference between `row => expr` +# and `expr`. +abstract type AbstractBin end - c, _ = to_exacore(moim; backend = backend, T = T) - return ExaModels.ExaModel(c; prod = prod) -end +# This is an `expr`. It gets added to the objective. +struct ObjectiveBin <:AbstractBin end -function to_exacore(moim::MOI.ModelLike; backend = nothing, T = Float64) - minimize = check_supported(T, moim) +# Things are passed through unchanged. +(::ObjectiveBin)(f) = f - c = ExaModels.ExaCore(T; backend = backend, minimize = minimize, concrete = Val(true)) +# Except objective constants are passed as `Null`. +(::ObjectiveBin)(f::Real) = ExaModels.Null(f) - c, var_to_idx = copy_variables!(c, moim, T) - c, con_to_idx = copy_constraints!(c, moim, var_to_idx, T) - c = copy_objective!(c, moim, var_to_idx) +# This is a `row => expr`. We keep the row in a closure. +struct ConstraintBin <: AbstractBin + row::Int +end - return c, (var_to_idx, con_to_idx) +# When passing, we convert to a pair. +(bin::ConstraintBin)(f) = bin.row => f + +# VariableIndices are handled directly. +function update_bin!( + bins::Vector{Bin}, + fn::AbstractBin, + f::Union{Real,MOI.VariableIndex}, +) + return update_bin!(bins, fn(f)) end -function fill_variable_bounds!(moim, lvar, uvar, var_to_idx, T) - for ci in - MOI.get(moim, MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.GreaterThan{T}}()) - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - lvar[var_to_idx[vi]] = MOI.get(moim, MOI.ConstraintSet(), ci).lower - end - for ci in - MOI.get(moim, MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.LessThan{T}}()) - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - uvar[var_to_idx[vi]] = MOI.get(moim, MOI.ConstraintSet(), ci).upper +# Add the additive terms separately, instead of creating a single +(args...) +# expression. +function update_bin!( + bins::Vector{Bin}, + fn::AbstractBin, + f::MOI.ScalarAffineFunction, +) + for term in f.terms + update_bin!(bins, fn(term)) end - for ci in MOI.get(moim, MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.EqualTo{T}}()) - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - fixed_val = MOI.get(moim, MOI.ConstraintSet(), ci).value - lvar[var_to_idx[vi]] = fixed_val - uvar[var_to_idx[vi]] = fixed_val + if !iszero(f.constant) + update_bin!(bins, fn(f.constant)) end + return bins end -function fill_variable_start!(moim, x0, param_vis) - var_to_idx = Dict{MOI.VariableIndex,Int}() - i = 0 - for vi in MOI.get(moim, MOI.ListOfVariableIndices()) - vi ∈ param_vis && continue - i += 1 - var_to_idx[vi] = i - start = if MOI.supports(moim, MOI.VariablePrimalStart(), typeof(vi)) - MOI.get(moim, MOI.VariablePrimalStart(), vi) - else - nothing - end - isnothing(start) && continue - x0[i] = start +# Add the additive terms separately, instead of creating a single +(args...) +# expression. +function update_bin!( + bins::Vector{Bin}, + fn::AbstractBin, + f::MOI.ScalarQuadraticFunction, +) + for term in f.affine_terms + update_bin!(bins, fn(term)) end - return var_to_idx -end - -function _get_parameters(moim::MOI.ModelLike, T) - cis = MOI.get(moim, MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{T}}()) - parameters = Vector{Tuple{MOI.VariableIndex,MOI.Parameter{T}}}() - for ci in cis - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - set = MOI.get(moim, MOI.ConstraintSet(), ci) - push!(parameters, (vi, set)) + for term in f.quadratic_terms + update_bin!(bins, fn(term)) end - sort!(parameters, by = x -> x[1].value) - return parameters + if !iszero(f.constant) + update_bin!(bins, fn(f.constant)) + end + return bins end +_is_zero(x::Real) = iszero(x) -function copy_variables!(c, moim, T) - nvarpar = MOI.get(moim, MOI.NumberOfVariables()) - parameters = _get_parameters(moim, T) - npar = length(parameters) - nvar = nvarpar - npar - - x0 = zeros(T, nvar) - var_to_idx = fill_variable_start!(moim, x0, first.(parameters)) +_is_zero(::Any) = false - lvar = fill(T(-Inf), nvar) - uvar = fill(T(Inf), nvar) - fill_variable_bounds!(moim, lvar, uvar, var_to_idx, T) - - c, _ = ExaModels.add_var(c, nvar; start = x0, lvar = lvar, uvar = uvar) - - varpar_to_idx = Dict() - for (vi, i) in var_to_idx - varpar_to_idx[vi] = (type = :variable, idx = i) +function update_bin!( + bins::Vector{Bin}, + fn::AbstractBin, + f::MOI.ScalarNonlinearFunction, +) + if f.head == :- && length(f.args) == 2 + # Optimization: :(x - y) -> :(+(x, -y)) + # This allows additive terms in the left-hand side to be added + # separately. This is a common case in JuMP because + # `@constraint(model, lhs <= rhs)` normalizes to `lhs - rhs <= 0`. + update_bin!(bins, fn, f.args[1]) + if !_is_zero(f.args[2]) + rhs = MOI.Utilities.operate(-, Float64, f.args[2]) + update_bin!(bins, fn(rhs)) + end + return bins + elseif f.head != :+ + return update_bin!(bins, fn(f)) end - - if npar > 0 - p0 = zeros(T, npar) - for (i, (vi, set)) in enumerate(parameters) - p0[i] = T(set.value) - varpar_to_idx[vi] = (type = :parameter, idx = i) + # Optimization: if the expression is a `:+`, add the child arguments as + # separate terms. This keeps the size of the expressions small for ExaModels. + constant = 0.0 + for arg in f.args + if arg isa MOI.ScalarAffineFunction + for term in arg.terms + update_bin!(bins, fn(term)) + end + constant += arg.constant + elseif arg isa MOI.ScalarQuadraticFunction + for term in arg.affine_terms + update_bin!(bins, fn(term)) + end + for term in arg.quadratic_terms + update_bin!(bins, fn(term)) + end + constant += arg.constant + else + # This is NOT fn(arg) here because we want to be able to lift any + # nested `+(+(args...), args....)`. + update_bin!(bins, fn, arg) end - c, _ = ExaModels.add_par(c, p0) end - - return c, varpar_to_idx + if !iszero(constant) + update_bin!(bins, fn(constant)) + end + return bins end -function copy_objective!(c, moim, var_to_idx) - obj_type = MOI.get(moim, MOI.ObjectiveFunctionType()) +# _exafy - bin = BinNull() - bin = exafy_obj(MOI.get(moim, MOI.ObjectiveFunction{obj_type}()), bin, var_to_idx) +# This method is used for objective constants. +_exafy(f::ExaModels.Null, data::Tuple, ::Any) = f, data - return build_objective!(c, bin) +# This method is used when a constant appears in a function. +function _exafy(f::Real, data::Tuple, ::Any) + e = ExaModels.DataIndexed(ExaModels.DataSource(), length(data) + 1) + return e, (data..., f) end -function copy_constraints!(c, moim, var_to_idx, T) - bin = BinNull() - offset = 0 - lcon = zeros(T, 0) - ucon = zeros(T, 0) - y0 = zeros(T, 0) - con_to_idx = Dict{MOI.ConstraintIndex,Int}() - - con_types = MOI.get(moim, MOI.ListOfConstraintTypesPresent()) - for (F, S) in con_types - cis = MOI.get(moim, MOI.ListOfConstraintIndices{F,S}()) - if F <: MOI.VariableIndex - for ci in cis - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - vartype, var_idx = var_to_idx[vi] - if vartype === :variable - con_to_idx[ci] = var_idx - end - end - continue - end - bin, offset = - exafy_con(moim, cis, bin, offset, lcon, ucon, y0, var_to_idx, con_to_idx) - end - c, cons = ExaModels.add_con(c, offset; start = y0, lcon = lcon, ucon = ucon) - c = build_constraint!(c, cons, bin) - - return c, con_to_idx -end - -function _exafy_con( - i, - c::C, - bin, - var_to_idx, - con_to_idx; - pos = true, -) where {C<:MOI.ScalarAffineFunction} - for mm in c.terms - e, p = _exafy(mm, var_to_idx) - e = pos ? e : -e - bin = update_bin!( - bin, - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) => e, - (p..., con_to_idx[i]), - ) # augment data with constraint index - end - bin = update_bin!(bin, ExaModels.Null(c.constant), (1,)) - return bin -end -function _exafy_con( - i, - c::C, - bin, - var_to_idx, - con_to_idx; - pos = true, -) where {C<:MOI.ScalarQuadraticFunction} - for mm in c.affine_terms - e, p = _exafy(mm, var_to_idx) - e = pos ? e : -e - bin = update_bin!( - bin, - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) => e, - (p..., con_to_idx[i]), - ) # augment data with constraint index - end - for mm in c.quadratic_terms - e, p = _exafy(mm, var_to_idx) - e = pos ? e : -e - bin = update_bin!( - bin, - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) => e, - (p..., con_to_idx[i]), - ) # augment data with constraint index - end - bin = update_bin!(bin, ExaModels.Null(c.constant), (1,)) - return bin -end -function _exafy_con( - i, - c::C, - bin, - var_to_idx, - con_to_idx; - pos = true, -) where {C<:MOI.ScalarNonlinearFunction} - if c.head == :+ - for mm in c.args - bin = _exafy_con(i, mm, bin, var_to_idx, con_to_idx) - end - # elseif c.head == :- - # bin, offset = _exafy_con(i, c.args[1], bin, offset) - # bin, offset = _exafy_con(i, c.args[2], bin, offset; pos = false) - else - e, p = _exafy(c, var_to_idx) - e = pos ? e : -e - bin = update_bin!( - bin, - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) => e, - (p..., con_to_idx[i]), - ) # augment data with constraint index - end - return bin -end -function _exafy_con(i, c::C, bin, var_to_idx, con_to_idx; pos = true) where {C<:Real} - e = - pos ? ExaModels.DataIndexed(ExaModels.DataSource(), 1) : - -ExaModels.DataIndexed(ExaModels.DataSource(), 1) - bin = update_bin!( - bin, - ExaModels.DataIndexed(ExaModels.DataSource(), 2) => 0 * ExaModels.Var(1) + e, - (c, con_to_idx[i]), - ) - - return bin -end - -function exafy_con( - moim, - cons::V, - bin, - offset, - lcon, - ucon, - y0, - var_to_idx, - con_to_idx, -) where {V<:Vector{<:MOI.ConstraintIndex}} - l = length(cons) - - resize!(lcon, offset + l) - resize!(ucon, offset + l) - resize!(y0, offset + l) - for (i, ci) in enumerate(cons) - func = MOI.get(moim, MOI.ConstraintFunction(), ci) - set = MOI.get(moim, MOI.ConstraintSet(), ci) - con_to_idx[ci] = offset + i - start = if MOI.supports( - moim, MOI.ConstraintPrimalStart(), typeof(ci) - ) - MOI.get(moim, MOI.ConstraintPrimalStart(), ci) - else - nothing +function _exafy( + f::MOI.VariableIndex, + data::Tuple, + var_to_data::Union{Nothing,Dict{Int,Int}}, +) + if f.value > PARAMETER_INDEX_THRESHOLD + e = ExaModels.DataIndexed(ExaModels.DataSource(), length(data) + 1) + idx = f.value - PARAMETER_INDEX_THRESHOLD + return ExaModels.ParameterNode(e), (data..., idx) + end + if var_to_data !== nothing + # An optimization: the variable `f` may already appear in the tuple + # `data`. If so, we want to re-use the slot instead of appending a new + # element to `data`. (ExaModels could be clever here and check for + # duplicates.) + # + # We don't have this optimization for ParameterNode's because the main + # problem with duplicates are they they show up as duplicated elements + # in the Jacobian and Hessian. + if (pidx = get(var_to_data, f.value, nothing)) !== nothing + p_cache = ExaModels.DataIndexed(ExaModels.DataSource(), pidx) + return ExaModels.Var(p_cache), data end - _exafy_con_update_start(ci, start, y0, con_to_idx) - _exafy_con_update_vector(ci, set, lcon, ucon, con_to_idx) - bin = _exafy_con(ci, func, bin, var_to_idx, con_to_idx) + var_to_data[f.value] = length(data) + 1 end - return bin, (offset += l) + e = ExaModels.DataIndexed(ExaModels.DataSource(), length(data) + 1) + return ExaModels.Var(e), (data..., f.value) end -function _exafy_con_update_start(i, start, y0, con_to_idx) - y0[con_to_idx[i]] = start +function _exafy( + f::MOI.ScalarAffineTerm, + data::Tuple, + var_to_data::Union{Nothing,Dict{Int,Int}}, +) + x_head, data = _exafy(f.variable, data, var_to_data) + c_head, data = _exafy(f.coefficient, data, var_to_data) + return c_head * x_head, data end -function _exafy_con_update_start(i, ::Nothing, y0, con_to_idx) - y0[con_to_idx[i]] = zero(eltype(y0)) +# This method is used when a ScalarAffineFunction appears inside a +# ScalarNonlinearFunction. For that reason we don't do anything clever with the +# additive terms. +function _exafy( + f::MOI.ScalarAffineFunction, + data::Tuple, + var_to_data::Union{Nothing,Dict{Int,Int}}, +) + head, data = _exafy(f.constant, data, var_to_data) + if !isempty(f.terms) + y = sum(begin + c1, data = _exafy(term, data, var_to_data) + c1 + end for term in f.terms) + head += y + end + return head, data end -function _exafy_con_update_vector(i, e::MOI.Interval{T}, lcon, ucon, con_to_idx) where {T} - lcon[con_to_idx[i]] = e.lower - ucon[con_to_idx[i]] = e.upper +function _exafy( + f::MOI.ScalarQuadraticTerm, + data::Tuple, + var_to_data::Union{Nothing,Dict{Int,Int}}, +) + if f.variable_1 == f.variable_2 + x_head, data = _exafy(f.variable_1, data, var_to_data) + c_head, data = _exafy(f.coefficient / 2, data, var_to_data) + return c_head * abs2(x_head), data + end + x1_head, data = _exafy(f.variable_1, data, var_to_data) + x2_head, data = _exafy(f.variable_2, data, var_to_data) + c_head, data = _exafy(f.coefficient, data, var_to_data) + return c_head * x1_head * x2_head, data end -function _exafy_con_update_vector(i, e::MOI.LessThan{T}, lcon, ucon, con_to_idx) where {T} - lcon[con_to_idx[i]] = -Inf - ucon[con_to_idx[i]] = e.upper +# This method is used when a ScalarQuadraticFunction appears inside a +# ScalarNonlinearFunction. For that reason we don't do anything clever with the +# additive terms. +function _exafy( + f::MOI.ScalarQuadraticFunction, + data::Tuple, + var_to_data::Union{Nothing,Dict{Int,Int}}, +) + head, data = _exafy(f.constant, data, var_to_data) + if !isempty(f.affine_terms) + head += sum(begin + c1, data = _exafy(term, data, var_to_data) + c1 + end for term in f.affine_terms) + end + if !isempty(f.quadratic_terms) + head += sum(begin + c1, data = _exafy(term, data, var_to_data) + c1 + end for term in f.quadratic_terms) + end + return head, data end -function _exafy_con_update_vector( - i, - e::MOI.GreaterThan{T}, - lcon, - ucon, - con_to_idx, -) where {T} - ucon[con_to_idx[i]] = Inf - lcon[con_to_idx[i]] = e.lower +function _exafy(f::MOI.ScalarNonlinearFunction, data::Tuple, ::Nothing) + # Replace the incoming `var_to_data === nothing` with a dictionary that maps + # the variable index with the element in `data`. This is used when there are + # repeated variable indices in `f`. See `_exafy(::VariableIndex, args...)`. + return _exafy(f, data, Dict{Int,Int}()) end -function _exafy_con_update_vector(i, e::MOI.EqualTo{T}, lcon, ucon, con_to_idx) where {T} - lcon[con_to_idx[i]] = e.value - ucon[con_to_idx[i]] = e.value +function _exafy( + f::MOI.ScalarNonlinearFunction, + data::Tuple, + var_to_data::Dict{Int,Int}, +) + # This assumes that we support only the default functions in `MOI.Nonlinear` + op = getfield(MOI.Nonlinear, f.head) + if length(f.args) == 1 + # A special case when there is one argument. + arg, data = _exafy(only(f.args), data, var_to_data) + return op(arg), data + elseif length(f.args) == 2 + # A special case when there are two arguments + arg1, data = _exafy(f.args[1], data, var_to_data) + arg2, data = _exafy(f.args[2], data, var_to_data) + return op(arg1, arg2), data + end + args = () + for arg in f.args + head, data = _exafy(arg, data, var_to_data) + args = (args..., head) + end + return op(args...), data end +""" + ExaModels.ExaModel( + src::MOI.ModelLike; + backend = nothing, + prod::Bool = false, + T = ExaModels.default_T(backend), + ) -function build_constraint!(c, cons, bin) - c = build_constraint!(c, cons, bin.inner) - c, _ = ExaModels.add_con!(c, cons, Base.Generator(_ -> bin.head, bin.data)) - return c +Convert `src` to an `ExaModel`. +""" +function ExaModels.ExaModel( + src::MOI.ModelLike; + backend = nothing, + prod = false, + T = ExaModels.default_T(backend), +) + dest = Optimizer{T}(nothing; backend) + MOI.copy_to(dest, src) + c = to_exacore(dest, backend) + return ExaModels.ExaModel(c; prod) end -function build_constraint!(c, cons, ::BinNull) - return c -end +# Now comes the MOI interface -function build_objective!(c, bin) - c = build_objective!(c, bin.inner) - c, _ = ExaModels.add_obj(c, bin.head, bin.data) - return c +""" + ExaModels.Optimizer(solver, backend = nothing) + +Create a new ExaModels.Optimizer object. + +## Examples + +```julia-repl +julia> import ExaModels, NLPModelsIpopt, KernelAbstractions + +julia> optimizer = () -> ExaModels.Optimizer(NLPModelsIpopt.ipopt); + +julia> optimizer = () -> ExaModels.Optimizer(NLPModelsIpopt.ipopt, KernelAbstractions.CPU()); +``` +""" +mutable struct Optimizer{T} <: MOI.AbstractOptimizer + solver::Any + backend::Any + result::Any + solve_time::Float64 + options::Dict{Symbol,Any} + # Problem cache + sense::MOI.OptimizationSense + lvar::Vector{T} + uvar::Vector{T} + startvar::Vector{T} + pstart::Vector{T} + lcon::Vector{T} + ucon::Vector{T} + objs::Vector{Bin} + cons::Vector{Bin} + + function Optimizer{T}(solver, backend = nothing; kwargs...) where {T} + return new( + solver, + backend, + nothing, + 0.0, + Dict{Symbol,Any}(kwargs...), + # Problem cache + MOI.FEASIBILITY_SENSE, + T[], + T[], + T[], + T[], + T[], + T[], + Bin[], + Bin[], + ) + end end -function build_objective!(c, ::BinNull) - return c +function Optimizer(solver, backend = nothing; kwargs...) + T = ExaModels.default_T(backend) + return Optimizer{T}(solver, backend; kwargs...) end -function exafy_obj(o::Nothing, bin, var_to_idx) - return bin +function MOI.empty!(model::ExaModelsMOI.Optimizer) + model.result = nothing + model.solve_time = 0.0 + model.sense = MOI.FEASIBILITY_SENSE + empty!(model.lvar) + empty!(model.uvar) + empty!(model.startvar) + empty!(model.pstart) + empty!(model.lcon) + empty!(model.ucon) + empty!(model.cons) + empty!(model.objs) + return end -function exafy_obj(o::MOI.VariableIndex, bin, var_to_idx) - e, p = _exafy(o, var_to_idx) - return update_bin!(bin, e, p) +function MOI.is_empty(model::Optimizer) + return isempty(model.lvar) && + isempty(model.pstart) && + isempty(model.cons) && + isempty(model.objs) end -function exafy_obj(o::MOI.ScalarQuadraticFunction{T}, bin, var_to_idx) where {T} - for m in o.affine_terms - e, p = _exafy(m, var_to_idx) - bin = update_bin!(bin, e, p) - end - for m in o.quadratic_terms - e, p = _exafy(m, var_to_idx) - bin = update_bin!(bin, e, p) - end +# MOI.ObjectiveSense - return update_bin!(bin, ExaModels.Null(o.constant), (1,)) -end +MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true -function exafy_obj(o::MOI.ScalarAffineFunction{T}, bin, var_to_idx) where {T} - for m in o.terms - e, p = _exafy(m, var_to_idx) - bin = update_bin!(bin, e, p) - end +MOI.get(model::Optimizer, ::MOI.ObjectiveSense) = model.sense - return update_bin!(bin, ExaModels.Null(o.constant), (1,)) +function MOI.set( + model::Optimizer, + ::MOI.ObjectiveSense, + sense::MOI.OptimizationSense, +) + model.sense = sense + if sense == MOI.FEASIBILITY_SENSE + empty!(model.objs) + end + return end -function exafy_obj(o::MOI.ScalarNonlinearFunction, bin, var_to_idx) - constant = 0.0 - if o.head == :+ - for m in o.args - if m isa MOI.ScalarAffineFunction - for mm in m.terms - e, p = _exafy(mm, var_to_idx) - bin = update_bin!(bin, e, p) - end - elseif m isa MOI.ScalarQuadraticFunction - for mm in m.affine_terms - e, p = _exafy(mm, var_to_idx) - bin = update_bin!(bin, e, p) - end - for mm in m.quadratic_terms - e, p = _exafy(mm, var_to_idx) - bin = update_bin!(bin, e, p) - end - constant += m.constant - else - e, p = _exafy(m, var_to_idx) - bin = update_bin!(bin, e, p) - end - end - else - e, p = _exafy(o, var_to_idx) - bin = update_bin!(bin, e, p) - end +# MOI.ObjectiveFunction + +function MOI.supports( + ::Optimizer{T}, + ::MOI.ObjectiveFunction{F}, +) where { + T, + F<:Union{ + MOI.VariableIndex, + MOI.ScalarAffineFunction{T}, + MOI.ScalarQuadraticFunction{T}, + MOI.ScalarNonlinearFunction, + }, +} + return true +end - return update_bin!(bin, ExaModels.Null(constant), (1,)) # TODO see if this can be empty tuple +function MOI.set( + model::Optimizer{T}, + ::MOI.ObjectiveFunction{F}, + f::F, +) where { + T, + F<:Union{ + MOI.VariableIndex, + MOI.ScalarAffineFunction{T}, + MOI.ScalarQuadraticFunction{T}, + MOI.ScalarNonlinearFunction, + }, +} + empty!(model.objs) + update_bin!(model.objs, ObjectiveBin(), f) + return end -function _exafy(v::MOI.VariableIndex, var_to_idx, p = ()) - i = ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) - vartype, idx = var_to_idx[v] - if vartype === :variable - return ExaModels.Var(i), (p..., idx) - elseif vartype === :parameter - return ExaModels.ParameterNode(i), (p..., idx) - else - error("Unknown variable type: $vartype") - end +# MOI.add_variable + +function MOI.add_variable(model::Optimizer{T}) where {T} + push!(model.lvar, typemin(T)) + push!(model.uvar, typemax(T)) + push!(model.startvar, zero(T)) + return MOI.VariableIndex(length(model.lvar)) end -function _exafy(i::R, var_to_idx, p) where {R<:Real} - return ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1), (p..., i) +# MOI.add_constrained_variable + +function MOI.supports_add_constrained_variable( + ::Optimizer{T}, + ::Type{MOI.Parameter{T}}, +) where {T} + return true end -function _exafy(e::MOI.ScalarNonlinearFunction, var_to_idx, p = ()) - return op(e.head)((begin - c, p = _exafy(e, var_to_idx, p) - c - end for e in e.args)...), p +function MOI.add_constrained_variable( + model::Optimizer{T}, + set::MOI.Parameter{T}, +) where {T} + push!(model.pstart, set.value) + index = PARAMETER_INDEX_THRESHOLD + length(model.pstart) + ci = MOI.ConstraintIndex{MOI.VariableIndex,typeof(set)}(index) + return MOI.VariableIndex(index), ci end -function _exafy(e::MOI.ScalarAffineFunction{T}, var_to_idx, p = ()) where {T} - ec = if !isempty(e.terms) - sum(begin - c1, p = _exafy(term, var_to_idx, p) - c1 - end for term in e.terms) + - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) - else - ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) - end +# VariableIndex-in-Set constraints - return ec, (p..., e.constant) +function MOI.supports_constraint( + ::Optimizer{T}, + ::Type{MOI.VariableIndex}, + ::Type{S}, +) where { + T, + S<:Union{ + MOI.GreaterThan{T}, + MOI.LessThan{T}, + MOI.EqualTo{T}, + MOI.Interval{T}, + }, +} + return true end -function _exafy(e::MOI.ScalarAffineTerm{T}, var_to_idx, p = ()) where {T} - c1, p = _exafy(e.variable, var_to_idx, p) - return *(c1, ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1)), - (p..., e.coefficient) +function _update_bound(model::Optimizer, col::Int, set::MOI.GreaterThan) + model.lvar[col] = set.lower + return end -function _exafy(e::MOI.ScalarQuadraticFunction{T}, var_to_idx, p = ()) where {T} - t = ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) - p = (p..., e.constant) - - if !isempty(e.affine_terms) - t += sum(begin - c1, p = _exafy(term, var_to_idx, p) - c1 - end for term in e.affine_terms) - end +function _update_bound(model::Optimizer, col::Int, set::MOI.LessThan) + model.uvar[col] = set.upper + return +end - if !isempty(e.quadratic_terms) - t += sum(begin - c1, p = _exafy(term, var_to_idx, p) - c1 - end for term in e.quadratic_terms) - end - - return t, p -end - -function _exafy(e::MOI.ScalarQuadraticTerm{T}, var_to_idx, p = ()) where {T} - - if e.variable_1 == e.variable_2 - v, p = _exafy(e.variable_1, var_to_idx, p) - return ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) * abs2(v), - (p..., e.coefficient / 2) # it seems that MOI assumes this by default - else - v1, p = _exafy(e.variable_1, var_to_idx, p) - v2, p = _exafy(e.variable_2, var_to_idx, p) - return ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) * v1 * v2, - (p..., e.coefficient) - end -end - -# eval can be a performance killer -- we want to explicitly include symbols for frequently used operations. -function op(s::Symbol) - # uni/multi - if s === :+ - return + - elseif s === :- - return - - # multi - elseif s === :* - return * - elseif s === :^ - return ^ - elseif s === :/ - return / - # uni - elseif s === :abs - return abs - elseif s === :sign - error("sign not supported") - elseif s === :sqrt - return sqrt - elseif s === :cbrt - return cbrt - elseif s === :abs2 - return abs2 - elseif s === :inv - return inv - elseif s === :log - return log - elseif s === :log10 - return log10 - elseif s === :log2 - return log2 - elseif s === :log1p - return log1p - elseif s === :exp - return exp - elseif s === :exp2 - return exp2 - elseif s === :expm1 - error("expm1 not supported") - # trig - elseif s === :sin - return sin - elseif s === :cos - return cos - elseif s === :tan - return tan - elseif s === :sec - return sec - elseif s === :csc - return csc - elseif s === :cot - return cot - elseif s === :sind - return sind - elseif s === :cosd - return cosd - elseif s === :tand - return tand - elseif s === :secd - return secd - elseif s === :cscd - return cscd - elseif s === :cotd - return cotd - elseif s === :asin - return asin - elseif s === :acos - return acos - elseif s === :atan - return atan - elseif s === :asec - error("asec not supported") - elseif s === :acsc - error("acsc not supported") - elseif s === :acot - return acot - elseif s === :asind - error("asind not supported") - elseif s === :acosd - error("acosd not supported") - elseif s === :atand - return atand - elseif s === :asecd - error("aced not supported") - elseif s === :acscd - error("acscd not supported") - elseif s === :acotd - return acotd - elseif s === :sinh - return sinh - elseif s === :cosh - return cosh - elseif s === :tanh - return tanh - elseif s === :sech - return sech - elseif s === :csch - return csch - elseif s === :coth - return coth - elseif s === :asinh - return asinh - elseif s === :acosh - return acosh - elseif s === :atanh - return atanh - elseif s === :asech - error("asech not supported") - elseif s === :acsch - error("acsch not supported") - elseif s === :acoth - return acoth - # special (commented will use `eval` which would succeed if SpecialFunctions is loaded) - elseif s === :deg2rad - error("deg2rad not supported") - elseif s === :rad2deg - error("rad2deg not supported") - # elseif s === :erf error("erf not supported") - # elseif s === :erfinv error("erfinv not supported") - # elseif s === :erfc error("erfc not supported") - # elseif s === :erfcinv error("erfcinv not supported") - # elseif s === :erfi error("erfi not supported") - # elseif s === :gamma error("gamma not supported") - elseif s === :lgamma - error("lgamma not supported") - # elseif s === :digamma error("digamma not supported") - # elseif s === :invdigamma error("invdigamma not supported") - # elseif s === :trigamma error("trigamma not supported") - # elseif s === :airyai error("airyai not supported") - # elseif s === :airybi error("airybi not supported") - # elseif s === :airyaiprime error("airyaiprime not supported") - # elseif s === :airybiprime error("airybiprime not supported") - # elseif s === :besselj0 error("besselj0 not supported") - # elseif s === :besselj1 error("besselj1 not supported") - # elseif s === :bessely0 error("bessely0 not supported") - # elseif s === :bessely1 error("bessely1 not supported") - # elseif s === :erfcx error("erfcx not supported") - # elseif s === :dawson error("dawson not supported") - - # not in MOI - elseif s === :exp10 - return exp10 - elseif s === :beta - return beta - elseif s === :logbeta - return logbeta - else - return eval(s) - end -end - - -# struct EmptyOptimizer{B} -# backend::B -# end -mutable struct Optimizer{B,S} <: MOI.AbstractOptimizer - solver::S - backend::B - model::Union{Nothing,ExaModels.ExaModel} - result::Any - solve_time::Float64 - options::Dict{Symbol,Any} +function _update_bound(model::Optimizer, col::Int, set::MOI.EqualTo) + model.lvar[col] = model.uvar[col] = set.value + return end -MOI.is_empty(model::Optimizer) = isnothing(model.model) +function _update_bound(model::Optimizer, col::Int, set::MOI.Interval) + model.lvar[col], model.uvar[col] = set.lower, set.upper + return +end -function MOI.supports_constraint( - ::Optimizer, - ::Type{<:SUPPORTED_FUNC_TYPE}, - ::Type{<:SUPPORTED_FUNC_SET_TYPE}, -) - return true +function MOI.add_constraint( + model::Optimizer{T}, + f::MOI.VariableIndex, + s::Union{ + MOI.GreaterThan{T}, + MOI.LessThan{T}, + MOI.EqualTo{T}, + MOI.Interval{T}, + }, +) where {T} + @assert f.value < PARAMETER_INDEX_THRESHOLD + _update_bound(model, f.value, s) + return MOI.ConstraintIndex{typeof(f),typeof(s)}(f.value) end -function MOI.supports_constraint( + +# MOI.VariablePrimalStart + +function MOI.supports( ::Optimizer, + ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}, - ::Type{<:SUPPORTED_VAR_SET_TYPE}, ) return true end -function MOI.supports(::Optimizer, ::MOI.ObjectiveSense) - return true -end -function MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{<:SUPPORTED_FUNC_TYPE_WITH_VAR}) - return true + +function MOI.set( + model::Optimizer{T}, + ::MOI.VariablePrimalStart, + x::MOI.VariableIndex, + value::Union{Nothing,T}, +) where {T} + @assert x.value < PARAMETER_INDEX_THRESHOLD + model.startvar[x.value] = something(value, zero(T)) + return end -function MOI.supports(::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}) + +# Function-in-Set constraints + +function MOI.supports_constraint( + ::Optimizer{T}, + ::Type{F}, + ::Type{S}, +) where { + T, + F<:Union{ + MOI.ScalarAffineFunction{T}, + MOI.ScalarQuadraticFunction{T}, + MOI.ScalarNonlinearFunction, + }, + S<:Union{ + MOI.GreaterThan{T}, + MOI.LessThan{T}, + MOI.EqualTo{T}, + MOI.Interval{T}, + }, +} return true end -function ExaModels.Optimizer(solver, backend = nothing; kwargs...) - return Optimizer(solver, backend, nothing, nothing, 0.0, Dict{Symbol,Any}(kwargs...)) +_bounds(s::MOI.Interval) = (s.lower, s.upper) + +_bounds(s::MOI.EqualTo) = (s.value, s.value) + +_bounds(s::MOI.GreaterThan{T}) where {T} = (s.lower, typemax(T)) + +_bounds(s::MOI.LessThan{T}) where {T} = (typemin(T), s.upper) + +function MOI.add_constraint( + model::Optimizer{T}, + f::Union{ + MOI.ScalarAffineFunction{T}, + MOI.ScalarQuadraticFunction{T}, + MOI.ScalarNonlinearFunction, + }, + s::Union{ + MOI.GreaterThan{T}, + MOI.LessThan{T}, + MOI.EqualTo{T}, + MOI.Interval{T}, + }, +) where {T} + row = length(model.lcon) + 1 + update_bin!(model.cons, ConstraintBin(row), f) + l, u = _bounds(s) + push!(model.lcon, l) + push!(model.ucon, u) + return MOI.ConstraintIndex{typeof(f),typeof(s)}(row) end -function MOI.empty!(model::ExaModelsMOI.Optimizer) - model.model = nothing +function to_exacore(model::Optimizer{T}, backend) where {T} + c = ExaModels.ExaCore( + T; + backend, + minimize = model.sense != MOI.MAX_SENSE, + concrete = Val(true), + ) + if !isempty(model.pstart) + c, _ = ExaModels.add_par(c, model.pstart) + end + c, _ = ExaModels.add_var( + c, + length(model.lvar); + start = model.startvar, + lvar = model.lvar, + uvar = model.uvar, + ) + if !isempty(model.cons) + c, cons = ExaModels.add_con(c, length(model.lcon); model.lcon, model.ucon) + for bin in model.cons + c, _ = ExaModels.add_con!(c, cons, (bin.head for _ in bin.data)) + end + end + for bin in model.objs + c, _ = ExaModels.add_obj(c, bin.head, bin.data) + end + return c end +# MOI.copy_to + +MOI.supports_incremental_interface(::Optimizer) = true + function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike) - core, maps = to_exacore(src; backend = dest.backend) - dest.model = ExaModels.ExaModel(core; prod = true) - - return _make_index_map(src, maps) -end - -function MOI.optimize!(optimizer::Optimizer) - optimizer.solve_time = @elapsed begin - result = optimizer.solver(optimizer.model; optimizer.options...) - optimizer.result = ( - objective = result.objective, - solution = Array(result.solution), - multipliers = Array(result.multipliers), - multipliers_L = Array(result.multipliers_L), - multipliers_U = Array(result.multipliers_U), - status = result.status, - ) - end + return MOI.Utilities.default_copy_to(dest, src) +end - return optimizer +# MOI.optimize! + +function MOI.optimize!(model::Optimizer) + core = to_exacore(model, model.backend) + exa_model = ExaModels.ExaModel(core; prod = true) + start_time = time() + result = model.solver(exa_model; model.options...) + model.result = ( + objective = result.objective, + solution = Array(result.solution), + multipliers = Array(result.multipliers), + multipliers_L = Array(result.multipliers_L), + multipliers_U = Array(result.multipliers_U), + status = result.status, + ) + model.solve_time = time() - start_time + return end +# MOI.TerminationStatus + # SolverCore returns a `Symbol` in `result.status` for any solver implementing # the NLPModels callable interface (e.g. `madnlp(::AbstractNLPModel)`, # `ipopt(::AbstractNLPModel)`). The vocabulary is defined by SolverCore.jl @@ -774,133 +723,420 @@ const _TERMINATION_STATUS_CODES = Dict{Symbol, MOI.TerminationStatusCode}( :first_order => MOI.LOCALLY_SOLVED, :acceptable => MOI.ALMOST_LOCALLY_SOLVED, :small_step => MOI.SLOW_PROGRESS, - :infeasible => MOI.INFEASIBLE_OR_UNBOUNDED, + :infeasible => MOI.INFEASIBLE, :max_iter => MOI.ITERATION_LIMIT, :max_time => MOI.TIME_LIMIT, :user => MOI.INTERRUPTED, :exception => MOI.OTHER_ERROR, ) + +MOI.get(model::Optimizer, ::MOI.RawStatusString) = string(model.result.status) + +function MOI.get(model::Optimizer, ::MOI.TerminationStatus) + if model.result === nothing + return MOI.OPTIMIZE_NOT_CALLED + end + return get(_TERMINATION_STATUS_CODES, model.result.status, MOI.OTHER_ERROR) +end + +# MOI.PrimalStatus, MOI.DualStatus + const _RESULT_STATUS_CODES = Dict{Symbol, MOI.ResultStatusCode}( :first_order => MOI.FEASIBLE_POINT, :acceptable => MOI.NEARLY_FEASIBLE_POINT, :infeasible => MOI.INFEASIBLE_POINT, ) -MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus) = - Base.get(_TERMINATION_STATUS_CODES, optimizer.result.status, MOI.OTHER_ERROR) -MOI.get(model::Optimizer, attr::Union{MOI.PrimalStatus,MOI.DualStatus}) = - Base.get(_RESULT_STATUS_CODES, model.result.status, MOI.UNKNOWN_RESULT_STATUS) +function MOI.get(model::Optimizer, attr::Union{MOI.PrimalStatus,MOI.DualStatus}) + if model.result === nothing || attr.result_index != 1 + return MOI.NO_SOLUTION + end + return get( + _RESULT_STATUS_CODES, + model.result.status, + MOI.UNKNOWN_RESULT_STATUS, + ) +end + +# MOI.VariablePrimal -function MOI.get(model::Optimizer, attr::MOI.VariablePrimal, vi::MOI.VariableIndex) +function MOI.get( + model::Optimizer, + attr::MOI.VariablePrimal, + vi::MOI.VariableIndex, +) MOI.check_result_index_bounds(model, attr) if vi.value > PARAMETER_INDEX_THRESHOLD - return model.model.θ[vi.value-PARAMETER_INDEX_THRESHOLD] - else - return model.result.solution[vi.value] + return model.pstart[vi.value] end + return model.result.solution[vi.value] +end + +# MOI.ConstraintDual + +function _scale(model::Optimizer{T}) where {T} + return model.sense == MOI.MAX_SENSE ? -one(T) : one(T) end function MOI.get( model::Optimizer, attr::MOI.ConstraintDual, - ci::MOI.ConstraintIndex{<:SUPPORTED_FUNC_TYPE,<:SUPPORTED_FUNC_SET_TYPE}, + ci::MOI.ConstraintIndex, ) MOI.check_result_index_bounds(model, attr) - # MOI.throw_if_not_valid(model, ci) - s = -1.0 - return s * model.result.multipliers[ci.value] + return -_scale(model) * model.result.multipliers[ci.value] end +function _reduced_cost(model, col) + return model.result.multipliers_L[col] - model.result.multipliers_U[col] +end function MOI.get( - model::Optimizer, + model::Optimizer{T}, attr::MOI.ConstraintDual, - ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{Float64}}, -) + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{T}}, +) where {T} MOI.check_result_index_bounds(model, attr) - # MOI.throw_if_not_valid(model, ci) - rc = model.result.multipliers_L[ci.value] - model.result.multipliers_U[ci.value] - return min(0.0, rc) + rc = _reduced_cost(model, ci.value) + return min(zero(rc), _scale(model) * rc) end function MOI.get( model::Optimizer, attr::MOI.ConstraintDual, - ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.GreaterThan{Float64}}, -) + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.GreaterThan{T}}, +) where {T} MOI.check_result_index_bounds(model, attr) - # MOI.throw_if_not_valid(model, ci) - rc = model.result.multipliers_L[ci.value] - model.result.multipliers_U[ci.value] - return max(0.0, rc) + rc = _reduced_cost(model, ci.value) + return max(zero(rc), _scale(model) * rc) end function MOI.get( - model::Optimizer, + model::Optimizer{T}, attr::MOI.ConstraintDual, - ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{Float64}}, -) + ci::MOI.ConstraintIndex{MOI.VariableIndex,S}, +) where {T,S<:Union{MOI.Interval{T},MOI.EqualTo{T}}} MOI.check_result_index_bounds(model, attr) - # MOI.throw_if_not_valid(model, ci) - rc = model.result.multipliers_L[ci.value] - model.result.multipliers_U[ci.value] - return rc + return _scale(model) * _reduced_cost(model, ci.value) end +# MOI.ResultCount -function MOI.get(model::Optimizer, ::MOI.ResultCount) - return (model.result !== nothing) ? 1 : 0 -end +MOI.get(model::Optimizer, ::MOI.ResultCount) = model.result !== nothing ? 1 : 0 + +# MOI.ObjectiveValue function MOI.get(model::Optimizer, attr::MOI.ObjectiveValue) MOI.check_result_index_bounds(model, attr) - # scale = (model.sense == MOI.MAX_SENSE) ? -1 : 1 - # return scale * model.result.objective return model.result.objective end +# MOI.SolveTimeSec + MOI.get(model::Optimizer, ::MOI.SolveTimeSec) = model.solve_time -MOI.get( - model::Optimizer, - ::MOI.SolverName, -) = "$(string(model.solver)) running with ExaModels" -function MOI.set(model::Optimizer, p::MOI.RawOptimizerAttribute, value) - model.options[Symbol(p.name)] = value +# MOI.SolverName + +function MOI.get(model::Optimizer, ::MOI.SolverName) + return "$(string(model.solver)) running with ExaModels" +end + +# MOI.RawOptimizerAttribute + +function MOI.set(model::Optimizer, attr::MOI.RawOptimizerAttribute, value) + model.options[Symbol(attr.name)] = value # No need to reset model.solver because this gets handled in optimize!. return end +# MOI.NLPBlock -_make_index_map(model::MOI.ModelLike, maps) = _make_index_map(model, maps[1], maps[2]) -function _make_index_map(model::MOI.ModelLike, var_to_idx, con_to_idx) - variables = MOI.get(model, MOI.ListOfVariableIndices()) - map = MOI.Utilities.IndexMap() - for x in variables - vartype, rawidx = var_to_idx[x] - if vartype === :variable - map[x] = typeof(x)(rawidx) - elseif vartype === :parameter - map[x] = typeof(x)(rawidx + PARAMETER_INDEX_THRESHOLD) - else - error("Unknown variable type $vartype") +function MOI.set(::Optimizer, ::MOI.NLPBlock, ::MOI.NLPBlockData) + return error( + """ + The legacy nonlinear model interface is not supported. + + Please use the new MOI-based interface. + """, + ) +end + +### +### ExaModels as an MOI.Nonlinear automatic-differentiation backend +### +### This is the reverse role of `ExaModels.Optimizer` above: instead of +### ExaModels pretending to be a solver, any MOI solver that supports +### `MOI.AutomaticDifferentiationBackend` can evaluate its nonlinear model +### through ExaModels by setting the backend to `ExaModels.SIMDMode()`. + +""" + SIMDMode(; device = nothing) + +An automatic-differentiation backend for `MOI.Nonlinear` that evaluates the +model with ExaModels' SIMD abstraction instead of +`MOI.Nonlinear.SparseReverseMode`. + +`device` is the `KernelAbstractions` device to evaluate on (`nothing` means +the CPU). + +Pass it to solvers via `MOI.AutomaticDifferentiationBackend()`. + +The variables of the model must be `MOI.VariableIndex.(1:n)`: the constraints +are translated to SIMD-grouped bins as they are added, referencing the +variables by their raw index, so `MOI.initialize` errors if the +`ordered_variables` of the evaluator are not the identity. +""" +struct SIMDMode{B} <: MOI.Nonlinear.AbstractAutomaticDifferentiation + device::B +end + +SIMDMode(; device = nothing) = SIMDMode(device) + +""" + SIMDNonlinearModel + +The nonlinear model built by `MOI.Nonlinear.model(::ExaModels.SIMDMode)`. + +The objective and the constraints are translated to SIMD-grouped bins as they +are added, with the same machinery as `ExaModels.Optimizer`; the `ExaCore` is +assembled during `MOI.initialize` of the evaluator, once the number of +variables is known. + +Unlike `MOI.Nonlinear.Model`, this model consumes `MOI.ScalarAffineFunction` +and `MOI.ScalarQuadraticFunction` objectives and constraints natively (their +terms are grouped into SIMD kernels), so it must not be wrapped in +`MOI.Nonlinear.ModelWithQuad`. +""" +mutable struct SIMDNonlinearModel + objs::Vector{Bin} + cons::Vector{Bin} + lcon::Vector{Float64} + ucon::Vector{Float64} + linearity::Vector{MOI.Nonlinear.Linearity} + objective_linearity::MOI.Nonlinear.Linearity + + function SIMDNonlinearModel() + return new( + Bin[], + Bin[], + Float64[], + Float64[], + MOI.Nonlinear.Linearity[], + MOI.Nonlinear.CONSTANT, + ) + end +end + +# ExaModels handles affine and quadratic functions natively, so only the +# oracle layer is stacked on top. +function MOI.Nonlinear.model(::SIMDMode) + return MOI.Nonlinear.ModelWithOracles(SIMDNonlinearModel()) +end + +MOI.Nonlinear.exploits_structure(::SIMDMode) = true + +_linearity(::MOI.VariableIndex) = MOI.Nonlinear.LINEAR +_linearity(::MOI.ScalarAffineFunction) = MOI.Nonlinear.LINEAR +_linearity(::MOI.ScalarQuadraticFunction) = MOI.Nonlinear.QUADRATIC +_linearity(::Any) = MOI.Nonlinear.NONLINEAR + +function MOI.Nonlinear.set_objective(model::SIMDNonlinearModel, obj) + empty!(model.objs) + model.objective_linearity = MOI.Nonlinear.CONSTANT + if obj !== nothing + update_bin!(model.objs, ObjectiveBin(), obj) + model.objective_linearity = _linearity(obj) + end + return +end + +function MOI.Nonlinear.add_constraint( + model::SIMDNonlinearModel, + f::Union{ + MOI.ScalarAffineFunction{Float64}, + MOI.ScalarQuadraticFunction{Float64}, + MOI.ScalarNonlinearFunction, + }, + s::Union{ + MOI.GreaterThan{Float64}, + MOI.LessThan{Float64}, + MOI.EqualTo{Float64}, + MOI.Interval{Float64}, + }, +) + row = length(model.lcon) + 1 + update_bin!(model.cons, ConstraintBin(row), f) + l, u = _bounds(s) + push!(model.lcon, l) + push!(model.ucon, u) + push!(model.linearity, _linearity(f)) + return MOI.Nonlinear.ConstraintIndex(row) +end + +function MOI.Nonlinear.register_operator( + ::SIMDNonlinearModel, + op::Symbol, + ::Int, + ::Function..., +) + return error( + "The operator `$op` cannot be registered: ExaModels does not " * + "support user-defined operators through `ExaModels.SIMDMode`.", + ) +end + +function MOI.Nonlinear.add_parameter(::SIMDNonlinearModel, ::Real) + return error( + "`MOI.Nonlinear` parameters are not supported by " * + "`ExaModels.SIMDMode`.", + ) +end + +function MOI.Nonlinear.add_expression(::SIMDNonlinearModel, expr) + return error( + "`MOI.Nonlinear` expressions are not supported by " * + "`ExaModels.SIMDMode`.", + ) +end + +mutable struct SIMDEvaluator{B} <: MOI.AbstractNLPEvaluator + model::SIMDNonlinearModel + mode::SIMDMode{B} + ordered_variables::Vector{MOI.VariableIndex} + # The `ExaModels.ExaModel`, built during `MOI.initialize`. + exa::Any +end + +function MOI.Nonlinear.Evaluator( + model::SIMDNonlinearModel, + mode::SIMDMode, + ordered_variables::Vector{MOI.VariableIndex}, +) + return SIMDEvaluator(model, mode, ordered_variables, nothing) +end + +function MOI.features_available(::SIMDEvaluator) + return [:Grad, :Jac, :JacVec, :Hess, :HessVec] +end + +function MOI.Nonlinear.num_constraints(d::SIMDEvaluator) + return length(d.model.lcon) +end + +function MOI.Nonlinear.constraint_bounds(d::SIMDEvaluator) + return MOI.NLPBoundsPair[ + MOI.NLPBoundsPair(l, u) for (l, u) in zip(d.model.lcon, d.model.ucon) + ] +end + +MOI.Nonlinear.constraint_linearity(d::SIMDEvaluator) = copy(d.model.linearity) + +MOI.Nonlinear.objective_linearity(d::SIMDEvaluator) = d.model.objective_linearity + +MOI.Nonlinear._has_objective(d::SIMDEvaluator) = !isempty(d.model.objs) + +function MOI.initialize(d::SIMDEvaluator, features::Vector{Symbol}) + T = Float64 + n = length(d.ordered_variables) + # The bins reference the variables by their raw index, so the columns of + # the evaluator must coincide with the variable indices. + if d.ordered_variables != MOI.VariableIndex.(1:n) + error( + "`ExaModels.SIMDMode` requires the variables of the model " * + "to be `MOI.VariableIndex.(1:n)`, in order.", + ) + end + c = ExaModels.ExaCore( + T; + backend = d.mode.device, + minimize = true, + concrete = Val(true), + ) + c, _ = ExaModels.add_var( + c, + n; + start = zeros(T, n), + lvar = fill(typemin(T), n), + uvar = fill(typemax(T), n), + ) + m = d.model + if !isempty(m.cons) + c, cons = + ExaModels.add_con(c, length(m.lcon); lcon = m.lcon, ucon = m.ucon) + for bin in m.cons + c, _ = ExaModels.add_con!(c, cons, (bin.head for _ in bin.data)) end end - for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent()) - _make_constraints_map(model, map.con_map[F, S], con_to_idx) + for bin in m.objs + c, _ = ExaModels.add_obj(c, bin.head, bin.data) end - return map + prod = :JacVec in features || :HessVec in features + d.exa = ExaModels.ExaModel(c; prod = prod) + return end -function _make_constraints_map( - model, - map::MOI.Utilities.DoubleDicts.IndexDoubleDictInner{F,S}, - con_to_idx, -) where {F,S} - for c in MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) - map[c] = typeof(c)(con_to_idx[c]) + +function _exa(d::SIMDEvaluator) + if d.exa === nothing + error("You must call `MOI.initialize` before evaluating.") end + return d.exa +end + +MOI.eval_objective(d::SIMDEvaluator, x) = ExaModels.NLPModels.obj(_exa(d), x) + +function MOI.eval_objective_gradient(d::SIMDEvaluator, grad, x) + ExaModels.NLPModels.grad!(_exa(d), x, grad) + return +end + +function MOI.eval_constraint(d::SIMDEvaluator, g, x) + ExaModels.NLPModels.cons!(_exa(d), x, g) + return +end + +function MOI.jacobian_structure(d::SIMDEvaluator) + exa = _exa(d) + nnzj = exa.meta.nnzj + rows, cols = Vector{Int}(undef, nnzj), Vector{Int}(undef, nnzj) + ExaModels.NLPModels.jac_structure!(exa, rows, cols) + return collect(zip(rows, cols)) +end + +function MOI.eval_constraint_jacobian(d::SIMDEvaluator, J, x) + ExaModels.NLPModels.jac_coord!(_exa(d), x, J) return end -function MOI.set(model::Optimizer, ::MOI.NLPBlock, nlp_data::MOI.NLPBlockData) - error("The legacy nonlinear model interface is not supported. Please use the new MOI-based interface.") +function MOI.hessian_lagrangian_structure(d::SIMDEvaluator) + exa = _exa(d) + nnzh = exa.meta.nnzh + rows, cols = Vector{Int}(undef, nnzh), Vector{Int}(undef, nnzh) + ExaModels.NLPModels.hess_structure!(exa, rows, cols) + return collect(zip(rows, cols)) +end + +function MOI.eval_hessian_lagrangian(d::SIMDEvaluator, H, x, σ, μ) + ExaModels.NLPModels.hess_coord!(_exa(d), x, μ, H; obj_weight = σ) + return +end + +function MOI.eval_constraint_jacobian_product(d::SIMDEvaluator, y, x, w) + ExaModels.NLPModels.jprod!(_exa(d), x, w, y) + return +end + +function MOI.eval_constraint_jacobian_transpose_product( + d::SIMDEvaluator, + y, + x, + w, +) + ExaModels.NLPModels.jtprod!(_exa(d), x, w, y) + return +end + +function MOI.eval_hessian_lagrangian_product(d::SIMDEvaluator, h, x, v, σ, μ) + ExaModels.NLPModels.hprod!(_exa(d), x, μ, v, h; obj_weight = σ) + return end end # module diff --git a/src/templates.jl b/src/templates.jl index ae09cb8f6..9c7f0ec73 100644 --- a/src/templates.jl +++ b/src/templates.jl @@ -36,5 +36,8 @@ end # to avoid type privacy sort!(array; kwargs...) = Base.sort!(array; kwargs...) -# MOI -function Optimizer end +# Placeholder for ExaModels.Optimizer +global Optimizer + +# Placeholder for ExaModels.SIMDMode +global SIMDMode diff --git a/test/JuMPTest/JuMPTest.jl b/test/JuMPTest/JuMPTest.jl index 22ef631b5..37485c28f 100644 --- a/test/JuMPTest/JuMPTest.jl +++ b/test/JuMPTest/JuMPTest.jl @@ -1,163 +1,221 @@ module JuMPTest -using Test, JuMP, ExaModels, PowerModels, NLPModelsIpopt, ..NLPTest - +using Test + +import ExaModels +import Ipopt +import JuMP +using JuMP: MOI +import NLPModels +import NLPModelsIpopt +import NLPModelsJuMP +import PowerModels + +import ..NLPTest import ..BACKENDS -import ..ad_tolerance, ..sol_tolerance, ..solver_tolerance - -const JUMP_INTERFACE_INSTANCES = [ - (:jump_luksan_vlcek_model, [3, 10]), - (:jump_ac_power_model, ["pglib_opf_case3_lmbd.m", "pglib_opf_case14_ieee.m"]), -] +import ..sol_tolerance +import ..solver_tolerance -function jump_luksan_vlcek_model(N) - jm = JuMP.Model() +function runtests() + is_test(name) = startswith("$name", "test_") + @testset "$name" for name in filter(is_test, names(@__MODULE__; all = true)) + getfield(@__MODULE__, name)() + end + return +end - JuMP.@variable(jm, x[i=1:N], start = mod(i, 2) == 1 ? -1.2 : 1.0) - JuMP.@constraint( - jm, - s[i=1:(N-2)], - 3x[i+1]^3 + 2x[i+2] - 5 + sin(x[i+1] - x[i+2])sin(x[i+1] + x[i+2]) + 4x[i+1] - - x[i]exp(x[i] - x[i+1]) - 3 == 0.0 +function test_moi_tests() + model = MOI.instantiate( + () -> ExaModels.Optimizer(NLPModelsIpopt.ipopt); + with_bridge_type = Float64, + with_cache_type = Float64, ) - JuMP.@objective(jm, Min, sum(100(x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2 for i = 2:N)) - - return jm + MOI.set(model, MOI.RawOptimizerAttribute("print_level"), 0) + MOI.Test.runtests( + model, + MOI.Test.Config(; + atol = 1e-4, + optimal_status = MOI.LOCALLY_SOLVED, + exclude = Any[ + MOI.DualObjectiveValue, + MOI.ObjectiveBound, + MOI.SolverVersion, + MOI.ConstraintBasisStatus, + MOI.VariableBasisStatus, + ], + ), + exclude = [ + # NLPModels doesn't detect unboundedness + r"^test_linear_DUAL_INFEASIBLE$", + r"^test_linear_DUAL_INFEASIBLE_2$", + r"^test_solve_TerminationStatus_DUAL_INFEASIBLE$", + # Returns INVALID_MODEL becuase of the empty row + r"^test_linear_VectorAffineFunction_empty_row$", + # Ipopt fails because of co-linear constraint and objective and + # redundant constraint. + r"^test_linear_transform$", + ], + ) + return end -function nlp_legacy_runtests() +function test_nlp_legacy() jm = JuMP.Model() - JuMP.@variable(jm, x[1:10]) JuMP.@NLobjective(jm, Min, sum(x[i] for i=1:10)) - - @test_throws ErrorException ExaModel(jm) - + @test_throws ErrorException ExaModels.ExaModel(jm) jm = JuMP.Model(() -> ExaModels.Optimizer(NLPModelsIpopt.ipopt)) - @test_throws ErrorException optimize!(jm) + @test_throws ErrorException JuMP.optimize!(jm) + return end - -function fixed_variable_e2etest() - N=5 - jm = JuMP.Model() +function test_fixed_variable_e2etest() + N = 5 + jm = JuMP.Model() JuMP.@variable(jm, x[1:N]) JuMP.fix(x[1], 1.0) JuMP.@constraint(jm, sum(x) == 1.0) JuMP.@objective(jm, Min, sum(2*x[i]^2 for i = 1:N)) - - em = ExaModel(jm) + em = ExaModels.ExaModel(jm) @test only(em.meta.lcon) == only(em.meta.ucon) == 1.0 - - # em.cons is a Tuple: (ConstraintAugmentation{Null}, ConstraintAugmentation{Pair}, Constraint{Null{Nothing}}) @test em.cons[1] isa ExaModels.ConstraintAugmentation - @test em.cons[1].f.f isa ExaModels.Null - - @test em.cons[2] isa ExaModels.ConstraintAugmentation - @test em.cons[2].f.f isa Pair - - @test typeof(em.cons[2].f.f.second) <: ExaModels.Node2{ + @test em.cons[1].f.f isa Pair + @test em.cons[1].f.f.second isa ExaModels.Node2{ typeof(*), - ExaModels.Var{T1}, - T2, - } where {T1<:ExaModels.DataIndexed,T2<:ExaModels.DataIndexed} - - @test em.cons[3] isa ExaModels.Constraint - @test em.cons[3].f.f isa ExaModels.Null{Nothing} - - @test em.objs[1].f.f isa ExaModels.Null - @test typeof(em.objs[2].f.f) <: ExaModels.Node2{ + <:ExaModels.DataIndexed, + <:ExaModels.Var{<:ExaModels.DataIndexed}, + } + @test em.cons[2] isa ExaModels.Constraint + @test em.cons[2].f.f isa ExaModels.Null{Nothing} + @test length(em.objs) == 1 + @test em.objs[1].f.f isa ExaModels.Node2{ typeof(*), - T1, - ExaModels.Node1{typeof(abs2),ExaModels.Var{T2}}, - } where {T1<:ExaModels.DataIndexed,T2<:ExaModels.DataIndexed} + <:ExaModels.DataIndexed, + <:ExaModels.Node1{typeof(abs2),<:ExaModels.Var{<:ExaModels.DataIndexed}}, + } + return +end +function test_parameter_e2etest() + N = 5 jm = JuMP.Model() - JuMP.@variable(jm, x[1:N]) JuMP.@variable(jm, p in JuMP.Parameter(1.0)) JuMP.@constraint(jm, sum(x) == p) JuMP.@objective(jm, Min, sum(x)) - - em = ExaModel(jm) + em = ExaModels.ExaModel(jm) @test only(em.meta.lcon) == only(em.meta.ucon) == 0.0 @test only(em.θ) == 1.0 - # em.cons: (ConstraintAugmentation{Null}, ConstraintAugmentation{Pair/Param}, ConstraintAugmentation{Pair/Var}, Constraint{Null{Nothing}}) @test em.cons[1] isa ExaModels.ConstraintAugmentation - @test em.cons[1].f.f isa ExaModels.Null + @test em.cons[1].f.f isa Pair + @test em.cons[1].f.f.second isa ExaModels.Node2{ + typeof(*), + <:ExaModels.DataIndexed, + <:ExaModels.ParameterNode{<:ExaModels.DataIndexed}, + } @test em.cons[2] isa ExaModels.ConstraintAugmentation @test em.cons[2].f.f isa Pair - @test typeof(em.cons[2].f.f.second) <: ExaModels.Node2{ - typeof(*), - ExaModels.ParameterNode{T1}, - T2, - } where {T1<:ExaModels.DataIndexed,T2<:ExaModels.DataIndexed} - @test em.cons[3] isa ExaModels.ConstraintAugmentation - @test em.cons[3].f.f isa Pair - @test typeof(em.cons[3].f.f.second) <: ExaModels.Node2{ + @test em.cons[2].f.f.second isa ExaModels.Node2{ typeof(*), - ExaModels.Var{T1}, - T2, - } where {T1<:ExaModels.DataIndexed,T2<:ExaModels.DataIndexed} - @test em.cons[4] isa ExaModels.Constraint - @test em.cons[4].f.f isa ExaModels.Null{Nothing} + <:ExaModels.DataIndexed, + <:ExaModels.Var{<:ExaModels.DataIndexed}, + } + @test em.cons[3] isa ExaModels.Constraint + @test em.cons[3].f.f isa ExaModels.Null{Nothing} jm = JuMP.Model() JuMP.@variable(jm, x) - @test_broken em = ExaModel(jm) # FIXME: support feasibility problems? - - return jm + @test ExaModels.ExaModel(jm) isa ExaModels.ExaModel + return end -function no_constraints_e2etest() - N=5 + +function test_no_constraints_e2etest() + N = 5 jm = JuMP.Model() JuMP.@variable(jm, x[1:N]) JuMP.@objective(jm, Max, sum(sin(x[i]) for i = 1:N)) + em = ExaModels.ExaModel(jm) + @test isempty(em.cons) + @test length(em.objs) == 1 + @test em.objs[1].f.f isa + ExaModels.Node1{typeof(sin),<:ExaModels.Var{<:ExaModels.DataIndexed}} + return +end - em = ExaModel(jm) - - @test length(em.cons) == 1 - @test em.cons[1] isa ExaModels.Constraint - - @test em.objs[1].f.f isa ExaModels.Null - @test typeof(em.objs[2].f.f) <: - ExaModels.Node1{typeof(sin),ExaModels.Var{T1}} where {T1<:ExaModels.DataIndexed} - - N=5 +function test_no_constraints_simd_failure() + N = 5 jm = JuMP.Model() JuMP.@variable(jm, x[1:N]) JuMP.@objective(jm, Max, sin(sum(x[i] for i = 1:N))) - - em = ExaModel(jm) - - @test length(em.cons) == 1 - @test em.cons[1] isa ExaModels.Constraint - - @test em.objs[1].f.f isa ExaModels.Null + em = ExaModels.ExaModel(jm) + @test isempty(em.cons) + @test length(em.objs) == 1 # broken since ExaMOI fails to detect SIMD in this case - @test_broken typeof(em.objs[2].f.f) <: - ExaModels.Node1{typeof(sin),ExaModels.Var{T1}} where {T1} + @test_broken em.objs[1].f.f isa ExaModels.Node1{typeof(sin),<:ExaModels.Var} + return end -function generic_e2etest() - N=5 + +function test_generic_e2etest() + N = 5 jm = JuMP.GenericModel{Float32}() JuMP.@variable(jm, x[1:N]) JuMP.@constraint(jm, sum(x) == 1.0f0) JuMP.@objective(jm, Min, sum(x[i]^2 for i = 1:N)) + em = ExaModels.ExaModel(jm) + @test typeof(em) <: ExaModels.ExaModel{Float32} + @test eltype(em.cons[1].itr) <: Tuple{Int,Float32,Int} + return +end - em = ExaModel(jm) +function _test_jump_interface(modelfunction, case) + jm = modelfunction(case) + JuMP.set_optimizer(jm, Ipopt.Optimizer) + JuMP.set_optimizer_attribute(jm, "print_level", 0) + JuMP.optimize!(jm) + sol = JuMP.value.(JuMP.all_variables(jm)) + dsol = JuMP.dual.(JuMP.all_constraints(jm, include_variable_in_set_constraints = true)) + JuMP.set_optimizer(jm, () -> ExaModels.Optimizer(NLPModelsIpopt.ipopt)) + JuMP.set_optimizer_attribute(jm, "print_level", 0) + JuMP.optimize!(jm) + sol2 = JuMP.value.(JuMP.all_variables(jm)) + dsol2 = JuMP.dual.(JuMP.all_constraints(jm, include_variable_in_set_constraints = true)) + @test sol ≈ sol2 atol = sol_tolerance(eltype(sol), eltype(sol2)) + @test dsol ≈ dsol2 atol = sol_tolerance(eltype(sol), eltype(sol2)) + @testset "$backend" for backend in BACKENDS + m = ExaModels.WrapperNLPModel(ExaModels.ExaModel(jm; backend)) + result = NLPModelsIpopt.ipopt( + m; + print_level = 0, + tol = solver_tolerance(eltype(m.inner.meta.x0)), + ) + @test sol ≈ result.solution atol = sol_tolerance(eltype(m.inner.meta.x0)) + end + return +end - @test typeof(em) <: ExaModel{Float32} - @test typeof(getindex.(em.cons[2].itr, 2)) <: Vector{Float32} +function jump_luksan_vlcek_model(N) + jm = JuMP.Model() + JuMP.@variable(jm, x[i=1:N], start = mod(i, 2) == 1 ? -1.2 : 1.0) + JuMP.@constraint( + jm, + s[i=1:(N-2)], + 3x[i+1]^3 + 2x[i+2] - 5 + sin(x[i+1] - x[i+2])sin(x[i+1] + x[i+2]) + 4x[i+1] - + x[i]exp(x[i] - x[i+1]) - 3 == 0.0 + ) + JuMP.@objective(jm, Min, sum(100(x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2 for i = 2:N)) + return jm end -function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") +function test_jump_luksan_vlcek() + @testset "$N" for N in [3, 10] + _test_jump_interface(jump_luksan_vlcek_model, N) + end + return +end +function jump_ac_power_model(filename::String) ref = NLPTest.get_power_data_ref(filename) - model = JuMP.Model() - #JuMP.set_optimizer_attribute(model, "print_level", 0) - JuMP.@variable(model, va[i in keys(ref[:bus])]) JuMP.@variable(model, will_delete) JuMP.@variable( @@ -165,7 +223,6 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") ref[:bus][i]["vmin"] <= vm[i in keys(ref[:bus])] <= ref[:bus][i]["vmax"], start = 1.0 ) - JuMP.@variable( model, ref[:gen][i]["pmin"] <= pg[i in keys(ref[:gen])] <= ref[:gen][i]["pmax"] @@ -174,7 +231,6 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") model, ref[:gen][i]["qmin"] <= qg[i in keys(ref[:gen])] <= ref[:gen][i]["qmax"] ) - JuMP.@variable( model, -ref[:branch][l]["rate_a"] <= @@ -187,7 +243,6 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") q[(l, i, j) in ref[:arcs]] <= ref[:branch][l]["rate_a"] ) - JuMP.@objective( model, Min, @@ -196,22 +251,18 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") (i, gen) in ref[:gen] ) ) - for (i, bus) in ref[:ref_buses] JuMP.@constraint(model, va[i] == 0) end - for (i, bus) in ref[:bus] bus_loads = [ref[:load][l] for l in ref[:bus_loads][i]] bus_shunts = [ref[:shunt][s] for s in ref[:bus_shunts][i]] - JuMP.@constraint( model, sum(p[a] for a in ref[:bus_arcs][i]) == sum(pg[g] for g in ref[:bus_gens][i]) - sum(load["pd"] for load in bus_loads) - sum(shunt["gs"] for shunt in bus_shunts) * vm[i]^2 ) - JuMP.@constraint( model, sum(q[a] for a in ref[:bus_arcs][i]) == @@ -219,22 +270,18 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") sum(shunt["bs"] for shunt in bus_shunts) * vm[i]^2 ) end - # Branch power flow physics and limit constraints for (i, branch) in ref[:branch] f_idx = (i, branch["f_bus"], branch["t_bus"]) t_idx = (i, branch["t_bus"], branch["f_bus"]) - p_fr = p[f_idx] q_fr = q[f_idx] p_to = p[t_idx] q_to = q[t_idx] - vm_fr = vm[branch["f_bus"]] vm_to = vm[branch["t_bus"]] va_fr = va[branch["f_bus"]] va_to = va[branch["t_bus"]] - g, b = PowerModels.calc_branch_y(branch) tr, ti = PowerModels.calc_branch_t(branch) ttm = tr^2 + ti^2 @@ -242,7 +289,6 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") b_fr = branch["b_fr"] g_to = branch["g_to"] b_to = branch["b_to"] - # From side of the branch flow JuMP.@constraint( model, @@ -258,7 +304,6 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") (-b * tr - g * ti) / ttm * (vm_fr * vm_to * cos(va_fr - va_to)) + (-g * tr + b * ti) / ttm * (vm_fr * vm_to * sin(va_fr - va_to)) ) - # To side of the branch flow JuMP.@constraint( model, @@ -274,63 +319,380 @@ function jump_ac_power_model(filename = "pglib_opf_case3_lmbd.m") (-b * tr + g * ti) / ttm * (vm_to * vm_fr * cos(va_to - va_fr)) + (-g * tr - b * ti) / ttm * (vm_to * vm_fr * sin(va_to - va_fr)) ) - # Voltage angle difference limit JuMP.@constraint(model, branch["angmin"] <= va_fr - va_to <= branch["angmax"]) - # Apparent power limit, from side and to side JuMP.@constraint(model, p_fr^2 + q_fr^2 <= branch["rate_a"]^2) JuMP.@constraint(model, p_to^2 + q_to^2 <= branch["rate_a"]^2) end - JuMP.delete(model, will_delete) - return model end -function runtests() - @testset "JuMP Interface test" begin - for (model, cases) in JUMP_INTERFACE_INSTANCES - for case in cases - @testset "$model $case" begin - modelfunction = getfield(@__MODULE__, model) - - # solve JuMP problem - jm = modelfunction(case) - set_optimizer(jm, NLPModelsIpopt.Ipopt.Optimizer) - set_optimizer_attribute(jm, "print_level", 0) - optimize!(jm) - sol = value.(all_variables(jm)) - dsol = dual.(all_constraints(jm, include_variable_in_set_constraints = true)) - - set_optimizer(jm, () -> ExaModels.Optimizer(ipopt)) - set_optimizer_attribute(jm, "print_level", 0) - optimize!(jm) - sol2 = value.(all_variables(jm)) - dsol2 = dual.(all_constraints(jm, include_variable_in_set_constraints = true)) - @test sol ≈ sol2 atol = sol_tolerance(eltype(sol), eltype(sol2)) - @test dsol ≈ dsol2 atol = sol_tolerance(eltype(sol), eltype(sol2)) - - for backend in BACKENDS - @testset "$backend" begin - m = WrapperNLPModel(ExaModel(jm; backend = backend)) - result = ipopt(m; print_level = 0, tol = solver_tolerance(eltype(m.inner.meta.x0))) - - @test sol ≈ result.solution atol = sol_tolerance(eltype(m.inner.meta.x0)) - end - end - end - end - end - @testset "E2E tests" begin - generic_e2etest() - fixed_variable_e2etest() - no_constraints_e2etest() - end - @testset "NLP legacy test" begin - nlp_legacy_runtests() +function test_jump_ac_power_model() + @testset "$file" for file in ["pglib_opf_case3_lmbd.m", "pglib_opf_case14_ieee.m"] + _test_jump_interface(jump_ac_power_model, file) + end + return +end + +function _jacobian_matrix(model, x) + rows = zeros(Int, model.meta.nnzj) + cols = zeros(Int, model.meta.nnzj) + values = zeros(eltype(x), model.meta.nnzj) + NLPModels.jac_structure!(model, rows, cols) + NLPModels.jac_coord!(model, x, values) + jacobian = zeros(eltype(x), model.meta.ncon, model.meta.nvar) + for k in eachindex(values) + jacobian[rows[k], cols[k]] += values[k] + end + return jacobian +end + +function _hessian_matrix(model, x, y; obj_weight) + rows = zeros(Int, model.meta.nnzh) + cols = zeros(Int, model.meta.nnzh) + values = zeros(eltype(x), model.meta.nnzh) + NLPModels.hess_structure!(model, rows, cols) + NLPModels.hess_coord!(model, x, y, values; obj_weight = obj_weight) + hessian = zeros(eltype(x), model.meta.nvar, model.meta.nvar) + for k in eachindex(values) + hessian[rows[k], cols[k]] += values[k] + if rows[k] != cols[k] + hessian[cols[k], rows[k]] += values[k] end end + return hessian +end + +function _test_callback_equivalence( + model::JuMP.GenericModel{T}, + points::Vector{Vector{T}}, +) where {T} + model_exa = ExaModels.ExaModel(model) + model_nlp = NLPModelsJuMP.MathOptNLPModel(model) + # Constraint rows may be ordered differently by the two adapters. Restrict + # this helper to models with at most one constraint, where direct callback + # comparison is unambiguous. + @assert model_exa.meta.ncon <= 1 + y = T(0.37) .* collect(T, 1:model_exa.meta.ncon) + obj_weight = T(0.61) + for x in points + @test NLPModels.obj(model_exa, x) ≈ NLPModels.obj(model_nlp, x) + @test NLPModels.cons(model_exa, x) ≈ NLPModels.cons(model_nlp, x) + @test NLPModels.grad(model_exa, x) ≈ NLPModels.grad(model_nlp, x) + @test _jacobian_matrix(model_exa, x) ≈ _jacobian_matrix(model_nlp, x) + @test _hessian_matrix(model_exa, x, y; obj_weight) ≈ + _hessian_matrix(model_nlp, x, y; obj_weight) + end + return +end + +function test_nonlinear_constraint_derivative_sparsity() + model = JuMP.Model() + JuMP.@variable(model, p) + JuMP.@variable(model, vmf) + JuMP.@variable(model, vmt) + JuMP.@variable(model, vaf) + JuMP.@variable(model, vat) + JuMP.@constraint( + model, + p - 1.2vmf^2 - 0.7vmf * vmt * cos(vaf - vat) - 0.3vmf * vmt * sin(vaf - vat) == 0.0, + ) + JuMP.@objective(model, Min, p) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 10 + @test model_exa.meta.nnzh == 21 + jacobian_rows = zeros(Int, model_exa.meta.nnzj) + jacobian_cols = zeros(Int, model_exa.meta.nnzj) + NLPModels.jac_structure!(model_exa, jacobian_rows, jacobian_cols) + @test length(unique(zip(jacobian_rows, jacobian_cols))) == 5 + # Hessian coordinates are unique here because this model has one + # constraint row. Different rows may legitimately repeat coordinates. + hessian_rows = zeros(Int, model_exa.meta.nnzh) + hessian_cols = zeros(Int, model_exa.meta.nnzh) + NLPModels.hess_structure!(model_exa, hessian_rows, hessian_cols) + @test length(unique(zip(hessian_rows, hessian_cols))) == 10 + _test_callback_equivalence( + model, + [[0.2, 1.0, 0.9, 0.1, -0.2], [-0.4, 1.1, 1.05, -0.3, 0.25]], + ) + return +end + +function test_nonlinear_constraint_repeated_variable() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@constraint(model, sin(x) + x^2 + cos(x - y) == 0.0) + JuMP.@objective(model, Min, x + y) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 4 + @test model_exa.meta.nnzh == 5 + _test_callback_equivalence(model, [[0.3, -0.7], [1.2, 0.4]]) + return +end + +function test_nonlinear_objective_derivative_sparsity() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@objective(model, Min, sin(x) + x^2 + cos(x - y) + 2.5) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzo == 2 + @test model_exa.meta.nnzh == 5 + _test_callback_equivalence(model, [[0.3, -0.7], [1.2, 0.4]]) + return +end + +function test_nonlinear_objective_nested_expr() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@objective(model, Min, exp(sin(x) + x^2 + cos(x - y))) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzo == 2 + @test model_exa.meta.nnzh == 3 + _test_callback_equivalence(model, [[0.3, -0.7], [1.2, 0.4]]) + return +end + +function test_nonlinear_objective_parameter() + model = JuMP.Model() + JuMP.@variable(model, x[1:8]) + JuMP.@variable(model, p in JuMP.Parameter(0.4)) + JuMP.@objective(model, Min, sum(sin(p * x[i]) for i = 1:8)) + model_exa = ExaModels.ExaModel(model) + @test length(model_exa.objs) == 1 + @test model_exa.meta.nnzo == 8 + @test model_exa.meta.nnzh == 8 + parameter_point = collect(range(-0.7, 0.7; length = 8)) + @test NLPModels.obj(model_exa, parameter_point) ≈ sum(sin.(0.4 .* parameter_point)) + @test NLPModels.grad(model_exa, parameter_point) ≈ 0.4 .* cos.(0.4 .* parameter_point) + return +end + +function test_nonlinear_objective_coupled() + N = 20 + model = JuMP.Model() + JuMP.@variable(model, x[1:N]) + JuMP.@objective(model, Min, sum(100(x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2 for i = 2:N)) + model_exa = ExaModels.ExaModel(model) + _test_callback_equivalence(model, [collect(range(-0.5, 0.5; length = N))]) + return +end + +function test_nonlinear_aliasing_and_batching() + model = JuMP.Model() + JuMP.@variable(model, z[1:4]) + JuMP.@constraint(model, sin(z[1] * z[2]) == 0.0) + JuMP.@constraint(model, sin(z[3] * z[3]) == 0.0) + JuMP.@constraint(model, sin(z[3] * z[4]) == 0.0) + JuMP.@objective(model, Min, sum(z)) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 5 + @test model_exa.meta.nnzh == 7 + p = [0.2, -0.4, 0.7, 1.1] + @test NLPModels.cons(model_exa, p) ≈ [sin(p[1] * p[2]), sin(p[3]^2), sin(p[3] * p[4])] + return +end + +function test_nonlinear_batching() + K = 4 + model = JuMP.Model() + JuMP.@variable(model, p[1:K]) + JuMP.@variable(model, vmf[1:K]) + JuMP.@variable(model, vmt[1:K]) + JuMP.@variable(model, vaf[1:K]) + JuMP.@variable(model, vat[1:K]) + JuMP.@constraint( + model, + [i = 1:K], + p[i] - + 1.2vmf[i]^2 - + 0.7vmf[i] * vmt[i] * cos(vaf[i] - vat[i]) - + 0.3vmf[i] * vmt[i] * sin(vaf[i] - vat[i]) == 0.0, + ) + JuMP.@objective(model, Min, sum(p)) + model_exa = ExaModels.ExaModel(model) + @test length(model_exa.cons) == 5 + @test model_exa.meta.nnzj == 40 + @test model_exa.meta.nnzh == 84 + batched_point = vcat( + collect(0.1:0.1:0.4), + fill(1.0, K), + fill(0.9, K), + collect(0.05:0.05:0.2), + collect(-0.2:0.05:-0.05), + ) + @test NLPModels.cons(model_exa, batched_point) ≈ + [ + batched_point[i] - + 1.2batched_point[K+i]^2 - + 0.7batched_point[K+i] * + batched_point[2K+i] * + cos(batched_point[3K+i] - batched_point[4K+i]) - + 0.3batched_point[K+i] * + batched_point[2K+i] * + sin(batched_point[3K+i] - batched_point[4K+i]) for i = 1:K + ] + return +end + +function test_nonlinear_parameters_and_nested_expressions() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, p in JuMP.Parameter(0.4)) + JuMP.@constraint(model, sin(p * x) + cos(p * x) + p == 0.0) + JuMP.@objective(model, Min, x) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 2 + @test NLPModels.cons(model_exa, [0.7]) ≈ [sin(0.4 * 0.7) + cos(0.4 * 0.7) + 0.4] + return +end + +function test_nonlinear_prameter_only() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, p in JuMP.Parameter(0.4)) + JuMP.@constraint(model, sin(p) + p^2 == 0.0) + JuMP.@objective(model, Min, x) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 0 + @test NLPModels.cons(model_exa, [0.7]) ≈ [sin(0.4) + 0.4^2] + return +end + +function test_nonlinear_nested_affine() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@constraint(model, sin(2x + 3y) + x == 0.0) + JuMP.@objective(model, Min, x + y) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 3 + @test model_exa.meta.nnzh == 3 + _test_callback_equivalence(model, [[0.3, -0.7], [1.2, 0.4]]) + return +end + +function test_nonlinear_nested_quadratic() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@constraint(model, sin(x^2 + x * y) + x == 0.0) + JuMP.@objective(model, Min, x + y) + model_exa = ExaModels.ExaModel(model) + @test model_exa.meta.nnzj == 3 + @test model_exa.meta.nnzh == 3 + _test_callback_equivalence(model, [[0.3, -0.7], [1.2, 0.4]]) + return +end + +function test_nonlinear_shapes() + model = JuMP.Model() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@variable(model, z) + JuMP.@constraint(model, sin(x) == 0.0) + JuMP.@constraint(model, exp(y + z) - 1.0 == 0.0) + JuMP.@constraint(model, cos(x * y) + z == 0.0) + JuMP.@objective(model, Min, x + y + z) + model_exa = ExaModels.ExaModel(model) + @test length(model_exa.cons) == 6 + @test model_exa.meta.nnzj == 6 + shapes_point = [0.2, -0.4, 0.7] + @test NLPModels.cons(model_exa, shapes_point) ≈ + [ + sin(shapes_point[1]), + exp(shapes_point[2] + shapes_point[3]) - 1.0, + cos(shapes_point[1] * shapes_point[2]) + shapes_point[3], + ] + return +end + +function test_nonlinear_float32() + model = JuMP.GenericModel{Float32}() + JuMP.@variable(model, x) + JuMP.@variable(model, y) + JuMP.@constraint(model, sin(x) + x^2 + cos(x - y) == 0.0f0) + JuMP.@objective(model, Min, x + y) + model_exa = ExaModels.ExaModel(model) + @test typeof(model_exa) <: ExaModels.ExaModel{Float32} + @test model_exa.meta.nnzj == 4 + @test model_exa.meta.nnzh == 5 + float_point = Float32[0.3, -0.7] + @test eltype(NLPModels.cons(model_exa, float_point)) == Float32 + difference = float_point[1] - float_point[2] + @test NLPModels.cons(model_exa, float_point) ≈ + Float32[sin(float_point[1]) + float_point[1]^2 + cos(difference)] + @test NLPModels.grad(model_exa, float_point) ≈ ones(Float32, 2) + @test _jacobian_matrix(model_exa, float_point) ≈ + Float32[(cos(float_point[1]) + 2float_point[1] - sin(difference)) sin(difference)] + expected_hessian = Float32[ + -sin(float_point[1])+2-cos(difference) cos(difference) + cos(difference) -cos(difference) + ] + @test _hessian_matrix( + model_exa, + float_point, + Float32[0.37]; + obj_weight = 0.61f0, + ) ≈ 0.37f0 .* expected_hessian + return +end + +function test_nonlinear_constraint_sum() + N = 20 + model = JuMP.Model() + JuMP.@variable(model, x[1:N]) + JuMP.@constraint(model, 0 <= sum(100(x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2 for i = 2:N) <= 1) + model_exa = ExaModels.ExaModel(model) + @test length(model_exa.cons) == 5 + return +end + +function test_sum_objective_decomposition() + model = JuMP.Model() + JuMP.@variable(model, x[1:3]) + JuMP.@expression(model, a, sum(x)) + JuMP.@expression(model, b, sum(x.^2)) + JuMP.@expression(model, c, sum(exp.(x))) + JuMP.@objective(model, Min, a + b + c) + model_exa = ExaModels.ExaModel(model) + # The five cons are the Constraint + + # i[2] * x[i[1]], i[2] * x[i[1]]^2, and exp(i[1]) + @test length(model_exa.objs) == 3 + return +end + +function test_sum_constraint_decomposition() + model = JuMP.Model() + JuMP.@variable(model, x[1:3]) + JuMP.@expression(model, a, sum(x)) + JuMP.@expression(model, b, sum(x.^2)) + JuMP.@expression(model, c, sum(exp.(x))) + JuMP.@constraint(model, a + b + c == 0) + model_exa = ExaModels.ExaModel(model) + # The four cons are the Constraint + + # i[2] * x[i[1]], i[2] * x[i[1]]^2, and exp(i[1]) + @test length(model_exa.cons) == 4 + return +end + +function test_sum_constraint_decomposition_multiple_rhs_terms() + model = JuMP.Model() + JuMP.@variable(model, x[1:3]) + JuMP.@expression(model, a, sum(x)) + JuMP.@expression(model, b, sum(x.^2)) + JuMP.@expression(model, c, sum(exp.(x))) + JuMP.@constraint(model, a == b + c - 2) + model_exa = ExaModels.ExaModel(model) + # TODO(odow): we'd like this one to be the same as + # test_sum_constraint_decomposition, but it requires fixing how we handle + # :(-(arg)) terms. + @test_broken length(model_exa.cons) == 4 + return end end # module diff --git a/test/Project.toml b/test/Project.toml index 043fb9ec8..61a2645a2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,6 +8,7 @@ JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" diff --git a/test/SIMDModeTest/SIMDModeTest.jl b/test/SIMDModeTest/SIMDModeTest.jl new file mode 100644 index 000000000..3652e4d9f --- /dev/null +++ b/test/SIMDModeTest/SIMDModeTest.jl @@ -0,0 +1,137 @@ +module SIMDModeTest + +using Test +import ExaModels +import MathOptInterface as MOI + +function runtests() + + + + + @testset "SIMDMode adapter" begin + x, y = MOI.VariableIndex(1), MOI.VariableIndex(2) + mode = ExaModels.SIMDMode() + @test mode isa MOI.Nonlinear.AbstractAutomaticDifferentiation + model = MOI.Nonlinear.model(mode) + @test model isa MOI.Nonlinear.ModelWithOracles + # Objective: x^2 (handled natively by ExaModels, no quad layer). + MOI.Nonlinear.set_objective( + model, + MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ), + ) + # Constraints, in row order: the oracle layer's rows come first. + oracle = MOI.VectorNonlinearOracle(; + dimension = 1, + l = [0.0], + u = [1.0], + eval_f = (ret, z) -> (ret[1] = z[1]^2), + jacobian_structure = [(1, 1)], + eval_jacobian = (ret, z) -> (ret[1] = 2.0 * z[1]), + hessian_lagrangian_structure = [(1, 1)], + eval_hessian_lagrangian = (ret, z, μ) -> (ret[1] = 2.0 * μ[1]), + ) + MOI.Nonlinear.add_constraint(model, MOI.VectorOfVariables([x]), oracle) + MOI.Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction( + [MOI.ScalarAffineTerm(2.0, x), MOI.ScalarAffineTerm(3.0, y)], + 0.0, + ), + MOI.LessThan(4.0), + ) + MOI.Nonlinear.add_constraint( + model, + MOI.ScalarQuadraticFunction( + [ + MOI.ScalarQuadraticTerm(2.0, x, x), + MOI.ScalarQuadraticTerm(1.0, x, y), + ], + [MOI.ScalarAffineTerm(1.0, y)], + 0.0, + ), + MOI.Interval(0.0, 1.0), + ) + sin_x = MOI.ScalarNonlinearFunction(:sin, Any[x]) + MOI.Nonlinear.add_constraint(model, sin_x, MOI.LessThan(0.5)) + d = MOI.Nonlinear.Evaluator(model, mode, [x, y]) + @test d isa MOI.Nonlinear.EvaluatorWithOracles + @test MOI.features_available(d) == [:Grad, :Jac, :Hess] + # Row queries work before MOI.initialize. + @test MOI.Nonlinear.num_constraints(d) == 4 + @test MOI.Nonlinear.constraint_bounds(d) == [ + MOI.NLPBoundsPair(0.0, 1.0), + MOI.NLPBoundsPair(-Inf, 4.0), + MOI.NLPBoundsPair(0.0, 1.0), + MOI.NLPBoundsPair(-Inf, 0.5), + ] + @test MOI.Nonlinear.constraint_linearity(d) == [ + MOI.Nonlinear.NONLINEAR, + MOI.Nonlinear.LINEAR, + MOI.Nonlinear.QUADRATIC, + MOI.Nonlinear.NONLINEAR, + ] + @test MOI.Nonlinear.objective_linearity(d) == MOI.Nonlinear.QUADRATIC + MOI.initialize(d, [:Grad, :Jac, :Hess]) + xv = [1.0, 2.0] + @test MOI.eval_objective(d, xv) == 1.0 + grad = fill(NaN, 2) + MOI.eval_objective_gradient(d, grad, xv) + @test grad == [2.0, 0.0] + g = fill(NaN, 4) + MOI.eval_constraint(d, g, xv) + @test g ≈ [1.0, 8.0, 5.0, sin(1.0)] + J_structure = MOI.jacobian_structure(d) + J_values = fill(NaN, length(J_structure)) + MOI.eval_constraint_jacobian(d, J_values, xv) + J = zeros(4, 2) + for ((row, col), value) in zip(J_structure, J_values) + J[row, col] += value + end + @test J ≈ [ + 2.0 0.0 + 2.0 3.0 + 4.0 2.0 + cos(1.0) 0.0 + ] + H_structure = MOI.hessian_lagrangian_structure(d) + σ, μ = 2.0, [10.0, 100.0, 1_000.0, 10_000.0] + H_values = fill(NaN, length(H_structure)) + MOI.eval_hessian_lagrangian(d, H_values, xv, σ, μ) + H = zeros(2, 2) + for ((row, col), value) in zip(H_structure, H_values) + H[row, col] += value + if row != col + H[col, row] += value + end + end + @test H[1, 1] ≈ 2σ + 2 * μ[1] + 2 * μ[3] - sin(1.0) * μ[4] + @test H[1, 2] ≈ μ[3] + @test H[2, 2] ≈ 0.0 + end + @testset "SIMDMode requires identity variable order" begin + x, y = MOI.VariableIndex(1), MOI.VariableIndex(2) + mode = ExaModels.SIMDMode() + model = MOI.Nonlinear.model(mode) + MOI.Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0), + MOI.LessThan(1.0), + ) + d = MOI.Nonlinear.Evaluator(model, mode, [y, x]) + @test_throws( + ErrorException( + "`ExaModels.SIMDMode` requires the variables of the model " * + "to be `MOI.VariableIndex.(1:n)`, in order.", + ), + MOI.initialize(d, [:Grad, :Jac]), + ) + end + return +end + +end # module diff --git a/test/runtests.jl b/test/runtests.jl index cc22b2dac..bbdc90fbf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -27,8 +27,12 @@ include("PrettyPrintTest.jl") include("ConcreteModeTest.jl") # include("OptimalControlTest/OptimalControlTest.jl") include("OracleTest/OracleTest.jl") +include("SIMDModeTest/SIMDModeTest.jl") @testset verbose = true "ExaModels test" begin + @info "Running SIMDMode Test" + SIMDModeTest.runtests() + @info "Running Argument Test" ArgumentTest.runtests()