New interface blocks(a) for array-of-arrays view (#117)

* Add blocks(a) view

* Run tests

* Fix setindex!

* Add blocks(::Adjoint) and  blocks(::Transpose)

* Test blocks(view(...))

* Add a special handling for SubArray

* Import `only` from Compat

* Trying to improve coverage

* Avoid shadowing `blocks` in doctest

Author: tkf

Project.toml

@@ -4,10 +4,12 @@ version = "0.12.7"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
 ArrayLayouts = "0.3"
+Compat = "2.2, 3"
 julia = "1.1"
 
 [extras]

docs/src/lib/public.md

@@ -36,6 +36,7 @@ blocksize
 blockfirsts
 blocklasts
 blocklengths
+blocks
 eachblock
 getblock
 getblock!

src/BlockArrays.jl

@@ -4,7 +4,7 @@ using LinearAlgebra, ArrayLayouts
 
 # AbstractBlockArray interface exports
 export AbstractBlockArray, AbstractBlockMatrix, AbstractBlockVector, AbstractBlockVecOrMat
-export Block, getblock, getblock!, setblock!, eachblock
+export Block, getblock, getblock!, setblock!, eachblock, blocks
 export blockaxes, blocksize, blocklength, blockcheckbounds, BlockBoundsError, BlockIndex
 export blocklengths, blocklasts, blockfirsts, blockisequal
 export BlockRange, blockedrange, BlockedUnitRange
@@ -39,13 +39,18 @@ import ArrayLayouts: _fill_lmul!, MatMulVecAdd, MatMulMatAdd, MatLmulVec, MatLdi
                         materialize!, MemoryLayout, sublayout, transposelayout, conjlayout, 
                         triangularlayout, triangulardata, _inv
 
+if !@isdefined(only)
+    using Compat: only
+end
+
 include("blockindices.jl")                        
 include("blockaxis.jl")
 include("abstractblockarray.jl")
 include("blockarray.jl")
 include("pseudo_blockarray.jl")
 include("views.jl")
 include("show.jl")
+include("blocks.jl")
 include("blockarrayinterface.jl")
 include("blockbroadcast.jl")
 include("blocklinalg.jl")

src/blockarray.jl

@@ -216,17 +216,19 @@ BlockMatrix{T}(λ::UniformScaling, block_sizes::Vararg{AbstractVector{Int},2}) w
 Construct a `BlockArray` from `blocks`.  `block_sizes` is computed from
 `blocks` if it is not given.
 
+This is an "inverse" of [`blocks`](@ref).
+
 # Examples
 ```jldoctest; setup = quote using BlockArrays end
-julia> blocks = permutedims(reshape([
+julia> arrays = permutedims(reshape([
                   1ones(1, 3), 2ones(1, 2),
                   3ones(2, 3), 4ones(2, 2),
               ], (2, 2)))
 2×2 Array{Array{Float64,2},2}:
  [1.0 1.0 1.0]               [2.0 2.0]
  [3.0 3.0 3.0; 3.0 3.0 3.0]  [4.0 4.0; 4.0 4.0]
 
-julia> mortar(blocks)
+julia> mortar(arrays)
 2×2-blocked 3×5 BlockArray{Float64,2}:
  1.0  1.0  1.0  │  2.0  2.0
  ───────────────┼──────────

src/blocks.jl

@@ -0,0 +1,105 @@
+"""
+    blocks(a::AbstractArray{T,N}) :: AbstractArray{<:AbstractArray{T,N},N}
+
+Return the array-of-arrays view to `a` such that
+
+```
+blocks(a)[i₁, i₂, ..., iₙ] == a[Block(i₁), Block(i₂), ..., Block(iₙ)]
+```
+
+This function does not copy the blocks and give a mutable viwe to the original
+array.  This is an "inverse" of [`mortar`](@ref).
+
+# Examples
+```jldoctest; setup = quote using BlockArrays end
+julia> bs1 = permutedims(reshape([
+               1ones(1, 3), 2ones(1, 2),
+               3ones(2, 3), 4ones(2, 2),
+           ], (2, 2)))
+2×2 Array{Array{Float64,2},2}:
+ [1.0 1.0 1.0]               [2.0 2.0]
+ [3.0 3.0 3.0; 3.0 3.0 3.0]  [4.0 4.0; 4.0 4.0]
+
+julia> a = mortar(bs1)
+2×2-blocked 3×5 BlockArray{Float64,2}:
+ 1.0  1.0  1.0  │  2.0  2.0
+ ───────────────┼──────────
+ 3.0  3.0  3.0  │  4.0  4.0
+ 3.0  3.0  3.0  │  4.0  4.0
+
+julia> bs2 = blocks(a)
+2×2 Array{Array{Float64,2},2}:
+ [1.0 1.0 1.0]               [2.0 2.0]
+ [3.0 3.0 3.0; 3.0 3.0 3.0]  [4.0 4.0; 4.0 4.0]
+
+julia> bs1 == bs2
+true
+
+julia> bs2[1, 1] .*= 100;
+
+julia> a  # in-place mutation is reflected to the block array
+2×2-blocked 3×5 BlockArray{Float64,2}:
+ 100.0  100.0  100.0  │  2.0  2.0
+ ─────────────────────┼──────────
+   3.0    3.0    3.0  │  4.0  4.0
+   3.0    3.0    3.0  │  4.0  4.0
+```
+"""
+blocks(a::AbstractArray) = blocks(PseudoBlockArray(a, axes(a)))
+blocks(a::AbstractBlockArray) = BlocksView(a)
+blocks(a::BlockArray) = a.blocks
+blocks(A::Adjoint) = adjoint(blocks(parent(A)))
+blocks(A::Transpose) = transpose(blocks(parent(A)))
+
+# convert a tuple of BlockRange to a tuple of `AbstractUnitRange{Int}`
+_blockrange2int() = ()
+_blockrange2int(A, B...) = tuple(Int.(A.block), _blockrange2int(B...)...)
+
+blocks(A::SubArray{<:Any,N,<:Any,<:NTuple{N,BlockSlice}}) where N =
+    view(blocks(parent(A)), _blockrange2int(parentindices(A)...)...)
+
+struct BlocksView{
+    S,                            # eltype(eltype(BlocksView(...)))
+    N,                            # ndims
+    T<:AbstractArray{S,N},        # eltype(BlocksView(...)), i.e., block type
+    B<:AbstractBlockArray{S,N},   # array to be wrapped
+} <: AbstractArray{T,N}
+    array::B
+end
+
+BlocksView(a::AbstractBlockArray{S,N}) where {S,N} =
+    BlocksView{S,N,AbstractArray{eltype(a),N},typeof(a)}(a)
+# Note: deciding concrete eltype of `BlocksView` requires some extra
+# interface for `AbstractBlockArray`.
+
+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))
+
+@propagate_inbounds _view(a::PseudoBlockArray, i::Block) = a[i]
+@propagate_inbounds _view(a::AbstractBlockArray, i::Block) = view(a, i)
+
+#=
+This is broken for now. See: https://github.com/JuliaArrays/BlockArrays.jl/issues/120
+# IndexLinear implementations
+@propagate_inbounds Base.getindex(a::BlocksView, i::Int) = _view(a.array, Block(i))
+@propagate_inbounds Base.setindex!(a::BlocksView, b, i::Int) = copyto!(a[i], b)
+=#
+
+# IndexCartesian implementations
+@propagate_inbounds Base.getindex(a::BlocksView{T,N}, i::Vararg{Int,N}) where {T,N} =
+    _view(a.array, Block(i...))
+@propagate_inbounds Base.setindex!(a::BlocksView{T,N}, b, i::Vararg{Int,N}) where {T,N} =
+    copyto!(a[i...], b)
+
+function Base.showarg(io::IO, a::BlocksView, toplevel::Bool)
+    if toplevel
+        print(io, "blocks of ")
+        Base.showarg(io, a.array, true)
+    else
+        print(io, "::BlocksView{…,")
+        Base.showarg(io, a.array, false)
+        print(io, '}')
+    end
+end

src/pseudo_blockarray.jl

@@ -302,6 +302,17 @@ Base.reshape(parent::PseudoBlockArray, shp::Tuple{Union{Integer,Base.OneTo}, Var
 Base.reshape(parent::PseudoBlockArray, dims::Tuple{Int,Vararg{Int}}) =
     Base._reshape(parent, dims)
 
+function Base.showarg(io::IO, A::PseudoBlockArray, toplevel::Bool)
+    if toplevel
+        print(io, "PseudoBlockArray of ")
+        Base.showarg(io, A.blocks, true)
+    else
+        print(io, "::PseudoBlockArray{…,")
+        Base.showarg(io, A.blocks, false)
+        print(io, '}')
+    end
+end
+
 
 ###########################
 # Strided Array interface #

test/runtests.jl

@@ -4,6 +4,7 @@ using BlockArrays, LinearAlgebra, Test
     include("test_blockindices.jl")
     include("test_blockarrays.jl")
     include("test_blockviews.jl")
+    include("test_blocks.jl")
     include("test_blockrange.jl")
     include("test_blockarrayinterface.jl")
     include("test_blockbroadcast.jl")

test/test_blocks.jl

@@ -0,0 +1,102 @@
+using Test, BlockArrays
+
+@testset "blocks(::BlockVector)" begin
+    vector_blocks = [[1, 2], [3, 4, 5], Int[]]
+    @test blocks(mortar(vector_blocks)) === vector_blocks
+    @test collect(blocks(mortar(vector_blocks))) == vector_blocks
+end
+
+@testset "blocks(::BlockMatrix)" begin
+    matrix_blocks = permutedims(reshape([
+        1ones(1, 3), 2ones(1, 2),
+        3ones(2, 3), 4ones(2, 2),
+    ], (2, 2)))
+    @test blocks(mortar(matrix_blocks)) === matrix_blocks
+end
+
+@testset "blocks(::PseudoBlockVector)" begin
+    v0 = rand(3)
+    vb = PseudoBlockArray(v0, [1, 2])
+    @test size(blocks(vb)) == (2,)
+    blocks(vb)[1] = [123]
+    @test v0[1] == 123
+    @test parent(blocks(vb)[1]) === v0
+
+    # toplevel = true:
+    str = sprint(show, "text/plain", blocks(vb))
+    @test occursin("blocks of PseudoBlockArray of", str)
+
+    # toplevel = false:
+    str = sprint(show, "text/plain", view(blocks(vb), 1:1))
+    @test occursin("::BlocksView{…,::PseudoBlockArray{…,", str)
+end
+
+@testset "blocks(::PseudoBlockMatrix)" begin
+    m0 = rand(2, 4)
+    mb = PseudoBlockArray(m0, [1, 1], [2, 1, 1])
+    @test size(blocks(mb)) == (2, 3)
+    blocks(mb)[1, 1] = [123 456]
+    @test m0[1, 1] == 123
+    @test m0[1, 2] == 456
+    @test parent(blocks(mb)[1, 1]) === m0
+
+    # linear indexing
+    @test blocks(mb)[1] == m0[1:1, 1:2]
+    blocks(mb)[1] = [111 222]
+    @test mb[Block(1, 1)] == [111 222]
+
+    # toplevel = true:
+    str = sprint(show, "text/plain", blocks(mb))
+    @test occursin("blocks of PseudoBlockArray of", str)
+
+    # toplevel = false:
+    str = sprint(show, "text/plain", view(blocks(mb), 1:1, 1:1))
+    @test occursin("::BlocksView{…,::PseudoBlockArray{…,", str)
+end
+
+@testset "blocks(::Vector)" begin
+    v = rand(3)
+    @test size(blocks(v)) == (1,)
+    blocks(v)[1][1] = 123
+    @test v[1] == 123
+    @test parent(blocks(v)[1]) === v
+end
+
+@testset "blocks(::Matrix)" begin
+    m = rand(2, 4)
+    @test size(blocks(m)) == (1, 1)
+    blocks(m)[1, 1][1, 1] = 123
+    @test m[1, 1] == 123
+    @test parent(blocks(m)[1, 1]) === m
+end
+
+@testset "blocks(::Adjoint|Transpose)" begin
+    m = BlockArray([rand(ComplexF64, 2, 2) for _ in 1:3, _ in 1:5], [1, 2], [2, 3])
+    @testset for i in 1:2, j in 1:2
+        @test blocks(m')[i, j] == m'[Block(i), Block(j)]
+        @test blocks(transpose(m))[i, j] == transpose(m)[Block(i), Block(j)]
+    end
+end
+
+@testset "blocks(::SubArray)" begin
+    vector_blocks = [[1, 2], [3, 4, 5], Int[]]
+    b = view(mortar(vector_blocks), Block(1):Block(2))
+    v = blocks(b)
+    @test v == vector_blocks[1:2]
+    v[1][1] = 111
+    @test b[1] == 111
+    @test parent(v) === parent(b).blocks  # special path works
+end
+
+@testset "blocks(::SubArray)" begin
+    matrix_blocks = permutedims(reshape([
+        1ones(1, 3), 2ones(1, 2),
+        3ones(2, 3), 4ones(2, 2),
+    ], (2, 2)))
+    b = view(mortar(matrix_blocks), Block(1):Block(2), Block(2):Block(2))
+    m = blocks(b)
+    @test m == matrix_blocks[1:2, 2:2]
+    m[1, 1][1, 1] = 111
+    @test b[1, 1] == 111
+    @test parent(m) === parent(b).blocks  # special path works
+end