Merge pull request #2070 from JuliaOpt/bl/@container

Create containers with map instead of for loops

Author: blegat

docs/src/constraints.md

@@ -251,7 +251,7 @@ following.
 One way of adding a group of constraints compactly is the following:
 ```jldoctest constraint_arrays; setup=:(model=Model(); @variable(model, x))
 julia> @constraint(model, con[i = 1:3], i * x <= i + 1)
-3-element Array{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,1}:
+3-element Array{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},1}:
  con[1] : x <= 2.0
  con[2] : 2 x <= 3.0
  con[3] : 3 x <= 4.0
@@ -264,7 +264,7 @@ julia> con[1]
 con[1] : x <= 2.0
 
 julia> con[2:3]
-2-element Array{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,1}:
+2-element Array{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},1}:
  con[2] : 2 x <= 3.0
  con[3] : 3 x <= 4.0
 ```
@@ -273,7 +273,7 @@ Anonymous containers can also be constructed by dropping the name (e.g. `con`)
 before the square brackets:
 ```jldoctest constraint_arrays
 julia> @constraint(model, [i = 1:2], i * x <= i + 1)
-2-element Array{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,1}:
+2-element Array{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},1}:
  x <= 2.0
  2 x <= 3.0
 ```
@@ -294,10 +294,10 @@ variables.
 
 ```jldoctest constraint_jumparrays; setup=:(model=Model(); @variable(model, x))
 julia> @constraint(model, con[i = 1:2, j = 2:3], i * x <= j + 1)
-2-dimensional DenseAxisArray{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,2,...} with index sets:
-    Dimension 1, 1:2
+2-dimensional DenseAxisArray{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},2,...} with index sets:
+    Dimension 1, Base.OneTo(2)
     Dimension 2, 2:3
-And data, a 2×2 Array{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,2}:
+And data, a 2×2 Array{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},2}:
  con[1,2] : x <= 3.0    con[1,3] : x <= 4.0
  con[2,2] : 2 x <= 3.0  con[2,3] : 2 x <= 4.0
 ```
@@ -311,7 +311,7 @@ similar to the [syntax for constructing](@ref variable_sparseaxisarrays) a
 
 ```jldoctest constraint_jumparrays; setup=:(model=Model(); @variable(model, x))
 julia> @constraint(model, con[i = 1:2, j = 1:2; i != j], i * x <= j + 1)
-JuMP.Containers.SparseAxisArray{ConstraintRef{Model,C,Shape} where Shape<:AbstractShape where C,2,Tuple{Any,Any}} with 2 entries:
+JuMP.Containers.SparseAxisArray{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},2,Tuple{Int64,Int64}} with 2 entries:
   [1, 2]  =  con[1,2] : x <= 3.0
   [2, 1]  =  con[2,1] : 2 x <= 2.0
 ```

docs/src/containers.md

@@ -21,10 +21,15 @@ JuMP.Containers.SparseAxisArray
 Containers in macros
 --------------------
 
-The `generate_container` function encodes the logic for how containers are
-constructed in JuMP's macros.
+The `container` function encodes the logic for how containers are
+constructed in JuMP's macros. The `@container` macro is available to create
+containers independently of any JuMP model.
 ```@docs
-JuMP.Containers.generate_container
+JuMP.Containers.container
+JuMP.Containers.default_container
+JuMP.Containers.VectorizedProductIterator
+JuMP.Containers.NestedIterator
+JuMP.Containers.@container
 ```
 
 In the [`@variable`](@ref) (resp. [`@constraint`](@ref)) macro, containers of
@@ -44,12 +49,12 @@ Each expression `index_set_i` can either be
   keyword arguments to be expressions depending on the `index_name`.
 
 The macro then creates the container using the
-[`JuMP.Containers.generate_container`](@ref) function with the following
+[`JuMP.Containers.container`](@ref) function with the following
 arguments:
 
-1. `VariableRef` for the [`@variable`](@ref) macro and `ConstraintRef` for the
-   [`@constraint`](@ref) macro.
-2. The index variables and arbitrary symbols for dimensions for which no
-   variable index is specified.
-3. The index sets specified.
-4. The value of the `keyword` argument if given or `:Auto`.
+1. A function taking as argument the value of the indices and returning the
+   value to be stored in the container, e.g. a variable for the
+   [`@variable`](@ref) macro and a constraint for the [`@constraint`](@ref)
+   macro.
+2. An iterator over the indices of the container.
+4. The value of the `container` keyword argument if given.

docs/src/variables.md

@@ -324,7 +324,7 @@ return a `DenseAxisArray`. For example:
 ```jldoctest variables_jump_arrays; setup=:(model=Model())
 julia> @variable(model, x[1:2, [:A,:B]])
 2-dimensional DenseAxisArray{VariableRef,2,...} with index sets:
-    Dimension 1, 1:2
+    Dimension 1, Base.OneTo(2)
     Dimension 2, Symbol[:A, :B]
 And data, a 2×2 Array{VariableRef,2}:
  x[1,A]  x[1,B]
@@ -371,7 +371,7 @@ For example, this applies when indices have a dependence upon previous
 indices (called *triangular indexing*). JuMP supports this as follows:
 ```jldoctest; setup=:(model=Model())
 julia> @variable(model, x[i=1:2, j=i:2])
