[BlockSparseArrays] Use new SparseArrayDOK type in BlockSparseArrays

NDTensors/src/arraystorage/blocksparsearray/storage/contract.jl

@@ -1,6 +1,5 @@
-# TODO: Change to:
-# using .SparseArrayDOKs: SparseArrayDOK
-using .BlockSparseArrays: SparseArray
+# TODO: Define in `SparseArrayInterface`.
+using ..SparseArrayDOKs: SparseArrayDOK
 
 # TODO: This is inefficient, need to optimize.
 # Look at `contract_labels`, `contract_blocks` and `maybe_contract_blocks!` in:
@@ -39,11 +38,11 @@ function default_contract_muladd(a1, labels1, a2, labels2, a_dest, labels_dest)
 end
 
 function contract!(
-  a_dest::SparseArray,
+  a_dest::SparseArrayDOK,
   labels_dest,
-  a1::SparseArray,
+  a1::SparseArrayDOK,
   labels1,
-  a2::SparseArray,
+  a2::SparseArrayDOK,
   labels2;
   muladd=default_contract_muladd,
 )

NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl

@@ -23,16 +23,17 @@ using Dictionaries: Dictionary, Indices, getindices, set! # TODO: Move to `Spars
 using LinearAlgebra: Hermitian
 using SplitApplyCombine: groupcount
 
-export BlockSparseArray, SparseArray
+export BlockSparseArray # , SparseArray
 
+include("defaults.jl")
 include("tensor_product.jl")
 include("base.jl")
 include("axes.jl")
 include("abstractarray.jl")
 include("permuteddimsarray.jl")
 include("blockarrays.jl")
 # TODO: Split off into `SparseArraysExtensions` module, rename to `SparseArrayDOK`.
-include("sparsearray.jl")
+# include("sparsearray.jl")
 include("blocksparsearray.jl")
 include("allocate_output.jl")
 include("subarray.jl")

NDTensors/src/lib/BlockSparseArrays/src/LinearAlgebraExt/LinearAlgebraExt.jl

@@ -1,7 +1,7 @@
 module LinearAlgebraExt
 using ...AlgorithmSelection: Algorithm, @Algorithm_str
 using BlockArrays: BlockArrays, blockedrange, blocks
-using ..BlockSparseArrays: SparseArray, nonzero_keys # TODO: Move to `SparseArraysExtensions` module, rename `SparseArrayDOK`.
+using ..BlockSparseArrays: nonzero_keys # TODO: Move to `SparseArraysExtensions` module, rename `SparseArrayDOK`.
 using ..BlockSparseArrays: BlockSparseArrays, BlockSparseArray, nonzero_blockkeys
 using LinearAlgebra: LinearAlgebra, Hermitian, Transpose, I, eigen, qr
 using SparseArrays: SparseArrays, SparseMatrixCSC, spzeros, sparse

NDTensors/src/lib/BlockSparseArrays/src/LinearAlgebraExt/qr.jl

@@ -1,3 +1,5 @@
+using ...SparseArrayDOKs: SparseArrayDOK
+
 # Check if the matrix has 1 or fewer entries
 # per row/column.
 function is_permutation_matrix(a::SparseMatrixCSC)
@@ -6,7 +8,7 @@ end
 
 # Check if the matrix has 1 or fewer entries
 # per row/column.
-function is_permutation_matrix(a::SparseArray{<:Any,2})
+function is_permutation_matrix(a::SparseArrayDOK{<:Any,2})
   keys = collect(Iterators.map(Tuple, nonzero_keys(a)))
   I = first.(keys)
   J = last.(keys)
@@ -17,7 +19,8 @@ function findnonzerorows(a::SparseMatrixCSC, col)
   return view(a.rowval, a.colptr[col]:(a.colptr[col + 1] - 1))
 end
 
-function SparseArrays.SparseMatrixCSC(a::SparseArray{<:Any,2})
+# TODO: Is this already defined?
+function SparseArrays.SparseMatrixCSC(a::SparseArrayDOK{<:Any,2})
   # Not defined:
   # a_csc = SparseMatrixCSC{eltype(a)}(size(a))
   a_csc = spzeros(eltype(a), size(a))
