Redefine blocksizes

.gitignore

@@ -9,4 +9,5 @@ src/.DS_Store
 Manifest.toml
 Manifest-v*.*.toml
 .DS_Store
+.*.swp
 build

src/blockaxis.jl

@@ -353,35 +353,6 @@ blocksize(A) = map(length, blockaxes(A))
 blocksize(A,i) = length(blockaxes(A,i))
 @inline blocklength(t) = prod(blocksize(t))
 
-"""
-    blocksizes(A::AbstractArray)
-    blocksizes(A::AbstractArray, i::Int)
-
-Return the tuple of the sizes of blocks along each
-dimension. See also size and blocksize.
-
-# 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> blocksizes(A)
-([2, 1], [1, 1, 1])
-
-julia> blocksizes(A,2)
-3-element Vector{Int64}:
- 1
- 1
- 1
-```
-"""
-blocksizes(A) = map(blocklengths, axes(A))
-blocksizes(A,i) = blocklengths(axes(A,i))
-
 axes(b::AbstractBlockedUnitRange) = (BlockedOneTo(blocklasts(b) .- (first(b)-oneunit(eltype(b)))),)
 unsafe_indices(b::AbstractBlockedUnitRange) = axes(b)
 # ::Integer works around case where blocklasts might return different type

src/blocks.jl

@@ -76,7 +76,7 @@ BlocksView(a::AbstractArray{S,N}) where {S,N} =
 Base.IteratorEltype(::Type{<:BlocksView}) = Base.EltypeUnknown()
 
 Base.size(a::BlocksView) = blocksize(a.array)
-Base.axes(a::BlocksView) = map(br -> only(br.indices), blockaxes(a.array))
+Base.axes(a::BlocksView) = map(br -> Int.(br), blockaxes(a.array))
 
 #=
 This is broken for now. See: https://github.com/JuliaArrays/BlockArrays.jl/issues/120
@@ -92,3 +92,48 @@ This is broken for now. See: https://github.com/JuliaArrays/BlockArrays.jl/issue
     copyto!(a[i...], b)
     a
 end
+
+"""
+    blocksizes(A::AbstractArray)
+    blocksizes(A::AbstractArray, d::Integer)
+
+Return an iterator over the sizes of each block.
+See also size and blocksize.
+
+# 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> 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, 1)  (2, 1)  (2, 1)
+ (1, 1)  (1, 1)  (1, 1)
+
+julia> blocksizes(A)[1,2]
+(2, 1)
+
+julia> blocksizes(A,2)
+3-element Vector{Int64}:
+ 1
+ 1
+ 1
+```
+"""
+blocksizes(A::AbstractArray) = BlockSizes(A)
+blocksizes(A::AbstractArray, d::Integer) = blocklengths(axes(A, d))
+
+struct BlockSizes{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)
+
+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)...))

test/test_blockarrayinterface.jl

@@ -15,8 +15,9 @@ end
     A = randn(5)
     @test blocksize(A) == (1,)
     @test blocksize(A, 1) == 1
-    @test blocksizes(A) == ([5],)
-    @test blocksizes(A, 1) == [5]
+    @test blocklengths.(axes(A)) == ([5],)
+    @test blocklengths(axes(A, 1)) == [5]
+    @test blocksizes(A) == [(5,)]
     @test A[Block(1)] == A
     view(A, Block(1))[1] = 2
     @test A[1] == 2

test/test_blockarrays.jl

@@ -109,7 +109,8 @@ end
 
             # test that BlockArrays may be created from immutable arrays
             B = BlockArray(reshape(1:9,3,3), [2,1], [2,1])
-            @test blocksizes(B) == ([2,1], [2,1])
+            @test blocklengths.(axes(B)) == ([2,1], [2,1])
+            @test blocksizes(B) == [(2,2) (2,1); (1,2) (1,1)]
             @test B == reshape([1:9;],3,3)
             @test blocks(B) isa Matrix{Matrix{Int}}
 

test/test_blockdeque.jl

@@ -167,7 +167,8 @@ using BlockArrays, Test
             @test pop!(B) == 6
             @test B == 1:5
             @test !any(isempty, blocks(B))
-            @test blocksizes(B,1) == [1,2,2]
+            @test blocklengths(axes(B,1)) == [1,2,2]
+            @test blocksizes(B) == [(1,), (2,), (2,)]
         end
     end
 
@@ -183,13 +184,15 @@ using BlockArrays, Test
         A = BlockArray([1:6;], [2,2,2])
         @test popfirst!(A) == 1
         @test A == 2:6
-        @test blocksizes(A,1) == [1,2,2]
+        @test blocklengths(axes(A,1)) == [1,2,2]
+        @test blocksizes(A) == [(1,), (2,), (2,)]
 
         @testset "empty blocks" begin
             B = BlockArray([1:6;], [0,0,1,2,3])
             @test popfirst!(B) == 1
             @test B == 2:6
-            @test blocksizes(B,1) == [2,3]
+            @test blocklengths(axes(B,1)) == [2,3]
+            @test blocksizes(B) == [(2,), (3,)]
             @test !any(isempty, blocks(B))
         end
     end

test/test_blockindices.jl

@@ -786,13 +786,6 @@ end
         @test eltype(blocklengths(r)) === UInt16
     end
 
-    @testset "blocksizes" begin
-        x = blockedrange(2:4)
-        @test blocksizes(x,1) === 2:4
-        y = blockedrange([2:4;])
-        @test blocksizes(x,1) == blocksizes(y,1)
-    end
-
     @testset "show" begin
         b = blockedrange([1,2])
         @test repr(b) == "$BlockedOneTo($([1,3]))"

test/test_blocklinalg.jl

@@ -101,12 +101,14 @@ bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
         v = BlockArray(rand(6), 1:3)
         w = A * v
         @test w isa BlockArray
-        @test blocksizes(w,1) == fill(2, 3)
+        @test blocklengths(axes(w,1)) == fill(2, 3)
+        @test blocksizes(w) == [(2,), (2,), (2,)]
         @test w ≈ Array(A) * v ≈ A * Array(v) ≈ Array(A) * Array(v)
 
         z = A * w
         @test z isa BlockArray
-        @test blocksizes(z,1) == fill(2, 3)
+        @test blocklengths(axes(z,1)) == fill(2, 3)
+        @test blocksizes(z) == [(2,), (2,), (2,)]
         @test z ≈ Array(A) * w ≈ A * Array(w) ≈ Array(A) * Array(w)
     end
 

test/test_blocks.jl

@@ -105,4 +105,28 @@ using Test, BlockArrays
     end
 end
 
+@testset "blocksizes" begin
+    @testset "blocksizes" begin
+        v = Array(reshape(1:20, (5, 4)))
+        A = BlockArray(v, [2, 3], [3, 1])
+        @test blocklengths.(axes(A)) == ([2, 3], [3, 1])
+        bs = @inferred(blocksizes(A))
+        @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 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]
+    end
+end
+
 end # module