Update for latest Derive, SparseArraysBase

Author: mtfishman

examples/README.jl

@@ -36,7 +36,7 @@ julia> Pkg.add(url="https://github.com/ITensor/BlockSparseArrays.jl")
 # ## Examples
 
 using BlockArrays: BlockArrays, BlockedVector, Block, blockedrange
-using BlockSparseArrays: BlockSparseArray, block_stored_length
+using BlockSparseArrays: BlockSparseArray, block_storedlength
 using Test: @test, @test_broken
 
 function main()
@@ -63,13 +63,13 @@ function main()
   ]
   b = BlockSparseArray(nz_blocks, d_blocks, i_axes)
 
-  @test block_stored_length(b) == 2
+  @test block_storedlength(b) == 2
 
   ## Blocks with discontiguous underlying data
   d_blocks = randn.(nz_block_sizes)
   b = BlockSparseArray(nz_blocks, d_blocks, i_axes)
 
-  @test block_stored_length(b) == 2
+  @test block_storedlength(b) == 2
 
   ## Access a block
   @test b[Block(1, 1)] == d_blocks[1]
@@ -93,7 +93,7 @@ function main()
   @test b + b ≈ Array(b) + Array(b)
   @test b + b isa BlockSparseArray
   ## TODO: Fix this, broken.
-  @test_broken block_stored_length(b + b) == 2
+  @test_broken block_storedlength(b + b) == 2
 
   scaled_b = 2b
   @test scaled_b ≈ 2Array(b)

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -21,7 +21,7 @@ using BlockArrays:
   findblockindex
 using Dictionaries: Dictionary, Indices
 using GradedUnitRanges: blockedunitrange_getindices, to_blockindices
-using SparseArraysBase: SparseArraysBase, stored_length, stored_indices
+using SparseArraysBase: SparseArraysBase, storedlength, eachstoredindex
 
 # A return type for `blocks(array)` when `array` isn't blocked.
 # Represents a vector with just that single block.
@@ -269,7 +269,7 @@ tuple_oneto(n) = ntuple(identity, n)
 function block_reshape(a::AbstractArray, axes::Tuple{Vararg{AbstractUnitRange}})
   reshaped_blocks_a = reshape(blocks(a), blocklength.(axes))
   reshaped_a = similar(a, axes)
-  for I in stored_indices(reshaped_blocks_a)
+  for I in eachstoredindex(reshaped_blocks_a)
     block_size_I = map(i -> length(axes[i][Block(I[i])]), tuple_oneto(length(axes)))
     # TODO: Better converter here.
     reshaped_a[Block(Tuple(I))] = reshape(reshaped_blocks_a[I], block_size_I)
@@ -465,8 +465,8 @@ function findblocks(axis::AbstractUnitRange, range::AbstractUnitRange)
   return findblock(axis, first(range)):findblock(axis, last(range))
 end
 
-function block_stored_indices(a::AbstractArray)
-  return Block.(Tuple.(stored_indices(blocks(a))))
+function block_eachstoredindex(a::AbstractArray)
+  return Block.(Tuple.(eachstoredindex(blocks(a))))
 end
 
 _block(indices) = block(indices)
@@ -533,13 +533,13 @@ function Base.setindex!(a::BlockView{<:Any,N}, value, index::Vararg{Int,N}) wher
   return a
 end
 
-function SparseArraysBase.stored_length(a::BlockView)
+function SparseArraysBase.storedlength(a::BlockView)
   # TODO: Store whether or not the block is stored already as
   # a Bool in `BlockView`.
   I = CartesianIndex(Int.(a.block))
-  # TODO: Use `block_stored_indices`.
-  if I ∈ stored_indices(blocks(a.array))
-    return stored_length(blocks(a.array)[I])
+  # TODO: Use `block_eachstoredindex`.
+  if I ∈ eachstoredindex(blocks(a.array))
+    return storedlength(blocks(a.array)[I])
   end
   return 0
 end

src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl

@@ -1,11 +1,11 @@
 using BlockArrays: AbstractBlockArray, BlocksView
-using SparseArraysBase: SparseArraysBase, stored_length
+using SparseArraysBase: SparseArraysBase, storedlength
 
-function SparseArraysBase.stored_length(a::AbstractBlockArray)
-  return sum(b -> stored_length(b), blocks(a); init=zero(Int))
+function SparseArraysBase.storedlength(a::AbstractBlockArray)
+  return sum(b -> storedlength(b), blocks(a); init=zero(Int))
 end
 
 # TODO: Handle `BlocksView` wrapping a sparse array?
-function SparseArraysBase.storage_indices(a::BlocksView)
+function SparseArraysBase.eachstoredindex(a::BlocksView)
   return CartesianIndices(a)
 end

src/abstractblocksparsearray/map.jl

@@ -62,7 +62,8 @@ end
 # is used to determine `union_stored_blocked_cartesianindices(...)`).
 # `reblock` is a partial solution to that, but a bit ad-hoc.
 # TODO: Move to `blocksparsearrayinterface/map.jl`.
