Broadcast a function returning multiple values

If I understand it correctly, a function that looks like returning multiple values actually returns a Tuple of the values and therefore the result of broadcasting on the function is an array of the Tuples:

function f1(x,y,c); return x+c*y; end
function f2(x,y,c); return x+c*y, x-c*y; end

xs = 1:3
ys = 7:9

r1 = f1.(xs,ys,2)
r2 = f2.(xs,ys,2) # Vector of Tuples
s2, t2 = f2.(xs,ys,2);
s2 # the 1st Tuple
t2 # the 2nd Tuple

But, I need s2 and t2 to be vectors of returned values. What’s the idiomatic way to get them?

During this little research, I learned that columntable() from Tables can do that

using Tables
s2, t2 = columntable(f2.(xs,ys,2));
s2 # Vector
t2 # Vector

columntable() seems to be originally designed for NamedTuples. It assigns some names to the elements of the Tuple and convert the Vector of the NamedTuples to a NamedTuple of Vectors. In the above sample code, the tuple deconstruction strips out the names.

But, “transposing” Vectors of Tuples or Vectors of Vectors is a common enough situation that there should be a common idiom, I imagine.

For this specific case, you can do

julia> first.(r2), last.(r2)
([15, 18, 21], [-13, -14, -15])

but for the general idea, there is invert from SplitApplyCombine.jl

julia> using SplitApplyCombine: invert

julia> invert(r2)
([15, 18, 21], [-13, -14, -15])
3 Likes

Just to note that for this specific linear function f2, there is no need to broadcast, as it will consume (abstract) arrays too: f2(xs,ys,2) will produce the desired output and Julia will add the corresponding method automatically to f2.

For a more general case, one could add a method manually:
f2(x::AbstractArray, y::AbstractArray, c)

1 Like

There’s also

julia> s2, t2 = eachrow(stack(r2))
1 Like

Another option would be to use StructArrays.jl:

julia> using StructArrays

julia> f2(t::Tuple) = f2(t...);

julia> sa = StructArray((xs, ys, fill(2, size(xs))))
3-element StructArray(::UnitRange{Int64}, ::UnitRange{Int64}, ::Vector{Int64}) with eltype Tuple{Int64, Int64, Int64}:
 (1, 7, 2)
 (2, 8, 2)
 (3, 9, 2)

julia> out = f2.(sa)
3-element StructArray(::Vector{Int64}, ::Vector{Int64}) with eltype Tuple{Int64, Int64}:
 (15, -13)
 (18, -14)
 (21, -15)

julia> s2, t2 = StructArrays.components(out)
([15, 18, 21], [-13, -14, -15])

julia> f2(t::Tuple, c) = f2(t..., c);

julia> StructArrays.components(f2.(StructArray((xs, ys)), 2))  # You don't have to include the 2s into the StructArray
([15, 18, 21], [-13, -14, -15])

It’s a bit less convenient than e.g. the invert above, but it will normally be more performant, as f2.(sa) directly creates s2 and t2 (i.e. the StructArray out), instead of as an extra processing step.

julia> @btime SplitApplyCombine.invert(f2.($xs, $ys, 2))
  66.496 ns (6 allocations: 272 bytes)
([15, 18, 21], [-13, -14, -15])

julia> @btime StructArrays.components(f2.($sa))
  45.187 ns (4 allocations: 160 bytes)
([15, 18, 21], [-13, -14, -15])

julia> @btime StructArrays.components(f2.(StructArray(($xs, $ys, fill(2, size($xs))))))
  68.033 ns (6 allocations: 240 bytes)
([15, 18, 21], [-13, -14, -15])

julia> @btime StructArrays.components(f2.(StructArray(($xs, $ys)), 2))
  48.734 ns (4 allocations: 160 bytes)
([15, 18, 21], [-13, -14, -15])
2 Likes

Isn’t there anything more Julian for going from a n array of d tuples to a d array of n arrays?

There is the getindex option:

n, d = 10, 3
nt = [Tuple(randn(d)) for i in 1:n]
dn = [getindex.(nt, i) for i in 1:d]

IMHO, TensorCast.jl provides the most clear syntax:

using TensorCast
@cast dn[j][i] := nt[i][j]
2 Likes

