Add ExaModels wrapper

Author: blegat

examples/quad.jl

@@ -0,0 +1,115 @@
+N = 3
+backend = nothing
+
+n = 9
+p = 4
+d(i, j, N) =
+    (j == 1 ? 1 * sin(2 * pi / N * i) : 0.0) +
+    (j == 3 ? 2 * sin(4 * pi / N * i) : 0.0) +
+    (j == 5 ? 2 * i / N : 0.0)
+dt = 1/N
+R = fill(1 / 10, 4)
+Q = [1, 0, 1, 0, 1, 0, 1, 1, 1]
+Qf = [1, 0, 1, 0, 1, 0, 1, 1, 1] / dt
+
+x0 = zeros(n)
+
+using JuMP
+model = Model()
+
+@variable(model, x[1:(N+1), 1:n])
+@variable(model, u[1:N, 1:p])
+
+using GenOpt
+container = ParametrizedArray
+
+@constraint(
+    model,
+    [i in 1:n],
+    x[1, i] == x0[i],
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 1] == x[i, 1] + (x[i, 2]) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 2] ==
+    x[i, 2] +
+    (
+        u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * cos(x[i, 9]) +
+        u[i, 1] * sin(x[i, 7]) * sin(x[i, 9])
+    ) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 3] == x[i, 3] + (x[i, 4]) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 4] ==
+    x[i, 4] +
+    (
+        u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * sin(x[i, 9]) -
+        u[i, 1] * sin(x[i, 7]) * cos(x[i, 9])
+    ) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 5] == x[i, 5] + (x[i, 6]) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 6] == x[i, 6] + (u[i, 1] * cos(x[i, 7]) * cos(x[i, 8]) - 9.8) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 7] ==
+    x[i, 7] +
+    (u[i, 2] * cos(x[i, 7]) / cos(x[i, 8]) + u[i, 3] * sin(x[i, 7]) / cos(x[i, 8])) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 8] == x[i, 8] + (-u[i, 2] * sin(x[i, 7]) + u[i, 3] * cos(x[i, 7])) * dt,
+    container = container,
+)
+@constraint(
+    model,
+    [i in 1:N],
+    x[i+1, 9] ==
+    x[i, 9] +
+    (
+        u[i, 2] * cos(x[i, 7]) * tan(x[i, 8]) +
+        u[i, 3] * sin(x[i, 7]) * tan(x[i, 8]) +
+        u[i, 4]
+    ) * dt,
+    container = container,
+)
+
+@objective(
+    model,
+    Min,
+    sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) +
+    sum(0.5 * Q[j] * (x[i, j] - d(i, j, N))^2 for i in 1:N, j in 1:n) +
+    sum(0.5 * Qf[j] * (x[N+1, j] - d(N + 1, j, N))^2 for j in 1:n),
+)
+
+using NLPModelsIpopt
+set_optimizer(model, () -> GenOpt.ExaOptimizer(ipopt))
+optimize!(model)

src/GenOpt.jl

@@ -7,6 +7,7 @@ module GenOpt
 
 include("MOI_wrapper.jl")
 include("JuMP_wrapper.jl")
+include("examodels.jl")
 
 # Copied from JuMP.jl:
 const _EXCLUDE_SYMBOLS = [Symbol(@__MODULE__), :eval, :include]

src/JuMP_wrapper.jl

@@ -15,7 +15,7 @@ include("operators.jl")
     end
 
 Iterator `iterators[index.value]`.
-"""
+""" # TODO remove
 struct IteratorInExpr
     iterators::Iterators
     index::IteratorIndex
@@ -27,6 +27,38 @@ JuMP._is_real(::Union{IteratorInExpr,IteratorIndex}) = true
 JuMP.moi_function(i::Union{IteratorInExpr,IteratorIndex}) = i
 JuMP.jump_function(_, i::Union{IteratorInExpr,IteratorIndex}) = i
 
+struct ArrayOfVariables{T,N} <: AbstractArray{JuMP.GenericVariableRef{T},N}
+    model::JuMP.GenericModel{T}
+    offset::Int64
+    size::NTuple{N,Int64}
+end
+
+Base.size(array::ArrayOfVariables) = array.size
+function Base.getindex(A::ArrayOfVariables{T}, I...) where {T}
+    index = A.offset + Base._to_linear_index(Base.CartesianIndices(A.size), I...)
+    return JuMP.GenericVariableRef{T}(A.model, MOI.VariableIndex(index))
+end
+
+JuMP._is_real(::ArrayOfVariables) = true
+JuMP.moi_function(array::ArrayOfVariables) = ContiguousArrayOfVariables(array.offset, array.size)
+function JuMP.jump_function(model::JuMP.GenericModel{T}, array::ContiguousArrayOfVariables{N}) where {T,N}
+    return ArrayOfVariables{T,N}(model, array.offset, array.size)
+end
+
+function Base.convert(::Type{ArrayOfVariables{T,N}}, array::Array{JuMP.GenericVariableRef{T},N}) where {T,N}
+    model = JuMP.owner_model(array[1])
+    offset = JuMP.index(array[1]).value - 1
+    for i in eachindex(array)
+        @assert JuMP.owner_model(array[i]) === model
+        @assert JuMP.index(array[i]).value == offset + i
+    end
+    return ArrayOfVariables{T,N}(model, offset, size(array))
+end
+
+function to_generator(array::Array{JuMP.GenericVariableRef{T},N}) where {T,N}
+    return convert(ArrayOfVariables{T,N}, array)
+end
+
 struct ExprGenerator{E,V<:JuMP.AbstractVariableRef} <:
        AbstractVector{JuMP.GenericNonlinearExpr{V}}
     expr::ExprTemplate{E,V}

src/MOI_wrapper.jl

@@ -5,6 +5,14 @@
 
 import MathOptInterface as MOI
 
+struct ContiguousArrayOfVariables{N} <: AbstractArray{MOI.VariableIndex,N}
+    offset::Int64
+    size::NTuple{N,Int64}
+end
+
+Base.copy(array::ContiguousArrayOfVariables) = array
+Base.size(array::ContiguousArrayOfVariables) = array.size
+
 """
     struct Iterator{T}
         values::Vector{T}
@@ -14,6 +22,8 @@ struct Iterator{T}
     values::Vector{T}
 end
 
+Iterator(values::AbstractArray) = Iterator(collect(values))
+
 struct IteratorIndex
     value::Int
 end
@@ -31,3 +41,10 @@ end
 function MOI.Utilities.is_canonical(f::FunctionGenerator)
     return MOI.Utilities.is_canonical(f.func)
 end
+function MOI.Utilities.map_indices(::MOI.Utilities.IndexMap, func::FunctionGenerator)
+    # TODO check it's identity
+    return func
+end
+function MOI.Utilities.is_coefficient_type(::Type{FunctionGenerator{E}}, ::Type{T}) where {E,T}
+    return MOI.Utilities.is_coefficient_type(E, T)
+end