Make blockwise map friendlier to abstract block types

Author: mtfishman

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -42,6 +42,15 @@ function eachblockstoredindex(a::AbstractArray)
   return Block.(Tuple.(eachstoredindex(blocks(a))))
 end
 
+function SparseArraysBase.isstored(
+  a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}
+) where {N}
+  return isstored(blocks(a), Int.(I)...)
+end
+function SparseArraysBase.isstored(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
+  return isstored(a, Tuple(I)...)
+end
+
 using DiagonalArrays: diagindices
 # Block version of `DiagonalArrays.diagindices`.
 function blockdiagindices(a::AbstractArray)

src/blocksparsearrayinterface/map.jl

@@ -18,6 +18,38 @@ function union_eachblockstoredindex(as::AbstractArray...)
   return ∪(map(eachblockstoredindex, as)...)
 end
 
+# Get a view of a block assuming it is stored.
+function viewblock_stored(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+  return blocks(a)[Int.(I)...]
+end
+function viewblock_stored(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
+  return viewblock_stored(a, Tuple(I)...)
+end
+
+using FillArrays: Zeros
+# Get a view of a block if it is stored, otherwise return a lazy zeros.
+function viewblock_or_zeros(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+  if isstored(a, I...)
+    return viewblock_stored(a, I...)
+  else
+    block_ax = map((ax, i) -> eachblockaxis(ax)[Int(i)], axes(a), I)
+    return Zeros{eltype(a)}(block_ax)
+  end
+end
+function viewblock_or_zeros(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
+  return viewblock_or_zeros(a, Tuple(I)...)
+end
+
+function map_block!(f, a_dest::AbstractArray, I::Block, a_srcs::AbstractArray...)
+  a_srcs_I = map(a_src -> viewblock_or_zeros(a_src, I), a_srcs)
+  if isstored(a_dest, I)
+    a_dest[I] .= f.(a_srcs_I...)
+  else
+    a_dest[I] = Broadcast.broadcast_preserving_zero_d(f, a_srcs_I...)
+  end
+  return a_dest
+end
+
 function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...)
   # TODO: This assumes element types are numbers, generalize this logic.
   f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest))
@@ -27,22 +59,7 @@ function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...)
     BlockRange(a_dest)
   end
   for I in Is
-    # TODO: Use:
-    # block_dest = @view a_dest[I]
-    # or:
-    # block_dest = @view! a_dest[I]
-    block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(I))...]
-    # TODO: Use:
-    # block_srcs = map(a_src -> @view(a_src[I]), a_srcs)
-    block_srcs = map(a_srcs) do a_src
-      return blocks_maybe_single(a_src)[Int.(Tuple(I))...]
-    end
-    # TODO: Use `map!!` to handle immutable blocks.
-    map!(f, block_dest, block_srcs...)
-    # Replace the entire block, handles initializing new blocks
-    # or if blocks are immutable.
-    # TODO: Use `a_dest[I] = block_dest`.
-    blocks(a_dest)[Int.(Tuple(I))...] = block_dest
+    map_block!(f, a_dest, I, a_srcs...)
   end
   return a_dest
 end

test/test_abstract_blocktype.jl

@@ -0,0 +1,44 @@
+using Adapt: adapt
+using BlockArrays: Block
+using BlockSparseArrays: BlockSparseMatrix, blockstoredlength
+using JLArrays: JLArray
+using SparseArraysBase: storedlength
+using Test: @test, @test_broken, @testset
+
+elts = (Float32, Float64, ComplexF32)
+arrayts = (Array, JLArray)
+@testset "Abstract block type (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts,
+  elt in elts
+
+  dev = adapt(arrayt)
+  a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3])
+  @test sprint(show, MIME"text/plain"(), a) isa String
+  @test iszero(storedlength(a))
+  @test iszero(blockstoredlength(a))
+  a[Block(1, 1)] = dev(randn(elt, 2, 2))
+  a[Block(2, 2)] = dev(randn(elt, 3, 3))
+  @test !iszero(a[Block(1, 1)])
+  @test a[Block(1, 1)] isa arrayt{elt,2}
+  @test !iszero(a[Block(2, 2)])
+  @test a[Block(2, 2)] isa arrayt{elt,2}
+  @test iszero(a[Block(2, 1)])
+  @test a[Block(2, 1)] isa Matrix{elt}
+  @test iszero(a[Block(1, 2)])
+  @test a[Block(1, 2)] isa Matrix{elt}
+
+  b = copy(a)
+  @test Array(b) ≈ Array(a)
+
+  b = a + a
+  @test Array(b) ≈ Array(a) + Array(a)
+
+  b = 3a
+  @test Array(b) ≈ 3Array(a)
+
+  if arrayt === Array
+    b = a * a
+    @test Array(b) ≈ Array(a) * Array(a)
+  else
+    @test_broken a * a
+  end
+end