[BlockSparseArrays] Towards block merging (#1512)

* [BlockSparseArrays] Towards block merging

* [NDTensors] Bump to v0.3.37

NDTensors/Project.toml

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

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -230,8 +230,26 @@ function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
   return r
 end
 
+# This handles changing the blocking, for example:
+# a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
+# I = blockedrange([4, 4])
+# a[I, I]
+# TODO: Generalize to `AbstractBlockedUnitRange`.
+function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockedOneTo{<:Integer})
+  # TODO: Probably this is incorrect and should be something like:
+  # return findblock(axis, first(r)):findblock(axis, last(r))
+  return only(blockaxes(r))
+end
+
+# This handles changing the blocking, for example:
+# a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
+# I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
+# a[I, I]
+# TODO: Generalize to `AbstractBlockedUnitRange` and `AbstractBlockVector`.
 function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockVector{<:Integer})
-  return error("Slicing not implemented for range of type `$(typeof(r))`.")
+  # TODO: Probably this is incorrect and should be something like:
+  # return findblock(axis, first(r)):findblock(axis, last(r))
+  return only(blockaxes(r))
 end
 
 using BlockArrays: BlockSlice

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

@@ -1,6 +1,7 @@
 using Adapt: Adapt, WrappedArray
 using BlockArrays:
   BlockArrays,
+  AbstractBlockVector,
   AbstractBlockedUnitRange,
   BlockIndexRange,
   BlockRange,
@@ -40,8 +41,9 @@ function Base.to_indices(
 end
 
 # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])]
+# a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])]
 function Base.to_indices(
-  a::BlockSparseArrayLike, inds, I::Tuple{BlockVector{<:Block{1}},Vararg{Any}}
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}}
 )
   return blocksparse_to_indices(a, inds, I)
 end

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

@@ -4,6 +4,7 @@ using BlockArrays:
   BlockIndex,
   BlockVector,
   BlockedUnitRange,
+  BlockedVector,
   block,
   blockcheckbounds,
   blocklengths,
@@ -46,6 +47,12 @@ function blocksparse_to_indices(a, inds, I::Tuple{BlockVector{<:Block{1}},Vararg
   return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
 end
 
+# TODO: Should this be combined with the version above?
+function blocksparse_to_indices(a, inds, I::Tuple{BlockedVector{<:Block{1}},Vararg{Any}})
+  I1 = blockedunitrange_getindices(inds[1], I[1])
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+
 # TODO: Need to implement this!
 function block_merge end
 
@@ -223,6 +230,9 @@ function Base.size(a::SparseSubArrayBlocks)
   return length.(axes(a))
 end
 function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N}
+  # TODO: Should this be defined as `@view a.array[Block(I)]` instead?
+  ## return @view a.array[Block(I)]
+
   parent_blocks = @view blocks(parent(a.array))[blockrange(a)...]
   parent_block = parent_blocks[I...]
   # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`.

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

@@ -7,6 +7,7 @@ using BlockArrays:
   BlockVector,
   BlockedOneTo,
   BlockedUnitRange,
+  BlockedVector,
   blockedrange,
   blocklength,
   blocklengths,
@@ -23,6 +24,24 @@ 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
+    a = BlockSparseArray{elt}([2, 2, 2, 2], [2, 2, 2, 2])
+    @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)]
+      a[I] = randn(elt, size(a[I]))
+    end
+
+    I = blockedrange([4, 4])
+    b = @view a[I, I]
+    @test_broken copy(b)
+
+    I = BlockedVector(Block.(1:4), [2, 2])
+    b = @view a[I, I]
+    @test_broken copy(b)
+
+    I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
+    b = @view a[I, I]
+    @test_broken copy(b)
+  end
   @testset "Basics" begin
     a = BlockSparseArray{elt}([2, 3], [2, 3])
     @test a == BlockSparseArray{elt}(blockedrange([2, 3]), blockedrange([2, 3]))

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

@@ -6,6 +6,7 @@ using BlockArrays:
   BlockRange,
   BlockSlice,
   BlockedUnitRange,
+  BlockedVector,
   block,
   blockindex,
   findblock,
@@ -70,6 +71,21 @@ function blockedunitrange_getindices(
   return blockedunitrange(indices .+ (first(a) - 1), blocklengths)
 end
 
+# TODO: Make sure this handles block labels (AbstractGradedUnitRange) correctly.
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::BlockedVector{<:Block{1},<:BlockRange{1}}
+)
+  blocklengths = map(bs -> sum(b -> length(a[b]), bs), blocks(indices))
+  return blockedrange(blocklengths)
+end
+
+# TODO: Make sure this handles block labels (AbstractGradedUnitRange) correctly.
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::BlockedVector{<:Block{1}}
+)
+  return mortar(map(bs -> mortar(map(b -> a[b], bs)), blocks(indices)))
+end
+
 # TODO: Move this to a `BlockArraysExtensions` library.
 function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockIndexRange)
   return a[block(indices)][only(indices.indices)]