-function SparseArraysBase.sparse_map!(
+## TODO: Make this an `@interface AbstractBlockSparseArray` function.
+function sparse_map!(
   ::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}
 )
   a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs)

src/abstractblocksparsearray/sparsearrayinterface.jl

@@ -1,5 +1,5 @@
 using BlockArrays: Block
-using SparseArraysBase: SparseArraysBase, sparse_storage, stored_indices
+using SparseArraysBase: SparseArraysBase, sparse_storage, eachstoredindex, storedlength
 
 # Structure storing the block sparse storage
 struct BlockSparseStorage{Arr<:AbstractBlockSparseArray}
@@ -11,7 +11,7 @@ function blockindex_to_cartesianindex(a::AbstractArray, blockindex)
 end
 
 function Base.keys(s::BlockSparseStorage)
-  stored_blockindices = Iterators.map(stored_indices(blocks(s.array))) do I
+  stored_blockindices = Iterators.map(eachstoredindex(blocks(s.array))) do I
     block_axes = axes(blocks(s.array)[I])
     blockindices = Block(Tuple(I))[block_axes...]
     return Iterators.map(
@@ -29,10 +29,11 @@ function Base.iterate(s::BlockSparseStorage, args...)
   return iterate(values(s), args...)
 end
 
-function SparseArraysBase.sparse_storage(a::AbstractBlockSparseArray)
-  return BlockSparseStorage(a)
-end
+## TODO: Delete this, define `getstoredindex`, etc.
+## function SparseArraysBase.sparse_storage(a::AbstractBlockSparseArray)
+##   return BlockSparseStorage(a)
+## end
 
-function SparseArraysBase.stored_length(a::AnyAbstractBlockSparseArray)
-  return sum(stored_length, sparse_storage(blocks(a)); init=zero(Int))
+function SparseArraysBase.storedlength(a::AnyAbstractBlockSparseArray)
+  return sum(storedlength, sparse_storage(blocks(a)); init=zero(Int))
 end

src/abstractblocksparsearray/views.jl

@@ -68,8 +68,8 @@ function BlockArrays.viewblock(
   a::AbstractBlockSparseArray{<:Any,N}, block::Vararg{Block{1},N}
 ) where {N}
   I = CartesianIndex(Int.(block))
-  # TODO: Use `block_stored_indices`.
-  if I ∈ stored_indices(blocks(a))
+  # TODO: Use `block_eachstoredindex`.
+  if I ∈ eachstoredindex(blocks(a))
     return blocks(a)[I]
   end
   return BlockView(a, block)
@@ -185,8 +185,8 @@ function BlockArrays.viewblock(
   block::Vararg{Block{1},N},
 ) where {T,N}
   I = CartesianIndex(Int.(block))
-  # TODO: Use `block_stored_indices`.
-  if I ∈ stored_indices(blocks(a))
+  # TODO: Use `block_eachstoredindex`.
+  if I ∈ eachstoredindex(blocks(a))
     return blocks(a)[I]
   end
   return BlockView(parent(a), Block.(Base.reindex(parentindices(blocks(a)), Tuple(I))))

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -13,7 +13,7 @@ using BlockArrays:
   blocks,
   findblockindex
 using LinearAlgebra: Adjoint, Transpose
-using SparseArraysBase: perm, iperm, stored_length, sparse_zero!
+using SparseArraysBase: perm, iperm, storedlength, sparse_zero!
 
 blocksparse_blocks(a::AbstractArray) = error("Not implemented")
 
@@ -136,8 +136,8 @@ function blocksparse_fill!(a::AbstractArray, value)
   return a
 end
 
-function block_stored_length(a::AbstractArray)
-  return stored_length(blocks(a))
+function block_storedlength(a::AbstractArray)
+  return storedlength(blocks(a))
 end
 
 # BlockArrays
@@ -169,18 +169,19 @@ function Base.getindex(
     blocks(parent(a.array))[_getindices(index, _invperm(a.array))...], _perm(a.array)
   )
 end
-function SparseArraysBase.stored_indices(a::SparsePermutedDimsArrayBlocks)
-  return map(I -> _getindices(I, _perm(a.array)), stored_indices(blocks(parent(a.array))))
+function SparseArraysBase.eachstoredindex(a::SparsePermutedDimsArrayBlocks)
+  return map(I -> _getindices(I, _perm(a.array)), eachstoredindex(blocks(parent(a.array))))
 end
 # TODO: Either make this the generic interface or define
 # `SparseArraysBase.sparse_storage`, which is used
 # to defined this.
-function SparseArraysBase.stored_length(a::SparsePermutedDimsArrayBlocks)
-  return length(stored_indices(a))
-end
-function SparseArraysBase.sparse_storage(a::SparsePermutedDimsArrayBlocks)
-  return error("Not implemented")
+function SparseArraysBase.storedlength(a::SparsePermutedDimsArrayBlocks)
+  return length(eachstoredindex(a))
 end
+## TODO: Delete.
+## function SparseArraysBase.sparse_storage(a::SparsePermutedDimsArrayBlocks)
+##   return error("Not implemented")
+## end
 
 reverse_index(index) = reverse(index)
 reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index)))