Don’t know if that’s more Julian, but the zip/splat trick for transposing lists of lists works here:

zip(nt...) .|> collect

Imho, in Julia working with multidimensional arrays instead of nested arrays is often easier/nicer, i.e., just stack it and work with that.

2 Likes

I have these snippets lying around that I’ve used for this before

function repack(x::AbstractArray{T}) where T<:Tuple
	# turn Array{Tuple{T1,T2...}} to Tuple{Array{T1},Array{T2},...}
	fT = ntuple(i->fieldtype(T,i),fieldcount(T)) # fieldtypes(T) would be the obvious choice but is type-unstable
	arrs = similar.(Ref(x),fT)
	for i in eachindex(x,arrs...)
		setindex!.(arrs,x[i],Ref(i))
	end
	return arrs
end

function repack(x::AbstractArray{NamedTuple{N,T}}) where {N,T<:Tuple}
	# turn Array{NamedTuple{N,Tuple{T1,T2...}}} to NamedTuple{N,Tuple{Array{T1},Array{T2},...}}
	fT = ntuple(i->fieldtype(T,i),fieldcount(T)) # fieldtypes(T) would be the obvious choice but is type-unstable
	arrs = similar.(Ref(x),fT)
	for i in eachindex(x,arrs...)
		setindex!.(arrs,Tuple(x[i]),Ref(i))
	end
	return NamedTuple{N}(arrs)
end

But there isn’t (as part of Base) a way to make a broadcast/map of a generic function return a tuple of collections (rather than collection of tuples) in the first place, everything here will revolve around some post-processing transformation (whether it is lazy like StructArrays.jl or eager like this snippet).

Ah, the elusive Base.unzip() · Issue #13942 · JuliaLang/julia · GitHub

Even for the simplest 2-tuple, there’s lots of considerations. Personally, I think StructArrays are great. You can also assign into a pre-allocated struct array with .= to more explicitly demand the vectors you want.

2 Likes

Thank you all who answered!

That’s exactly what I had in mind (and more)! Looking at the Introdution of the doc, I get the impression that it also covers (tries to cover) cases the Tables package covers.


But, where do people talk about this kind of stuff? ChatGPT didn’t find this obvious solution. Instead, it recommended Tables.columntable! Do people discuss these things where search engines can’t reach? Or simply, is this thread the first public discussion?

Here in Discourse, and many times. Just one example.

You could consider benchmarking it against @sgaure’s solution eachrow(stack()) — to me, that looked the best.

I didn’t see that. Thanks. But is it straightforward to determine whether that’s the case or not? I can see that the + function has Array-Array versions and the * version has scalar-Array versions and x + c*y can be applied to xs + c*ys without broadcast, but I don’t know what other functions behave that way. For example, when I was new to Julia, I naïvely thought sin(xs) would work with an array xs as it does in Fortran. What about the ternary operator? Etc.

And then, in matrixA * matrixB, the * function does something different from matrixA .* matrixB !

So, what happens to my f2(xs, ys, c) if I use a vector as c ?

Also, while scalar * array behaves like scalar .* array, scalar / array doesn’t work! So, my f2 would stop working when f2(x,y,c) = (x + c/y, x - c/y) for f2(xs,ys,c).

I understand that the operators *, +, and - are “extended” for linear algebra: vec + vec isn’t “broadcast” in spirit but the standard addition of two vectors in linear algebra.

Conversely, if you mean broadcast, it’s better to write vec .+ vec even if the result is the same as vec + vec. Conversely, if I meant my f2 to be a linear-algebraic function, I would have written it as f2(v1, v2, c) = (v1 + c*v2, v1 - c*v2).

In that case, why not implement invert as that?

I want to learn as generic and as simple a solution as possible and stick to it, until I find that is the bottleneck of the performance my program. I don’t want to be looking for optimal solutions for each case.

Does that mean that this Discourse isn’t scanned by ChatGPT? I’ve been struggling to find right information about Julia.

I did search this Discourse but my keywords weren’t good enough. Each and every time I ask a question here, some kind, knowlegeable persons like yourself point out past threads or documentations.

The search engine of Discourse isn’t nearly as “intelligent”(?) as Google search or ChatGPT.

I don’t know anything about AI, but I imagine that humans still learn and organize knowledge differently than machines do.

1 Like