Create containers with map instead of for loops

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.

src/Containers/container.jl

@@ -1,5 +1,70 @@
-const ArrayIndices{N} = VectorizedProductIterator{NTuple{N, Base.OneTo{Int}}}
+#  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)

src/Containers/macro.jl

@@ -1,5 +1,21 @@
+#  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/.
+
 using Base.Meta
 
+_get_name(c::Symbol) = c
+_get_name(c::Nothing) = ()
+_get_name(c::AbstractString) = c
+function _get_name(c::Expr)
+    if c.head == :string
+        return c
+    else
+        return c.args[1]
+    end
+end
+
 """
     _extract_kw_args(args)
 
@@ -65,8 +81,8 @@ function _parse_ref_sets(_error::Function, expr::Expr)
     c = copy(expr)
     idxvars = Any[]
     idxsets = Any[]
-    # On 0.7, :(t[i;j]) is a :ref, while t[i,j;j] is a :typed_vcat.
-    # In both cases :t is the first arg.
+    # `:(t[i,j;k])` is a :ref, while `:(t[i;j])` is a :typed_vcat.
+    # In both cases `:t` is the first arg.
     if isexpr(c, :typed_vcat) || isexpr(c, :ref)
         popfirst!(c.args)
     end
@@ -176,12 +192,71 @@ function parse_container(_error, var, value, requested_container)
     return container_code(idxvars, indices, value, requested_container)
 end
 
+"""
+    @container([i=..., j=..., ...], expr)
+
+Create a container with indices `i`, `j`, ... and values given by `expr` that
+may depend on the value of the indices.
+
+    @container(ref[i=..., j=..., ...], expr)
+
+Same as above but the container is assigned to the variable of name `ref`.
+
+The type of container can be controlled by the `container` keyword. See
+[Containers in macros](@ref). Note that when the index set is explicitly
+given as `1:n` for any expression `n`, it is transformed to `Base.OneTo(n)`
+before being given to [`container`](@ref).
+
+## Examples
+
+```jldoctest
+julia> Containers.@container([i = 1:3, j = 1:3], i + j)
+3×3 Array{Int64,2}:
+ 2  3  4
+ 3  4  5
+ 4  5  6
+
+julia> I = 1:3
+1:3
+
+julia> Containers.@container([i = I, j = I], i + j)
+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 = 2:3, j = 1:3], i + j)
+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 = 1:3, j = 1:3; i <= j], i + j)
+JuMP.Containers.SparseAxisArray{Int64,2,Tuple{Int64,Int64}} with 6 entries:
+  [1, 2]  =  3
+  [2, 3]  =  5
+  [3, 3]  =  6
+  [2, 2]  =  4
+  [1, 1]  =  2
+  [1, 3]  =  4
+```
+"""
 macro container(args...)
     args, kw_args, requested_container = _extract_kw_args(args)
     @assert length(args) == 2
     @assert isempty(kw_args)
     var, value = args
-    name = var.args[1]
     code = parse_container(error, var, esc(value), requested_container)
-    return :($(esc(name)) = $code)
+    anonvar = isexpr(var, :vect) || isexpr(var, :vcat)
+    if anonvar
+        return code
+    else
+        name = Containers._get_name(var)
+        return :($(esc(name)) = $code)
+    end
 end

src/Containers/nested_iterator.jl

@@ -1,3 +1,8 @@
+#  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/.
+
 """
     struct NestedIterator{T}
         iterators::T # Tuple of functions

src/Containers/vectorized_product_iterator.jl

