Slicing, trailing indices

Author: mtfishman

examples/README.jl

@@ -39,6 +39,8 @@ julia> Pkg.add("SparseArraysBase")
 
 using SparseArraysBase:
   SparseArrayDOK,
+  SparseMatrixDOK,
+  SparseVectorDOK,
   eachstoredindex,
   getstoredindex,
   getunstoredindex,
@@ -81,15 +83,24 @@ using Dictionaries: IndexError
 # AbstractArray functionality:
 
 b = a .+ 2 .* a'
-@test b isa SparseArrayDOK{Float64}
+@test b isa SparseMatrixDOK{Float64}
 @test b == [0 12; 24 0]
 @test storedlength(b) == 2
 
 b = permutedims(a, (2, 1))
-@test b isa SparseArrayDOK{Float64}
+@test b isa SparseMatrixDOK{Float64}
 @test b[1, 1] == a[1, 1]
 @test b[2, 1] == a[1, 2]
 @test b[1, 2] == a[2, 1]
 @test b[2, 2] == a[2, 2]
 
-a * a'
+b = a * a'
+@test b isa SparseMatrixDOK{Float64}
+@test b == [144 0; 0 0]
+@test storedlength(b) == 1
+
+# Second column.
+b = a[1:2, 2]
+@test b isa SparseVectorDOK{Float64}
+@test b == [12, 0]
+@test storedlength(b) == 1

src/abstractsparsearray.jl

@@ -22,5 +22,5 @@ using LinearAlgebra: LinearAlgebra
 # TODO: Define `AbstractMatrixOps` and overload for
 # `AnyAbstractSparseMatrix` and `AnyAbstractSparseVector`,
 # which is where matrix multiplication and factorizations
-# shoudl go.
+# should go.
 @derive AnyAbstractSparseArray AbstractArrayOps

src/abstractsparsearrayinterface.jl

@@ -1,44 +1,66 @@
 # Minimal interface for `SparseArrayInterface`.
 # TODO: Define default definitions for these based
 # on the dense case.
+# TODO: Define as `MethodError`.
 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()
 
+# TODO: Use `Base.to_indices`?
+isstored(a::AbstractArray, I::CartesianIndex) = isstored(a, Tuple(I)...)
+getstoredindex(a::AbstractArray, I::CartesianIndex) = getstoredindex(a, Tuple(I)...)
+getunstoredindex(a::AbstractArray, I::CartesianIndex) = getunstoredindex(a, Tuple(I)...)
+function setstoredindex!(a::AbstractArray, value, I::CartesianIndex)
+  return setstoredindex!(a, value, Tuple(I)...)
+end
+function setunstoredindex!(a::AbstractArray, value, I::CartesianIndex)
+  return setunstoredindex!(a, value, Tuple(I)...)
+end
+
 # 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))
+storedlength(a::AbstractArray) = length(storedvalues(a))
+storedpairs(a::AbstractArray) = map(I -> I => getstoredindex(a, I), eachstoredindex(a))
+function storedvalues(a::AbstractArray)
+  return @view a[collect(eachstoredindex(a))]
+end
 
 function eachstoredindex(a1, a2, a_rest...)
   # TODO: Make this more customizable, say with a function
   # `combine/promote_storedindices(a1, a2)`.
   return union(eachstoredindex.((a1, a2, a_rest...))...)
 end
 
-using Derive: Derive, @interface, AbstractArrayInterface
+using Derive: Derive, @derive, @interface, AbstractArrayInterface
 
 # 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
 
-# TODO: Use `ArrayLayouts.layout_getindex`, `ArrayLayouts.sub_materialize`
-# to handle slicing (implemented by copying SubArray).
-@interface AbstractSparseArrayInterface function Base.getindex(a, I::Int...)
+# We restrict to `I::Vararg{Int,N}` to allow more general functions to handle trailing
+# indices and linear indices.
+@interface ::AbstractSparseArrayInterface function Base.getindex(
+  a::AbstractArray{<:Any,N}, I::Vararg{Int,N}
+) where {N}
   !isstored(a, I...) && return getunstoredindex(a, I...)
   return getstoredindex(a, I...)
 end
 
