Filling an array with a generator that returns tuples

General Usage
, , , ,
1

I’m trying to create a function that takes a number of iterables and fills an array columnwise with each combination of those. The following function has the correct output, but the number of allocations is unexpected.

using Base.Iterators
using BenchmarkTools

function product_stack(iters...)
    output = zeros(length(iters), (*).(length.(iters)...))
    for (index, combination) in enumerate(product(iters...))
        output[:, index] .= combination
    end
    return output
end

julia> @btime product_stack(1:3, 1:3)
  7.512 μs (91 allocations: 3.69 KiB)
2×9 Matrix{Float64}:
 1.0  2.0  3.0  1.0  2.0  3.0  1.0  2.0  3.0
 1.0  1.0  1.0  2.0  2.0  2.0  3.0  3.0  3.0

Is it possible to do this with a single allocation?

Extra question:
I was hoping something like this might work, but unfortunately it does not:

function product_stack(iters...)
    output = zeros(length(iters), (*).(length.(iters)...))
    output .= product(iters...)
    return output
end

Why not?

2

This is much faster:

function product_stack2(iters...)
    p = flatten(product(iters...))
    reshape(collect(p), length(iters), :)
end

@btime product_stack2(1:3,1:3)
  101.695 ns (3 allocations: 320 bytes)
2×9 Matrix{Int64}:
 1  2  3  1  2  3  1  2  3
 1  1  1  2  2  2  3  3  3

The second example doesn’t work because when you say output .= x, this replaces all elements of output with the same x element. In your case, x is an iterator and each element of output is a Float64, they can’t match. You can modify it as follows but it won’t be as fast as the version above:

function product_stack(iters...)
    output = zeros(Int64, length(iters), (*).(length.(iters)...))
    output[1:end] .= flatten(product(iters...))
    return output
end
@btime product_stack(1:3,1:3)
  1.560 μs (17 allocations: 1.02 KiB)
2 Likes
3

Wow, thank you for your answer!
That broadened my horizon regarding iterators and especially reshape.

4

I stumbled upon this old post and fished out some scripts I had made long ago for a similar goal.
Stimulated by @Seif_Shebl 's brilliant solution, I tried to find a new one that has comparable performance

julia> @btime hcat(collect.(Base.Iterators.product(1:3,1:3))...)
       @btime mapreduce(collect, hcat, Base.Iterators.product(1:3,1:3))
       @btime reduce(hcat, collect.(Base.Iterators.product(1:3,1:3)));
  772.642 ns (16 allocations: 1.55 KiB)
  516.667 ns (17 allocations: 1.89 KiB)
  610.286 ns (19 allocations: 2.22 KiB)

julia> @btime reshape(collect(Iterators.flatten(Base.Iterators.product(1:3,1:3))),2,:);
  94.280 ns (3 allocations: 304 bytes)

julia>
julia> @btime  reinterpret(reshape,Int,foldr(prod_tuples, (1:3,1:3)))
  34.980 ns (1 allocation: 208 bytes)
2×9 reinterpret(reshape, Int64, ::Vector{Tuple{Int64, Int64}}) with eltype Int64:
 1  1  1  2  2  2  3  3  3
 1  2  3  1  2  3  1  2  3

#with


function prod_tuples(itrs1,itrs2)
l1,l2 = length(itrs1), length(itrs2)
n=length(itrs1[1])+length(itrs2[1])
m=Array{NTuple{n,Int64},1}(undef,l1*l2)
for (i,ei) in enumerate(itrs1)
    for (j,ej) in enumerate(itrs2)
        m[j+l2*(i-1)]=(ei,ej...)
    end
end
m
end

Another solution can be obtained by exploiting the possibilities of Cartesian indices

1 Like
5 
       function prod_itrs(itrs)
           foldr((itr1,itr2)->((e1,e2...) for e2 in itr2 for e1 in itr1), itrs)
       end

julia> @btime prod_itrs((1:3,1:3,:5:7))  0.001 ns (0 allocations: 0 bytes)
2 Likes