New design for handling unstored values (#65)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseArraysBase"
 uuid = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.6.0"
+version = "0.7.0"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

docs/Project.toml

@@ -8,4 +8,4 @@ SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 Dictionaries = "0.4.4"
 Documenter = "1.8.1"
 Literate = "2.20.1"
-SparseArraysBase = "0.6.0"
+SparseArraysBase = "0.7.0"

examples/Project.toml

@@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Dictionaries = "0.4.4"
-SparseArraysBase = "0.6.0"
+SparseArraysBase = "0.7.0"
 Test = "<0.0.1, 1"

src/abstractsparsearray.jl

@@ -11,6 +11,54 @@ function DerivableInterfaces.interface(::Type{<:AbstractSparseArray})
   return SparseArrayInterface()
 end
 
+function Base.copy(a::AnyAbstractSparseArray)
+  return copyto!(similar(a), a)
+end
+
+function similar_sparsearray(a::AnyAbstractSparseArray, unstored::Unstored)
+  return SparseArrayDOK(unstored)
+end
+function similar_sparsearray(a::AnyAbstractSparseArray, T::Type, ax::Tuple)
+  return similar_sparsearray(a, Unstored(unstoredsimilar(unstored(a), T, ax)))
+end
+function similar_sparsearray(a::AnyAbstractSparseArray, T::Type)
+  return similar_sparsearray(a, Unstored(unstoredsimilar(unstored(a), T)))
+end
+function similar_sparsearray(a::AnyAbstractSparseArray, ax::Tuple)
+  return similar_sparsearray(a, Unstored(unstoredsimilar(unstored(a), ax)))
+end
+function similar_sparsearray(a::AnyAbstractSparseArray)
+  return similar_sparsearray(a, Unstored(unstored(a)))
+end
+
+function Base.similar(a::AnyAbstractSparseArray, unstored::Unstored)
+  return similar_sparsearray(a, unstored)
+end
+function Base.similar(a::AnyAbstractSparseArray)
+  return similar_sparsearray(a)
+end
+function Base.similar(a::AnyAbstractSparseArray, T::Type)
+  return similar_sparsearray(a, T)
+end
+function Base.similar(a::AnyAbstractSparseArray, ax::Tuple)
+  return similar_sparsearray(a, ax)
+end
+function Base.similar(a::AnyAbstractSparseArray, T::Type, ax::Tuple)
+  return similar_sparsearray(a, T, ax)
+end
+# Fix ambiguity error.
+function Base.similar(a::AnyAbstractSparseArray, T::Type, ax::Tuple{Int,Vararg{Int}})
+  return similar_sparsearray(a, T, ax)
+end
+# Fix ambiguity error.
+function Base.similar(
+  a::AnyAbstractSparseArray,
+  T::Type,
+  ax::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}},
+)
+  return similar_sparsearray(a, T, ax)
+end
+
 using DerivableInterfaces: @derive
 
 # TODO: These need to be loaded since `AbstractArrayOps`
@@ -20,12 +68,30 @@ using DerivableInterfaces: @derive
 using ArrayLayouts: ArrayLayouts
 using LinearAlgebra: LinearAlgebra
 
-# DerivableInterfaces `Base.getindex`, `Base.setindex!`, etc.
-# TODO: Define `AbstractMatrixOps` and overload for
-# `AnyAbstractSparseMatrix` and `AnyAbstractSparseVector`,
-# which is where matrix multiplication and factorizations
-# should go.
-@derive AnyAbstractSparseArray AbstractArrayOps
+@derive (T=AnyAbstractSparseArray,) begin
+  Base.getindex(::T, ::Any...)
+  Base.getindex(::T, ::Int...)
+  Base.setindex!(::T, ::Any, ::Any...)
+  Base.setindex!(::T, ::Any, ::Int...)
+  Base.copy!(::AbstractArray, ::T)
+  Base.copyto!(::AbstractArray, ::T)
+  Base.map(::Any, ::T...)
+  Base.map!(::Any, ::AbstractArray, ::T...)
+  Base.mapreduce(::Any, ::Any, ::T...; kwargs...)
+  Base.reduce(::Any, ::T...; kwargs...)
+  Base.all(::Function, ::T)
+  Base.all(::T)
+  Base.iszero(::T)
+  Base.real(::T)
+  Base.fill!(::T, ::Any)
+  DerivableInterfaces.zero!(::T)
+  Base.zero(::T)
+  Base.permutedims!(::Any, ::T, ::Any)
+  Broadcast.BroadcastStyle(::Type{<:T})
+  Base.copyto!(::T, ::Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{0}})
+  ArrayLayouts.MemoryLayout(::Type{<:T})
+  LinearAlgebra.mul!(::AbstractMatrix, ::T, ::T, ::Number, ::Number)
+end
 
 using DerivableInterfaces.Concatenate: concatenate
 # We overload `Base._cat` instead of `Base.cat` since it