-@interface AbstractSparseArrayInterface function Base.setindex!(a, value, I::Int...)
+# We restrict to `I::Vararg{Int,N}` to allow more general functions to handle trailing
+# indices and linear indices.
+@interface ::AbstractSparseArrayInterface function Base.setindex!(
+  a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}
+) where {N}
   iszero(value) && return a
   if !isstored(a, I...)
     setunstoredindex!(a, value, I...)
@@ -50,14 +72,16 @@ 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, T::Type, size::Tuple{Vararg{Int}}
+@interface ::AbstractSparseArrayInterface function Base.similar(
+  a::AbstractArray, T::Type, size::Tuple{Vararg{Int}}
 )
   # TODO: Define `default_similartype` or something like that?
   return SparseArrayDOK{T}(size...)
 end
 
-@interface AbstractSparseArrayInterface function Base.map!(f, dest, as...)
+@interface ::AbstractSparseArrayInterface function Base.map!(
+  f, dest::AbstractArray, as::AbstractArray...
+)
   # Check `f` preserves zeros.
   # Define as `map_stored!`.
   # Define `eachstoredindex` promotion.
@@ -67,26 +91,26 @@ end
   return dest
 end
 
-# TODO: Make this a subtype of `Derive.AbstractArrayStyle{N}` instead.
-using Derive: Derive
-abstract type AbstractSparseArrayStyle{N} <: Derive.AbstractArrayStyle{N} end
+abstract type AbstractSparseArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
+
+@derive AbstractSparseArrayStyle AbstractArrayStyleOps
 
 struct SparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
 
 SparseArrayStyle{M}(::Val{N}) where {M,N} = SparseArrayStyle{N}()
 
-@interface AbstractSparseArrayInterface function Broadcast.BroadcastStyle(type::Type)
+@interface ::AbstractSparseArrayInterface function Broadcast.BroadcastStyle(type::Type)
   return SparseArrayStyle{ndims(type)}()
 end
 
 using ArrayLayouts: ArrayLayouts, MatMulMatAdd
 
 abstract type AbstractSparseLayout <: ArrayLayouts.MemoryLayout end
 
-struct SparseLayout <: AbstractSparseLayout end
-
-@interface AbstractSparseArrayInterface function ArrayLayouts.MemoryLayout(type::Type)
-  return SparseLayout()
+function ArrayLayouts.sub_materialize(::AbstractSparseLayout, a::AbstractArray, axes::Tuple)
+  a_dest = similar(a)
+  a_dest .= a
+  return a_dest
 end
 
 function mul_indices(I1::CartesianIndex{2}, I2::CartesianIndex{2})
