Skip to content

Commit

Permalink
Row and Column Iterator for Matrices (#29749)
Browse files Browse the repository at this point in the history
* first pass, sans tests

* add test

* refactor a bit

* fix eltype to show views

* fix whitespace

* fix typo

* rename [ci skip]

* typofix [ci skip]

* add eachslice() [ci skip]

* eachcolumn => eachcol and a bit of error handling [ci skip]

* feedback

* fix whitespace

* refactor eachslice; type-stable now?

* git feedback

* more git feedback

* remove default for eachslice

* add tests

* prune eachrow eachcol exports and newline fix?

* typo in exports

* update docs for AbstractArray [ci skip]

* Update test/arrayops.jl

Co-Authored-By: arnavs <soodarnav01@gmail.com>

* Update base/abstractarraymath.jl

Co-Authored-By: arnavs <soodarnav01@gmail.com>

* bikeshedding

* Add missing backticks and improve message

[ci skip]

* add exports for eachcol, eachrow

[ci skip]
  • Loading branch information
Arnav Sood authored and mbauman committed Dec 3, 2018
1 parent 15ce5ef commit 6b04291
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 0 deletions.
37 changes: 37 additions & 0 deletions base/abstractarraymath.jl
Original file line number Diff line number Diff line change
Expand Up @@ -407,3 +407,40 @@ _reperr(s, n, N) = throw(ArgumentError("number of " * s * " repetitions " *

return R
end

"""
eachrow(A::AbstractVecOrMat)
Create a generator that iterates over the first dimension of vector or matrix `A`,
returning the rows as views.
See also [`eachcol`](@ref) and [`eachslice`](@ref).
"""
eachrow(A::AbstractVecOrMat) = (view(A, i, :) for i in axes(A, 1))


"""
eachcol(A::AbstractVecOrMat)
Create a generator that iterates over the second dimension of matrix `A`, returning the
columns as views.
See also [`eachrow`](@ref) and [`eachslice`](@ref).
"""
eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2))

"""
eachslice(A::AbstractArray; dims)
Create a generator that iterates over dimensions `dims` of `A`, returning views that select all
the data from the other dimensions in `A`.
Only a single dimension in `dims` is currently supported. Equivalent to `(view(A,:,:,...,i,:,:
...)) for i in axes(A, dims))`, where `i` is in position `dims`.
See also [`eachrow`](@ref), [`eachcol`](@ref), and [`selectdim`](@ref).
"""
@inline function eachslice(A::AbstractArray; dims)
length(dims) == 1 || throw(ArgumentError("only single dimensions are supported"))
dim = first(dims)
dim <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions"))
idx1, idx2 = ntuple(d->(:), dim-1), ntuple(d->(:), ndims(A)-dim)
return (view(A, idx1..., i, idx2...) for i in axes(A, dim))
end
3 changes: 3 additions & 0 deletions base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,10 @@ export
cumsum!,
accumulate,
accumulate!,
eachcol,
eachindex,
eachrow,
eachslice,
extrema,
fill!,
fill,
Expand Down
16 changes: 16 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1971,6 +1971,22 @@ end
@test IndexStyle(selectdim(A, 3, 1)) == IndexStyle(view(A, :, :, 1)) == IndexLinear()
end

# row/column/slice iterator tests
using Base: eachrow, eachcol
@testset "row/column/slice iterators" begin
# Simple ones
M = [1 2 3; 4 5 6; 7 8 9]
@test collect(eachrow(M)) == collect(eachslice(M, dims = 1)) == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
@test collect(eachcol(M)) == collect(eachslice(M, dims = 2)) == [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
@test_throws DimensionMismatch eachslice(M, dims = 4)

# Higher-dimensional case
M = reshape([(1:16)...], 2, 2, 2, 2)
@test_throws MethodError collect(eachrow(M))
@test_throws MethodError collect(eachcol(M))
@test collect(eachslice(M, dims = 1))[1][:, :, 1] == [1 5; 3 7]
end

###
### IndexCartesian workout
###
Expand Down

0 comments on commit 6b04291

Please sign in to comment.