A julia function to stack two arrays of size (m, n), to get one of size (2, m, n)

I want to stack two arrays of size (m, n) to get an array of size (2, m, n), respectively (m, n, 2).

With numpy.stack, np.stack((x,y)) returns an array of size (2, m, n)
while
np.stack((x,y), axis=-1)- an array of size (m, n, 2).

I found a similar question, here https://discourse.julialang.org/t/function-like-numpy-stack/39394,
but the given answer doesn’t help in my case.

Could be this, probably there are other alternatives:

julia> x = rand(3,2) ; y = rand(3,2);

julia> reshape(hcat(x,y),3,2,2)
3×2×2 Array{Float64,3}:
[:, :, 1] =
 0.322737  0.808556
 0.644822  0.926724
 0.570159  0.548639

[:, :, 2] =
 0.740272  0.31467
 0.947123  0.843563
 0.952541  0.0188235


cat and permutedims should do the trick

julia> a = rand(1:10, 4, 3)
4×3 Array{Int64,2}:
 10   6  5
  9  10  9
  7   8  2
  1   4  2

julia> b = rand(11:20, 4, 3)
4×3 Array{Int64,2}:
 11  19  11
 17  16  19
 13  18  13
 20  16  16

julia> cat(a,b; dims=3)
4×3×2 Array{Int64,3}:
[:, :, 1] =
 10   6  5
  9  10  9
  7   8  2
  1   4  2

[:, :, 2] =
 11  19  11
 17  16  19
 13  18  13
 20  16  16

julia> permutedims(ans, [3, 1, 2])
2×4×3 Array{Int64,3}:
[:, :, 1] =
 10   9   7   1
 11  17  13  20

[:, :, 2] =
  6  10   8   4
 19  16  18  16

[:, :, 3] =
  5   9   2   2
 11  19  13  16
1 Like

I think this might be a bit clearer:

const newaxis = [CartesianIndex()]
vcat(x[newaxis,:,:], y[newaxis,:,:])
# or:
cat(x[newaxis,:,:], y[newaxis,:,:], dims=3)

or to avoid unecessary allocations:

vcat(view(x,newaxis,:,:), view(y,newaxis,:,:))

As mentioned in the linked post, stacking them to size (2,m,n) is probably not a good idea. If you’re doing that, it’s probably better to rearrange your data some other way.

You can do it of course, but it’s normally not optimal to just duplicate numpy-funcionality without considering the differences between the data structures.

4 Likes

Thank you @lmiq, @MarcMush, @yha, and @DNF for your answers. I selected that given by @MarcMush, because of its simplicity. It is an intuitive definition.