[WIP] Rewrite for latest Derive.jl and SparseArraysBase.jl

mtfishman · mtfishman · commit 25ae13732c8d · 2024-12-08T11:29:49.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BroadcastMapConversion = "4a4adec5-520f-4750-bb37-d5e66b4ddeb2"
+Derive = "a07dfc7f-7d04-4eb5-84cc-a97f051f655a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
@@ -21,19 +22,17 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [sources]
-BroadcastMapConversion = {url = "https://github.com/ITensor/BroadcastMapConversion.jl"}
 GradedUnitRanges = {url = "https://github.com/ITensor/GradedUnitRanges.jl"}
 LabelledNumbers = {url = "https://github.com/ITensor/LabelledNumbers.jl"}
 NestedPermutedDimsArrays = {url = "https://github.com/ITensor/NestedPermutedDimsArrays.jl"}
-SparseArraysBase = {url = "https://github.com/ITensor/SparseArraysBase.jl"}
 TensorAlgebra = {url = "https://github.com/ITensor/TensorAlgebra.jl"}
-TypeParameterAccessors = {url = "https://github.com/ITensor/TypeParameterAccessors.jl"}
 
 [compat]
 Adapt = "4.1.1"
 Aqua = "0.8.9"
 ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
+Derive = "0.3.0"
 Dictionaries = "0.4.3"
 GPUArraysCore = "0.1.0"
 LinearAlgebra = "1.10"
diff --git a/TODO.md b/TODO.md
@@ -1,3 +1,9 @@
+- Rename `block_storedlength` to `blockstoredlength`.
+- Rename `BlockZero` to `getunstoredblock`.
+- Use `Derive.@array_aliases`.
+- Use `@derive AbstractArrayInterface` to define `abstractblocksparsearray/map.jl` functions.
+
+
 - Add Aqua tests.
 - Turn the package extensions into actual package extensions:
   - BlockSparseArraysAdaptExt
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -12,10 +12,7 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [sources]
-BroadcastMapConversion = {url = "https://github.com/ITensor/BroadcastMapConversion.jl"}
 GradedUnitRanges = {url = "https://github.com/ITensor/GradedUnitRanges.jl"}
 LabelledNumbers = {url = "https://github.com/ITensor/LabelledNumbers.jl"}
 NestedPermutedDimsArrays = {url = "https://github.com/ITensor/NestedPermutedDimsArrays.jl"}
-SparseArraysBase = {url = "https://github.com/ITensor/SparseArraysBase.jl"}
 TensorAlgebra = {url = "https://github.com/ITensor/TensorAlgebra.jl"}
-TypeParameterAccessors = {url = "https://github.com/ITensor/TypeParameterAccessors.jl"}
diff --git a/examples/Project.toml b/examples/Project.toml
@@ -13,10 +13,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [sources]
-BroadcastMapConversion = {url = "https://github.com/ITensor/BroadcastMapConversion.jl"}
 GradedUnitRanges = {url = "https://github.com/ITensor/GradedUnitRanges.jl"}
 LabelledNumbers = {url = "https://github.com/ITensor/LabelledNumbers.jl"}
 NestedPermutedDimsArrays = {url = "https://github.com/ITensor/NestedPermutedDimsArrays.jl"}
-SparseArraysBase = {url = "https://github.com/ITensor/SparseArraysBase.jl"}
 TensorAlgebra = {url = "https://github.com/ITensor/TensorAlgebra.jl"}
