[BlockSparseArrays] Redesign and fix slicing with unit ranges (#1487)

* [BlockSparseArrays] Redesign and fix slicing with unit ranges

* [NDTensors] Bump to v0.3.22

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.21"
+version = "0.3.22"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl

@@ -3,7 +3,8 @@ using Compat: Returns
 using Test: @test, @testset, @test_broken
 using BlockArrays: Block, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
-using NDTensors.GradedAxes: GradedAxes, GradedUnitRange, UnitRangeDual, dual, gradedrange
+using NDTensors.GradedAxes:
+  GradedAxes, GradedUnitRange, UnitRangeDual, blocklabels, dual, gradedrange
 using NDTensors.LabelledNumbers: label
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
@@ -68,13 +69,26 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     a = BlockSparseArray{elt}(d1, d2, d1, d2)
     blockdiagonal!(randn!, a)
     m = fusedims(a, (1, 2), (3, 4))
-    @test axes(m, 1) isa GradedUnitRange
-    @test axes(m, 2) isa GradedUnitRange
+    # TODO: Once block merging is implemented, this should
+    # be the real test.
+    for ax in axes(m)
+      @test ax isa GradedUnitRange
+      @test_broken blocklabels(ax) == [U1(0), U1(1), U1(2)]
+      @test blocklabels(ax) == [U1(0), U1(1), U1(1), U1(2)]
+    end
+    for I in CartesianIndices(m)
+      if I ∈ CartesianIndex.([(1, 1), (4, 4)])
+        @test !iszero(m[I])
+      else
+        @test iszero(m[I])
+      end
+    end
     @test a[1, 1, 1, 1] == m[1, 1]
     @test a[2, 2, 2, 2] == m[4, 4]
     # TODO: Current `fusedims` doesn't merge
     # common sectors, need to fix.
     @test_broken blocksize(m) == (3, 3)
+    @test blocksize(m) == (4, 4)
     @test a == splitdims(m, (d1, d2), (d1, d2))
   end
   @testset "dual axes" begin

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl

@@ -14,10 +14,9 @@ using ..SparseArrayInterface:
 
 # Returns `Vector{<:CartesianIndices}`
 function union_stored_blocked_cartesianindices(as::Vararg{AbstractArray})
+  combined_axes = combine_axes(axes.(as)...)
   stored_blocked_cartesianindices_as = map(as) do a
-    return blocked_cartesianindices(
-      axes(a), combine_axes(axes.(as)...), block_stored_indices(a)
-    )
+    return blocked_cartesianindices(axes(a), combined_axes, block_stored_indices(a))
   end
   return ∪(stored_blocked_cartesianindices_as...)
 end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -1,5 +1,5 @@
 using Adapt: Adapt, WrappedArray
-using BlockArrays: BlockArrays, BlockedUnitRange, blockedrange, unblock
+using BlockArrays: BlockArrays, BlockedUnitRange, BlockRange, blockedrange, unblock
 using SplitApplyCombine: groupcount
 
 const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
@@ -11,28 +11,76 @@ const BlockSparseArrayLike{T,N} = Union{
   <:AbstractBlockSparseArray{T,N},<:WrappedAbstractBlockSparseArray{T,N}
 }
 
-# AbstractArray interface
-# TODO: Use `BlockSparseArrayLike`.
-# TODO: Need to handle block indexing.
-function Base.axes(a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray})
-  return ntuple(i -> sub_axis(axes(parent(a), i), a.indices[i]), ndims(a))
+# Used when making views.
+# TODO: Move to blocksparsearrayinterface.
+function blocksparse_to_indices(a, inds, I)
+  return (unblock(a, inds, I), to_indices(a, BlockArrays._maybetail(inds), Base.tail(I))...)
+end
+
+# TODO: Move to blocksparsearrayinterface.
+function blocksparse_to_indices(a, I)
+  return to_indices(a, axes(a), I)
 end
 
 # Used when making views.
 function Base.to_indices(
-  a::BlockSparseArrayLike, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}}
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractVector{<:Block{1}},Vararg{Any}}
 )
-  return (unblock(a, inds, I), to_indices(a, BlockArrays._maybetail(inds), Base.tail(I))...)
+  return blocksparse_to_indices(a, inds, I)
 end
-function Base.to_indices(a::BlockSparseArrayLike, I::Tuple{Vector{<:Block{1}},Vararg{Any}})
-  return to_indices(a, axes(a), I)
+
+function Base.to_indices(
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractUnitRange{<:Integer},Vararg{Any}}
+)
+  return blocksparse_to_indices(a, inds, I)
+end
+
+# Fixes ambiguity error with BlockArrays.
+function Base.to_indices(a::BlockSparseArrayLike, inds, I::Tuple{BlockRange{1},Vararg{Any}})
+  return blocksparse_to_indices(a, inds, I)
+end
+
+function Base.to_indices(
+  a::BlockSparseArrayLike, I::Tuple{AbstractVector{<:Block{1}},Vararg{Any}}
+)
+  return blocksparse_to_indices(a, 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)
+end
+
+function Base.to_indices(
+  a::BlockSparseArrayLike, I::Tuple{AbstractUnitRange{<:Integer},Vararg{Any}}
+)
+  return blocksparse_to_indices(a, I)
 end
 
 # Used inside `Base.to_indices` when making views.
