Bl/tmp

Project.toml

@@ -6,10 +6,12 @@ authors = ["Benoît Legat <benoit.legat@gmail.com>"]
 [deps]
 ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 
 [compat]
 ExaModels = "0.9"
 JuMP = "1.29"
+LinearAlgebra = "1.12.0"
 MathOptInterface = "1.46"
 julia = "1.10"

examples/quadrotor.jl

@@ -7,10 +7,11 @@ N = 3
 
 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)
+function d(i, j, N)
+    return (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)
+end
 dt = 1/N
 R = fill(1 / 10, 4)
 Q = [1, 0, 1, 0, 1, 0, 1, 1, 1]
@@ -27,12 +28,7 @@ model = Model()
 using GenOpt
 container = ParametrizedArray
 
-@constraint(
-    model,
-    [i in 1:n],
-    x[1, i] == x0[i],
-    container = container,
-)
+@constraint(model, [i in 1:n], x[1, i] == x0[i], container = container,)
 @constraint(
     model,
     [i in 1:N],
@@ -84,13 +80,17 @@ container = ParametrizedArray
     [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,
+    (
+        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,
+    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(

src/GenOpt.jl

@@ -5,8 +5,11 @@
 
 module GenOpt
 
+import LinearAlgebra
+
 include("MOI_wrapper.jl")
 include("JuMP_wrapper.jl")
+include("array_of_variables.jl")
 include("examodels.jl")
 
 # Copied from JuMP.jl:

src/JuMP_wrapper.jl

@@ -27,38 +27,6 @@ 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}
@@ -106,7 +74,8 @@ end
 
 function Base.getindex(expr::ExprGenerator, i::Integer)
     idx = CartesianIndices(Base.OneTo.(_size(expr)))[i]
-    values = [expr.iterators[i].values[idx[i]] for i in eachindex(expr.iterators)]
+    values =
+        [expr.iterators[i].values[idx[i]] for i in eachindex(expr.iterators)]
     return index_iterators(expr.expr.expr, values)
 end
 

src/MOI_wrapper.jl

@@ -41,7 +41,10 @@ end
 function MOI.Utilities.is_canonical(f::FunctionGenerator)
     return MOI.Utilities.is_canonical(f.func)
 end
-function MOI.Utilities.is_coefficient_type(::Type{FunctionGenerator{E}}, ::Type{T}) where {E,T}
+function MOI.Utilities.is_coefficient_type(
+    ::Type{FunctionGenerator{E}},
+    ::Type{T},
+) where {E,T}
     return MOI.Utilities.is_coefficient_type(E, T)
 end
 
@@ -54,12 +57,18 @@ function Base.copy(f::SumGenerator{F}) where {F}
     return SumGenerator{F}(copy(f.func), f.iterators)
 end
 
-function MOI.Utilities.map_indices(::MOI.Utilities.IndexMap, func::Union{FunctionGenerator,SumGenerator})
+function MOI.Utilities.map_indices(
+    ::MOI.Utilities.IndexMap,
+    func::Union{FunctionGenerator,SumGenerator},
+)
     # TODO check it's identity
     return func
 end
 
-function MOI.Utilities.map_indices(::Function, func::Union{FunctionGenerator,SumGenerator})
+function MOI.Utilities.map_indices(
+    ::Function,
+    func::Union{FunctionGenerator,SumGenerator},
+)
     # TODO check it's identity
     return func
 end

src/array_of_variables.jl

@@ -0,0 +1,60 @@
+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
+
+function JuMP.Containers.container(
+    f::Function,
+    indices::JuMP.Containers.VectorizedProductIterator{NTuple{N,Base.OneTo{Int}}},
+    ::Type{ArrayOfVariables},
+) where {N}
+    return to_generator(JuMP.Containers.container(f, indices, Array))
+end
+
+JuMP._is_real(::ArrayOfVariables) = true
+function JuMP.moi_function(array::ArrayOfVariables)
+    return ContiguousArrayOfVariables(array.offset, array.size)
+end
+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
+
+function LinearAlgebra.mul(A::ArrayOfVariables, B::Array)
+    return JuMP.NonlinearExpr(:*, Any[A, B])
+end
+
+function Base.show(io::IO, ::MIME"text/plain", v::ArrayOfVariables)
+    return println(io, Base.summary(v), " with offset ", v.offset)
+end
+
+function Base.show(io::IO, v::ArrayOfVariables)
+    return show(io, MIME"text/plain"(), v)
+end