@@ -27,8 +30,9 @@ function SparseArrays.SparseMatrixCSC(a::SparseArray{<:Any,2})
   return a_csc
 end
 
+# TODO: Is this already defined?
 # Get the sparse structure of a SparseArray as a SparseMatrixCSC.
-function sparse_structure(structure_type::Type{<:SparseMatrixCSC}, a::SparseArray{<:Any,2})
+function sparse_structure(structure_type::Type{<:SparseMatrixCSC}, a::SparseArrayDOK{<:Any,2})
   # Idealy would work but a bit too complicated for `map` right now:
   # return SparseMatrixCSC(map(x -> iszero(x) ? false : true, a))
   # TODO: Change to `spzeros(Bool, size(a))`.

NDTensors/src/lib/BlockSparseArrays/src/allocate_output.jl

@@ -89,6 +89,12 @@ end
 # SparseArray
 #############################################################################
 
+using ..SparseArrayDOKs: SparseArrayDOK
+
+const SparseArrayLike{T,N,Zero} = Union{
+  SparseArrayDOK{T,N,Zero},PermutedDimsArray{T,N,<:Any,<:Any,SparseArrayDOK{T,N,Zero}}
+}
+
 # TODO: Maybe store nonzero locations?
 # TODO: Make this backwards compatible.
 # TODO: Add a default for `eltype` and `zero`.
@@ -100,7 +106,7 @@ Base.@kwdef struct SparseArrayStructure{ElType,Axes,Zero} <:
   zero::Zero
 end
 
-function allocate(arraytype::Type{<:SparseArray}, structure::SparseArrayStructure)
+function allocate(arraytype::Type{<:SparseArrayDOK}, structure::SparseArrayStructure)
   # TODO: Use `set_eltype`.
   return arraytype{structure.eltype}(structure.axes, structure.zero)
 end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray.jl

@@ -1,8 +1,9 @@
 using BlockArrays: block
+## using ..SparseArrayDOKs: SparseArrayDOK
 
 # Also add a version with contiguous underlying data.
 struct BlockSparseArray{
-  T,N,A<:AbstractArray{T,N},Blocks<:SparseArray{A,N},Axes<:NTuple{N,AbstractUnitRange{Int}}
+  T,N,A<:AbstractArray{T,N},Blocks<:AbstractArray{A,N},Axes<:NTuple{N,AbstractUnitRange{Int}}
 } <: AbstractBlockArray{T,N}
   blocks::Blocks
   axes::Axes
@@ -26,7 +27,7 @@ function Base.reshape(a::BlockSparseArray, ax::Tuple{Vararg{AbstractUnitRange}})
   ## blocks_reshaped = reshape(blocks(a), blocklength.(ax))
   blocktype_reshaped = set_ndims(blocktype(a), length(ax))
   # TODO: Some other way of getting `zero` function?
-  blocks_reshaped = SparseArray(
+  blocks_reshaped = default_sparsearray_type()(
     Dictionary{CartesianIndex{length(ax)},blocktype_reshaped}(),
     blocklength.(ax),
     BlockZero(ax),
@@ -113,7 +114,7 @@ end
 ##   axes::Axes
 ## end
 
-function BlockSparseArray(a::SparseArray, axes::Tuple{Vararg{AbstractUnitRange}})
+function BlockSparseArray(a::AbstractArray, axes::Tuple{Vararg{AbstractUnitRange}})
   A = eltype(a)
   T = eltype(A)
   N = ndims(a)
@@ -132,7 +133,9 @@ function BlockSparseArray(
     map(block -> CartesianIndex(inttuple(block)), blocks)
   end
   cartesiandata = Dictionary(cartesianblocks, blockdata)
-  block_storage = SparseArray(cartesiandata, blocklength.(axes), BlockZero(axes))
+  block_storage = default_sparsearray_type()(
+    cartesiandata, blocklength.(axes), BlockZero(axes)
+  )
   return BlockSparseArray(block_storage, axes)
 end
 

NDTensors/src/lib/BlockSparseArrays/src/defaults.jl

@@ -0,0 +1,3 @@
+using ..SparseArrayDOKs: SparseArrayDOK
+
+default_sparsearray_type() = SparseArrayDOK

NDTensors/src/lib/BlockSparseArrays/test/Project.toml

@@ -1,3 +1,4 @@
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"