@@ -139,3 +163,9 @@ function ArrayLayouts.materialize!(
   sparse_mul!(m.C, m.A, m.B, m.α, m.β)
   return m.C
 end
+
+struct SparseLayout <: AbstractSparseLayout end
+
+@interface ::AbstractSparseArrayInterface function ArrayLayouts.MemoryLayout(type::Type)
+  return SparseLayout()
+end

src/sparsearraydok.jl

@@ -10,12 +10,20 @@ struct SparseArrayDOK{T,N,F} <: AbstractSparseArray{T,N}
   getunstoredindex::F
 end
 
+using Derive: Derive
+# This defines the destination type of various operations in Derive.jl.
+Derive.arraytype(::AbstractSparseArrayInterface, T::Type) = SparseArrayDOK{T}
+
 function SparseArrayDOK{T,N}(size::Vararg{Int,N}) where {T,N}
   getunstoredindex = default_getunstoredindex
   F = typeof(getunstoredindex)
   return SparseArrayDOK{T,N,F}(Dictionary{CartesianIndex{N},T}(), size, getunstoredindex)
 end
 
+function SparseArrayDOK{T}(::UndefInitializer, size::Tuple{Vararg{Int}}) where {T}
+  return SparseArrayDOK{T,length(size)}(size...)
+end
+
 function SparseArrayDOK{T}(size::Int...) where {T}
   return SparseArrayDOK{T,length(size)}(size...)
 end

src/sparsearrayinterface.jl

@@ -8,4 +8,4 @@ struct SparseArrayInterface <: AbstractSparseArrayInterface end
 # version of `map`.
 # const sparse = SparseArrayInterface()
 
-Derive.interface(::AbstractSparseArrayStyle) = SparseArrayInterface()
+Derive.interface(::Type{<:AbstractSparseArrayStyle}) = SparseArrayInterface()

src/wrappers.jl

@@ -1,70 +1,102 @@
+function cartesianindex_reverse(I::CartesianIndex)
+  return CartesianIndex(reverse(Tuple(I)))
+end
+tuple_oneto(n) = ntuple(identity, n)
+
+# TODO: Use `Base.PermutedDimsArrays.genperm` or
+# https://github.com/jipolanco/StaticPermutations.jl?
+genperm(v, perm) = map(j -> v[j], perm)
+
+## TODO: Use this and something similar for `Dictionary` to make a faster
+## implementation of `storedvalues(::SubArray)`.
+## function valuesview(d::Dict, keys)
+##   return @view d.vals[[Base.ht_keyindex(d, key) for key in keys]]
+## end
+
+function eachstoredparentindex(a::SubArray)
+  return filter(eachstoredindex(parent(a))) do I
+    return all(d -> I[d] ∈ parentindices(a)[d], 1:ndims(parent(a)))
+  end
+end
+function storedvalues(a::SubArray)
+  return @view parent(a)[collect(eachstoredparentindex(a))]
+end
+function isstored(a::SubArray, I::Int...)
+  return isstored(parent(a), Base.reindex(parentindices(a), I)...)
+end
+function getstoredindex(a::SubArray, I::Int...)
+  return getstoredindex(parent(a), Base.reindex(parentindices(a), I)...)
+end
+function getunstoredindex(a::SubArray, I::Int...)
+  return getunstoredindex(parent(a), Base.reindex(parentindices(a), I)...)
+end
+function eachstoredindex(a::SubArray)
+  nonscalardims = filter(tuple_oneto(ndims(parent(a)))) do d
+    return !(parentindices(a)[d] isa Real)
+  end
+  return collect((
+    CartesianIndex(
+      map(nonscalardims) do d
+        return findfirst(==(I[d]), parentindices(a)[d])
+      end,
+    ) for I in eachstoredparentindex(a)
+  ))
+end
+
+perm(::PermutedDimsArray{<:Any,<:Any,p}) where {p} = p
+iperm(::PermutedDimsArray{<:Any,<:Any,<:Any,ip}) where {ip} = ip
+
 storedvalues(a::PermutedDimsArray) = storedvalues(parent(a))
 function isstored(a::PermutedDimsArray, I::Int...)
-  return isstored(parent(a), reverse(I)...)
-end
-function eachstoredindex(a::PermutedDimsArray)
-  # TODO: Make lazy with `Iterators.map`.
-  return map(CartesianIndex ∘ reverse ∘ Tuple, collect(eachstoredindex(parent(a))))
+  return isstored(parent(a), genperm(I, iperm(a))...)
 end
 function getstoredindex(a::PermutedDimsArray, I::Int...)
-  return getstoredindex(parent(a), reverse(I)...)
+  return getstoredindex(parent(a), genperm(I, iperm(a))...)
 end
 function getunstoredindex(a::PermutedDimsArray, I::Int...)
-  return getunstoredindex(parent(a), reverse(I)...)
+  return getunstoredindex(parent(a), genperm(I, iperm(a))...)
 end
 function setstoredindex!(a::PermutedDimsArray, value, I::Int...)
-  setstoredindex!(parent(a), value, reverse(I)...)
+  # TODO: Should this be `iperm(a)`?
+  setstoredindex!(parent(a), value, genperm(I, perm(a))...)
   return a
 end
 function setunstoredindex!(a::PermutedDimsArray, value, I::Int...)
-  setunstoredindex!(parent(a), value, reverse(I)...)
+  # TODO: Should this be `iperm(a)`?
+  setunstoredindex!(parent(a), value, genperm(I, perm(a))...)
   return a
 end
-
-using LinearAlgebra: Adjoint
-storedvalues(a::Adjoint) = storedvalues(parent(a))
-function isstored(a::Adjoint, i::Int, j::Int)
-  return isstored(parent(a), j, i)
-end
-function eachstoredindex(a::Adjoint)
+function eachstoredindex(a::PermutedDimsArray)
   # TODO: Make lazy with `Iterators.map`.
-  return map(CartesianIndex ∘ reverse ∘ Tuple, collect(eachstoredindex(parent(a))))
-end
-function getstoredindex(a::Adjoint, i::Int, j::Int)
-  return getstoredindex(parent(a), j, i)'
-end
-function getunstoredindex(a::Adjoint, i::Int, j::Int)
-  return getunstoredindex(parent(a), j, i)'
-end
-function setstoredindex!(a::Adjoint, value, i::Int, j::Int)
-  setstoredindex!(parent(a), value', j, i)
-  return a
-end
-function setunstoredindex!(a::Adjoint, value, i::Int, j::Int)
-  setunstoredindex!(parent(a), value', j, i)
-  return a
+  return map(collect(eachstoredindex(parent(a)))) do I
+    return CartesianIndex(genperm(I, perm(a)))
+  end
 end
 
-using LinearAlgebra: Transpose
-storedvalues(a::Transpose) = storedvalues(parent(a))
-function isstored(a::Transpose, i::Int, j::Int)
-  return isstored(parent(a), j, i)
-end
-function eachstoredindex(a::Transpose)
-  # TODO: Make lazy with `Iterators.map`.
-  return map(CartesianIndex ∘ reverse ∘ Tuple, collect(eachstoredindex(parent(a))))
-end
-function getstoredindex(a::Transpose, i::Int, j::Int)
-  return transpose(getstoredindex(parent(a), j, i))
-end
-function getunstoredindex(a::Transpose, i::Int, j::Int)
-  return transpose(getunstoredindex(parent(a), j, i))
-end
-function setstoredindex!(a::Transpose, value, i::Int, j::Int)
-  setstoredindex!(parent(a), transpose(value), j, i)
-  return a
-end
-function setunstoredindex!(a::Transpose, value, i::Int, j::Int)
-  setunstoredindex!(parent(a), transpose(value), j, i)
-  return a
+for (type, func) in ((:Adjoint, :adjoint), (:Transpose, :transpose))
+  @eval begin
+    using LinearAlgebra: $type
+    storedvalues(a::$type) = storedvalues(parent(a))
+    function isstored(a::$type, i::Int, j::Int)
+      return isstored(parent(a), j, i)
+    end
+    function eachstoredindex(a::$type)
+      # TODO: Make lazy with `Iterators.map`.
+      return map(cartesianindex_reverse, collect(eachstoredindex(parent(a))))
+    end
+    function getstoredindex(a::$type, i::Int, j::Int)
+      return $func(getstoredindex(parent(a), j, i))
+    end
+    function getunstoredindex(a::$type, i::Int, j::Int)
+      return $func(getunstoredindex(parent(a), j, i))
+    end
+    function setstoredindex!(a::$type, value, i::Int, j::Int)
+      setstoredindex!(parent(a), $func(value), j, i)
+      return a
+    end
+    function setunstoredindex!(a::$type, value, i::Int, j::Int)
+      setunstoredindex!(parent(a), $func(value), j, i)
+      return a
+    end
+  end
 end