@@ -1,10 +1,16 @@
+#  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/.
+
 """
     struct VectorizedProductIterator{T}
         prod::Iterators.ProductIterator{T}
     end
 
-Same as `Base.Iterators.ProductIterator` except that it is independent
-on the `IteratorSize` of the elements of `prod.iterators`.
+Cartesian product of the iterators `prod.iterators`. It is the same iterator as
+`Base.Iterators.ProductIterator` except that it is independent of the
+`IteratorSize` of the elements of `prod.iterators`.
 For instance:
 * `size(Iterators.product(1, 2))` is `tuple()` while
   `size(VectorizedProductIterator(1, 2))` is `(1, 1)`.

src/macros.jl

@@ -41,17 +41,6 @@ function _add_kw_args(call, kw_args)
     end
 end
 
-_get_name(c::Symbol) = c
-_get_name(c::Nothing) = ()
-_get_name(c::AbstractString) = c
-function _get_name(c::Expr)
-    if c.head == :string
-        return c
-    else
-        return c.args[1]
-    end
-end
-
 _valid_model(m::AbstractModel, name) = nothing
 _valid_model(m, name) = error("Expected $name to be a JuMP model, but it has type ", typeof(m))
 
@@ -383,7 +372,7 @@ function _constraint_macro(args, macro_name::Symbol, parsefun::Function)
 
     anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(extra) != 1
     variable = gensym()
-    name = _get_name(c)
+    name = Containers._get_name(c)
     base_name = anonvar ? "" : string(name)
     # TODO: support the base_name keyword argument
 
@@ -841,7 +830,7 @@ macro expression(args...)
     if anonvar
         macro_code = _macro_return(code)
     else
-        macro_code = _macro_assign_and_return(code, variable, _get_name(c),
+        macro_code = _macro_assign_and_return(code, variable, Containers._get_name(c),
                                               model_for_registering = m)
     end
     return _assert_valid_model(m, macro_code)
@@ -1167,7 +1156,7 @@ macro variable(args...)
     anonvar && explicit_comparison && _error("Cannot use explicit bounds via >=, <= with an anonymous variable")
     variable = gensym()
     # TODO: Should we generate non-empty default names for variables?
-    name = _get_name(var)
+    name = Containers._get_name(var)
     if isempty(base_name_kw_args)
         base_name = anonvar ? "" : string(name)
     else
@@ -1348,7 +1337,7 @@ macro NLconstraint(m, x, args...)
         macro_code = _macro_return(creation_code)
     else
         macro_code = _macro_assign_and_return(creation_code, variable,
-                                              _get_name(c),
+                                              Containers._get_name(c),
                                               model_for_registering = esc_m)
     end
     return _assert_valid_model(esc_m, macro_code)
@@ -1397,7 +1386,7 @@ macro NLexpression(args...)
         macro_code = _macro_return(creation_code)
     else
         macro_code = _macro_assign_and_return(creation_code, variable,
-                                              _get_name(c),
+                                              Containers._get_name(c),
                                               model_for_registering = esc(m))
     end
     return _assert_valid_model(esc(m), macro_code)
@@ -1466,6 +1455,6 @@ macro NLparameter(m, ex, extra...)
     # TODO: NLparameters are not registered in the model because we don't yet
     # have an anonymous version.
     macro_code = _macro_assign_and_return(creation_code, variable,
-                                          _get_name(c))
+                                          Containers._get_name(c))
     return _assert_valid_model(esc_m, macro_code)
 end

test/Containers/Containers.jl

@@ -6,5 +6,7 @@ using JuMP.Containers
     include("DenseAxisArray.jl")
     include("SparseAxisArray.jl")
     include("generate_container.jl")
+    include("nested_iterator.jl")
+    include("nested_iterator.jl")
     include("macro.jl")
 end

test/Containers/macro.jl

@@ -6,7 +6,7 @@ using JuMP.Containers
     @testset "Array" begin
         Containers.@container(x[i = 1:3], i^2)
         @test x isa Vector{Int}
-        Containers.@container(x[i = 1:3, j = 1:3], i^2)
+        x = Containers.@container([i = 1:3, j = 1:3], i^2)
         @test x isa Matrix{Int}
     end
     @testset "DenseAxisArray" begin

test/Containers/nested_iterator.jl

@@ -0,0 +1,14 @@
+@testset "Nested Iterator" begin
+    iterators = (() -> 1:3, i -> 1:i)
+    condition(i, j) = j > i
+    @test isempty(Containers.nested(iterators..., condition=condition))
+    @test isempty(collect(Containers.nested(iterators..., condition=condition)))
+    condition(i, j) = isodd(i) || isodd(j)
+    @test collect(Containers.nested(iterators..., condition=condition)) == [
+        (1, 1)
+        (2, 1)
+        (3, 1)
+        (3, 2)
+        (3, 3)
+    ]
+end

test/Containers/vectorized_product_iterator.jl

@@ -0,0 +1,9 @@
+@testset "Vectorized Product Iterator" begin
+    I = [1 2
+         3 4]
+    @test isempty(Containers.vectorized_product(2, I, 1:0))
+    @test isempty(collect(Containers.vectorized_product(2, I, 1:0)))
+    @test collect(Containers.vectorized_product(2, I)) == [
+        (2, 1) (2, 3) (2, 2), (2, 4)
+    ]
+end