[BlockSparseArrays] Sub-slices of multiple blocks (#1489)

* [BlockSparseArrays] Sub-slices of multiple blocks

* [NDTensors] Bump to v0.3.24

Author: mtfishman

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -159,8 +159,8 @@ function blockrange(axis::AbstractUnitRange, r::UnitRange)
 end
 
 function blockrange(axis::AbstractUnitRange, r::Int)
-  error("Slicing with integer values isn't supported.")
-  return findblock(axis, r)
+  ## return findblock(axis, r)
+  return error("Slicing with integer values isn't supported.")
 end
 
 function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
@@ -187,6 +187,24 @@ function blockrange(axis::AbstractUnitRange, r::BlockIndexRange)
   return Block(r):Block(r)
 end
 
+function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:BlockIndexRange{1}})
+  return error("Slicing not implemented for range of type `$(typeof(r))`.")
+end
+
+function blockrange(
+  axis::AbstractUnitRange,
+  r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
+)
+  return map(b -> Block(b), blocks(r))
+end
+
+# This handles slicing with `:`/`Colon()`.
+function blockrange(axis::AbstractUnitRange, r::Base.Slice)
+  # TODO: Maybe use `BlockRange`, but that doesn't output
+  # the same thing.
+  return only(blockaxes(axis))
+end
+
 function blockrange(axis::AbstractUnitRange, r)
   return error("Slicing not implemented for range of type `$(typeof(r))`.")
 end
@@ -228,6 +246,22 @@ function blockindices(a::AbstractUnitRange, b::Block, r::BlockIndices)
   return blockindices(a, b, r.blocks)
 end
 
+function blockindices(
+  a::AbstractUnitRange,
+  b::Block,
+  r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
+)
+  # TODO: Change to iterate over `BlockRange(r)`
+  # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404
+  # is fixed.
+  for bl in blocks(r)
+    if b == Block(bl)
+      return only(bl.indices)
+    end
+  end
+  return error("Block not found.")
+end
+
 function cartesianindices(a::AbstractArray, b::Block)
   return cartesianindices(axes(a), b)
 end

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -1,5 +1,6 @@
 using Adapt: Adapt, WrappedArray
-using BlockArrays: BlockArrays, BlockedUnitRange, BlockRange, blockedrange, unblock
+using BlockArrays:
+  BlockArrays, BlockedUnitRange, BlockIndexRange, BlockRange, blockedrange, mortar, unblock
 using SplitApplyCombine: groupcount
 
 const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
@@ -46,6 +47,15 @@ function Base.to_indices(
   return blocksparse_to_indices(a, I)
 end
 
+# Handle case of indexing with `[Block(1)[1:2], Block(2)[1:2]]`
+# by converting it to a `BlockVector` with
+# `mortar([Block(1)[1:2], Block(2)[1:2]])`.
+function Base.to_indices(
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractVector{<:BlockIndexRange{1}},Vararg{Any}}
+)
+  return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
+end
+
 # Fixes ambiguity error with BlockArrays.
 function Base.to_indices(a::BlockSparseArrayLike, I::Tuple{BlockRange{1},Vararg{Any}})
   return blocksparse_to_indices(a, I)
@@ -126,14 +136,25 @@ function Base.getindex(
   return blocksparse_getindex(a, block...)
 end
 
-# TODO: Define `issasigned(a, ::Block{N})`.
+# TODO: Define `blocksparse_isassigned`.
 function Base.isassigned(
   a::BlockSparseArrayLike{<:Any,N}, index::Vararg{Block{1},N}
 ) where {N}
-  # TODO: Define `blocksparse_isassigned`.
   return isassigned(blocks(a), Int.(index)...)
 end
 
+function Base.isassigned(a::BlockSparseArrayLike{<:Any,N}, index::Block{N}) where {N}
+  return isassigned(a, Tuple(index)...)
+end
+
+# TODO: Define `blocksparse_isassigned`.
+function Base.isassigned(
+  a::BlockSparseArrayLike{<:Any,N}, index::Vararg{BlockIndex{1},N}
+) where {N}
+  b = block.(index)
+  return isassigned(a, b...) && isassigned(@view(a[b...]), blockindex.(index)...)
+end
+
 function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::BlockIndex{N}) where {N}
   blocksparse_setindex!(a, value, I)
   return a

