export and test BlockIndex (#16)

Author: KristofferC

docs/src/lib/internals.md

@@ -19,7 +19,6 @@ Pages = ["internals.md"]
 ## Internals
 
 ```@docs
-BlockIndex
 blockindex2global
 global2blockindex
 ```

docs/src/lib/public.md

@@ -29,6 +29,7 @@ This sections defines the functions a subtype of `AbstractBlockArray` should def
 AbstractBlockArray
 BlockBoundsError
 Block
+BlockIndex
 nblocks
 blocksize
 getblock

docs/src/man/pseudoblockarrays.md

@@ -19,7 +19,7 @@ a direct solver using `full`.
 
 Creating a `PseudoBlockArray` works in the same way as a `BlockArray`.
 
-```jldoctest
+```jldoctest A
 julia> pseudo = PseudoBlockArray(rand(3,3), [1,2], [2,1])
 2×2-blocked 3×3 BlockArrays.PseudoBlockArray{Float64,2,Array{Float64,2}}:
  0.590845  0.460085  │  0.200586
@@ -35,7 +35,7 @@ This "takes ownership" of the passed in array so no copy of the array is made.
 Setting and getting blocks uses the same API as `BlockArrays`. The difference here is that setting a block will update the block in place and getting a block
 will extract a copy of the block and return it. For `PseudoBlockArrays` there is a mutating block getter called `getblock!` which updates a passed in array to avoid a copy:
 
-```jldoctest
+```jldoctest A
 julia> A = zeros(2,2)
 2×2 Array{Float64,2}:
  0.0  0.0
@@ -49,4 +49,11 @@ julia> A
  0.566237  0.854147
 ```
 
+It is sometimes convenient to access an index in a certain block. We could of course write this as `A[Block(I,J)][i,j]` but the problem is that `A[Block(I,J)]` allocates its output so this type of indexing will be inefficient. Instead, it is possible to use the `A[BlockIndex((I,J), (i,j))]` indexing. Using the same block matrix `A` as above:
+
+```jldoctest A
+julia> pseudo[BlockIndex((2,1), (2,2))]
+0.8541465903790502
+```
+
 The underlying array is accessed with `Array` just like for `BlockArray`.

src/BlockArrays.jl

@@ -4,7 +4,7 @@ module BlockArrays
 
 # AbstractBlockArray interface exports
 export AbstractBlockArray, AbstractBlockMatrix, AbstractBlockVector, AbstractBlockVecOrMat
-export Block, getblock, getblock!, setblock!, nblocks, blocksize, blockcheckbounds, BlockBoundsError
+export Block, getblock, getblock!, setblock!, nblocks, blocksize, blockcheckbounds, BlockBoundsError, BlockIndex
 
 export BlockArray, BlockMatrix, BlockVector, BlockVecOrMat
 export PseudoBlockArray, PseudoBlockMatrix, PseudoBlockVector, PseudoBlockVecOrMat

src/blockarray.jl

@@ -148,7 +148,7 @@ end
     return block_arr
 end
 
-@inline function Base.setindex!{T,N}(block_array::BlockArray{T, N}, v, block_index::BlockIndex{N})
+@propagate_inbounds function Base.setindex!{T,N}(block_array::BlockArray{T, N}, v, block_index::BlockIndex{N})
     getblock(block_array, block_index.I...)[block_index.α...] = v
 end
 

src/blockindices.jl

@@ -3,12 +3,46 @@
 
 A `BlockIndex` is an index which stores a global index in two parts: the block
 and the offset index into the block.
+
+It can be used to index into `BlockArrays` in the following manner:
+
+```jldoctest
+julia> arr = Array(reshape(1:25, (5,5)));
+
+julia> a = PseudoBlockArray(arr, [3,2], [1,4])
+2×2-blocked 5×5 BlockArrays.PseudoBlockArray{Int64,2,Array{Int64,2}}:
+ 1  │   6  11  16  21
+ 2  │   7  12  17  22
+ 3  │   8  13  18  23
+ ───┼────────────────
+ 4  │   9  14  19  24
+ 5  │  10  15  20  25
+
+julia> a[BlockIndex((1,2), (1,2))]
+11
+
+julia> a[BlockIndex((2,2), (2,3))]
+20
+```
 """
 immutable BlockIndex{N}
     I::NTuple{N, Int}
     α::NTuple{N, Int}
 end
 
+@inline BlockIndex(a::Int, b::Int) = BlockIndex((a,), (b,))
+@inline BlockIndex(a::NTuple, b::Int) = BlockIndex(a, (b,))
+@inline BlockIndex(a::Int, b::NTuple) = BlockIndex((a,), b)
+
+@generated function BlockIndex{M, N}(I::NTuple{N, Int}, α::NTuple{M, Int})
+    @assert M < N
+    α_ex = Expr(:tuple, [k <= M ? :(α[$k]) : :(1) for k = 1:N]...)
+    return quote
+        $Expr(:meta, :inline)
+        @inbounds α2 = $α_ex
+        BlockIndex(I, α2)
+    end
+end
 
 """
     global2blockindex{N}(block_sizes::BlockSizes{N}, inds...) -> BlockIndex{N}
@@ -20,6 +54,7 @@ Converts from global indices `inds` to a `BlockIndex`.
     I_ex = Expr(:tuple, [:(block_index[$k][1]) for k = 1:N]...)
     α_ex = Expr(:tuple, [:(block_index[$k][2]) for k = 1:N]...)
     return quote
+        $Expr(:meta, :inline)
         @inbounds block_index = $block_index_ex
         @inbounds I = $I_ex
         @inbounds α = $α_ex
@@ -35,6 +70,7 @@ Converts from a block index to a tuple containing the global indices
 @generated function blockindex2global{N}(block_sizes::BlockSizes{N}, block_index::BlockIndex{N})
     ex = Expr(:tuple, [:(block_sizes[$k, block_index.I[$k]] + block_index.α[$k] - 1) for k = 1:N]...)
     return quote
+        $Expr(:meta, :inline)
         @inbounds v = $ex
         return $ex
     end

src/blocksizes.jl

@@ -93,6 +93,7 @@ end
     #start_indices_ex = Expr(:tuple, [:(1 + _cumsum(block_sizes[$i], block_index[$i]-1)) for i=1:N]...)
     indices_ex = Expr(:tuple, [:(block_sizes[$i, block_index[$i]]:block_sizes[$i, block_index[$i] + 1] - 1) for i = 1:N]...)
     return quote
+        $Expr(:meta, :inline)
         @inbounds inds = $indices_ex
         return inds
     end

src/pseudo_blockarray.jl

@@ -98,6 +98,12 @@ end
 # Indexing #
 ############
 
+@inline function Base.getindex{T, N}(block_arr::PseudoBlockArray{T,N}, blockindex::BlockIndex{N})
+    I = blockindex2global(block_arr.block_sizes, blockindex)
+    @boundscheck checkbounds(block_arr.blocks, I...)
+    @inbounds v = block_arr.blocks[I...]
+    return v
+end
 
 @inline function getblock{T,N}(block_arr::PseudoBlockArray{T,N}, block::Vararg{Int, N})
     range = globalrange(block_arr.block_sizes, block)
@@ -129,6 +135,13 @@ end
     end
 end
 
+@inline function Base.setindex!{T, N}(block_arr::PseudoBlockArray{T,N}, v, blockindex::BlockIndex{N})
+    I = blockindex2global(block_arr.block_sizes, blockindex)
+    @boundscheck checkbounds(block_arr.blocks, I...)
+    @inbounds block_arr.blocks[I...] = v
+    return block_arr
+end
+
 function _check_setblock!{T, N}(blockrange, x, block_arr::PseudoBlockArray{T,N}, block::NTuple{N, Int})
     blocksizes = blocksize(block_arr, block...)
     for i in 1:N
@@ -138,7 +151,6 @@ function _check_setblock!{T, N}(blockrange, x, block_arr::PseudoBlockArray{T,N},
     end
 end
 
-
 @generated function setblock!{T, N}(block_arr::PseudoBlockArray{T, N}, x, block::Vararg{Int, N})
     return quote
         blockrange = globalrange(block_arr.block_sizes, block)
@@ -153,6 +165,8 @@ end
     end
 end
 
+
+
 ########
 # Misc #
 ########

test/test_blockarrays.jl

@@ -3,6 +3,8 @@
     BA_1 = BlockArray(Vector{Float64}, [1,2,3])
     a_1 = rand(2)
     BA_1[Block(2)] = a_1
+    @test BA_1[BlockIndex(2, 1)] == a_1[1]
+    @test BA_1[BlockIndex(2, 2)] == a_1[2]
     @test BA_1[Block(2)] == a_1
     @test BA_1[2] == a_1[1]
     @test_throws DimensionMismatch BA_1[Block(3)] = rand(4)
@@ -23,12 +25,13 @@ end
     for BlockType in (BlockArray, PseudoBlockArray)
         a_1 = rand(6)
         BA_1 = BlockType(a_1, [1,2,3])
-        @test full(BA_1) == a_1
+        @test Array(BA_1) == a_1
         @test nblocks(BA_1) == (3,)
         @test nblocks(BA_1,1) == 3
         @test eltype(similar(BA_1, Float32)) == Float32
         q = rand(1)
         BA_1[Block(1)] = q
+        BA_1[BlockIndex(3, 2)] = a_1[5]
         @test BA_1[Block(1)] == q
         if BlockType == PseudoBlockArray
             q2 = zeros(q)
@@ -47,17 +50,18 @@ end
 
         a_1_sparse = sprand(6, 0.9)
         BA_1_sparse = BlockType(a_1_sparse, [1,2,3])
-        @test full(BA_1_sparse) == a_1_sparse
+        @test Array(BA_1_sparse) == a_1_sparse
         BA_1_sparse[4] = 3.0
         @test BA_1_sparse[4] == 3.0
 
 
         a_2 = rand(3, 7)
         BA_2 = BlockType(a_2, [1,2], [3,4])
-        @test full(BA_2) == a_2
+        @test Array(BA_2) == a_2
         @test nblocks(BA_2) == (2,2)
         @test nblocks(BA_2, 1) == 2
         @test nblocks(BA_2, 2, 1) == (2, 2)
+        BA_2[BlockIndex((2,1), (2,2))] = a_2[3,2]
         @test eltype(similar(BA_2, Float32)) == Float32
         q = rand(1,4)
         BA_2[Block(1,2)] = q
@@ -77,17 +81,18 @@ end
 
         a_2_sparse = sprand(3, 7, 0.9)
         BA_2_sparse = BlockType(a_2_sparse, [1,2], [3,4])
-        @test full(BA_2_sparse) == a_2_sparse
+        @test Array(BA_2_sparse) == a_2_sparse
         BA_2_sparse[1,2] = 3.0
         @test BA_2_sparse[1,2] == 3.0
 
         a_3 = rand(3, 7,4)
         BA_3 = BlockType(a_3, [1,2], [3,4], [1,2,1])
-        @test full(BA_3) == a_3
+        @test Array(BA_3) == a_3
         @test nblocks(BA_3) == (2,2,3)
         @test nblocks(BA_3, 1) == 2
         @test nblocks(BA_3, 3, 1) == (3, 2)
         @test nblocks(BA_3, 3) == 3
+        BA_3[BlockIndex((1,1,1), (1,1,1))] = a_3[1,1,1]
         @test eltype(similar(BA_3, Float32)) == Float32
         q = rand(1,4,2)
         BA_3[Block(1,2,2)] = q