-TypeParameterAccessors = {url = "https://github.com/ITensor/TypeParameterAccessors.jl"}
diff --git a/examples/README.jl b/examples/README.jl
@@ -15,18 +15,12 @@
 ```julia
 julia> using Pkg: Pkg
 
-julia> Pkg.add(url="https://github.com/ITensor/BroadcastMapConversion.jl")
-
 julia> Pkg.add(url="https://github.com/ITensor/NestedPermutedDimsArrays.jl")
 
-julia> Pkg.add(url="https://github.com/ITensor/TypeParameterAccessors.jl")
-
 julia> Pkg.add(url="https://github.com/ITensor/LabelledNumbers.jl")
 
 julia> Pkg.add(url="https://github.com/ITensor/GradedUnitRanges.jl")
 
-julia> Pkg.add(url="https://github.com/ITensor/SparseArraysBase.jl")
-
 julia> Pkg.add(url="https://github.com/ITensor/TensorAlgebra.jl")
 
 julia> Pkg.add(url="https://github.com/ITensor/BlockSparseArrays.jl")
@@ -36,7 +30,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 +57,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 +87,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)
diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl 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.
@@ -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])
+  if I ∈ eachstoredindex(blocks(a.array))
+    return storedlength(blocks(a.array)[I])
   end
   return 0
 end
diff --git a/src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl b/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
diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl
@@ -1,5 +1,6 @@
 module BlockSparseArrays
 include("BlockArraysExtensions/BlockArraysExtensions.jl")
+include("blocksparsearrayinterface/abstractblocksparsearrayinterface.jl")
 include("blocksparsearrayinterface/blocksparsearrayinterface.jl")
 include("blocksparsearrayinterface/linearalgebra.jl")
 include("blocksparsearrayinterface/blockzero.jl")
@@ -16,7 +17,7 @@ include("abstractblocksparsearray/views.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/sparsearrayinterface.jl")
 include("abstractblocksparsearray/broadcast.jl")
-include("abstractblocksparsearray/map.jl")
+## include("abstractblocksparsearray/map.jl")
 include("abstractblocksparsearray/linearalgebra.jl")
 include("abstractblocksparsearray/cat.jl")
 include("blocksparsearray/defaults.jl")
diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl
@@ -8,6 +8,31 @@ function nonzero_keys end
 
 abstract type AbstractBlockSparseArray{T,N} <: AbstractBlockArray{T,N} end
 
+using Derive: @array_aliases
+# Define AbstractSparseVector, AnyAbstractSparseArray, etc.
+@array_aliases AbstractBlockSparseArray
+
+using Derive: Derive
+function Derive.interface(::Type{<:AbstractBlockSparseArray})
+  return BlockSparseArrayInterface()
+end
+
+using Derive: @derive
+
+# TODO: These need to be loaded since `AbstractArrayOps`
+# includes overloads of functions from these modules.
+# Ideally that wouldn't be needed and can be circumvented
+# with `GlobalRef`.
+using ArrayLayouts: ArrayLayouts
+using LinearAlgebra: LinearAlgebra
+
+# Derive `Base.getindex`, `Base.setindex!`, etc.
+# TODO: Define `AbstractMatrixOps` and overload for
+# `AnyAbstractSparseMatrix` and `AnyAbstractSparseVector`,
+# which is where matrix multiplication and factorizations
+# shoudl go.
+@derive AnyAbstractBlockSparseArray AbstractArrayOps
+
 ## Base `AbstractArray` interface
 
 Base.axes(::AbstractBlockSparseArray) = error("Not implemented")
diff --git a/src/abstractblocksparsearray/arraylayouts.jl b/src/abstractblocksparsearray/arraylayouts.jl
@@ -3,11 +3,12 @@ using BlockArrays: BlockLayout
 using SparseArraysBase: SparseLayout
 using TypeParameterAccessors: parenttype, similartype
 
-function ArrayLayouts.MemoryLayout(arraytype::Type{<:AnyAbstractBlockSparseArray})
-  outer_layout = typeof(MemoryLayout(blockstype(arraytype)))
-  inner_layout = typeof(MemoryLayout(blocktype(arraytype)))
-  return BlockLayout{outer_layout,inner_layout}()
-end
+## TODO: Bring this back.
+## function ArrayLayouts.MemoryLayout(arraytype::Type{<:AnyAbstractBlockSparseArray})
+##   outer_layout = typeof(MemoryLayout(blockstype(arraytype)))
+##   inner_layout = typeof(MemoryLayout(blocktype(arraytype)))
+##   return BlockLayout{outer_layout,inner_layout}()
+## end
 
 # TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`.
 function ArrayLayouts.MemoryLayout(
diff --git a/src/abstractblocksparsearray/map.jl b/src/abstractblocksparsearray/map.jl
@@ -57,13 +57,14 @@ function reblock(
   return @view parent(a)[map(I -> Vector(I.blocks), parentindices(a))...]
 end
 
+# TODO: Define as `@interface BlockSparseArrayInterface Base.map!(...)`.
 # TODO: Rewrite this so that it takes the blocking structure
 # made by combining the blocking of the axes (i.e. the blocking that
 # 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!(
-  ::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}
+function blocksparse_map!(
+  f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}
 )
   a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs)
   for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...)
@@ -92,59 +93,71 @@ end
 # function SparseArraysBase.sparse_mapreduce(::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray})
 # end
 