@@ -240,25 +241,32 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe
   # TODO: Implement this properly.
   return true
 end
-function SparseArraysBase.stored_indices(a::SparseSubArrayBlocks)
-  return stored_indices(view(blocks(parent(a.array)), blockrange(a)...))
+function SparseArraysBase.eachstoredindex(a::SparseSubArrayBlocks)
+  return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...))
 end
 # TODO: Either make this the generic interface or define
 # `SparseArraysBase.sparse_storage`, which is used
 # to defined this.
-SparseArraysBase.stored_length(a::SparseSubArrayBlocks) = length(stored_indices(a))
+SparseArraysBase.storedlength(a::SparseSubArrayBlocks) = length(eachstoredindex(a))
 
 ## struct SparseSubArrayBlocksStorage{Array<:SparseSubArrayBlocks}
 ##   array::Array
 ## end
-function SparseArraysBase.sparse_storage(a::SparseSubArrayBlocks)
-  return map(I -> a[I], stored_indices(a))
-end
 
-function SparseArraysBase.getindex_zero_function(a::SparseSubArrayBlocks)
-  # TODO: Base it off of `getindex_zero_function(blocks(parent(a.array))`, but replace the
-  # axes with `axes(a.array)`.
-  return BlockZero(axes(a.array))
+## TODO: Delete.
+## function SparseArraysBase.sparse_storage(a::SparseSubArrayBlocks)
+##   return map(I -> a[I], eachstoredindex(a))
+## end
+
+## TODO: Delete.
+## function SparseArraysBase.getindex_zero_function(a::SparseSubArrayBlocks)
+##   # TODO: Base it off of `getindex_zero_function(blocks(parent(a.array))`, but replace the
+##   # axes with `axes(a.array)`.
+##   return BlockZero(axes(a.array))
+## end
+
+function SparseArraysBase.getunstoredindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N}
+  error("Not implemented.")
 end
 
 to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block)
@@ -271,4 +279,4 @@ function blocksparse_blocks(
 end
 
 using BlockArrays: BlocksView
-SparseArraysBase.stored_length(a::BlocksView) = length(a)
+SparseArraysBase.storedlength(a::BlocksView) = length(a)

src/blocksparsearrayinterface/blockzero.jl

@@ -18,28 +18,32 @@ struct BlockZero{Axes}
   axes::Axes
 end
 
-function (f::BlockZero)(a::AbstractArray, I)
-  return f(eltype(a), I)
+function (f::BlockZero)(a::AbstractArray, I...)
+  return f(eltype(a), I...)
 end
 
-function (f::BlockZero)(arraytype::Type{<:SubArray{<:Any,<:Any,P}}, I) where {P}
-  return f(P, I)
+function (f::BlockZero)(arraytype::Type{<:SubArray{<:Any,<:Any,P}}, I...) where {P}
+  return f(P, I...)
 end
 
-function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I)
+function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I::CartesianIndex)
+  return f(arraytype, Tuple(I)...)
+end
+
+function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I::Int...)
   # TODO: Make sure this works for sparse or block sparse blocks, immutable
   # blocks, diagonal blocks, etc.!
-  blck_size = block_size(f.axes, Block(Tuple(I)))
+  blck_size = block_size(f.axes, Block(I))
   blck_type = similartype(arraytype, blck_size)
   return fill!(blck_type(undef, blck_size), false)
 end
 
 # Fallback so that `SparseArray` with scalar elements works.
-function (f::BlockZero)(blocktype::Type{<:Number}, I)
+function (f::BlockZero)(blocktype::Type{<:Number}, I...)
   return zero(blocktype)
 end
 
 # Fallback to Array if it is abstract
-function (f::BlockZero)(arraytype::Type{AbstractArray{T,N}}, I) where {T,N}
-  return f(Array{T,N}, I)
+function (f::BlockZero)(arraytype::Type{AbstractArray{T,N}}, I...) where {T,N}
+  return f(Array{T,N}, I...)
 end

src/blocksparsearrayinterface/cat.jl

@@ -3,7 +3,9 @@ using SparseArraysBase: SparseArraysBase, allocate_cat_output, sparse_cat!
 
 # TODO: Maybe move to `SparseArraysBaseBlockArraysExt`.
 # TODO: Handle dual graded unit ranges, for example in a new `SparseArraysBaseGradedUnitRangesExt`.
-function SparseArraysBase.axis_cat(
+## TODO: Add this back.
+## function SparseArraysBase.axis_cat(
+function axis_cat(
   a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange
 )
   return blockedrange(vcat(blocklengths(a1), blocklengths(a2)))