Introduce eachblockaxes, more general blocklengths (#476)

Introduces and exports `eachblockaxes`, which is analogous to the
redefinition of `blocksizes` introduced in #399. It outputs a (lazy)
array with a size `blocksize(A)` that when you index into it it returns
the axes of the corresponding block of `A`. It also defines a private
function `BlockArrays.eachblockaxes1` analogous to `Base.axes1` for
getting the axes of the blocks in the first dimension.

The motivation for this is related to the discussion in #446, this helps
with use cases where the blocks of a block array have non-trivial axes,
for example they themselves might be blocked, have a Kronecker
structure, have extra information like symmetry sector information
(which is an application of [graded vector
spaces](https://en.wikipedia.org/wiki/Graded_vector_space)), have labels
that need to be preserved, etc.

It also introduces a generalization of
[`blocklengths`](https://juliaarrays.github.io/BlockArrays.jl/v1.6/lib/public/#BlockArrays.blocklengths)
beyond `AbstractUnitRange{<:Integer}` to any `AbstractArray` by
outputting a (lazy) array that when you index into it returns the length
of the corresponding block.

This PR also changes the implementation of `blocksizes`. Before, it was
implemented as a lazy version of `size.(blocks(A))`. Now, it is based on
a new `ProductArray` type, where the `blocklengths` of the axes in each
dimension are stored and accessed. Then, `ProductArray` can be shared
with and used for the implementation of `eachblockaxes`. I think that
design is easier to reason about and in the case of `eachblockaxes` it
is easier to make the element type concrete.

Related to #369 and #446.

Author: mtfishman

docs/src/lib/public.md

@@ -37,6 +37,7 @@ blockfirsts
 blocklasts
 blocklengths
 blocksizes
+eachblockaxes
 blocks
 eachblock
 blockcheckbounds

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

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)))

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