-JuMP.Containers.SparseAxisArray{VariableRef,2,Tuple{Any,Any}} with 3 entries:
+JuMP.Containers.SparseAxisArray{VariableRef,2,Tuple{Int64,Int64}} with 3 entries:
   [1, 2]  =  x[1,2]
   [2, 2]  =  x[2,2]
   [1, 1]  =  x[1,1]
@@ -382,7 +382,7 @@ syntax appends a comparison check that depends upon the named indices and is
 separated from the indices by a semi-colon (`;`). For example:
 ```jldoctest; setup=:(model=Model())
 julia> @variable(model, x[i=1:4; mod(i, 2)==0])
-JuMP.Containers.SparseAxisArray{VariableRef,1,Tuple{Any}} with 2 entries:
+JuMP.Containers.SparseAxisArray{VariableRef,1,Tuple{Int64}} with 2 entries:
   [4]  =  x[4]
   [2]  =  x[2]
 ```

src/Containers/Containers.jl

@@ -26,5 +26,10 @@ export DenseAxisArray, SparseAxisArray
 include("DenseAxisArray.jl")
 include("SparseAxisArray.jl")
 include("generate_container.jl")
+include("vectorized_product_iterator.jl")
+include("nested_iterator.jl")
+include("no_duplicate_dict.jl")
+include("container.jl")
+include("macro.jl")
 
 end

src/Containers/container.jl

@@ -0,0 +1,95 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+"""
+    default_container(indices)
+
+If `indices` is a [`NestedIterator`](@ref), return a
+[`SparseAxisArray`](@ref). Otherwise, `indices` should be
+a `VectorizedProductIterator` and the function returns
+`Array` if all iterators of the product are `Base.OneTo` and retunrs
+[`DenseAxisArray`](@ref) otherwise.
+"""
+function default_container end
+
+"""
+    container(f::Function, indices, ::Type{C})
+
+Create a container of type `C` with indices `indices` and values at given
+indices given by `f`.
+
+    container(f::Function, indices)
+
+Create a container with indices `indices` and values at given indices given by
+`f`. The type of container used is determined by [`default_container`](@ref).
+
+## Examples
+
+```@jldoctest
+julia> Containers.container((i, j) -> i + j, Containers.vectorized_product(Base.OneTo(3), Base.OneTo(3)))
+3×3 Array{Int64,2}:
+ 2  3  4
+ 3  4  5
+ 4  5  6
+
+julia> Containers.container((i, j) -> i + j, Containers.vectorized_product(1:3, 1:3))
+2-dimensional DenseAxisArray{Int64,2,...} with index sets:
+    Dimension 1, 1:3
+    Dimension 2, 1:3
+And data, a 3×3 Array{Int64,2}:
+ 2  3  4
+ 3  4  5
+ 4  5  6
+
+julia> Containers.container((i, j) -> i + j, Containers.vectorized_product(2:3, Base.OneTo(3)))
+2-dimensional DenseAxisArray{Int64,2,...} with index sets:
+    Dimension 1, 2:3
+    Dimension 2, Base.OneTo(3)
+And data, a 2×3 Array{Int64,2}:
+ 3  4  5
+ 4  5  6
+
+julia> Containers.container((i, j) -> i + j, Containers.nested(() -> 1:3, i -> i:3, condition = (i, j) -> isodd(i) || isodd(j)))
+SparseAxisArray{Int64,2,Tuple{Int64,Int64}} with 5 entries:
+  [1, 2]  =  3
+  [2, 3]  =  5
+  [3, 3]  =  6
+  [1, 1]  =  2
+  [1, 3]  =  4
+```
+"""
+function container end
+
+container(f::Function, indices) = container(f, indices, default_container(indices))
+
+const ArrayIndices{N} = VectorizedProductIterator{NTuple{N, Base.OneTo{Int}}}
+default_container(::ArrayIndices) = Array
+function container(f::Function, indices::ArrayIndices, ::Type{Array})
+    return map(I -> f(I...), indices)
+end
+function _oneto(indices)
+    if indices isa UnitRange{Int} && indices == 1:length(indices)
+        return Base.OneTo(length(indices))
+    end
+    error("Index set for array is not one-based interval.")
+end
+function container(f::Function, indices::VectorizedProductIterator,
+                   ::Type{Array})
+    container(f, vectorized_product(_oneto.(indices.prod.iterators)...), Array)
+end
+default_container(::VectorizedProductIterator) = DenseAxisArray
+function container(f::Function, indices::VectorizedProductIterator,
+                   ::Type{DenseAxisArray})
+    return DenseAxisArray(map(I -> f(I...), indices), indices.prod.iterators...)
+end
+default_container(::NestedIterator) = SparseAxisArray
+function container(f::Function, indices,
+                   ::Type{SparseAxisArray})
+    # Same as `map` but does not allocate the resulting vector.
+    mappings = Base.Generator(I -> I => f(I...), indices)
+    # Same as `Dict(mapping)` but it will error if two indices are the same.
+    data = NoDuplicateDict(mappings)
+    return SparseAxisArray(data.dict)
+end