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
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.
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.
Clearer, maybe, but it is about a factor 10x slower that MLUtils unsqueeze method.
using MLUtils: unsqueeze
x = rand(2,3)
@benchmark ts[:, newaxis, :] # 576ns
@benchmark unsqueeze(ts, 2) # 50ns
You are comparing making a copy to a function that create a view (and also including the access to a global variable in your benchmark).
Creating a view view(x,:,newaxis,:) is basically free, but reading it later is indeed slower than reading the result of unsqueeze (which is doing reshape).
So much Julia magic going on in that solution 
I believe this is perfect use case for:
For other library solutions, you might also want to consider
using JuliennedArrays
Align([x, y], 2, 3) # Inner dimensions at 2 and 3, i.e., result 2 x 3 x 2
Align([x, y], 1, 2) # Inner dimensions at 1 and 2, i.e., result 3 x 2 x 2
or
using SplitApplyCombine
combinedims([x, y]) # outer dimension is last, i.e., result 3 x 2 x 2
combinedimsview([x, y]) # in case you want a view instead
And for completeness, combinedims([x, y], 1) if you want to put the outer dimension first.
Thanks, had missed that as it does not show up in the docstring … should have checked the actual methods instead.
[x;;;y]
[[x;;;y][:,1,:]' ;;;[x;;;y][:,2,:]']
:edited
I remembered a discussion in which the subject of the use of the semicolon was dealt with.
Using the macro Meta. @ Lower I discovered the hvncat function from whose very concise help (*) I could derive the following use. [hvncat((2,3,2),true, [x;y]...)]
julia> Meta.@lower([x;;;y])
:($(Expr(:thunk, CodeInfo(
@ none within `top-level scope`
1 ─ %1 = Base.hvncat(3, x, y)
└── return %1
))))
julia> x
3×2 reshape(::UnitRange{Int64}, 3, 2) with eltype Int64:
1 4
2 5
3 6
julia> y
3×2 reshape(::UnitRange{Int64}, 3, 2) with eltype Int64:
7 10
8 11
9 12
julia> hvncat((2,3,2),true, [x;y]...)
2×3×2 Array{Int64, 3}:
[:, :, 1] =
1 2 3
7 8 9
[:, :, 2] =
4 5 6
10 11 12
julia> [[x;;;y][:,1,:]' ;;;[x;;;y][:,2,:]']
2×3×2 Array{Int64, 3}:
[:, :, 1] =
1 2 3
7 8 9
[:, :, 2] =
4 5 6
10 11 12
(*) It would be helpful if someone who knows it thoroughly would provide clearer examples of the third method. Explaining where and how it can be useful to use it.
Fyi, using LazyStack.jl seems to be faster than any of the solutions above for the example provided in the accepted solution:
a = rand(1:10, 4, 3)
b = rand(11:20, 4, 3)
using LazyStack # PermutedDimsArray() gets: 2×4×3
PermutedDimsArray(stack(a,b), (3,1,2)) # 372 ns (5 allocs: 192 bytes)
Note that a simple function like this should perform as the best alternatives:
function stack(m1::Matrix{T},m2::Matrix{T}) where T
M = Array{T,3}(undef,size(m1)...,2)
M[:,:,1] .= m1
M[:,:,2] .= m2
return M
end
(didn’t test - phone). And this can be easily generalized.
It is good to know that one can write our own solutions instead of having necessarily to depend on packages or functions that not always are easy to find and combine.
Your function is like one order of magnitude faster.
LazyStack would be faster for the lazy view stack(a,b) which has dimensions 4x3x2. It is PermutedDimsArray() that screws it up to get 2x4x3. Maybe there is another way to reinterpret this.
I tried to grapple with a possible generalization, but the running time goes from about 70ns to about 2us.
what is the right way to generalize and not lose the performance of the specific case?
function mystack(m1::Matrix{T},m2::Matrix{T}, d=ndims(m2)+1) where T
sz=insert!([size(a)...],d,2)
M = Array{T,3}(undef,sz...)
idx1=Union{Colon,Int}[:,:,:]
idx2=Union{Colon,Int}[:,:,:]
idx1[d]=1
idx2[d]=2
M[idx1...] .= m1
M[idx2...] .= m2
return M
end
this lasts 1us
function mystack(m1::Matrix{T},m2::Matrix{T}, d=ndims(m2)+1) where T
sz=insert!([size(a)...],d,2)
M = Array{Int64,3}(undef,sz...)
idx1=ntuple(x->x==d ? 1 : :, ndims(m1)+1)
idx2=ntuple(x->x==d ? 2 : :, ndims(m1)+1)
M[idx1...] .= m1
M[idx2...] .= m2
return M
end
There’s no reason why a general combine-dims function couldn’t have the same performance as the specialized stack() by @lmiq. Currently, combinedims from SplitApplyCombine has some overhead and unnecessary allocations that make it ~4 times slower on small arrays. However, this overhead is easily fixable: reduce combinedims overhead by aplavin · Pull Request #55 · JuliaData/SplitApplyCombine.jl · GitHub. Thanks to the VS Code profiler, that single line immediately stood out.
With this change I get:
julia> @btime combinedims($([a, b]), 1)
110.627 ns (1 allocation: 256 bytes)
2×4×3 Array{Int64, 3}:
julia> @btime stack($a, $b)
108.349 ns (1 allocation: 256 bytes)
2×4×3 Array{Int64, 3}:
Also a hcat followed by a reshape/reinterpret (I never remember the difference) should be fast, as the memory layout is the same.