-function BlockArrays.unblock(a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}})
+# TODO: Move to blocksparsearrayinterface.
+# TODO: Make a special definition for `BlockedVector{<:Block{1}}` in order
+# to merge blocks.
+function blocksparse_unblock(a, inds, I::Tuple{AbstractVector{<:Block{1}},Vararg{Any}})
   return BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1]))
 end
 
+# TODO: Move to blocksparsearrayinterface.
+function blocksparse_unblock(a, inds, I::Tuple{AbstractUnitRange{<:Integer},Vararg{Any}})
+  bs = blockrange(inds[1], I[1])
+  return BlockSlice(bs, blockedunitrange_getindices(inds[1], I[1]))
+end
+
+function BlockArrays.unblock(a, inds, I::Tuple{AbstractVector{<:Block{1}},Vararg{Any}})
+  return blocksparse_unblock(a, inds, I)
+end
+
+function BlockArrays.unblock(
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractUnitRange{<:Integer},Vararg{Any}}
+)
+  return blocksparse_unblock(a, inds, I)
+end
+
 # BlockArrays `AbstractBlockArray` interface
 BlockArrays.blocks(a::BlockSparseArrayLike) = blocksparse_blocks(a)
 

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -267,7 +267,7 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe
   return true
 end
 function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks)
-  return stored_indices(view(blocks(parent(a.array)), axes(a)...))
+  return stored_indices(view(blocks(parent(a.array)), blockrange(a)...))
 end
 # TODO: Either make this the generic interface or define
 # `SparseArrayInterface.sparse_storage`, which is used

NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl

@@ -1,5 +1,6 @@
 @eval module $(gensym())
-using BlockArrays: Block, BlockRange, BlockedUnitRange, blockedrange, blocklength, blocksize
+using BlockArrays:
+  Block, BlockRange, BlockedUnitRange, BlockVector, blockedrange, blocklength, blocksize
 using LinearAlgebra: mul!
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored, block_reshape
 using NDTensors.SparseArrayInterface: nstored
@@ -8,6 +9,27 @@ using Test: @test, @test_broken, @test_throws, @testset
 include("TestBlockSparseArraysUtils.jl")
 @testset "BlockSparseArrays (eltype=$elt)" for elt in
                                                (Float32, Float64, ComplexF32, ComplexF64)
+  @testset "Broken" begin
+    # TODO: These are broken, need to fix.
+    a = BlockSparseArray{elt}([2, 3], [2, 3])
+    for I in (Block.(1:2), [Block(1), Block(2)])
+      b = @view a[I, I]
+      x = randn(elt, 2, 2)
+      b[Block(1, 1)] = x
+      # These outputs a block of zeros,
+      # for some reason the block
+      # is not getting set.
+      # I think the issue is that:
+      # ```julia
+      # @view(@view(a[I, I]))[Block(1, 1)]
+      # ```
+      # creates a doubly-wrapped SubArray
+      # instead of flattening down to a
+      # single SubArray wrapper.
+      @test_broken a[Block(1, 1)] == x
+      @test_broken b[Block(1, 1)] == x
+    end
+  end
   @testset "Basics" begin
     a = BlockSparseArray{elt}([2, 3], [2, 3])
     @test a == BlockSparseArray{elt}(blockedrange([2, 3]), blockedrange([2, 3]))
@@ -255,58 +277,67 @@ include("TestBlockSparseArraysUtils.jl")
     a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
     x = randn(elt, 1, 2)
     @view(a[Block(2, 2)])[1:1, 1:2] = x
-    @test @view(a[Block(2, 2)])[1:1, 1:2] == x
     @test a[Block(2, 2)][1:1, 1:2] == x
-
-    # TODO: This is broken, fix!
-    @test_broken a[3:3, 4:5] == x
+    @test @view(a[Block(2, 2)])[1:1, 1:2] == x
+    @test a[3:3, 4:5] == x
 
     a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
     x = randn(elt, 1, 2)
     @views a[Block(2, 2)][1:1, 1:2] = x
-    @test @view(a[Block(2, 2)])[1:1, 1:2] == x
     @test a[Block(2, 2)][1:1, 1:2] == x