+# TODO: @derive AbstractArrayInterface
 function Base.map!(f, a_dest::AbstractArray, a_srcs::Vararg{AnyAbstractBlockSparseArray})
   sparse_map!(f, a_dest, a_srcs...)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.map(f, as::Vararg{AnyAbstractBlockSparseArray})
   return f.(as...)
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.copy!(a_dest::AbstractArray, a_src::AnyAbstractBlockSparseArray)
   sparse_copy!(a_dest, a_src)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.copyto!(a_dest::AbstractArray, a_src::AnyAbstractBlockSparseArray)
   sparse_copyto!(a_dest, a_src)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 # Fix ambiguity error
 function Base.copyto!(a_dest::LayoutArray, a_src::AnyAbstractBlockSparseArray)
   sparse_copyto!(a_dest, a_src)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.copyto!(
   a_dest::AbstractMatrix, a_src::Transpose{T,<:AbstractBlockSparseMatrix{T}}
 ) where {T}
   sparse_copyto!(a_dest, a_src)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.copyto!(
   a_dest::AbstractMatrix, a_src::Adjoint{T,<:AbstractBlockSparseMatrix{T}}
 ) where {T}
   sparse_copyto!(a_dest, a_src)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
+# TODO: Define as `a_dest .= PermutedDimsArray(a_src, perm)`
 function Base.permutedims!(a_dest, a_src::AnyAbstractBlockSparseArray, perm)
   sparse_permutedims!(a_dest, a_src, perm)
   return a_dest
 end
 
+# TODO: @derive AbstractArrayInterface
 function Base.mapreduce(f, op, as::Vararg{AnyAbstractBlockSparseArray}; kwargs...)
   return sparse_mapreduce(f, op, as...; kwargs...)
 end
 
+# TODO: @derive AbstractArrayInterface
 # TODO: Why isn't this calling `mapreduce` already?
 function Base.iszero(a::AnyAbstractBlockSparseArray)
   return sparse_iszero(blocks(a))
 end
 
+# TODO: @derive AbstractArrayInterface
 # TODO: Why isn't this calling `mapreduce` already?
 function Base.isreal(a::AnyAbstractBlockSparseArray)
   return sparse_isreal(blocks(a))
diff --git a/src/abstractblocksparsearray/sparsearrayinterface.jl b/src/abstractblocksparsearray/sparsearrayinterface.jl
@@ -1,5 +1,5 @@
 using BlockArrays: Block
-using SparseArraysBase: SparseArraysBase, sparse_storage, stored_indices
+using SparseArraysBase: SparseArraysBase, eachstoredindex, storedvalues
 
 # Structure storing the block sparse storage
 struct BlockSparseStorage{Arr<:AbstractBlockSparseArray}
@@ -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: Use `SparseArraysBase.getstoredindex`, `storedvalues`, 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, storedvalues(blocks(a)); init=zero(Int))
 end
diff --git a/src/abstractblocksparsearray/views.jl b/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 `eachstoredblock`.
+  if I ∈ eachstoredindex(blocks(a))
     return blocks(a)[I]
   end
   return BlockView(a, block)
diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl
@@ -11,14 +11,14 @@ using BlockArrays:
 using SplitApplyCombine: groupcount
 using TypeParameterAccessors: similartype
 
-const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
-  T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N}
-}
-
-# TODO: Rename `AnyBlockSparseArray`.
-const AnyAbstractBlockSparseArray{T,N} = Union{
-  <:AbstractBlockSparseArray{T,N},<:WrappedAbstractBlockSparseArray{T,N}
-}
+## const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
+##   T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N}
+## }
+
+## # TODO: Rename `AnyBlockSparseArray`.
+## const AnyAbstractBlockSparseArray{T,N} = Union{
+##   <:AbstractBlockSparseArray{T,N},<:WrappedAbstractBlockSparseArray{T,N}
+## }
 
 # a[1:2, 1:2]
 function Base.to_indices(
diff --git a/src/blocksparsearrayinterface/abstractblocksparsearrayinterface.jl b/src/blocksparsearrayinterface/abstractblocksparsearrayinterface.jl
@@ -0,0 +1,15 @@
+using Derive: Derive, @interface
+using SparseArraysBase: AbstractSparseArrayInterface
+
+using SparseArraysBase: AbstractSparseArrayStyle
+
+abstract type AbstractBlockSparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
+
+struct BlockSparseArrayStyle{N} <: AbstractBlockSparseArrayStyle{N} end
+
+# TODO: Add `ndims` type parameter.
+# TODO: This isn't used to define interface functions right now.
+# Currently, `@interface` expects an instance, probably it should take a
+# type instead so fallback functions can use abstract types.
+abstract type AbstractBlockSparseArrayInterface <: AbstractSparseArrayInterface end
+
diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
diff --git a/src/blocksparsearrayinterface/blockzero.jl b/src/blocksparsearrayinterface/blockzero.jl
diff --git a/src/blocksparsearrayinterface/broadcast.jl b/src/blocksparsearrayinterface/broadcast.jl
diff --git a/src/blocksparsearrayinterface/cat.jl b/src/blocksparsearrayinterface/cat.jl
diff --git a/test/Project.toml b/test/Project.toml
