[BlockSparseArrays] Fix initializing blocks using broadcasting (#1506)

* [BlockSparseArrays] Fix initializing blocks using broadcasting

* [NDTensors] Bump to v0.3.34

Author: mtfishman

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -185,7 +185,22 @@ function Base.setindex!(a::BlockSparseArrayLike{<:Any,1}, value, I::Block{1})
   return a
 end
 
+function Base.fill!(a::AbstractBlockSparseArray, value)
+  if iszero(value)
+    # This drops all of the blocks.
+    sparse_zero!(blocks(a))
+    return a
+  end
+  blocksparse_fill!(a, value)
+  return a
+end
+
 function Base.fill!(a::BlockSparseArrayLike, value)
+  # TODO: Even if `iszero(value)`, this doesn't drop
+  # blocks from `a`, and additionally allocates
+  # new blocks filled with zeros, unlike
+  # `fill!(a::AbstractBlockSparseArray, value)`.
+  # Consider changing that behavior when possible.
   blocksparse_fill!(a, value)
   return a
 end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -96,11 +96,6 @@ function blocksparse_setindex!(
 end
 
 function blocksparse_fill!(a::AbstractArray, value)
-  if iszero(value)
-    # This drops all of the blocks.
-    sparse_zero!(blocks(a))
-    return a
-  end
   for b in BlockRange(a)
     # We can't use:
     # ```julia
@@ -284,12 +279,14 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) w
   return a[Tuple(I)...]
 end
 function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N}
-  parent_blocks = view(blocks(parent(a.array)), axes(a)...)
+  parent_blocks = @view blocks(parent(a.array))[blockrange(a)...]
   # TODO: The following line is required to instantiate
   # uninstantiated blocks, maybe use `@view!` instead,
   # or some other code pattern.
   parent_blocks[I...] = parent_blocks[I...]
-  return parent_blocks[I...][blockindices(parent(a.array), Block(I), a.array.indices)...] =
+  # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`.
+  block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a)))
+  return parent_blocks[I...][blockindices(parent(a.array), block, a.array.indices)...] =
     value
 end
 function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N}

test/backup/runtests.jl

@@ -13,6 +13,7 @@ using BlockArrays:
   blocksize,
   blocksizes,
   mortar
+using Compat: @compat
 using LinearAlgebra: mul!
 using NDTensors.BlockSparseArrays:
   @view!, BlockSparseArray, block_nstored, block_reshape, view!
@@ -23,22 +24,21 @@ include("TestBlockSparseArraysUtils.jl")
 @testset "BlockSparseArrays (eltype=$elt)" for elt in
                                                (Float32, Float64, ComplexF32, ComplexF64)
   @testset "Broken" begin
-    a = BlockSparseArray{elt}([2, 3], [3, 4])
-    @test_broken a[Block(1, 2)] .= 1
-
     a = BlockSparseArray{elt}([2, 3], [3, 4])
     b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
+    @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
     @test_broken b[2:4, 2:4]
 
+    a = BlockSparseArray{elt}([2, 3], [3, 4])
+    b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(1, 1)]
+    @test_broken b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
+
     a = BlockSparseArray{elt}([2, 3], [3, 4])
     b = @views a[Block(1, 1)][1:2, 1:1]
+    @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
     for i in parentindices(b)
       @test_broken i isa BlockSlice{<:BlockIndexRange{1}}
     end
-
-    a = BlockSparseArray{elt}([2, 3], [3, 4])
-    b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
-    @test_broken b[Block(1, 1)] = randn(3, 3)
   end
   @testset "Basics" begin
     a = BlockSparseArray{elt}([2, 3], [2, 3])
@@ -82,6 +82,26 @@ include("TestBlockSparseArraysUtils.jl")
     @test block_nstored(a) == 2
     @test nstored(a) == 2 * 4 + 3 * 3
 
+    a = BlockSparseArray{elt}([2, 3], [3, 4])
+    a[Block(1, 2)] .= 2
+    @test eltype(a) == elt
+    @test all(==(2), a[Block(1, 2)])
+    @test iszero(a[Block(1, 1)])
+    @test iszero(a[Block(2, 1)])
+    @test iszero(a[Block(2, 2)])
+    @test block_nstored(a) == 1
+    @test nstored(a) == 2 * 4
+
+    a = BlockSparseArray{elt}([2, 3], [3, 4])
+    a[Block(1, 2)] .= 0
+    @test eltype(a) == elt
+    @test iszero(a[Block(1, 1)])
+    @test iszero(a[Block(2, 1)])
+    @test iszero(a[Block(1, 2)])
+    @test iszero(a[Block(2, 2)])
+    @test block_nstored(a) == 1
+    @test nstored(a) == 2 * 4
+
     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]))
@@ -490,13 +510,45 @@ include("TestBlockSparseArraysUtils.jl")
       @test b[Block(2, 2)] == x
     end
 
+    function f1()
+      a = BlockSparseArray{elt}([2, 3], [3, 4])
+      b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
+      x = randn(elt, 3, 4)
+      b[Block(1, 1)] .= x
+      return (; a, b, x)
+    end
+    function f2()
+      a = BlockSparseArray{elt}([2, 3], [3, 4])
+      b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
+      x = randn(elt, 3, 4)
+      b[Block(1, 1)] = x
+      return (; a, b, x)
+    end
+    for abx in (f1(), f2())
+      @compat (; a, b, x) = abx
+      @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
+      @test block_nstored(b) == 1
+      @test b[Block(1, 1)] == x
+      for blck in [Block(2, 1), Block(1, 2), Block(2, 2)]
+        @test iszero(b[blck])
+      end
+      @test block_nstored(a) == 1
+      @test a[Block(2, 2)] == x
+      for blck in [Block(1, 1), Block(2, 1), Block(1, 2)]
+        @test iszero(a[blck])
+      end
+      @test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3)
+    end
+
     a = BlockSparseArray{elt}([2, 3], [3, 4])
-    b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
-    x = randn(elt, 3, 4)
-    b[Block(1, 1)] .= x
-    @test b[Block(1, 1)] == x
-    @test a[Block(2, 2)] == x
-    @test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3)
+    b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(2, 1)]
+    @test iszero(b)
+    @test size(b) == (2, 4)
+    x = randn(elt, 2, 4)
+    b .= x
+    @test b == x
+    @test a[Block(1, 2)] == x
+    @test block_nstored(a) == 1
 
     a = BlockSparseArray{elt}([2, 3], [3, 4])
     b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]