[BlockSparseArrays] Use new SparseArrayDOK type in BlockSparseArrays (#1272)

Author: mtfishman

src/DiagonalArrays.jl

@@ -1,11 +1,12 @@
 module DiagonalArrays
-include("defaults.jl")
-include("diaginterface.jl")
-include("diagindex.jl")
-include("diagindices.jl")
-include("diagonalarray.jl")
-include("sparsearrayinterface.jl")
-include("diagonalarraydiaginterface.jl")
-include("diagonalmatrix.jl")
-include("diagonalvector.jl")
+include("diaginterface/defaults.jl")
+include("diaginterface/diaginterface.jl")
+include("diaginterface/diagindex.jl")
+include("diaginterface/diagindices.jl")
+include("abstractdiagonalarray/abstractdiagonalarray.jl")
+include("abstractdiagonalarray/sparsearrayinterface.jl")
+include("abstractdiagonalarray/diagonalarraydiaginterface.jl")
+include("diagonalarray/diagonalarray.jl")
+include("diagonalarray/diagonalmatrix.jl")
+include("diagonalarray/diagonalvector.jl")
 end

src/abstractdiagonalarray/abstractdiagonalarray.jl

@@ -0,0 +1,3 @@
+using ..SparseArrayInterface: AbstractSparseArray
+
+abstract type AbstractDiagonalArray{T,N} <: AbstractSparseArray{T,N} end

src/diagonalarraydiaginterface.jl renamed to src/abstractdiagonalarray/diagonalarraydiaginterface.jl

@@ -2,22 +2,22 @@ using ..SparseArrayInterface: SparseArrayInterface, StorageIndex, StorageIndices
 
 SparseArrayInterface.StorageIndex(i::DiagIndex) = StorageIndex(index(i))
 
-function Base.getindex(a::DiagonalArray, i::DiagIndex)
+function Base.getindex(a::AbstractDiagonalArray, i::DiagIndex)
   return a[StorageIndex(i)]
 end
 
-function Base.setindex!(a::DiagonalArray, value, i::DiagIndex)
+function Base.setindex!(a::AbstractDiagonalArray, value, i::DiagIndex)
   a[StorageIndex(i)] = value
   return a
 end
 
 SparseArrayInterface.StorageIndices(i::DiagIndices) = StorageIndices(indices(i))
 
-function Base.getindex(a::DiagonalArray, i::DiagIndices)
+function Base.getindex(a::AbstractDiagonalArray, i::DiagIndices)
   return a[StorageIndices(i)]
 end
 
-function Base.setindex!(a::DiagonalArray, value, i::DiagIndices)
+function Base.setindex!(a::AbstractDiagonalArray, value, i::DiagIndices)
   a[StorageIndices(i)] = value
   return a
 end

src/abstractdiagonalarray/sparsearrayinterface.jl

@@ -0,0 +1,22 @@
+using Compat: Returns, allequal
+using ..SparseArrayInterface: SparseArrayInterface
+
+# `SparseArrayInterface` interface
+function SparseArrayInterface.index_to_storage_index(
+  a::AbstractDiagonalArray{<:Any,N}, I::CartesianIndex{N}
+) where {N}
+  !allequal(Tuple(I)) && return nothing
+  return first(Tuple(I))
+end
+
+function SparseArrayInterface.storage_index_to_index(a::AbstractDiagonalArray, I)
+  return CartesianIndex(ntuple(Returns(I), ndims(a)))
+end
+
+## # 1-dimensional case can be `AbstractDiagonalArray`.
+## function SparseArrayInterface.sparse_similar(
+##   a::AbstractDiagonalArray, elt::Type, dims::Tuple{Int}
+## )
+##   # TODO: Handle preserving zero element function.
+##   return similar(a, elt, dims)
+## end

src/defaults.jl renamed to src/diaginterface/defaults.jl

@@ -1,6 +1,8 @@
-using ..SparseArrayInterface: Zero
+using ..SparseArrayInterface: Zero, getindex_zero_function
+# TODO: Put into `DiagonalArraysSparseArrayDOKsExt`?
+using ..SparseArrayDOKs: SparseArrayDOKs, SparseArrayDOK
 
-struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractArray{T,N}
+struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractDiagonalArray{T,N}
   diag::Diag
   dims::NTuple{N,Int}
   zero::Zero
@@ -50,18 +52,48 @@ function DiagonalArray{<:Any,N}(diag::AbstractVector{T}) where {T,N}
 end
 
 # undef
-function DiagonalArray{T,N}(::UndefInitializer, d::Tuple{Vararg{Int,N}}) where {T,N}
-  return DiagonalArray{T,N}(Vector{T}(undef, minimum(d)), d)
+function DiagonalArray{T,N}(
+  ::UndefInitializer, d::Tuple{Vararg{Int,N}}, zero=Zero()
+) where {T,N}
+  return DiagonalArray{T,N}(Vector{T}(undef, minimum(d)), d, zero)
 end
 
 function DiagonalArray{T,N}(::UndefInitializer, d::Vararg{Int,N}) where {T,N}
   return DiagonalArray{T,N}(undef, d)
 end
 
-function DiagonalArray{T}(::UndefInitializer, d::Tuple{Vararg{Int,N}}) where {T,N}
-  return DiagonalArray{T,N}(undef, d)
+function DiagonalArray{T}(
+  ::UndefInitializer, d::Tuple{Vararg{Int,N}}, zero=Zero()
+) where {T,N}
+  return DiagonalArray{T,N}(undef, d, zero)
 end
 
 function DiagonalArray{T}(::UndefInitializer, d::Vararg{Int,N}) where {T,N}
   return DiagonalArray{T,N}(undef, d)
 end
+
+# Minimal `AbstractArray` interface
+Base.size(a::DiagonalArray) = a.dims
+
+function Base.similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}})
+  # TODO: Preserve zero element function.
+  return DiagonalArray{elt}(undef, dims, getindex_zero_function(a))
+end
+
+# Minimal `SparseArrayInterface` interface
+SparseArrayInterface.sparse_storage(a::DiagonalArray) = a.diag
+
+# `SparseArrayInterface`
+# Defines similar when the output can't be `DiagonalArray`,
+# such as in `reshape`.
+# TODO: Put into `DiagonalArraysSparseArrayDOKsExt`?
+# TODO: Special case 2D to output `SparseMatrixCSC`?
+function SparseArrayInterface.sparse_similar(
+  a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}
+)
+  return SparseArrayDOK{elt}(undef, dims, getindex_zero_function(a))
+end
+
+function SparseArrayInterface.getindex_zero_function(a::DiagonalArray)
+  return a.zero
+end