Improve testing and organization

ITensor · mtfishman · Dec 2, 2023 · Nov 25, 2023 · Nov 26, 2023 · Nov 26, 2023
commit 32d375af44ac8ed20e330740f1e7b0b9793a8d36
diff --git a/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl
@@ -10,12 +10,14 @@ include("sparsearrayinterface/broadcast.jl")
 include("sparsearrayinterface/conversion.jl")
 include("sparsearrayinterface/wrappers.jl")
 include("sparsearrayinterface/zero.jl")
+include("sparsearrayinterface/SparseArrayInterfaceLinearAlgebraExt.jl")
 include("abstractsparsearray/abstractsparsearray.jl")
 include("abstractsparsearray/sparsearrayinterface.jl")
 include("abstractsparsearray/base.jl")
 include("abstractsparsearray/broadcast.jl")
 include("abstractsparsearray/map.jl")
 include("abstractsparsearray/baseinterface.jl")
 include("abstractsparsearray/convert.jl")
-include("abstractsparsearray/SparseArrayDOKsSparseArraysExt.jl")
+include("abstractsparsearray/SparseArrayInterfaceSparseArraysExt.jl")
+include("abstractsparsearray/SparseArrayInterfaceLinearAlgebraExt.jl")
 end
diff --git a/.../lib/SparseArrayInterface/src/abstractsparsearray/SparseArrayInterfaceLinearAlgebraExt.jl b/.../lib/SparseArrayInterface/src/abstractsparsearray/SparseArrayInterfaceLinearAlgebraExt.jl
@@ -0,0 +1,3 @@
+using LinearAlgebra: LinearAlgebra
+
+LinearAlgebra.norm(a::AbstractSparseArray, p::Real=2) = sparse_norm(a, p)
diff --git a/...searray/SparseArrayDOKsSparseArraysExt.jl → ...ay/SparseArrayInterfaceSparseArraysExt.jl b/...searray/SparseArrayDOKsSparseArraysExt.jl → ...ay/SparseArrayInterfaceSparseArraysExt.jl
diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl
@@ -3,30 +3,13 @@ using ..SparseArrayInterface: SparseArrayInterface
 
 SparseArrayInterface.sparse_storage(::AbstractSparseArray) = error("Not implemented")
 
-# TODO: Move to `SparseArrayInterface`?
-getindex_zero_function(::AbstractSparseArray) = error("Not implemented")
-
 function SparseArrayInterface.index_to_storage_index(
   a::AbstractSparseArray{<:Any,N}, I::CartesianIndex{N}
 ) where {N}
   !isassigned(SparseArrayInterface.sparse_storage(a), I) && return nothing
   return I
 end
 
-function SparseArrayInterface.getindex_notstored(
-  a::AbstractSparseArray{<:Any,N}, I::CartesianIndex{N}
-) where {N}
-  return getindex_zero_function(a)(a, I)
-end
-
-## # TODO: Generalize with `is_wrapped_array` trait.
-## function SparseArrayInterface.getindex_notstored(
-##   a::PermutedDimsArray{<:Any,N,<:AbstractSparseArray}, I::CartesianIndex{N}
-## ) where {N}
-##   # TODO: Need to permute `I`?
-##   return getindex_zero_function(a)(a, I)
-## end
-
 function SparseArrayInterface.setindex_notstored!(
   a::AbstractSparseArray{<:Any,N}, value, I::CartesianIndex{N}
 ) where {N}

diff --git a/...lib/SparseArrayInterface/src/sparsearrayinterface/SparseArrayInterfaceLinearAlgebraExt.jl b/...lib/SparseArrayInterface/src/sparsearrayinterface/SparseArrayInterfaceLinearAlgebraExt.jl
@@ -0,0 +1,3 @@
+using LinearAlgebra: norm
+
+sparse_norm(a::AbstractArray, p::Real=2) = norm(sparse_storage(a))
diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl
@@ -1,7 +1,9 @@
 # Optional interface.
 # Access a zero value.
+getindex_zero_function(::AbstractArray) = Zero()
+
 function getindex_notstored(a::AbstractArray, I)
-  return zero(eltype(a))
+  return getindex_zero_function(a)(a, I)
 end
 
 # Optional interface.

