diff --git a/NEWS.md b/NEWS.md index 36660f0078c76..89d61c58081ca 100644 --- a/NEWS.md +++ b/NEWS.md @@ -59,6 +59,8 @@ Library changes * `Dict` can be now shrunk manually by `sizehint!` ([#45004]). * `@time` now separates out % time spent recompiling invalidated methods ([#45015]). * `@time_imports` now shows any compilation and recompilation time percentages per import ([#45064]). +* `eachslice` now works over multiple dimensions; `eachslice`, `eachrow` and `eachcol` return + a `Slices` object, which allows dispatching to provide more efficient methods ([#32310]). Standard library changes ------------------------ diff --git a/base/Base.jl b/base/Base.jl index e3fec462215ef..d50228de198a1 100644 --- a/base/Base.jl +++ b/base/Base.jl @@ -181,6 +181,7 @@ include("multinverses.jl") using .MultiplicativeInverses include("abstractarraymath.jl") include("arraymath.jl") +include("slicearray.jl") # SIMD loops sizeof(s::String) = Core.sizeof(s) # needed by gensym as called from simdloop diff --git a/base/abstractarraymath.jl b/base/abstractarraymath.jl index 9690fc0f2e4c4..70c304d9060c1 100644 --- a/base/abstractarraymath.jl +++ b/base/abstractarraymath.jl @@ -516,113 +516,3 @@ function repeat_inner(arr, inner) end end#module - -""" - eachrow(A::AbstractVecOrMat) - -Create a generator that iterates over the first dimension of vector or matrix `A`, -returning the rows as `AbstractVector` views. - -See also [`eachcol`](@ref), [`eachslice`](@ref), [`mapslices`](@ref). - -!!! compat "Julia 1.1" - This function requires at least Julia 1.1. - -# Example - -```jldoctest -julia> a = [1 2; 3 4] -2×2 Matrix{Int64}: - 1 2 - 3 4 - -julia> first(eachrow(a)) -2-element view(::Matrix{Int64}, 1, :) with eltype Int64: - 1 - 2 - -julia> collect(eachrow(a)) -2-element Vector{SubArray{Int64, 1, Matrix{Int64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}: - [1, 2] - [3, 4] -``` -""" -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 `AbstractVector` views. - -See also [`eachrow`](@ref) and [`eachslice`](@ref). - -!!! compat "Julia 1.1" - This function requires at least Julia 1.1. - -# Example - -```jldoctest -julia> a = [1 2; 3 4] -2×2 Matrix{Int64}: - 1 2 - 3 4 - -julia> first(eachcol(a)) -2-element view(::Matrix{Int64}, :, 1) with eltype Int64: - 1 - 3 - -julia> collect(eachcol(a)) -2-element Vector{SubArray{Int64, 1, Matrix{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}: - [1, 3] - [2, 4] -``` -""" -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), [`mapslices`](@ref), and [`selectdim`](@ref). - -!!! compat "Julia 1.1" - This function requires at least Julia 1.1. - -# Example - -```jldoctest -julia> M = [1 2 3; 4 5 6; 7 8 9] -3×3 Matrix{Int64}: - 1 2 3 - 4 5 6 - 7 8 9 - -julia> first(eachslice(M, dims=1)) -3-element view(::Matrix{Int64}, 1, :) with eltype Int64: - 1 - 2 - 3 - -julia> collect(eachslice(M, dims=2)) -3-element Vector{SubArray{Int64, 1, Matrix{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}: - [1, 4, 7] - [2, 5, 8] - [3, 6, 9] -``` -""" -@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")) - inds_before = ntuple(Returns(:), dim-1) - inds_after = ntuple(Returns(:), ndims(A)-dim) - return (view(A, inds_before..., i, inds_after...) for i in axes(A, dim)) -end diff --git a/base/exports.jl b/base/exports.jl index a8c8ff8dcda33..9843aa060d3b5 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -18,6 +18,7 @@ export AbstractMatrix, AbstractRange, AbstractSet, + AbstractSlices, AbstractUnitRange, AbstractVector, AbstractVecOrMat, @@ -41,6 +42,7 @@ export ComplexF32, ComplexF16, ComposedFunction, + ColumnSlices, DenseMatrix, DenseVecOrMat, DenseVector, @@ -80,8 +82,10 @@ export RoundNearestTiesUp, RoundToZero, RoundUp, + RowSlices, Set, Some, + Slices, StepRange, StepRangeLen, StridedArray, diff --git a/base/slicearray.jl b/base/slicearray.jl new file mode 100644 index 0000000000000..c9371622f6aff --- /dev/null +++ b/base/slicearray.jl @@ -0,0 +1,234 @@ +""" + AbstractSlices{S,N} <: AbstractArray{S,N} + +Supertype for arrays of slices into a parent array over some dimension(s), +returning views that select all the data from the other dimensions. + +`parent` will return the parent array. +""" +abstract type AbstractSlices{T,N} <: AbstractArray{T,N} end + +""" + Slices{P,SM,AX,S,N} <: AbstractSlices{S,N} + +An `AbstractArray` of slices into a parent array over specified dimension(s), +returning views that select all the data from the other dimension(s). + +These should typically be constructed by [`eachslice`](@ref), [`eachcol`](@ref) or +[`eachrow`](@ref). + +[`parent(s::Slices)`](@ref) will return the parent array. +""" +struct Slices{P,SM,AX,S,N} <: AbstractSlices{S,N} + """ + Parent array + """ + parent::P + """ + A tuple of length `ndims(parent)`, denoting how each dimension should be handled: + - an integer `i`: this is the `i`th dimension of the outer `Slices` object. + - `:`: an "inner" dimension + """ + slicemap::SM + """ + A tuple of length `N` containing the [`axes`](@ref) of the `Slices` object. + """ + axes::AX +end + +unitaxis(::AbstractArray) = Base.OneTo(1) + +function Slices(A::P, slicemap::SM, ax::AX) where {P,SM,AX} + N = length(ax) + S = Base._return_type(view, Tuple{P, map((a,l) -> l === (:) ? Colon : eltype(a), axes(A), slicemap)...}) + Slices{P,SM,AX,S,N}(A, slicemap, ax) +end + +_slice_check_dims(N) = nothing +function _slice_check_dims(N, dim, dims...) + 1 <= dim <= N || throw(DimensionMismatch("Invalid dimension $dim")) + dim in dims && throw(DimensionMismatch("Dimensions $dims are not unique")) + _slice_check_dims(N,dims...) +end + +@inline function _eachslice(A::AbstractArray{T,N}, dims::NTuple{M,Integer}, drop::Bool) where {T,N,M} + _slice_check_dims(N,dims...) + if drop + # if N = 4, dims = (3,1) then + # axes = (axes(A,3), axes(A,1)) + # slicemap = (2, :, 1, :) + ax = map(dim -> axes(A,dim), dims) + slicemap = ntuple(dim -> something(findfirst(isequal(dim), dims), (:)), N) + return Slices(A, slicemap, ax) + else + # if N = 4, dims = (3,1) then + # axes = (axes(A,1), OneTo(1), axes(A,3), OneTo(1)) + # slicemap = (1, :, 3, :) + ax = ntuple(dim -> dim in dims ? axes(A,dim) : unitaxis(A), N) + slicemap = ntuple(dim -> dim in dims ? dim : (:), N) + return Slices(A, slicemap, ax) + end +end +@inline function _eachslice(A::AbstractArray, dim::Integer, drop::Bool) + _eachslice(A, (dim,), drop) +end + +""" + eachslice(A::AbstractArray; dims, drop=true) + +Create a [`Slices`](@ref) object that is an array of slices over dimensions `dims` of `A`, returning +views that select all the data from the other dimensions in `A`. `dims` can either by an +integer or a tuple of integers. + +If `drop = true` (the default), the outer `Slices` will drop the inner dimensions, and +the ordering of the dimensions will match those in `dims`. If `drop = false`, then the +`Slices` will have the same dimensionality as the underlying array, with inner +dimensions having size 1. + +See also [`eachrow`](@ref), [`eachcol`](@ref), [`mapslices`](@ref) and [`selectdim`](@ref). + +!!! compat "Julia 1.1" + This function requires at least Julia 1.1. + +!!! compat "Julia 1.9" + Prior to Julia 1.9, this returned an iterator, and only a single dimension `dims` was supported. + +# Example + +```jldoctest +julia> m = [1 2 3; 4 5 6; 7 8 9] +3×3 Matrix{Int64}: + 1 2 3 + 4 5 6 + 7 8 9 + +julia> s = eachslice(m, dims=1) +3-element RowSlices{Matrix{Int64}, Tuple{Base.OneTo{Int64}}, SubArray{Int64, 1, Matrix{Int64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}: + [1, 2, 3] + [4, 5, 6] + [7, 8, 9] + +julia> s[1] +3-element view(::Matrix{Int64}, 1, :) with eltype Int64: + 1 + 2 + 3 + +julia> eachslice(m, dims=1, drop=false) +3×1 Slices{Matrix{Int64}, Tuple{Int64, Colon}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, SubArray{Int64, 1, Matrix{Int64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, 2}: + [1, 2, 3] + [4, 5, 6] + [7, 8, 9] +``` +""" +@inline function eachslice(A; dims, drop=true) + _eachslice(A, dims, drop) +end + +""" + eachrow(A::AbstractVecOrMat) <: AbstractVector + +Create a [`RowSlices`](@ref) object that is a vector of rows of matrix or vector `A`. +Row slices are returned as `AbstractVector` views of `A`. + +See also [`eachcol`](@ref), [`eachslice`](@ref) and [`mapslices`](@ref). + +!!! compat "Julia 1.1" + This function requires at least Julia 1.1. + +!!! compat "Julia 1.9" + Prior to Julia 1.9, this returned an iterator. + +# Example + +```jldoctest +julia> a = [1 2; 3 4] +2×2 Matrix{Int64}: + 1 2 + 3 4 + +julia> s = eachrow(a) +2-element RowSlices{Matrix{Int64}, Tuple{Base.OneTo{Int64}}, SubArray{Int64, 1, Matrix{Int64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}: + [1, 2] + [3, 4] + +julia> s[1] +2-element view(::Matrix{Int64}, 1, :) with eltype Int64: + 1 + 2 +``` +""" +eachrow(A::AbstractMatrix) = _eachslice(A, (1,), true) +eachrow(A::AbstractVector) = eachrow(reshape(A, size(A,1), 1)) + +""" + eachcol(A::AbstractVecOrMat) <: AbstractVector + +Create a [`ColumnSlices`](@ref) object that is a vector of columns of matrix or vector `A`. +Column slices are returned as `AbstractVector` views of `A`. + +See also [`eachrow`](@ref), [`eachslice`](@ref) and [`mapslices`](@ref). + +!!! compat "Julia 1.1" + This function requires at least Julia 1.1. + +!!! compat "Julia 1.9" + Prior to Julia 1.9, this returned an iterator. + +# Example + +```jldoctest +julia> a = [1 2; 3 4] +2×2 Matrix{Int64}: + 1 2 + 3 4 + +julia> s = eachcol(a) +2-element ColumnSlices{Matrix{Int64}, Tuple{Base.OneTo{Int64}}, SubArray{Int64, 1, Matrix{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}: + [1, 3] + [2, 4] + +julia> s[1] +2-element view(::Matrix{Int64}, :, 1) with eltype Int64: + 1 + 3 +``` +""" +eachcol(A::AbstractMatrix) = _eachslice(A, (2,), true) +eachcol(A::AbstractVector) = eachcol(reshape(A, size(A, 1), 1)) + +""" + RowSlices{M,AX,S} + +A special case of [`Slices`](@ref) that is a vector of row slices of a matrix, as +constructed by [`eachrow`](@ref). + +[`parent`](@ref) can be used to get the underlying matrix. +""" +const RowSlices{P<:AbstractMatrix,AX,S<:AbstractVector} = Slices{P,Tuple{Int,Colon},AX,S,1} + +""" + ColumnSlices{M,AX,S} + +A special case of [`Slices`](@ref) that is a vector of column slices of a matrix, as +constructed by [`eachcol`](@ref). + +[`parent`](@ref) can be used to get the underlying matrix. +""" +const ColumnSlices{P<:AbstractMatrix,AX,S<:AbstractVector} = Slices{P,Tuple{Colon,Int},AX,S,1} + + +IteratorSize(::Type{Slices{P,SM,AX,S,N}}) where {P,SM,AX,S,N} = HasShape{N}() +axes(s::Slices) = s.axes +size(s::Slices) = map(length, s.axes) + +@inline function _slice_index(s::Slices, c...) + return map(l -> l === (:) ? (:) : c[l], s.slicemap) +end + +Base.@propagate_inbounds getindex(s::Slices{P,SM,AX,S,N}, I::Vararg{Int,N}) where {P,SM,AX,S,N} = + view(s.parent, _slice_index(s, I...)...) +Base.@propagate_inbounds setindex!(s::Slices{P,SM,AX,S,N}, val, I::Vararg{Int,N}) where {P,SM,AX,S,N} = + s.parent[_slice_index(s, I...)...] = val + +parent(s::Slices) = s.parent diff --git a/doc/src/base/arrays.md b/doc/src/base/arrays.md index 1dc2d8ed926af..853e4c7a4ec1b 100644 --- a/doc/src/base/arrays.md +++ b/doc/src/base/arrays.md @@ -30,6 +30,9 @@ Base.StridedArray Base.StridedVector Base.StridedMatrix Base.StridedVecOrMat +Base.Slices +Base.RowSlices +Base.ColumnSlices Base.getindex(::Type, ::Any...) Base.zeros Base.ones diff --git a/test/arrayops.jl b/test/arrayops.jl index e289f6d87d889..9639281f60b04 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -2164,19 +2164,67 @@ end end # row/column/slice iterator tests -using Base: eachrow, eachcol @testset "row/column/slice iterators" begin + # check type aliases + @test RowSlices <: AbstractSlices{<:AbstractVector, 1} <: AbstractVector{<:AbstractVector} + @test eachrow(ones(3)) isa RowSlices + @test eachrow(ones(3,3)) isa RowSlices + @test ColumnSlices <: AbstractSlices{<:AbstractVector, 1} <: AbstractVector{<:AbstractVector} + @test eachcol(ones(3)) isa ColumnSlices + @test eachcol(ones(3,3)) isa ColumnSlices + # 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 eachrow(M) == eachslice(M, dims = 1) == [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + @test eachcol(M) == 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) + SR = @inferred eachrow(M) + @test SR[2] isa eltype(SR) + SR[2] = [14,15,16] + @test SR[2] == M[2,:] == [14,15,16] + @test parent(SR) === M + + SC = @inferred eachcol(M) + @test SC[3] isa eltype(SC) + SC[3] = [23,26,29] + @test SC[3] == M[:,3] == [23,26,29] + @test parent(SC) === M + + # Higher-dimensional cases + M = reshape(collect(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] + + S1 = eachslice(M, dims = 1) + @test S1 isa AbstractSlices{<:AbstractArray{Int, 3}, 1} + @test size(S1) == (2,) + @test S1[1] == M[1,:,:,:] + + S1K = eachslice(M, dims = 1, drop=false) + @test S1K isa AbstractSlices{<:AbstractArray{Int, 3}, 4} + @test size(S1K) == (2,1,1,1) + @test S1K[1,1,1,1] == M[1,:,:,:] + + S23 = eachslice(M, dims = (2,3)) + @test S23 isa AbstractSlices{<:AbstractArray{Int, 2}, 2} + @test size(S23) == (2,2) + @test S23[2,1] == M[:,2,1,:] + + S23K = eachslice(M, dims = (2,3), drop=false) + @test S23K isa AbstractSlices{<:AbstractArray{Int, 2}, 4} + @test size(S23K) == (1,2,2,1) + @test S23K[1,2,1,1] == M[:,2,1,:] + + S32 = eachslice(M, dims = (3,2)) + @test S32 isa AbstractSlices{<:AbstractArray{Int, 2}, 2} + @test size(S32) == (2,2) + @test S32[2,1] == M[:,1,2,:] + + S32K = eachslice(M, dims = (3,2), drop=false) + @test S32K isa AbstractSlices{<:AbstractArray{Int, 2}, 4} + @test size(S32K) == (1,2,2,1) + @test S32K[1,2,1,1] == M[:,2,1,:] end ###