I often run into this situation and have yet to find a solution other than writing a for-loop.
Is there a way to vectorize the MWE below using features in Base?
Suppose I have a function accepting a 2d array and two integers that returns a transformed array:
function foo(a::Array{Float64,2}, b::Int, c::Int)
#some transform
return newa
end
I want to be able to generalize the usage of foo() with a broadcast:
a = ones(5,5)
b = [1,2,3]
c = [1,2,3]
result = zeros(5,5,3)
result .= foo.(a, b, c)
In this case I would intuit that a gets expanded on its 3rd dimension. However, broadcast can’t handle that. The closest one-liner solution I can come up with is:
result = foo.([a],b,c)
…but this returns an Array of Arrays which obviously isn’t great either.
I had thought the fairly recent addition of eachslice() might help but it doesn’t so much in this situation.
Any ideas? Is it possible to provide a clue to broadcast that indicates how the first Array should be expanded? Or is it “just write a for loop”???
Oh, I missed the signature of foo(a::Array{Float64,2}, b::Int, c::Int). My answer assumes that foo is truly a scalar operation and operates element wise across all three arrays. You want to operate element wise across b and c, but treat a as an array, and concatenate the answers into a third dimension. That’s simply not what broadcast does.
Yeah, I assumed it was pushing things into odd territory. But seeing as I often pleasantly find clever (often undocumented) tricks, worth an ask. I will likely just write an extra dispatch that envelops the original function in a for loop. Thanks!
It’s the concatenating part that’s the challenge to express with broadcast alone.
Using Ref is the same as doing [a], which was tried in the original post and returns an array of arrays. You can always concatenate them together afterwards, though.