diff --git a/...s/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl b/...s/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl
@@ -0,0 +1,50 @@
+module AbstractSparseArrays
+using NDTensors.SparseArrayInterface: SparseArrayInterface, AbstractSparseArray
+
+struct SparseArray{T,N} <: AbstractSparseArray{T,N}
+  data::Vector{T}
+  dims::Tuple{Vararg{Int,N}}
+  index_to_dataindex::Dict{CartesianIndex{N},Int}
+  dataindex_to_index::Vector{CartesianIndex{N}}
+end
+function SparseArray{T,N}(dims::Tuple{Vararg{Int,N}}) where {T,N}
+  return SparseArray{T,N}(
+    T[], dims, Dict{CartesianIndex{N},Int}(), Vector{CartesianIndex{N}}()
+  )
+end
+SparseArray{T,N}(dims::Vararg{Int,N}) where {T,N} = SparseArray{T,N}(dims)
+SparseArray{T}(dims::Tuple{Vararg{Int}}) where {T} = SparseArray{T,length(dims)}(dims)
+SparseArray{T}(dims::Vararg{Int}) where {T} = SparseArray{T}(dims)
+
+# AbstractArray interface
+Base.size(a::SparseArray) = a.dims
+function Base.similar(a::SparseArray, elt::Type, dims::Tuple{Vararg{Int}})
+  return SparseArray{elt}(dims)
+end
+
+# Minimal interface
+SparseArrayInterface.sparse_storage(a::SparseArray) = a.data
+function SparseArrayInterface.index_to_storage_index(
+  a::SparseArray{<:Any,N}, I::CartesianIndex{N}
+) where {N}
+  return get(a.index_to_dataindex, I, nothing)
+end
+SparseArrayInterface.storage_index_to_index(a::SparseArray, I) = a.dataindex_to_index[I]
+function SparseArrayInterface.setindex_notstored!(
+  a::SparseArray{<:Any,N}, value, I::CartesianIndex{N}
+) where {N}
+  push!(a.data, value)
+  push!(a.dataindex_to_index, I)
+  a.index_to_dataindex[I] = length(a.data)
+  return a
+end
+
+# Empty the storage, helps with efficiency in `map!` to drop
+# zeros.
+function SparseArrayInterface.dropall!(a::SparseArray)
+  empty!(a.data)
+  empty!(a.index_to_dataindex)
+  empty!(a.dataindex_to_index)
+  return a
+end
+end
diff --git a/.../SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl b/.../SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl
@@ -1,4 +1,5 @@
 module SparseArrayInterfaceTestUtils
-include("SparseArrays.jl")
+include("AbstractSparseArrays.jl")
 include("DiagonalArrays.jl")
+include("SparseArrays.jl")
 end
diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl
@@ -18,15 +18,16 @@ SparseArray{T}(dims::Vararg{Int}) where {T} = SparseArray{T}(dims)
 
 # AbstractArray interface
 Base.size(a::SparseArray) = a.dims
+function Base.similar(a::SparseArray, elt::Type, dims::Tuple{Vararg{Int}})
+  return SparseArray{elt}(dims)
+end
+
 function Base.getindex(a::SparseArray, I...)
   return SparseArrayInterface.sparse_getindex(a, I...)
 end
 function Base.setindex!(a::SparseArray, I...)
   return SparseArrayInterface.sparse_setindex!(a, I...)
 end
-function Base.similar(a::SparseArray, elt::Type, dims::Tuple{Vararg{Int}})
-  return SparseArray{elt}(dims)
-end
 
 # Minimal interface
 SparseArrayInterface.sparse_storage(a::SparseArray) = a.data

diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl
@@ -1,5 +1,5 @@
 @eval module $(gensym())
-include("test_array.jl")
-include("test_diagonalarray.jl")
-include("test_sparsearray.jl")
+for filename in ["abstractsparsearray", "array", "diagonalarray"]
+  include("test_$filename.jl")
+end
 end
diff --git a/...seArrayInterface/test/test_sparsearray.jl → ...nterface/test/test_abstractsparsearray.jl b/...seArrayInterface/test/test_sparsearray.jl → ...nterface/test/test_abstractsparsearray.jl
@@ -2,9 +2,14 @@
 using LinearAlgebra: norm
 using NDTensors.SparseArrayInterface: SparseArrayInterface
 include("SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl")
-using .SparseArrayInterfaceTestUtils.SparseArrays: SparseArray
+using .SparseArrayInterfaceTestUtils.AbstractSparseArrays: AbstractSparseArrays
+using .SparseArrayInterfaceTestUtils.SparseArrays: SparseArrays
 using Test: @test, @testset
-@testset "SparseArray (eltype=$elt)" for elt in (Float32, ComplexF32, Float64, ComplexF64)
+@testset "AbstractSparseArray (arraytype=$SparseArray, eltype=$elt)" for SparseArray in (
+    AbstractSparseArrays.SparseArray, SparseArrays.SparseArray
+  ),
+  elt in (Float32, ComplexF32, Float64, ComplexF64)
+
   a = SparseArray{elt}(2, 3)
   @test size(a) == (2, 3)
   @test axes(a) == (1:2, 1:3)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		using LinearAlgebra: LinearAlgebra

		LinearAlgebra.norm(a::AbstractSparseArray, p::Real=2) = sparse_norm(a, p)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		using LinearAlgebra: norm

		sparse_norm(a::AbstractArray, p::Real=2) = norm(sparse_storage(a))