Skip to content

Commit

Permalink
stack(f, iter)
Browse files Browse the repository at this point in the history
  • Loading branch information
mcabbott committed Dec 5, 2021
1 parent c6aeed6 commit 87b514f
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 0 deletions.
37 changes: 37 additions & 0 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2562,6 +2562,43 @@ true
stack(itr) = _stack_iter(IteratorSize(itr), itr)
stack(A::AbstractArray{<:AbstractArray}) = _typed_stack(mapreduce(eltype, promote_type, A), A)

"""
stack(f, args)
Apply `f` to each element of `args`, and `stack` the result.
See also [`mapslices`](@ref), [`mapreduce`](@ref).
# Examples
```jldoctest
julia> stack("julia") do c
(c, c-32)
end
2×5 Matrix{Char}:
'j' 'u' 'l' 'i' 'a'
'J' 'U' 'L' 'I' 'A'
julia> ans == mapreduce(c -> [c, c-32], hcat, "julia")
true
julia> stack(x -> x*x', eachcol([1 2; 10 20; 100 200]))
3×3×2 Array{Int64, 3}:
[:, :, 1] =
1 10 100
10 100 1000
100 1000 10000
[:, :, 2] =
4 40 400
40 400 4000
400 4000 40000
julia> ans == cat([1,10,100] * [1,10,100]', [2,20,200] * [2,20,200]'; dims=3)
true
```
"""
stack(f, itr) = stack(Iterators.map(f, itr))

function _stack_iter(::HasShape, itr)
w, val = _vstack_plus(itr)
reshape(w, axes(val)..., axes(itr)...)
Expand Down
4 changes: 4 additions & 0 deletions test/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1577,6 +1577,10 @@ end
@test size(stack(Iterators.product(1:3, 1:4))) == (2,3,4)
@test stack([('a', 'b'), ('c', 'd')]) == ['a' 'c'; 'b' 'd']

# stack(f, iter)
@test stack(x -> [x, 2x], 3:5) == [3 4 5; 6 8 10]
@test stack(x -> x*x'/2, [1:2, 3:4]) == [0.5 1.0; 1.0 2.0;;; 4.5 6.0; 6.0 8.0]

# Mismatched sizes
@test_throws DimensionMismatch stack([1:2, 1:3])
@test_throws DimensionMismatch stack(x for x in [1:2, 1:3])
Expand Down

0 comments on commit 87b514f

Please sign in to comment.