@@ -35,7 +101,12 @@ function Base._cat(dims, a::AnyAbstractSparseArray...)
   return concatenate(dims, a...)
 end
 
+# TODO: Use `map(WeakPreserving(f), a)` instead.
+# Currently that has trouble with type unstable maps, since
+# the element type becomes abstract and therefore the zero/unstored
+# values are not well defined.
 function map_stored(f, a::AnyAbstractSparseArray)
+  iszero(storedlength(a)) && return a
   kvs = storedpairs(a)
   # `collect` to convert to `Vector`, since otherwise
   # if it stays as `Dictionary` we might hit issues like
@@ -52,6 +123,10 @@ end
 
 using Adapt: adapt
 function Base.print_array(io::IO, a::AnyAbstractSparseArray)
+  # TODO: Use `map(WeakPreserving(adapt(Array)), a)` instead.
+  # Currently that has trouble with type unstable maps, since
+  # the element type becomes abstract and therefore the zero/unstored
+  # values are not well defined.
   a′ = map_stored(adapt(Array), a)
   return @invoke Base.print_array(io::typeof(io), a′::AbstractArray{<:Any,ndims(a)})
 end
@@ -75,27 +150,30 @@ from the input indices.
 This constructor does not take ownership of the supplied storage, and will result in an
 independent container.
 """
-sparse(::Union{AbstractDict,AbstractDictionary}, dims...; kwargs...)
+sparse(::Union{AbstractDict,AbstractDictionary}, dims...)
 
 const AbstractDictOrDictionary = Union{AbstractDict,AbstractDictionary}
 # checked constructor from data: use `setindex!` to validate/convert input
-function sparse(storage::AbstractDictOrDictionary, dims::Dims; kwargs...)
-  A = SparseArrayDOK{valtype(storage)}(undef, dims; kwargs...)
+function sparse(storage::AbstractDictOrDictionary, unstored::AbstractArray)
+  A = SparseArrayDOK(Unstored(unstored))
   for (i, v) in pairs(storage)
     A[i] = v
   end
   return A
 end
-function sparse(storage::AbstractDictOrDictionary, dims::Int...; kwargs...)
-  return sparse(storage, dims; kwargs...)
+function sparse(storage::AbstractDictOrDictionary, ax::Tuple)
+  return sparse(storage, Zeros{valtype(storage)}(ax))
+end
+function sparse(storage::AbstractDictOrDictionary, dims::Int...)
+  return sparse(storage, dims)
 end
 # Determine the size automatically.
-function sparse(storage::AbstractDictOrDictionary; kwargs...)
+function sparse(storage::AbstractDictOrDictionary)
   dims = ntuple(Returns(0), length(keytype(storage)))
   for I in keys(storage)
     dims = map(max, dims, Tuple(I))
   end
-  return sparse(storage, dims; kwargs...)
+  return sparse(storage, dims)
 end
 
 using Random: Random, AbstractRNG, default_rng
@@ -107,12 +185,18 @@ Create an empty size `dims` sparse array.
 The optional `T` argument specifies the element type, which defaults to `Float64`.
 """ sparsezeros
 