test/test_basics.jl

@@ -1,6 +1,14 @@
 @eval module $(gensym())
 using BlockArrays:
-  Block, BlockRange, BlockedUnitRange, BlockVector, blockedrange, blocklength, blocksize
+  Block,
+  BlockRange,
+  BlockedUnitRange,
+  BlockVector,
+  blockedrange,
+  blocklength,
+  blocklengths,
+  blocksize,
+  mortar
 using LinearAlgebra: mul!
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored, block_reshape
 using NDTensors.SparseArrayInterface: nstored
@@ -331,13 +339,69 @@ include("TestBlockSparseArraysUtils.jl")
       @test b[4, 4] == 44
     end
 
-    # This is outputting only zero blocks.
     a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
-    a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)])))
-    a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)])))
+    @views for b in [Block(1, 2), Block(2, 1)]
+      a[b] = randn(elt, size(a[b]))
+    end
     b = a[Block(2):Block(2), Block(1):Block(2)]
     @test block_nstored(b) == 1
     @test b == Array(a)[3:5, 1:end]
+
+    a = BlockSparseArray{elt}(undef, ([2, 3, 4], [2, 3, 4]))
+    # TODO: Define `block_diagindices`.
+    @views for b in [Block(1, 1), Block(2, 2), Block(3, 3)]
+      a[b] = randn(elt, size(a[b]))
+    end
+    for (I1, I2) in (
+      (mortar([Block(2)[2:3], Block(3)[1:3]]), mortar([Block(2)[2:3], Block(3)[2:3]])),
+      ([Block(2)[2:3], Block(3)[1:3]], [Block(2)[2:3], Block(3)[2:3]]),
+    )
+      for b in (a[I1, I2], @view(a[I1, I2]))
+        # TODO: Rename `block_stored_length`.
+        @test block_nstored(b) == 2
+        @test b[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]]
+        @test b[Block(2, 2)] == a[Block(3, 3)[1:3, 2:3]]
+      end
+    end
+
+    a = BlockSparseArray{elt}(undef, ([3, 3], [3, 3]))
+    # TODO: Define `block_diagindices`.
+    @views for b in [Block(1, 1), Block(2, 2)]
+      a[b] = randn(elt, size(a[b]))
+    end
+    I = mortar([Block(1)[1:2], Block(2)[1:2]])
+    b = a[:, I]
+    @test b[Block(1, 1)] == a[Block(1, 1)][:, 1:2]
+    @test b[Block(2, 1)] == a[Block(2, 1)][:, 1:2]
+    @test b[Block(1, 2)] == a[Block(1, 2)][:, 1:2]
+    @test b[Block(2, 2)] == a[Block(2, 2)][:, 1:2]
+    @test blocklengths.(axes(b)) == ([3, 3], [2, 2])
+    # TODO: Rename `block_stored_length`.
+    @test blocksize(b) == (2, 2)
+    @test block_nstored(b) == 2
+
+    a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
+    @views for b in [Block(1, 2), Block(2, 1)]
+      a[b] = randn(elt, size(a[b]))
+    end
+    @test isassigned(a, 1, 1)
+    @test isassigned(a, 5, 7)
+    @test !isassigned(a, 0, 1)
+    @test !isassigned(a, 5, 8)
+    @test isassigned(a, Block(1), Block(1))
+    @test isassigned(a, Block(2), Block(2))
+    @test !isassigned(a, Block(1), Block(0))
+    @test !isassigned(a, Block(3), Block(2))
+    @test isassigned(a, Block(1, 1))
+    @test isassigned(a, Block(2, 2))
+    @test !isassigned(a, Block(1, 0))
+    @test !isassigned(a, Block(3, 2))
+    @test isassigned(a, Block(1)[1], Block(1)[1])
+    @test isassigned(a, Block(2)[3], Block(2)[4])
+    @test !isassigned(a, Block(1)[0], Block(1)[1])
+    @test !isassigned(a, Block(2)[3], Block(2)[5])
+    @test !isassigned(a, Block(1)[1], Block(0)[1])
+    @test !isassigned(a, Block(3)[3], Block(2)[4])
   end
   @testset "LinearAlgebra" begin
     a1 = BlockSparseArray{elt}([2, 3], [2, 3])