diff --git a/docs/src/lib/public.md b/docs/src/lib/public.md index bc7caded..7bcf7c85 100644 --- a/docs/src/lib/public.md +++ b/docs/src/lib/public.md @@ -37,6 +37,7 @@ blockfirsts blocklasts blocklengths blocksizes +eachblockaxes blocks eachblock blockcheckbounds diff --git a/src/BlockArrays.jl b/src/BlockArrays.jl index 2820b6d6..e60fa7e6 100644 --- a/src/BlockArrays.jl +++ b/src/BlockArrays.jl @@ -5,7 +5,7 @@ using LinearAlgebra, ArrayLayouts, FillArrays export AbstractBlockArray, AbstractBlockMatrix, AbstractBlockVector, AbstractBlockVecOrMat export Block, getblock, getblock!, setblock!, eachblock, blocks export blockaxes, blocksize, blocklength, blockcheckbounds, BlockBoundsError, BlockIndex, BlockIndexRange -export blocksizes, blocklengths, blocklasts, blockfirsts, blockisequal +export blocksizes, blocklengths, eachblockaxes, blocklasts, blockfirsts, blockisequal export BlockRange, blockedrange, BlockedUnitRange, BlockedOneTo export BlockArray, BlockMatrix, BlockVector, BlockVecOrMat, mortar diff --git a/src/blocks.jl b/src/blocks.jl index ca8610c7..747709e9 100644 --- a/src/blocks.jl +++ b/src/blocks.jl @@ -93,6 +93,21 @@ This is broken for now. See: https://github.com/JuliaArrays/BlockArrays.jl/issue a end +# AbstractArray version of `Iterators.product`. +# https://en.wikipedia.org/wiki/Cartesian_product +# https://github.com/lazyLibraries/ProductArrays.jl +# https://github.com/JuliaData/SplitApplyCombine.jl#productviewf-a-b +# https://github.com/JuliaArrays/MappedArrays.jl/pull/42 +struct ProductArray{T,N,V<:Tuple{Vararg{AbstractVector,N}}} <: AbstractArray{T,N} + vectors::V +end +ProductArray(vectors::Vararg{AbstractVector,N}) where {N} = + ProductArray{Tuple{map(eltype, vectors)...},N,typeof(vectors)}(vectors) +Base.size(p::ProductArray) = map(length, p.vectors) +Base.axes(p::ProductArray) = map(Base.axes1, p.vectors) +@propagate_inbounds getindex(p::ProductArray{T,N}, I::Vararg{Int,N}) where {T,N} = + map((v, i) -> v[i], p.vectors, I) + """ blocksizes(A::AbstractArray) blocksizes(A::AbstractArray, d::Integer) @@ -110,7 +125,7 @@ julia> A = BlockArray(ones(3,3),[2,1],[1,1,1]) 1.0 │ 1.0 │ 1.0 julia> blocksizes(A) -2×3 BlockArrays.BlockSizes{Tuple{Int64, Int64}, 2, BlockMatrix{Float64, Matrix{Matrix{Float64}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}}: +2×3 BlockArrays.ProductArray{Tuple{Int64, Int64}, 2, Tuple{Vector{Int64}, Vector{Int64}}}: (2, 1) (2, 1) (2, 1) (1, 1) (1, 1) (1, 1) @@ -124,16 +139,80 @@ julia> blocksizes(A,2) 1 ``` """ -blocksizes(A::AbstractArray) = BlockSizes(A) +blocksizes(A::AbstractArray) = ProductArray(map(blocklengths, axes(A))...) @inline blocksizes(A::AbstractArray, d::Integer) = blocklengths(axes(A, d)) -struct BlockSizes{T,N,A<:AbstractArray{<:Any,N}} <: AbstractArray{T,N} +""" + blocklengths(A::AbstractArray) + +Return an iterator over the lengths of each block. +See also blocksizes. + +# Examples +```jldoctest +julia> A = BlockArray(ones(3,3),[2,1],[1,1,1]) +2×3-blocked 3×3 BlockMatrix{Float64}: + 1.0 │ 1.0 │ 1.0 + 1.0 │ 1.0 │ 1.0 + ─────┼───────┼───── + 1.0 │ 1.0 │ 1.0 + +julia> blocklengths(A) +2×3 BlockArrays.BlockLengths{Int64, 2, BlockMatrix{Float64, Matrix{Matrix{Float64}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}}: + 2 2 2 + 1 1 1 + +julia> blocklengths(A)[1,2] +2 +``` +""" +blocklengths(A::AbstractArray) = BlockLengths(A) +blocklengths(A::AbstractVector) = map(length, blocks(A)) + +struct BlockLengths{T,N,A<:AbstractArray{<:Any,N}} <: AbstractArray{T,N} array::A end -BlockSizes(a::AbstractArray{<:Any,N}) where {N} = - BlockSizes{Tuple{eltype.(axes(a))...},N,typeof(a)}(a) +BlockLengths(a::AbstractArray{<:Any,N}) where {N} = + BlockLengths{typeof(length(a)),N,typeof(a)}(a) -size(bs::BlockSizes) = blocksize(bs.array) -axes(bs::BlockSizes) = map(br -> Int.(br), blockaxes(bs.array)) -@propagate_inbounds getindex(a::BlockSizes{T,N}, i::Vararg{Int,N}) where {T,N} = - size(view(a.array, Block.(i)...)) +size(bs::BlockLengths) = blocksize(bs.array) +axes(bs::BlockLengths) = map(br -> Int.(br), blockaxes(bs.array)) +@propagate_inbounds getindex(a::BlockLengths{T,N}, i::Vararg{Int,N}) where {T,N} = + length(view(a.array, Block.(i)...)) + +""" + eachblockaxes(A::AbstractArray) + eachblockaxes(A::AbstractArray, d::Integer) + +Return an iterator over the axes of each block. +See also blocksizes and blocklengths. + +# Examples +```jldoctest +julia> A = BlockArray(ones(3,3),[2,1],[1,1,1]) +2×3-blocked 3×3 BlockMatrix{Float64}: + 1.0 │ 1.0 │ 1.0 + 1.0 │ 1.0 │ 1.0 + ─────┼───────┼───── + 1.0 │ 1.0 │ 1.0 + +julia> eachblockaxes(A) +2×3 BlockArrays.ProductArray{Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, 2, Tuple{Vector{Base.OneTo{Int64}}, Vector{Base.OneTo{Int64}}}}: + (Base.OneTo(2), Base.OneTo(1)) … (Base.OneTo(2), Base.OneTo(1)) + (Base.OneTo(1), Base.OneTo(1)) (Base.OneTo(1), Base.OneTo(1)) + +julia> eachblockaxes(A)[1,2] +(Base.OneTo(2), Base.OneTo(1)) + +julia> eachblockaxes(A,2) +3-element Vector{Base.OneTo{Int64}}: + Base.OneTo(1) + Base.OneTo(1) + Base.OneTo(1) +``` +""" +eachblockaxes(A::AbstractArray) = + ProductArray(map(ax -> map(Base.axes1, blocks(ax)), axes(A))...) +eachblockaxes(A::AbstractVector) = map(axes, blocks(Base.axes1(A))) +eachblockaxes(A::AbstractArray, d::Integer) = map(Base.axes1, blocks(axes(A, d))) +eachblockaxes1(A::AbstractArray) = map(Base.axes1, blocks(Base.axes1(A))) diff --git a/test/test_blocks.jl b/test/test_blocks.jl index 3e5e4c9a..a8636c3b 100644 --- a/test/test_blocks.jl +++ b/test/test_blocks.jl @@ -1,6 +1,7 @@ module TestBlocks using Test, BlockArrays +import BlockArrays: eachblockaxes1 @testset "blocks" begin @testset "blocks(::BlockVector)" begin @@ -143,18 +144,18 @@ end @test @inferred(size(bs)) == (2, 2) @test @inferred(length(bs)) == 4 @test @inferred(axes(bs)) == (1:2, 1:2) - @test @inferred(eltype(bs)) === Tuple{Int,Int} + @test @inferred(eltype(bs)) ≡ Tuple{Int,Int} @test bs == [(2, 3) (2, 1); (3, 3) (3, 1)] - @test @inferred(bs[1, 1]) == (2, 3) - @test @inferred(bs[2, 1]) == (3, 3) - @test @inferred(bs[1, 2]) == (2, 1) - @test @inferred(bs[2, 2]) == (3, 1) - @test @inferred(bs[1]) == (2, 3) - @test @inferred(bs[2]) == (3, 3) - @test @inferred(bs[3]) == (2, 1) - @test @inferred(bs[4]) == (3, 1) - @test blocksizes(A, 1) == [2, 3] - @test blocksizes(A, 2) == [3, 1] + @test @inferred(bs[1, 1]) ≡ (2, 3) + @test @inferred(bs[2, 1]) ≡ (3, 3) + @test @inferred(bs[1, 2]) ≡ (2, 1) + @test @inferred(bs[2, 2]) ≡ (3, 1) + @test @inferred(bs[1]) ≡ (2, 3) + @test @inferred(bs[2]) ≡ (3, 3) + @test @inferred(bs[3]) ≡ (2, 1) + @test @inferred(bs[4]) ≡ (3, 1) + @test @inferred((x -> blocksizes(x, 1))(A)) == [2, 3] + @test @inferred((x -> blocksizes(x, 2))(A)) == [3, 1] end @testset "Inference: issue #425" begin @@ -166,4 +167,69 @@ end end end +@testset "blocklengths" begin + v = Array(reshape(1:20, (5, 4))) + A = BlockArray(v, [2, 3], [3, 1]) + bls = @inferred(blocklengths(A)) + @test bls == [6 2; 9 3] + @test @inferred(length(bls)) ≡ 4 + @test @inferred(size(bls)) ≡ (2, 2) + @test @inferred(eltype(bls)) ≡ Int + @test @inferred(bls[1, 1]) ≡ 6 + @test @inferred(bls[2, 1]) ≡ 9 + @test @inferred(bls[1, 2]) ≡ 2 + @test @inferred(bls[2, 2]) ≡ 3 + @test @inferred(bls[1]) ≡ 6 + @test @inferred(bls[2]) ≡ 9 + @test @inferred(bls[3]) ≡ 2 + @test @inferred(bls[4]) ≡ 3 +end + +@testset "eachblockaxes" begin + v = Array(reshape(1:20, (5, 4))) + A = BlockArray(v, [2, 3], [3, 1]) + bas = @inferred(eachblockaxes(A)) + @test bas == [(Base.OneTo(2), Base.OneTo(3)) (Base.OneTo(2), Base.OneTo(1)); (Base.OneTo(3), Base.OneTo(3)) (Base.OneTo(3), Base.OneTo(1))] + @test @inferred(length(bas)) ≡ 4 + @test @inferred(size(bas)) ≡ (2, 2) + @test @inferred(eltype(bas)) ≡ Tuple{Base.OneTo{Int},Base.OneTo{Int}} + @test @inferred(bas[1, 1]) ≡ (Base.OneTo(2), Base.OneTo(3)) + @test @inferred(bas[2, 1]) ≡ (Base.OneTo(3), Base.OneTo(3)) + @test @inferred(bas[1, 2]) ≡ (Base.OneTo(2), Base.OneTo(1)) + @test @inferred(bas[2, 2]) ≡ (Base.OneTo(3), Base.OneTo(1)) + @test @inferred(bas[1]) ≡ (Base.OneTo(2), Base.OneTo(3)) + @test @inferred(bas[2]) ≡ (Base.OneTo(3), Base.OneTo(3)) + @test @inferred(bas[3]) ≡ (Base.OneTo(2), Base.OneTo(1)) + @test @inferred(bas[4]) ≡ (Base.OneTo(3), Base.OneTo(1)) + + bas2 = @inferred (x -> eachblockaxes(x, 2))(A) + @test bas2 == [Base.OneTo(3), Base.OneTo(1)] + @test length(bas2) ≡ 2 + @test size(bas2) ≡ (2,) + @test @inferred(eltype(bas2)) ≡ Base.OneTo{Int} + @test @inferred(bas2[1]) ≡ Base.OneTo(3) + @test @inferred(bas2[2]) ≡ Base.OneTo(1) + @test @inferred((x -> eachblockaxes(x, 3))(A)) == [Base.OneTo(1)] + + V = mortar([[2, 3], [4, 5, 6]]) + @test @inferred(eachblockaxes(V)) == [(Base.OneTo(2),), (Base.OneTo(3),)] + @test @inferred((x -> eachblockaxes(x, 1))(V)) == [Base.OneTo(2), Base.OneTo(3)] + @test @inferred((x -> eachblockaxes(x, 2))(V)) == [Base.OneTo(1)] +end + +@testset "eachblockaxes1" begin + v = Array(reshape(1:20, (5, 4))) + A = BlockArray(v, [2, 3], [3, 1]) + bas = @inferred(eachblockaxes1(A)) + @test bas == [Base.OneTo(2), Base.OneTo(3)] + @test @inferred(length(bas)) ≡ 2 + @test @inferred(size(bas)) ≡ (2,) + @test @inferred(eltype(bas)) ≡ Base.OneTo{Int} + @test @inferred(bas[1]) ≡ Base.OneTo(2) + @test @inferred(bas[2]) ≡ Base.OneTo(3) + + @test @inferred(eachblockaxes1(mortar([[2, 3], [4, 5, 6]]))) == [Base.OneTo(2), Base.OneTo(3)] + @test @inferred(eachblockaxes1(fill(2))) == [Base.OneTo(1)] +end + end # module