-
-    # TODO: This is broken, fix!
-    @test_broken a[3:3, 4:5] == x
+    @test @view(a[Block(2, 2)])[1:1, 1:2] == x
+    @test a[3:3, 4:5] == x
 
     a = BlockSparseArray{elt}([2, 3], [2, 3])
     @views for b in [Block(1, 1), Block(2, 2)]
       # TODO: Use `blocksizes(a)[Int.(Tuple(b))...]` once available.
       a[b] = randn(elt, size(a[b]))
     end
-    b = @view a[[Block(1), Block(2)], [Block(1), Block(2)]]
-    for I in CartesianIndices(a)
-      @test b[I] == a[I]
-    end
-    for block in BlockRange(a)
-      @test b[block] == a[block]
+    for I in (
+      Block.(1:2),
+      [Block(1), Block(2)],
+      BlockVector([Block(1), Block(2)], [1, 1]),
+      # TODO: This should merge blocks.
+      BlockVector([Block(1), Block(2)], [2]),
+    )
+      b = @view a[I, I]
+      for I in CartesianIndices(a)
+        @test b[I] == a[I]
+      end
+      for block in BlockRange(a)
+        @test b[block] == a[block]
+      end
     end
 
     a = BlockSparseArray{elt}([2, 3], [2, 3])
     @views for b in [Block(1, 1), Block(2, 2)]
       # TODO: Use `blocksizes(a)[Int.(Tuple(b))...]` once available.
       a[b] = randn(elt, size(a[b]))
     end
-    b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
-    @test b[Block(1, 1)] == a[Block(2, 2)]
-    @test b[Block(2, 1)] == a[Block(1, 2)]
-    @test b[Block(1, 2)] == a[Block(2, 1)]
-    @test b[Block(2, 2)] == a[Block(1, 1)]
-    @test b[1, 1] == a[3, 3]
-    @test b[4, 4] == a[1, 1]
-    b[4, 4] = 44
-    @test b[4, 4] == 44
-
-    ## Broken, need to fix.
+    for I in (
+      [Block(2), Block(1)],
+      BlockVector([Block(2), Block(1)], [1, 1]),
+      # TODO: This should merge blocks.
+      BlockVector([Block(2), Block(1)], [2]),
+    )
+      b = @view a[I, I]
+      @test b[Block(1, 1)] == a[Block(2, 2)]
+      @test b[Block(2, 1)] == a[Block(1, 2)]
+      @test b[Block(1, 2)] == a[Block(2, 1)]
+      @test b[Block(2, 2)] == a[Block(1, 1)]
+      @test b[1, 1] == a[3, 3]
+      @test b[4, 4] == a[1, 1]
+      b[4, 4] = 44
+      @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)])))
     b = a[Block(2):Block(2), Block(1):Block(2)]
-    @test_broken block_nstored(b) == 1
-    @test_broken b == Array(a)[3:5, 1:end]
+    @test block_nstored(b) == 1
+    @test b == Array(a)[3:5, 1:end]
   end
   @testset "LinearAlgebra" begin
     a1 = BlockSparseArray{elt}([2, 3], [2, 3])

NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl

@@ -36,8 +36,15 @@ function blockedunitrange_findblockindex(a::BlockedUnitRange, index::Integer)
   return @inbounds findblockindex(a, index + first(a) - 1)
 end
 
+function blockedunitrange_getindices(a::AbstractUnitRange, indices)
+  return a[indices]
+end
+
 # TODO: Move this to a `BlockArraysExtensions` library.
 # Like `a[indices]` but preserves block structure.
+# TODO: Consider calling this something else, for example
+# `blocked_getindex`. See the discussion here:
+# https://github.com/JuliaArrays/BlockArrays.jl/issues/347
 function blockedunitrange_getindices(
   a::BlockedUnitRange, indices::AbstractUnitRange{<:Integer}
 )
@@ -72,10 +79,21 @@ function blockedunitrange_getindices(a::BlockedUnitRange, indices::Vector{<:Inte
 end
 
 # TODO: Move this to a `BlockArraysExtensions` library.
+# TODO: Make a special definition for `BlockedVector{<:Block{1}}` in order
+# to merge blocks.
 function blockedunitrange_getindices(
-  a::BlockedUnitRange, indices::Vector{<:Union{Block{1},BlockIndexRange{1}}}
+  a::BlockedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
 )
-  return mortar(map(index -> a[index], indices))
+  # Without converting `indices` to `Vector`,
+  # mapping `indices` outputs a `BlockVector`
+  # which is harder to reason about.
+  blocks = map(index -> a[index], Vector(indices))
+  # We pass `length.(blocks)` to `mortar` in order
+  # to pass block labels to the axes of the output,
+  # if they exist. This makes it so that
+  # `only(axes(a[indices])) isa `GradedUnitRange`
+  # if `a isa `GradedUnitRange`, for example.
+  return mortar(blocks, length.(blocks))
 end
 
 # TODO: Move this to a `BlockArraysExtensions` library.