Construct BlockSparseArray when slicing with graded unit ranges

Author: mtfishman

Project.toml

@@ -25,6 +25,7 @@ LabelledNumbers = "f856a3a6-4152-4ec4-b2a7-02c1a55d7993"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
 [extensions]
+BlockSparseArraysGradedUnitRangesExt = "GradedUnitRanges"
 BlockSparseArraysTensorAlgebraExt = ["LabelledNumbers", "TensorAlgebra"]
 
 [compat]

ext/BlockSparseArraysGradedUnitRangesExt/BlockSparseArraysGradedUnitRangesExt.jl

@@ -0,0 +1,44 @@
+module BlockSparseArraysGradedUnitRangesExt
+
+using BlockSparseArrays: BlockSparseArray
+using GradedUnitRanges: AbstractGradedUnitRange
+
+# A block spare array similar to the input (dense) array.
+# TODO: Make `BlockSparseArrays.blocksparse_similar` more general and use that,
+# and also turn it into an DerivableInterfaces.jl-based interface function.
+function similar_blocksparse(
+  a::AbstractArray,
+  elt::Type,
+  axes::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}},
+)
+  # TODO: Probably need to unwrap the type of `a` in certain cases
+  # to make a proper block type.
+  return BlockSparseArray{elt,length(axes),typeof(a)}(axes)
+end
+
+function Base.similar(
+  a::AbstractArray,
+  elt::Type,
+  axes::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}},
+)
+  return similar_blocksparse(a, elt, axes)
+end
+
+# Fix ambiguity error with `BlockArrays.jl`.
+function Base.similar(
+  a::StridedArray,
+  elt::Type,
+  axes::Tuple{
+    AbstractGradedUnitRange,AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}
+  },
+)
+  return similar_blocksparse(a, elt, axes)
+end
+
+function Base.getindex(a::AbstractArray, I::AbstractGradedUnitRange...)
+  a′ = similar(a, only.(axes.(I))...)
+  a′ .= a
+  return a′
+end
+
+end

src/abstractblocksparsearray/map.jl

@@ -69,11 +69,19 @@ end
 @interface interface::AbstractBlockSparseArrayInterface function Base.map!(
   f, a_dest::AbstractArray, a_srcs::AbstractArray...
 )
+  if isempty(a_srcs)
+    # Broadcast expressions of the form `a .= 2`.
+    error("Not implemented.")
+  end
   if iszero(ndims(a_dest))
     @interface interface map_zero_dim!(f, a_dest, a_srcs...)
     return a_dest
   end
 
+  # TODO: This assumes element types are numbers, generalize this logic.
+  elt = promote_type(eltype.(a_srcs)...)
+  f_preserves_zeros = f(zero(elt)) == zero(elt)
+
   a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs)
   for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...)
     BI_dest = blockindexrange(a_dest, I)
@@ -88,11 +96,13 @@ end
     end
     subblock_dest = @view block_dest[BI_dest.indices...]
     subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs))
-    # TODO: Use `map!!` to handle immutable blocks.
-    map!(f, subblock_dest, subblock_srcs...)
-    # Replace the entire block, handles initializing new blocks
-    # or if blocks are immutable.
-    blocks(a_dest)[Int.(Tuple(_block(BI_dest)))...] = block_dest
+    if f_preserves_zeros && any(!iszero, subblock_srcs)
+      # TODO: Use `map!!` to handle immutable blocks.
+      map!(f, subblock_dest, subblock_srcs...)
+      # Replace the entire block, handles initializing new blocks
+      # or if blocks are immutable.
+      blocks(a_dest)[Int.(Tuple(_block(BI_dest)))...] = block_dest
+    end
   end
   return a_dest
 end
@@ -120,7 +130,17 @@ end
 end
 
 function Base.map!(f, a_dest::AbstractArray, a_srcs::AnyAbstractBlockSparseArray...)
-  @interface interface(a_srcs...) map!(f, a_dest, a_srcs...)
+  @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...)
+  return a_dest
+end
+function Base.map!(f, a_dest::AnyAbstractBlockSparseArray, a_srcs::AbstractArray...)
+  @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...)
+  return a_dest
+end
+function Base.map!(
+  f, a_dest::AnyAbstractBlockSparseArray, a_srcs::AnyAbstractBlockSparseArray...
+)
+  @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...)
   return a_dest
 end
 

src/blocksparsearrayinterface/broadcast.jl

@@ -44,6 +44,46 @@ function Base.copyto!(
   # convert to map
   # flatten and only keep the AbstractArray arguments
   m = Mapped(bc)
-  @interface interface(bc) map!(m.f, dest, m.args...)
+  @interface interface(dest, bc) map!(m.f, dest, m.args...)
+  return dest
+end
+
+# Broadcasting implementation
+# TODO: Delete this in favor of `DerivableInterfaces` version.
+function Base.copyto!(dest::AnyAbstractBlockSparseArray, bc::Broadcasted)
+  # convert to map
+  # flatten and only keep the AbstractArray arguments
+  m = Mapped(bc)
+  # TODO: Include `bc` when determining interface, currently
+  # `interface(::Type{<:Base.Broadcast.DefaultArrayStyle})`
+  # isn't defined.
+  @interface interface(dest) map!(m.f, dest, m.args...)
+  return dest
+end
+
+# Broadcasting implementation
+# TODO: Delete this in favor of `DerivableInterfaces` version.
+function Base.copyto!(
+  dest::AnyAbstractBlockSparseArray, bc::Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}}
+)
+  # convert to map
+  # flatten and only keep the AbstractArray arguments
+  m = Mapped(bc)
+  # TODO: Include `bc` when determining interface, currently
+  # `interface(::Type{<:Base.Broadcast.DefaultArrayStyle})`
+  # isn't defined.
+  @interface interface(dest) map!(m.f, dest, m.args...)
+  return dest
+end
+
+# Broadcasting implementation
+# TODO: Delete this in favor of `DerivableInterfaces` version.
+function Base.copyto!(
+  dest::AnyAbstractBlockSparseArray{<:Any,N}, bc::Broadcasted{BlockSparseArrayStyle{N}}
+) where {N}
+  # convert to map
+  # flatten and only keep the AbstractArray arguments
+  m = Mapped(bc)
+  @interface interface(dest, bc) map!(m.f, dest, m.args...)
   return dest
 end