Row and Column Iterator for Matrices (#29749)

* 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]
JuliaLang · Dec 3, 2018 · 6b04291 · 6b04291
1 parent 15ce5ef
commit 6b04291
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 0 deletions.
diff --git a/base/abstractarraymath.jl b/base/abstractarraymath.jl
@@ -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
diff --git a/base/exports.jl b/base/exports.jl
@@ -365,7 +365,10 @@ export
     cumsum!,
     accumulate,
     accumulate!,
+    eachcol,
     eachindex,
+    eachrow,
+    eachslice,
     extrema,
     fill!,
     fill,

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -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
 ###