Create a generator which returns two values at once in Julia

user_17884397 0
Viewed 135

Given a generator:

myVec1 = rand(0:4, 2)
myVec2 = rand(0:4, 8)

myGen = (val1 + val2 for val1 in myVec1, val2 in myVec2)

This is basically a matrix with 2 columns. It can be seen by using collect(myGen).

How can I create a generator which yields two values per call (basically a column)?

Conceptually, something equivalent of:

for myCol in eachcol(collect(myGen))
    @show myCol;
end

Just without any explicit allocation of the matrix.

Can I wrap myGen for the following case:

for value1, value2 in myGen
  dosomethingelse1(value1, value2)
end

In other words, I am after a way to create a generator which returns 2 (or more?) consecutive values at once and can be used in a loop to do so.

So basically, we create a 2D array in the generator and I'd like to access the whole slice at once. I could do it with eachcol and eachrow for actual array, but what about the generator?

Here is a test case:

myVec1 = rand(0:4, 2);
myVec2 = rand(0:4, 800);

@btime begin
    myMat = [val1 + val2 for val1 in myVec1, val2 in myVec2];
    outVec = [sum(myCol) for myCol in eachcol(myMat)];
end

@btime begin
    myGen = (val1 + val2 for val1 in myVec1, val2 in myVec2);
    outVec = [sum(myCol) for myCol in Iterators.partition(myGen, 2)];
end

The solution by @Bogumił Kamiński indeed works, yet in practice, for some reason, it creates more allocations while the motivation was to reduce it.

julia generator
4 Answers

This is supported directly by Julia syntax. Iterating over generators of tuples works like interating over atomic values.

For an example you could try:

for (a,b) in ((x, 3x) for x in 1:4)
    println("a=$a, b=$b")
end
user_99577100

I assume you want something like:

julia> for (x1, x2) in Iterators.partition(1:10, 2)
           @show x1, x2
       end
(x1, x2) = (1, 2)
(x1, x2) = (3, 4)
(x1, x2) = (5, 6)
(x1, x2) = (7, 8)
(x1, x2) = (9, 10)

If this is what you want then Iterators.partition is a function you can use.

Edit: If you have two streams of sources use zip:

julia> for (x1, x2) in zip(1:5, 6:10)
           @show x1, x2
       end
(x1, x2) = (1, 6)
(x1, x2) = (2, 7)
(x1, x2) = (3, 8)
(x1, x2) = (4, 9)
(x1, x2) = (5, 10)

Edit 2: my first solution works already for your case:

julia> collect(myGen)
2×8 Matrix{Int64}:
 3  7  5  4  6  3  4  5
 1  5  3  2  4  1  2  3

julia> for (x1, x2) in Iterators.partition(myGen, 2)
           @show x1, x2
       end
(x1, x2) = (3, 1)
(x1, x2) = (7, 5)
(x1, x2) = (5, 3)
(x1, x2) = (4, 2)
(x1, x2) = (6, 4)
(x1, x2) = (3, 1)
(x1, x2) = (4, 2)
(x1, x2) = (5, 3)
user_12695670

Although the other answers are in some ways more general based on the information added by OP in their edits a more memory efficient option would be using nested generators. Something like:

function solution_nested(v1, v2)
    myGen = ((val1 + val2 for val1 in v1) for val2 in v2)
    [sum(myCol) for myCol in myGen]
end

When you test the solutions you should avoid using global variables and preferably wrap the solution in a function so you give Julia sufficient opportunity to optimize the code.

This solution gives the expected result of only one allocation:

julia> @btime solution_nested(myVec1, myVec2);
  1.856 μs (1 allocation: 6.38 KiB)

So while this solution does not quite fit the title it seems to fit what you are describing. We use a lazy sequence of lazy columns. The reason Iterators.partition is slow and memory inefficient is that it actually allocates the intermediate vectors of values in a partition: https://github.com/JuliaLang/julia/blob/dacf9d65aff4668b8fff25957d9aaa2cf03868c8/base/iterators.jl#L1232 .

user_152812240

You're basically missing brackets when destructuring the tuple of values in the second loop. To be more detailed, you can just return two values (a tuple) in your dosomething function. For example:

function dosomething(element)
    secondElement = element^2
    element, secondElement
end

And then you can use the loop by destructuring the return value, like so:

for (value1, value2) in myGen
    dosomethingelse(value1, value2)
end

If you want a full working example:

myArray = [1, 2, 3]

function dosomething(element)
    secondElement = element^2
    element, secondElement
end

myGen = (dosomething(myElement) for myElement in myArray)

function dosomethingelse(value1, value2)
    println("Value 1: $value1 \nValue 2: $value2 \n")
end

for (value1, value2) in myGen
    dosomethingelse(value1, value2)
end
user_144671320
Related