More functionality for SparseArrayDOK (#4)

Author: mtfishman

README.md

@@ -52,6 +52,7 @@ AbstractArray interface:
 
 ````julia
 a[1, 2] = 12
+@test a == [0 12; 0 0]
 @test a[1, 1] == 0
 @test a[2, 1] == 0
 @test a[1, 2] == 12
@@ -76,6 +77,17 @@ SparseArraysBase interface:
 @test issetequal(storedvalues(a), [12])
 ````
 
+AbstractArray functionality:
+
+````julia
+b = a .+ 2 .* a'
+@test b == [0 12; 24 0]
+@test storedlength(b) == 2
+@test b isa SparseArrayDOK{Float64}
+
+a * a'
+````
+
 ---
 
 *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*

examples/README.jl

@@ -55,6 +55,7 @@ a = SparseArrayDOK{Float64}(2, 2)
 # AbstractArray interface:
 
 a[1, 2] = 12
+@test a == [0 12; 0 0]
 @test a[1, 1] == 0
 @test a[2, 1] == 0
 @test a[1, 2] == 12
@@ -75,3 +76,12 @@ a[1, 2] = 12
 @test storedlength(a) == 1
 @test issetequal(storedpairs(a), [CartesianIndex(1, 2) => 12])
 @test issetequal(storedvalues(a), [12])
+
+# AbstractArray functionality:
+
+b = a .+ 2 .* a'
+@test b == [0 12; 24 0]
+@test storedlength(b) == 2
+@test b isa SparseArrayDOK{Float64}
+
+a * a'

src/SparseArraysBase.jl

@@ -2,6 +2,7 @@ module SparseArraysBase
 
 include("sparsearrayinterface.jl")
 include("wrappers.jl")
+include("abstractsparsearray.jl")
 include("sparsearraydok.jl")
 
 end

src/abstractsparsearray.jl

@@ -0,0 +1,24 @@
+abstract type AbstractSparseArray{T,N} <: AbstractArray{T,N} end
+
+using Derive: @array_aliases
+# Define AbstractSparseVector, AnyAbstractSparseArray, etc.
+@array_aliases AbstractSparseArray
+
+using Derive: Derive
+function Derive.interface(::Type{<:AbstractSparseArray})
+  return SparseArrayInterface()
+end
+
+using Derive: @derive
+# Derive `Base.getindex`, `Base.setindex!`, etc.
+@derive AnyAbstractSparseArray AbstractArrayOps
+
+using LinearAlgebra: LinearAlgebra
+@derive (T=AnyAbstractSparseVecOrMat,) begin
+  LinearAlgebra.mul!(::AbstractMatrix, ::T, ::T, ::Number, ::Number)
+end
+
+using ArrayLayouts: ArrayLayouts
+@derive (T=AnyAbstractSparseArray,) begin
+  ArrayLayouts.MemoryLayout(::Type{<:T})
+end

src/sparsearraydok.jl

@@ -1,25 +1,23 @@
 # TODO: Define `AbstractSparseArray`, make this a subtype.
-struct SparseArrayDOK{T,N} <: AbstractArray{T,N}
+struct SparseArrayDOK{T,N} <: AbstractSparseArray{T,N}
   storage::Dict{CartesianIndex{N},T}
   size::NTuple{N,Int}
 end
 
-function SparseArrayDOK{T}(size::Int...) where {T}
-  N = length(size)
+const SparseMatrixDOK{T} = SparseArrayDOK{T,2}
+const SparseVectorDOK{T} = SparseArrayDOK{T,1}
+
+function SparseArrayDOK{T,N}(size::Vararg{Int,N}) where {T,N}
   return SparseArrayDOK{T,N}(Dict{CartesianIndex{N},T}(), size)
 end
 
-using Derive: @wrappedtype
-# Define `WrappedSparseArrayDOK` and `AnySparseArrayDOK`.
-@wrappedtype SparseArrayDOK
-
-using Derive: Derive
-function Derive.interface(::Type{<:SparseArrayDOK})
-  return SparseArrayInterface()
+function SparseArrayDOK{T}(size::Int...) where {T}
+  return SparseArrayDOK{T,length(size)}(size...)
 end
 
-using Derive: @derive
-@derive AnySparseArrayDOK AbstractArrayOps
+using Derive: @array_aliases
+# Define `SparseMatrixDOK`, `AnySparseArrayDOK`, etc.
+@array_aliases SparseArrayDOK
 
 storage(a::SparseArrayDOK) = a.storage
 Base.size(a::SparseArrayDOK) = a.size

src/sparsearrayinterface.jl

@@ -5,10 +5,14 @@ storedvalues(a) = error()
 isstored(a, I::Int...) = error()
 eachstoredindex(a) = error()
 getstoredindex(a, I::Int...) = error()
-getunstoredindex(a, I::Int...) = error()
 setstoredindex!(a, value, I::Int...) = error()
 setunstoredindex!(a, value, I::Int...) = error()
 
+# Interface defaults.
+# TODO: Have a fallback that handles element types
+# that don't define `zero(::Type)`.
+getunstoredindex(a, I::Int...) = zero(eltype(a))
+
 # Derived interface.
 storedlength(a) = length(storedvalues(a))
 storedpairs(a) = map(I -> I => getstoredindex(a, I), eachstoredindex(a))
@@ -19,22 +23,28 @@ function eachstoredindex(a1, a2, a_rest...)
   return union(eachstoredindex.((a1, a2, a_rest...))...)
 end
 
-# TODO: Add `ndims` type parameter.
-# TODO: Define `AbstractSparseArrayInterface`, make this a subtype.
 using Derive: Derive, @interface, AbstractArrayInterface
-struct SparseArrayInterface <: AbstractArrayInterface 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 AbstractSparseArrayInterface <: AbstractArrayInterface end
+
+struct SparseArrayInterface <: AbstractSparseArrayInterface end
 
 # Convenient shorthand to refer to the sparse interface.
+# TODO: Define this as `InterfaceFunction(AbstractSparseArrayInterface)`
 const sparse = SparseArrayInterface()
 
 # TODO: Use `ArrayLayouts.layout_getindex`, `ArrayLayouts.sub_materialize`
 # to handle slicing (implemented by copying SubArray).
-@interface sparse function Base.getindex(a, I::Int...)
+@interface AbstractSparseArrayInterface function Base.getindex(a, I::Int...)
   !isstored(a, I...) && return getunstoredindex(a, I...)
   return getstoredindex(a, I...)
 end
 
-@interface sparse function Base.setindex!(a, value, I::Int...)
+@interface AbstractSparseArrayInterface function Base.setindex!(a, value, I::Int...)
   iszero(value) && return a
   if !isstored(a, I...)
     setunstoredindex!(a, value, I...)
@@ -46,23 +56,25 @@ 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 sparse function Base.similar(a, T::Type, size::Tuple{Vararg{Int}})
+@interface AbstractSparseArrayInterface function Base.similar(
+  a, T::Type, size::Tuple{Vararg{Int}}
+)
   return SparseArrayDOK{T}(size...)
 end
 
 ## TODO: Make this more general, handle mixtures of integers and ranges.
 ## TODO: Make this logic generic to all `similar(::AbstractInterface, ...)`.
-## @interface sparse function Base.similar(a, T::Type, dims::Tuple{Vararg{Base.OneTo}})
+## @interface  AbstractSparseArrayInterface function Base.similar(a, T::Type, dims::Tuple{Vararg{Base.OneTo}})
 ##   return sparse(similar)(interface, a, T, Base.to_shape(dims))
 ## end
 
-@interface sparse function Base.map(f, as...)
+@interface AbstractSparseArrayInterface function Base.map(f, as...)
   # This is defined in this way so we can rely on the Broadcast logic
   # for determining the destination of the operation (element type, shape, etc.).
   return f.(as...)
 end
 
-@interface sparse function Base.map!(f, dest, as...)
+@interface AbstractSparseArrayInterface function Base.map!(f, dest, as...)
   # Check `f` preserves zeros.
   # Define as `map_stored!`.
   # Define `eachstoredindex` promotion.
@@ -72,12 +84,13 @@ end
   return dest
 end
 
-# TODO: Define `AbstractSparseArrayStyle`, make this a subtype.
-struct SparseArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
+abstract type AbstractSparseArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
+
+struct SparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
 
 SparseArrayStyle{M}(::Val{N}) where {M,N} = SparseArrayStyle{N}()
 
-@interface sparse function Broadcast.BroadcastStyle(type::Type)
+@interface AbstractSparseArrayInterface function Broadcast.BroadcastStyle(type::Type)
   return SparseArrayStyle{ndims(type)}()
 end
 
@@ -100,12 +113,12 @@ abstract type AbstractSparseLayout <: ArrayLayouts.MemoryLayout end
 
 struct SparseLayout <: AbstractSparseLayout end
 
-@interface sparse function ArrayLayouts.MemoryLayout(type::Type)
+@interface AbstractSparseArrayInterface function ArrayLayouts.MemoryLayout(type::Type)
   return SparseLayout()
 end
 
 using LinearAlgebra: LinearAlgebra
-@interface sparse function LinearAlgebra.mul!(a_dest, a1, a2, α, β)
+@interface AbstractSparseArrayInterface function LinearAlgebra.mul!(a_dest, a1, a2, α, β)
   return ArrayLayouts.mul!(a_dest, a1, a2, α, β)
 end