-function sparsezeros(::Type{T}, dims::Dims; kwargs...) where {T}
-  return SparseArrayDOK{T}(undef, dims; kwargs...)
+function sparsezeros(::Type{T}, unstored::AbstractArray{<:Any,N}) where {T,N}
+  return SparseArrayDOK{T,N}(Unstored(unstored))
+end
+function sparsezeros(unstored::AbstractArray{T,N}) where {T,N}
+  return SparseArrayDOK{T,N}(Unstored(unstored))
+end
+function sparsezeros(::Type{T}, dims::Dims) where {T}
+  return sparsezeros(T, Zeros{T}(dims))
 end
-sparsezeros(::Type{T}, dims::Int...; kwargs...) where {T} = sparsezeros(T, dims; kwargs...)
-sparsezeros(dims::Dims; kwargs...) = sparsezeros(Float64, dims; kwargs...)
-sparsezeros(dims::Int...; kwargs...) = sparsezeros(Float64, dims; kwargs...)
+sparsezeros(::Type{T}, dims::Int...) where {T} = sparsezeros(T, dims)
+sparsezeros(dims::Dims) = sparsezeros(Float64, dims)
+sparsezeros(dims::Int...) = sparsezeros(Float64, dims)
 
 @doc """
     sparserand([rng], [T::Type], dims; density::Real=0.5, randfun::Function=rand) -> A::SparseArrayDOK{T}

src/abstractsparsearrayinterface.jl

@@ -1,7 +1,9 @@
 using Base: @_propagate_inbounds_meta
 using DerivableInterfaces:
   DerivableInterfaces, @derive, @interface, AbstractArrayInterface, zero!
+using FillArrays: Zeros
 
+function unstored end
 function eachstoredindex end
 function getstoredindex end
 function getunstoredindex end
@@ -12,9 +14,25 @@ function storedlength end
 function storedpairs end
 function storedvalues end
 
-# Replace the function for accessing
-# unstored values.
-function set_getunstoredindex end
+# Indicates that the array should be interpreted
+# as the unstored values of a sparse array.
+struct Unstored{T,N,P<:AbstractArray{T,N}} <: AbstractArray{T,N}
+  parent::P
+end
+Base.parent(a::Unstored) = a.parent
+
+unstored(a::AbstractArray) = Zeros{eltype(a)}(axes(a))
+
+function unstoredsimilar(a::AbstractArray, T::Type, ax::Tuple)
+  return Zeros{T}(ax)
+end
+function unstoredsimilar(a::AbstractArray, ax::Tuple)
+  return unstoredsimilar(a, eltype(a), ax)
+end
+function unstoredsimilar(a::AbstractArray, T::Type)
+  return AbstractArray{T}(a)
+end
+unstoredsimilar(a::AbstractArray) = a
 
 # Generic functionality for converting to a
 # dense array, trying to preserve information
@@ -84,14 +102,6 @@ Base.size(a::StoredValues) = size(a.storedindices)
   return setindex!(a.array, value, a.storedindices[I])
 end
 
-# TODO: This may need to be defined in `sparsearraydok.jl`, after `SparseArrayDOK`
-# is defined. And/or define `default_type(::SparseArrayStyle, T::Type) = SparseArrayDOK{T}`.
-@interface ::AbstractSparseArrayInterface function Base.similar(
-  a::AbstractArray, T::Type, size::Tuple{Vararg{Int}}
-)
-  # TODO: Define `default_similartype` or something like that?
-  return SparseArrayDOK{T}(undef, size)
-end
 using DerivableInterfaces: DerivableInterfaces, zero!
 
 # `zero!` isn't defined in `Base`, but it is defined in `ArrayLayouts`
@@ -136,7 +146,10 @@ end
 
 abstract type AbstractSparseArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
 
-@derive AbstractSparseArrayStyle AbstractArrayStyleOps
+@derive (T=AbstractSparseArrayStyle,) begin
+  Base.similar(::Broadcast.Broadcasted{<:T}, ::Type, ::Tuple)
+  Base.copyto!(::AbstractArray, ::Broadcast.Broadcasted{<:T})
+end
 
 struct SparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
 

src/map.jl

@@ -48,20 +48,20 @@ function ZeroPreserving(f, T::Type, Ts::Type...)
     return NonPreserving(f)
   end
 end
+ZeroPreserving(f::ZeroPreserving, T::Type, Ts::Type...) = f
 
-const _WEAK_FUNCTIONS = (:+, :-)
-for f in _WEAK_FUNCTIONS
+for F in (:(typeof(+)), :(typeof(-)), :(typeof(identity)))
   @eval begin
-    ZeroPreserving(::typeof($f), ::Type{<:Number}, ::Type{<:Number}...) = WeakPreserving($f)
+    ZeroPreserving(f::$F, ::Type, ::Type...) = WeakPreserving(f)
   end
 end
 
-const _STRONG_FUNCTIONS = (:*,)
-for f in _STRONG_FUNCTIONS
+using MapBroadcast: MapFunction
+for F in (:(typeof(*)), :(MapFunction{typeof(*)}))
   @eval begin
-    ZeroPreserving(::typeof($f), ::Type{<:Number}, ::Type{<:Number}...) = StrongPreserving(
-      $f
-    )
+    function ZeroPreserving(f::$F, ::Type, ::Type...)
+      return StrongPreserving(f)
+    end
   end
 end
 
@@ -71,29 +71,35 @@ end
   f, A::AbstractArray, Bs::AbstractArray...
 )
   f_pres = ZeroPreserving(f, A, Bs...)
-  return @interface I map(f_pres, A, Bs...)
+  return map_sparsearray(f_pres, A, Bs...)
 end
-@interface I::AbstractSparseArrayInterface function Base.map(
-  f::ZeroPreserving, A::AbstractArray, Bs::AbstractArray...
-)
+
+# This isn't an overload of `Base.map` since that leads to ambiguity errors. 
+function map_sparsearray(f::ZeroPreserving, A::AbstractArray, Bs::AbstractArray...)
   T = Base.Broadcast.combine_eltypes(f.f, (A, Bs...))
-  C = similar(I, T, size(A))
-  return @interface I map!(f, C, A, Bs...)
+  C = similar(A, T)
+  # TODO: Instead use:
+  # ```julia
+  # U = map(f.f, map(unstored, (A, Bs...))...)
+  # C = similar(A, Unstored(U))
+  # ```
+  # though right now `map` doesn't preserve `Zeros` or `BlockZeros`.
+  return map_sparsearray!(f, C, A, Bs...)
 end
 
 @interface I::AbstractSparseArrayInterface function Base.map!(
   f, C::AbstractArray, A::AbstractArray, Bs::AbstractArray...
 )
   f_pres = ZeroPreserving(f, A, Bs...)
-  return @interface I map!(f_pres, C, A, Bs...)
+  return map_sparsearray!(f_pres, C, A, Bs...)
 end
 
-@interface ::AbstractSparseArrayInterface function Base.map!(
+# This isn't an overload of `Base.map!` since that leads to ambiguity errors. 
+function map_sparsearray!(
   f::ZeroPreserving, C::AbstractArray, A::AbstractArray, Bs::AbstractArray...
 )
   checkshape(C, A, Bs...)
   unaliased = map(Base.Fix1(Base.unalias, C), (A, Bs...))
-
   if f isa StrongPreserving
     style = IndexStyle(C, unaliased...)
     inds = intersect(eachstoredindex.(Ref(style), unaliased)...)
@@ -107,19 +113,20 @@ end
   else
     error(lazy"unknown zero-preserving type $(typeof(f))")
   end
-
   @inbounds for I in inds
     C[I] = f.f(ith_all(I, unaliased)...)
   end
-
   return C
 end
 
 # Derived functions
 # -----------------
-@interface I::AbstractSparseArrayInterface Base.copyto!(C::AbstractArray, A::AbstractArray) = @interface I map!(
-  identity, C, A
+@interface I::AbstractSparseArrayInterface function Base.copyto!(
+  dest::AbstractArray, src::AbstractArray
 )
+  @interface I map!(identity, dest, src)
+  return dest
+end
 
 # Only map the stored values of the inputs.
 function map_stored! end