From 6401807f237d937b2e3908ee2b07d94261f6a540 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sat, 25 Nov 2023 17:32:42 -0500 Subject: [PATCH 01/17] [NDTensors] Start SparseArrayDOKs module --- NDTensors/src/NDTensors.jl | 46 +++--- NDTensors/src/abstractarray/similar.jl | 2 + .../abstractarray/tensoralgebra/contract.jl | 3 + NDTensors/src/array/permutedims.jl | 5 +- NDTensors/src/array/set_types.jl | 2 + .../arraystorage/storage/arraystorage.jl | 3 + .../blocksparsearray/storage/contract.jl | 4 + NDTensors/src/imports.jl | 7 + .../SparseArrayDOKs/src/SparseArrayDOKs.jl | 156 ++++++++++++++++++ .../src/lib/SparseArrayDOKs/test/runtests.jl | 43 +++++ NDTensors/src/lib/Unwrap/README.md | 6 + NDTensors/test/lib/runtests.jl | 1 + src/imports.jl | 4 +- test/Project.toml | 1 + 14 files changed, 254 insertions(+), 29 deletions(-) create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index c052cbffe5..8267ceb574 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -19,33 +19,25 @@ using Strided using TimerOutputs using TupleTools -# TODO: List types, macros, and functions being used. -include("lib/AlgorithmSelection/src/AlgorithmSelection.jl") -using .AlgorithmSelection: AlgorithmSelection -include("lib/BaseExtensions/src/BaseExtensions.jl") -using .BaseExtensions: BaseExtensions -include("lib/SetParameters/src/SetParameters.jl") -using .SetParameters -include("lib/BroadcastMapConversion/src/BroadcastMapConversion.jl") -using .BroadcastMapConversion: BroadcastMapConversion -include("lib/Unwrap/src/Unwrap.jl") -using .Unwrap -include("lib/RankFactorization/src/RankFactorization.jl") -using .RankFactorization: RankFactorization -include("lib/TensorAlgebra/src/TensorAlgebra.jl") -using .TensorAlgebra: TensorAlgebra -include("lib/DiagonalArrays/src/DiagonalArrays.jl") -using .DiagonalArrays -include("lib/BlockSparseArrays/src/BlockSparseArrays.jl") -using .BlockSparseArrays -include("lib/NamedDimsArrays/src/NamedDimsArrays.jl") -using .NamedDimsArrays: NamedDimsArrays -include("lib/SmallVectors/src/SmallVectors.jl") -using .SmallVectors -include("lib/SortedSets/src/SortedSets.jl") -using .SortedSets -include("lib/TagSets/src/TagSets.jl") -using .TagSets +for lib in [ + :AlgorithmSelection, + :BaseExtensions, + :SetParameters, + :BroadcastMapConversion, + :Unwrap, + :RankFactorization, + :TensorAlgebra, + :DiagonalArrays, + :SparseArrayDOKs, + :BlockSparseArrays, + :NamedDimsArrays, + :SmallVectors, + :SortedSets, + :TagSets, +] + include("lib/$(lib)/src/$(lib).jl") + @eval using .$lib: $lib +end using Base: @propagate_inbounds, ReshapedArray, DimOrInd, OneTo diff --git a/NDTensors/src/abstractarray/similar.jl b/NDTensors/src/abstractarray/similar.jl index dc0c6062ac..4a39212b75 100644 --- a/NDTensors/src/abstractarray/similar.jl +++ b/NDTensors/src/abstractarray/similar.jl @@ -1,3 +1,5 @@ +using .Unwrap: IsWrappedArray + ## Custom `NDTensors.similar` implementation. ## More extensive than `Base.similar`. diff --git a/NDTensors/src/abstractarray/tensoralgebra/contract.jl b/NDTensors/src/abstractarray/tensoralgebra/contract.jl index 269c8c39b2..139c9b9c0a 100644 --- a/NDTensors/src/abstractarray/tensoralgebra/contract.jl +++ b/NDTensors/src/abstractarray/tensoralgebra/contract.jl @@ -1,4 +1,7 @@ using LinearAlgebra: BlasFloat +using .Unwrap: expose + +# TODO: Delete these exports export backend_auto, backend_blas, backend_generic @eval struct GemmBackend{T} diff --git a/NDTensors/src/array/permutedims.jl b/NDTensors/src/array/permutedims.jl index 12a5cb9d2f..4151d03099 100644 --- a/NDTensors/src/array/permutedims.jl +++ b/NDTensors/src/array/permutedims.jl @@ -1,4 +1,7 @@ -## Create the Exposed version of Base.permutedims +using .Unwrap: Exposed, unexpose + +# TODO: Move to `Unwrap` module. +# Create the Exposed version of Base.permutedims function permutedims(E::Exposed{<:Array}, perm) ## Creating Mperm here to evaluate the permutation and ## avoid returning a Stridedview diff --git a/NDTensors/src/array/set_types.jl b/NDTensors/src/array/set_types.jl index 57b396f931..739c562dfc 100644 --- a/NDTensors/src/array/set_types.jl +++ b/NDTensors/src/array/set_types.jl @@ -1,3 +1,5 @@ +using .SetParameters: Position, set_parameters + """ TODO: Use `Accessors.jl` notation: ```julia diff --git a/NDTensors/src/arraystorage/arraystorage/storage/arraystorage.jl b/NDTensors/src/arraystorage/arraystorage/storage/arraystorage.jl index 2fa606be72..60a48bda3e 100644 --- a/NDTensors/src/arraystorage/arraystorage/storage/arraystorage.jl +++ b/NDTensors/src/arraystorage/arraystorage/storage/arraystorage.jl @@ -1,3 +1,6 @@ +using .BlockSparseArrays: BlockSparseArray +using .DiagonalArrays: DiagonalArray + # Used for dispatch to distinguish from Tensors wrapping TensorStorage. # Remove once TensorStorage is removed. const ArrayStorage{T,N} = Union{ diff --git a/NDTensors/src/arraystorage/blocksparsearray/storage/contract.jl b/NDTensors/src/arraystorage/blocksparsearray/storage/contract.jl index 3e790ae29a..7007f45d22 100644 --- a/NDTensors/src/arraystorage/blocksparsearray/storage/contract.jl +++ b/NDTensors/src/arraystorage/blocksparsearray/storage/contract.jl @@ -1,3 +1,7 @@ +# TODO: Change to: +# using .SparseArrayDOKs: SparseArrayDOK +using .BlockSparseArrays: SparseArray + # TODO: This is inefficient, need to optimize. # Look at `contract_labels`, `contract_blocks` and `maybe_contract_blocks!` in: # src/blocksparse/contract_utilities.jl diff --git a/NDTensors/src/imports.jl b/NDTensors/src/imports.jl index 4ad941515e..df1d71cbde 100644 --- a/NDTensors/src/imports.jl +++ b/NDTensors/src/imports.jl @@ -1,3 +1,10 @@ +# Makes `cpu` available as `NDTensors.cpu`. +# TODO: Define `cpu`, `cu`, etc. in a module `DeviceAbstractions`, +# similar to: +# https://github.com/JuliaGPU/KernelAbstractions.jl +# https://github.com/oschulz/HeterogeneousComputing.jl +using .Unwrap: cpu + import Base: # Types AbstractFloat, diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl new file mode 100644 index 0000000000..759514e949 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl @@ -0,0 +1,156 @@ +module SparseArrayDOKs +using Dictionaries: Dictionary, set! +using SparseArrays: SparseArrays, AbstractSparseArray + +# Also look into: +# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ + +# Required `SparseArrayInterface` interface. +# https://github.com/Jutho/SparseArrayKit.jl interface functions +nonzero_keys(a::AbstractArray) = error("Not implemented") +nonzero_values(a::AbstractArray) = error("Not implemented") +nonzero_pairs(a::AbstractArray) = error("Not implemented") + +# A dictionary-like structure +# TODO: Rename `nonzeros`, `structural_nonzeros`, etc.? +nonzero_structure(a::AbstractArray) = error("Not implemented") + +# Derived `SparseArrayInterface` interface. +nonzero_length(a::AbstractArray) = length(nonzero_keys(a)) +is_structural_nonzero(a::AbstractArray, I) = I ∈ nonzero_keys(a) + +# Overload if zero value is index dependent or +# doesn't match element type. +zerotype(a::AbstractArray) = eltype(a) +getindex_nonzero(a::AbstractArray, I) = nonzero_structure(a)[I] +getindex_zero(a::AbstractArray, I) = zero(zerotype(a)) +function setindex_zero!(a::AbstractArray, value, I) + # TODO: This may need to be modified. + nonzero_structure(a)[I] = value + return a +end +function setindex_nonzero!(a::AbstractArray, value, I) + nonzero_structure(a)[I] = value + return a +end + +struct Zero +end +(::Zero)(type, I) = zero(type) + +default_zero_type(type::Type) = type +default_zero() = Zero() # (eltype, I) -> zero(eltype) +default_keytype(ndims::Int) = CartesianIndex{ndims} +default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() + +# TODO: Define a constructor with a default `zero`. +struct SparseArrayDOK{T,N,ZT,Zero} <: AbstractSparseArray{T,Int,N} + data::Dictionary{CartesianIndex{N},T} + dims::NTuple{N,Int} + zero::Zero +end + +# Constructors +function SparseArrayDOK{T,N,ZT,Zero}(dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT,Zero} + return SparseArrayDOK{T,N,ZT,Zero}(default_data(T, N), dims, zero) +end + +function SparseArrayDOK{T,N,ZT}(dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT} + return SparseArrayDOK{T,N,ZT,typeof(zero)}(dims, default_zero()) +end + +function SparseArrayDOK{T,N,ZT}(dims::Tuple{Vararg{Int}}) where {T,N,ZT} + return SparseArrayDOK{T,N,ZT}(dims, default_zero()) +end + +function SparseArrayDOK{T,N}(dims::Tuple{Vararg{Int}}) where {T,N} + return SparseArrayDOK{T,N,default_zero_type(T)}(dims) +end + +function SparseArrayDOK{T}(dims::Tuple{Vararg{Int}}) where {T} + return SparseArrayDOK{T,length(dims)}(dims) +end + +function SparseArrayDOK{T}(dims::Int...) where {T} + return SparseArrayDOK{T}(dims) +end + +# undef +function SparseArrayDOK{T,N,ZT,Zero}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT,Zero} + return SparseArrayDOK{T,N,ZT,Zero}(dims, zero) +end + +function SparseArrayDOK{T,N,ZT}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT} + return SparseArrayDOK{T,N,ZT}(dims, zero) +end + +function SparseArrayDOK{T,N,ZT}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T,N,ZT} + return SparseArrayDOK{T,N,ZT}(dims) +end + +function SparseArrayDOK{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T,N} + return SparseArrayDOK{T,N}(dims) +end + +function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} + return SparseArrayDOK{T}(dims) +end + +function SparseArrayDOK{T}(::UndefInitializer, dims::Int...) where {T} + return SparseArrayDOK{T}(dims...) +end + +# Required `SparseArrayInterface` interface +nonzero_structure(a::SparseArrayDOK) = a.data +# TODO: Make this a generic function. +nonzero_keys(a::SparseArrayDOK) = keys(nonzero_structure(a)) + +# Julia Base `AbstractSparseArray` interface +SparseArrays.nnz(a::SparseArrayDOK) = nonzero_length(a) + +# Optional SparseArray interface +# TODO: Use `SetParameters`. +zerotype(a::SparseArrayDOK{<:Any,<:Any,ZT}) where {ZT} = ZT +zero_value(a::SparseArrayDOK, I) = a.zero(zerotype(a), I) + +# Accessors +Base.size(a::SparseArrayDOK) = a.dims + +function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N} + return a[CartesianIndex(I)] +end + +function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N}) where {N} + if !is_structural_nonzero(a, I) + return getindex_zero(a, I) + end + return getindex_nonzero(a, I) +end + +# `SparseArrayInterface` interface +function setindex_zero!(a::SparseArrayDOK, value, I) + # TODO: This is specific to the `Dictionaries.jl` + # interface, make more generic? + set!(nonzero_structure(a), I, value) + return a +end + + +function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N} + a[CartesianIndex(I)] = value + return a +end + +function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N}) where {N} + if !is_structural_nonzero(a, I) + setindex_zero!(a, value, I) + end + setindex_nonzero!(a, value, I) + return a +end + +# similar +# TODO: How does this deal with the converting the zero type? +Base.similar(a::SparseArrayDOK{T,N,ZT,Zero}) where {T,N,ZT,Zero} = SparseArrayDOK{T,N,ZT,Zero}(undef, size(a), a.zero) + +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl new file mode 100644 index 0000000000..de0c6b8f5e --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl @@ -0,0 +1,43 @@ +@eval module $(gensym()) +using Test: @test, @testset +using NDTensors.SparseArrayDOKs: SparseArrayDOK, nonzero_keys, nonzero_length +using SparseArrays: nnz +@testset "SparseArrayDOK (eltype=$elt)" for elt in ( + Float32, ComplexF32, Float64, ComplexF64 +) + a = SparseArrayDOK{elt}(3, 4) + @test a == SparseArrayDOK{elt}((3, 4)) + @test a == SparseArrayDOK{elt}(undef, 3, 4) + @test a == SparseArrayDOK{elt}(undef, (3, 4)) + @test iszero(a) + @test iszero(nnz(a)) + @test nonzero_length(a) == nnz(a) + @test size(a) == (3, 4) + @test eltype(a) == elt + for I in eachindex(a) + @test iszero(a[I]) + @test a[I] isa elt + end + @test isempty(nonzero_keys(a)) + + x12 = randn(elt) + x23 = randn(elt) + b = copy(a) + @test b isa SparseArrayDOK{elt} + @test iszero(b) + b[1, 2] = x12 + b[2, 3] = x23 + @test iszero(a) + @test !iszero(b) + @test b[1, 2] == x12 + @test b[2, 3] == x23 + + # To test: + # reshape + # zero (PermutedDimsArray) + # map[!] + # broadcast + # Custom zero type + # conversion to `SparseMatrixCSC` +end +end diff --git a/NDTensors/src/lib/Unwrap/README.md b/NDTensors/src/lib/Unwrap/README.md index 04f8c355db..d5d1b36d72 100644 --- a/NDTensors/src/lib/Unwrap/README.md +++ b/NDTensors/src/lib/Unwrap/README.md @@ -1,3 +1,9 @@ # Unwrap A module to unwrap complex array types to assist in the generic programming of array-type based functions. + +Related: +- https://juliaarrays.github.io/ArrayInterface.jl/stable/wrapping/ +- https://github.com/JuliaGPU/Adapt.jl +- https://github.com/chengchingwen/StructWalk.jl +- https://github.com/FluxML/Functors.jl diff --git a/NDTensors/test/lib/runtests.jl b/NDTensors/test/lib/runtests.jl index 102c83bad7..8f0fbbdfc8 100644 --- a/NDTensors/test/lib/runtests.jl +++ b/NDTensors/test/lib/runtests.jl @@ -10,6 +10,7 @@ using Test: @testset "SetParameters", "SmallVectors", "SortedSets", + "SparseArrayDOKs", "TagSets", "TensorAlgebra", "Unwrap", diff --git a/src/imports.jl b/src/imports.jl index 86a7d9b756..bea3457d4d 100644 --- a/src/imports.jl +++ b/src/imports.jl @@ -110,6 +110,9 @@ import LinearAlgebra: tr, transpose +using ITensors.NDTensors.Unwrap: + cpu + using ITensors.NDTensors: Algorithm, @Algorithm_str, @@ -117,7 +120,6 @@ using ITensors.NDTensors: _Tuple, _NTuple, blas_get_num_threads, - cpu, cu, disable_auto_fermion, double_precision, diff --git a/test/Project.toml b/test/Project.toml index 90bb006414..201f9a3bdf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,6 +9,7 @@ ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" From 9e4142c49225252128fda459495f769b96f6a673 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sun, 26 Nov 2023 15:13:26 -0500 Subject: [PATCH 02/17] Reorganization --- NDTensors/src/NDTensors.jl | 28 ++-- .../SparseArrayDOKs/src/SparseArrayDOKs.jl | 157 +----------------- .../src/SparseArrayDOKsSparseArraysExt.jl | 5 + .../lib/SparseArrayDOKs/src/sparsearraydok.jl | 104 ++++++++++++ .../src/sparsearrayinterface.jl | 39 +++++ .../src/lib/SparseArrayDOKs/test/runtests.jl | 82 +++++---- src/imports.jl | 3 +- 7 files changed, 212 insertions(+), 206 deletions(-) create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index 8267ceb574..a7f4efa37a 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -20,20 +20,20 @@ using TimerOutputs using TupleTools for lib in [ - :AlgorithmSelection, - :BaseExtensions, - :SetParameters, - :BroadcastMapConversion, - :Unwrap, - :RankFactorization, - :TensorAlgebra, - :DiagonalArrays, - :SparseArrayDOKs, - :BlockSparseArrays, - :NamedDimsArrays, - :SmallVectors, - :SortedSets, - :TagSets, + :AlgorithmSelection, + :BaseExtensions, + :SetParameters, + :BroadcastMapConversion, + :Unwrap, + :RankFactorization, + :TensorAlgebra, + :DiagonalArrays, + :SparseArrayDOKs, + :BlockSparseArrays, + :NamedDimsArrays, + :SmallVectors, + :SortedSets, + :TagSets, ] include("lib/$(lib)/src/$(lib).jl") @eval using .$lib: $lib diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl index 759514e949..d1e9a4ebd3 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl @@ -1,156 +1,5 @@ module SparseArrayDOKs -using Dictionaries: Dictionary, set! -using SparseArrays: SparseArrays, AbstractSparseArray - -# Also look into: -# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ - -# Required `SparseArrayInterface` interface. -# https://github.com/Jutho/SparseArrayKit.jl interface functions -nonzero_keys(a::AbstractArray) = error("Not implemented") -nonzero_values(a::AbstractArray) = error("Not implemented") -nonzero_pairs(a::AbstractArray) = error("Not implemented") - -# A dictionary-like structure -# TODO: Rename `nonzeros`, `structural_nonzeros`, etc.? -nonzero_structure(a::AbstractArray) = error("Not implemented") - -# Derived `SparseArrayInterface` interface. -nonzero_length(a::AbstractArray) = length(nonzero_keys(a)) -is_structural_nonzero(a::AbstractArray, I) = I ∈ nonzero_keys(a) - -# Overload if zero value is index dependent or -# doesn't match element type. -zerotype(a::AbstractArray) = eltype(a) -getindex_nonzero(a::AbstractArray, I) = nonzero_structure(a)[I] -getindex_zero(a::AbstractArray, I) = zero(zerotype(a)) -function setindex_zero!(a::AbstractArray, value, I) - # TODO: This may need to be modified. - nonzero_structure(a)[I] = value - return a -end -function setindex_nonzero!(a::AbstractArray, value, I) - nonzero_structure(a)[I] = value - return a -end - -struct Zero -end -(::Zero)(type, I) = zero(type) - -default_zero_type(type::Type) = type -default_zero() = Zero() # (eltype, I) -> zero(eltype) -default_keytype(ndims::Int) = CartesianIndex{ndims} -default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() - -# TODO: Define a constructor with a default `zero`. -struct SparseArrayDOK{T,N,ZT,Zero} <: AbstractSparseArray{T,Int,N} - data::Dictionary{CartesianIndex{N},T} - dims::NTuple{N,Int} - zero::Zero -end - -# Constructors -function SparseArrayDOK{T,N,ZT,Zero}(dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT,Zero} - return SparseArrayDOK{T,N,ZT,Zero}(default_data(T, N), dims, zero) -end - -function SparseArrayDOK{T,N,ZT}(dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT} - return SparseArrayDOK{T,N,ZT,typeof(zero)}(dims, default_zero()) -end - -function SparseArrayDOK{T,N,ZT}(dims::Tuple{Vararg{Int}}) where {T,N,ZT} - return SparseArrayDOK{T,N,ZT}(dims, default_zero()) -end - -function SparseArrayDOK{T,N}(dims::Tuple{Vararg{Int}}) where {T,N} - return SparseArrayDOK{T,N,default_zero_type(T)}(dims) -end - -function SparseArrayDOK{T}(dims::Tuple{Vararg{Int}}) where {T} - return SparseArrayDOK{T,length(dims)}(dims) -end - -function SparseArrayDOK{T}(dims::Int...) where {T} - return SparseArrayDOK{T}(dims) -end - -# undef -function SparseArrayDOK{T,N,ZT,Zero}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT,Zero} - return SparseArrayDOK{T,N,ZT,Zero}(dims, zero) -end - -function SparseArrayDOK{T,N,ZT}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T,N,ZT} - return SparseArrayDOK{T,N,ZT}(dims, zero) -end - -function SparseArrayDOK{T,N,ZT}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T,N,ZT} - return SparseArrayDOK{T,N,ZT}(dims) -end - -function SparseArrayDOK{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T,N} - return SparseArrayDOK{T,N}(dims) -end - -function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} - return SparseArrayDOK{T}(dims) -end - -function SparseArrayDOK{T}(::UndefInitializer, dims::Int...) where {T} - return SparseArrayDOK{T}(dims...) -end - -# Required `SparseArrayInterface` interface -nonzero_structure(a::SparseArrayDOK) = a.data -# TODO: Make this a generic function. -nonzero_keys(a::SparseArrayDOK) = keys(nonzero_structure(a)) - -# Julia Base `AbstractSparseArray` interface -SparseArrays.nnz(a::SparseArrayDOK) = nonzero_length(a) - -# Optional SparseArray interface -# TODO: Use `SetParameters`. -zerotype(a::SparseArrayDOK{<:Any,<:Any,ZT}) where {ZT} = ZT -zero_value(a::SparseArrayDOK, I) = a.zero(zerotype(a), I) - -# Accessors -Base.size(a::SparseArrayDOK) = a.dims - -function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N} - return a[CartesianIndex(I)] -end - -function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N}) where {N} - if !is_structural_nonzero(a, I) - return getindex_zero(a, I) - end - return getindex_nonzero(a, I) -end - -# `SparseArrayInterface` interface -function setindex_zero!(a::SparseArrayDOK, value, I) - # TODO: This is specific to the `Dictionaries.jl` - # interface, make more generic? - set!(nonzero_structure(a), I, value) - return a -end - - -function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N} - a[CartesianIndex(I)] = value - return a -end - -function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N}) where {N} - if !is_structural_nonzero(a, I) - setindex_zero!(a, value, I) - end - setindex_nonzero!(a, value, I) - return a -end - -# similar -# TODO: How does this deal with the converting the zero type? -Base.similar(a::SparseArrayDOK{T,N,ZT,Zero}) where {T,N,ZT,Zero} = SparseArrayDOK{T,N,ZT,Zero}(undef, size(a), a.zero) - +include("sparsearrayinterface.jl") +include("sparsearraydok.jl") +include("SparseArrayDOKsSparseArraysExt.jl") end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl new file mode 100644 index 0000000000..8bbb6121fd --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl @@ -0,0 +1,5 @@ +using SparseArrays: SparseArrays + +# Julia Base `AbstractSparseArray` interface +SparseArrays.nnz(a::SparseArrayDOK) = nonzero_length(a) + diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl new file mode 100644 index 0000000000..454fc09a7d --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl @@ -0,0 +1,104 @@ +using Dictionaries: Dictionary, set! + +# TODO: Define a constructor with a default `zero`. +struct SparseArrayDOK{T,N,Zero} <: AbstractArray{T,N} + data::Dictionary{CartesianIndex{N},T} + dims::NTuple{N,Int} + zero::Zero +end + +# Constructors +function SparseArrayDOK{T,N,Zero}(dims::Tuple{Vararg{Int}}, zero) where {T,N,Zero} + return SparseArrayDOK{T,N,Zero}(default_data(T, N), dims, zero) +end + +function SparseArrayDOK{T,N}(dims::Tuple{Vararg{Int}}, zero) where {T,N} + return SparseArrayDOK{T,N,typeof(zero)}(dims, zero) +end + +function SparseArrayDOK{T,N}(dims::Tuple{Vararg{Int}}) where {T,N} + return SparseArrayDOK{T,N}(dims, default_zero()) +end + +function SparseArrayDOK{T}(dims::Tuple{Vararg{Int}}) where {T} + return SparseArrayDOK{T,length(dims)}(dims) +end + +function SparseArrayDOK{T}(dims::Int...) where {T} + return SparseArrayDOK{T}(dims) +end + +# undef +function SparseArrayDOK{T,N,Zero}( + ::UndefInitializer, dims::Tuple{Vararg{Int}}, zero +) where {T,N,Zero} + return SparseArrayDOK{T,N,Zero}(dims, zero) +end + +function SparseArrayDOK{T,N}( + ::UndefInitializer, dims::Tuple{Vararg{Int}}, zero +) where {T,N} + return SparseArrayDOK{T,N}(dims, zero) +end + +function SparseArrayDOK{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T,N} + return SparseArrayDOK{T,N}(dims) +end + +function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} + return SparseArrayDOK{T}(dims) +end + +function SparseArrayDOK{T}(::UndefInitializer, dims::Int...) where {T} + return SparseArrayDOK{T}(dims...) +end + +# Required `SparseArrayInterface` interface +nonzero_structure(a::SparseArrayDOK) = a.data +# TODO: Make this a generic function. +nonzero_keys(a::SparseArrayDOK) = keys(nonzero_structure(a)) + +# Optional `SparseArrayInterface` interface +# TODO: Use `SetParameters`. +zero_value(a::SparseArrayDOK, I) = a.zero(eltype(a), I) + +# Accessors +Base.size(a::SparseArrayDOK) = a.dims + +function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N} + return a[CartesianIndex(I)] +end + +function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N}) where {N} + if !is_structural_nonzero(a, I) + return getindex_zero(a, I) + end + return getindex_nonzero(a, I) +end + +# `SparseArrayInterface` interface +function setindex_zero!(a::SparseArrayDOK, value, I) + # TODO: This is specific to the `Dictionaries.jl` + # interface, make more generic? + set!(nonzero_structure(a), I, value) + return a +end + +function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N} + a[CartesianIndex(I)] = value + return a +end + +function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N}) where {N} + if !is_structural_nonzero(a, I) + setindex_zero!(a, value, I) + end + setindex_nonzero!(a, value, I) + return a +end + +# similar +# TODO: How does this deal with the converting the zero type? +function Base.similar(a::SparseArrayDOK{T,N,Zero}) where {T,N,Zero} + return SparseArrayDOK{T,N,Zero}(undef, size(a), a.zero) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl new file mode 100644 index 0000000000..41f6f8f20b --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl @@ -0,0 +1,39 @@ +using Dictionaries: Dictionary + +# Also look into: +# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ + +# Required `SparseArrayInterface` interface. +# https://github.com/Jutho/SparseArrayKit.jl interface functions +nonzero_keys(a::AbstractArray) = error("Not implemented") +nonzero_values(a::AbstractArray) = error("Not implemented") +nonzero_pairs(a::AbstractArray) = error("Not implemented") + +# A dictionary-like structure +# TODO: Rename `nonzeros`, `structural_nonzeros`, etc.? +nonzero_structure(a::AbstractArray) = error("Not implemented") + +# Derived `SparseArrayInterface` interface. +nonzero_length(a::AbstractArray) = length(nonzero_keys(a)) +is_structural_nonzero(a::AbstractArray, I) = I ∈ nonzero_keys(a) + +# Overload if zero value is index dependent or +# doesn't match element type. +getindex_nonzero(a::AbstractArray, I) = nonzero_structure(a)[I] +getindex_zero(a::AbstractArray, I) = zero(eltype(a)) +function setindex_zero!(a::AbstractArray, value, I) + # TODO: This may need to be modified. + nonzero_structure(a)[I] = value + return a +end +function setindex_nonzero!(a::AbstractArray, value, I) + nonzero_structure(a)[I] = value + return a +end + +struct Zero end +(::Zero)(type, I) = zero(type) + +default_zero() = Zero() # (eltype, I) -> zero(eltype) +default_keytype(ndims::Int) = CartesianIndex{ndims} +default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() diff --git a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl index de0c6b8f5e..621f2e4bde 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl @@ -1,43 +1,53 @@ @eval module $(gensym()) -using Test: @test, @testset +using Test: @test, @testset, @test_broken using NDTensors.SparseArrayDOKs: SparseArrayDOK, nonzero_keys, nonzero_length using SparseArrays: nnz -@testset "SparseArrayDOK (eltype=$elt)" for elt in ( - Float32, ComplexF32, Float64, ComplexF64 -) - a = SparseArrayDOK{elt}(3, 4) - @test a == SparseArrayDOK{elt}((3, 4)) - @test a == SparseArrayDOK{elt}(undef, 3, 4) - @test a == SparseArrayDOK{elt}(undef, (3, 4)) - @test iszero(a) - @test iszero(nnz(a)) - @test nonzero_length(a) == nnz(a) - @test size(a) == (3, 4) - @test eltype(a) == elt - for I in eachindex(a) - @test iszero(a[I]) - @test a[I] isa elt - end - @test isempty(nonzero_keys(a)) +@testset "SparseArrayDOK (eltype=$elt)" for elt in + (Float32, ComplexF32, Float64, ComplexF64) + @testset "Basics" begin + a = SparseArrayDOK{elt}(3, 4) + @test a == SparseArrayDOK{elt}((3, 4)) + @test a == SparseArrayDOK{elt}(undef, 3, 4) + @test a == SparseArrayDOK{elt}(undef, (3, 4)) + @test iszero(a) + @test iszero(nnz(a)) + @test nonzero_length(a) == nnz(a) + @test size(a) == (3, 4) + @test eltype(a) == elt + for I in eachindex(a) + @test iszero(a[I]) + @test a[I] isa elt + end + @test isempty(nonzero_keys(a)) - x12 = randn(elt) - x23 = randn(elt) - b = copy(a) - @test b isa SparseArrayDOK{elt} - @test iszero(b) - b[1, 2] = x12 - b[2, 3] = x23 - @test iszero(a) - @test !iszero(b) - @test b[1, 2] == x12 - @test b[2, 3] == x23 + x12 = randn(elt) + x23 = randn(elt) + b = copy(a) + @test b isa SparseArrayDOK{elt} + @test iszero(b) + b[1, 2] = x12 + b[2, 3] = x23 + @test iszero(a) + @test !iszero(b) + @test b[1, 2] == x12 + @test b[2, 3] == x23 + @test iszero(nonzero_length(a)) + @test_broken nonzero_length(b) == 2 - # To test: - # reshape - # zero (PermutedDimsArray) - # map[!] - # broadcast - # Custom zero type - # conversion to `SparseMatrixCSC` + # To test: + # reshape + # zero (PermutedDimsArray) + # map[!] + # broadcast + # Custom zero type + # conversion to `SparseMatrixCSC` + end + @testset "map/broadcast" begin + a = SparseArrayDOK{elt}(3, 4) + a[1, 1] = 11 + a[3, 4] = 34 + @test nonzero_length(a) == 2 + 2 * a + end end end diff --git a/src/imports.jl b/src/imports.jl index bea3457d4d..57a1fb73b3 100644 --- a/src/imports.jl +++ b/src/imports.jl @@ -110,8 +110,7 @@ import LinearAlgebra: tr, transpose -using ITensors.NDTensors.Unwrap: - cpu +using ITensors.NDTensors.Unwrap: cpu using ITensors.NDTensors: Algorithm, From 0677dfab0baf5e6184c488488242a8a14abc6b6e Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sun, 26 Nov 2023 15:13:56 -0500 Subject: [PATCH 03/17] Format --- .../lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl | 1 - NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl | 4 +--- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl index 8bbb6121fd..93a1104bf0 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl @@ -2,4 +2,3 @@ using SparseArrays: SparseArrays # Julia Base `AbstractSparseArray` interface SparseArrays.nnz(a::SparseArrayDOK) = nonzero_length(a) - diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl index 454fc09a7d..4e972117f1 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl @@ -35,9 +35,7 @@ function SparseArrayDOK{T,N,Zero}( return SparseArrayDOK{T,N,Zero}(dims, zero) end -function SparseArrayDOK{T,N}( - ::UndefInitializer, dims::Tuple{Vararg{Int}}, zero -) where {T,N} +function SparseArrayDOK{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T,N} return SparseArrayDOK{T,N}(dims, zero) end From 995c3f9788ed33d6e10ba9db67a64c563d73b892 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 08:42:49 -0500 Subject: [PATCH 04/17] Start making SparseArrayInterface module, to be used by SparseArrayDOK and DiagonalArray --- NDTensors/src/NDTensors.jl | 1 + .../lib/DiagonalArrays/src/DiagonalArrays.jl | 11 + .../SparseArrayDOKs/src/SparseArrayDOKs.jl | 3 + .../src/lib/SparseArrayDOKs/src/broadcast.jl | 38 ++ .../src/lib/SparseArrayDOKs/src/copyto.jl | 12 + NDTensors/src/lib/SparseArrayDOKs/src/map.jl | 12 + .../lib/SparseArrayDOKs/src/sparsearraydok.jl | 24 +- .../src/lib/SparseArrayInterface/README.md | 59 +++ .../src/SparseArrayInterface.jl | 11 + .../SparseArrayInterfaceSparseArraysExt.jl | 10 + .../SparseArrayInterface/src/backup/basics.jl | 45 ++ .../src/backup/broadcast.jl | 38 ++ .../src/lib/SparseArrayInterface/src/base.jl | 101 +++++ .../lib/SparseArrayInterface/src/broadcast.jl | 37 ++ .../lib/SparseArrayInterface/src/copyto.jl | 15 + .../lib/SparseArrayInterface/src/indexing.jl | 57 +++ .../lib/SparseArrayInterface/src/interface.jl | 17 + .../src/interface_optional.jl | 12 + .../src/lib/SparseArrayInterface/src/map.jl | 12 + .../lib/SparseArrayInterface/src/wrappers.jl | 15 + .../SparseArrayInterface/test/Project.toml | 4 + .../lib/SparseArrayInterface/test/runtests.jl | 401 ++++++++++++++++++ .../src/lib/SparseArrayInterface/test/temp.jl | 105 +++++ 23 files changed, 1037 insertions(+), 3 deletions(-) create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/map.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/README.md create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterfaceSparseArraysExt.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/backup/basics.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/backup/broadcast.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/base.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/broadcast.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/copyto.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/indexing.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/interface.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/map.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/Project.toml create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/runtests.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/temp.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index a7f4efa37a..6d2284387b 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -27,6 +27,7 @@ for lib in [ :Unwrap, :RankFactorization, :TensorAlgebra, + :SparseArrayInterface, :DiagonalArrays, :SparseArrayDOKs, :BlockSparseArrays, diff --git a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl index 936ddf8843..2df6e33e4d 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl @@ -7,12 +7,15 @@ export DiagonalArray, DiagonalMatrix, DiagonalVector, DiagIndex, DiagIndices, de include("diagview.jl") +# TODO: Make use of `Zero` type in `SparseArrayInterface`. struct DefaultZero end function (::DefaultZero)(eltype::Type, I::CartesianIndex) return zero(eltype) end +# TODO: Rename `diag` to `nonzero_values`, to make more +# generic? struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractArray{T,N} diag::Diag dims::NTuple{N,Int} @@ -20,6 +23,7 @@ struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractArray{T,N} end # TODO: Use `Accessors.jl`. +# TODO: Make more generic, call `set_nonzero_values`. set_diag(a::DiagonalArray, diag) = DiagonalArray(diag, size(a), a.zero) Base.copy(a::DiagonalArray) = set_diag(a, copy(a[DiagIndices()])) @@ -27,6 +31,7 @@ Base.similar(a::DiagonalArray) = set_diag(a, similar(a[DiagIndices()])) Base.size(a::DiagonalArray) = a.dims +# TODO: Rename `nonzer_values`. diagview(a::DiagonalArray) = a.diag LinearAlgebra.diag(a::DiagonalArray) = copy(diagview(a)) @@ -92,6 +97,7 @@ function DiagonalArray{T}(::UndefInitializer, d::Vararg{Int,N}) where {T,N} return DiagonalArray{T,N}(undef, d) end +# TODO: Use `sparse_getindex` from `SparseArrayInterface`. function Base.getindex(a::DiagonalArray{T,N}, I::CartesianIndex{N}) where {T,N} i = diagindex(a, I) isnothing(i) && return a.zero(T, I) @@ -102,6 +108,7 @@ function Base.getindex(a::DiagonalArray{T,N}, I::Vararg{Int,N}) where {T,N} return a[CartesianIndex(I)] end +# TODO: Use `sparse_setindex!` from `SparseArrayInterface`. function Base.setindex!(a::DiagonalArray{T,N}, v, I::CartesianIndex{N}) where {T,N} i = diagindex(a, I) isnothing(i) && return error("Can't set off-diagonal element of DiagonalArray") @@ -119,11 +126,14 @@ function densearray(a::DiagonalArray) # TODO: Check this works on GPU. # TODO: Make use of `a.zero`? d = similar(diagview(a), size(a)) + # TODO: Use `getindex_zero(a)` or `zero_value(a)`. fill!(d, zero(eltype(a))) diagcopyto!(d, a) return d end +# TODO: Make generic in `SparseArrayInterface`, use a `SparseIndices(:)` +# iterator? function Base.permutedims!(a_dest::AbstractArray, a_src::DiagonalArray, perm) @assert ndims(a_src) == ndims(a_dest) == length(perm) a_dest[DiagIndices()] = a_src[DiagIndices()] @@ -151,6 +161,7 @@ function Base.map(f, a::DiagonalArray) return map_nonzeros(f, a) end +# TODO: Make this generic in `SparseArrayInterface`. function map_diag(f, a::DiagonalArray) return set_diag(a, map(f, a[DiagIndices()])) end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl index d1e9a4ebd3..3c648f7cba 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl @@ -1,5 +1,8 @@ module SparseArrayDOKs include("sparsearrayinterface.jl") include("sparsearraydok.jl") +include("broadcast.jl") +include("map.jl") +include("copyto.jl") include("SparseArrayDOKsSparseArraysExt.jl") end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl b/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl new file mode 100644 index 0000000000..c539aeb6b3 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl @@ -0,0 +1,38 @@ +using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted +using ..BroadcastMapConversion: map_function, map_args + +struct SparseArrayDOKStyle{N} <: AbstractArrayStyle{N} end + +function Broadcast.BroadcastStyle(::Type{<:SparseArrayDOK{<:Any,N}}) where {N} + return SparseArrayDOKStyle{N}() +end + +SparseArrayDOKStyle(::Val{N}) where {N} = SparseArrayDOKStyle{N}() +SparseArrayDOKStyle{M}(::Val{N}) where {M,N} = SparseArrayDOKStyle{N}() + +Broadcast.BroadcastStyle(a::SparseArrayDOKStyle, ::DefaultArrayStyle{0}) = a +function Broadcast.BroadcastStyle(::SparseArrayDOKStyle{N}, a::DefaultArrayStyle) where {N} + return BroadcastStyle(DefaultArrayStyle{N}(), a) +end +function Broadcast.BroadcastStyle( + ::SparseArrayDOKStyle{N}, ::Broadcast.Style{Tuple} +) where {N} + return DefaultArrayStyle{N}() +end + +# TODO: Use `allocate_output`, share logic with `map`. +function Base.similar(bc::Broadcasted{<:SparseArrayDOKStyle}, elt::Type) + # TODO: Is this a good definition? Probably should check that + # they have consistent axes. + return similar(first(map_args(bc)), elt) +end + +# Broadcasting implementation +function Base.copyto!( + dest::SparseArrayDOK{<:Any,N}, bc::Broadcasted{SparseArrayDOKStyle{N}} +) where {N} + # convert to map + # flatten and only keep the AbstractArray arguments + map!(map_function(bc), dest, map_args(bc)...) + return dest +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl b/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl new file mode 100644 index 0000000000..61c3c038b3 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl @@ -0,0 +1,12 @@ +# TODO: Rename `copy_nonzeros!`, move to `sparsearrayinterface.jl`. +function Base.copy!(dest::AbstractArray, src::SparseArrayDOK) + @assert axes(dest) == axes(src) + map!(identity, dest, src) + return dest +end + +# TODO: Rename `copyto_nonzeros!`, move to `sparsearrayinterface.jl`. +function Base.copyto!(dest::AbstractArray, src::SparseArrayDOK) + copy!(dest, src) + return dest +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/map.jl b/NDTensors/src/lib/SparseArrayDOKs/src/map.jl new file mode 100644 index 0000000000..673fb7fee8 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/map.jl @@ -0,0 +1,12 @@ +# TODO: Rename the preserve zero case to `map_nonzeros!`, +# move to `sparsearrayinterface.jl`. +function Base.map!(f, a_dest::AbstractArray, as::SparseArrayDOK...) + # TODO: Handle nonzero case, fill all values. + # Also, handle arrays of arrays better, using `getindex_zero(a, first(eachindex(a)))`. + @assert iszero(f(map(a -> getindex_zero(a, first(eachindex(a))), as)...)) + Is = union(nonzero_keys.(as)...) + for I in Is + a_dest[I] = f(map(a -> a[I], as)...) + end + return a_dest +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl index 4e972117f1..a0ded0cc9c 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl @@ -28,6 +28,11 @@ function SparseArrayDOK{T}(dims::Int...) where {T} return SparseArrayDOK{T}(dims) end +# Specify zero function +function SparseArrayDOK{T}(dims::Tuple{Vararg{Int}}, zero) where {T} + return SparseArrayDOK{T,length(dims)}(dims, zero) +end + # undef function SparseArrayDOK{T,N,Zero}( ::UndefInitializer, dims::Tuple{Vararg{Int}}, zero @@ -51,6 +56,10 @@ function SparseArrayDOK{T}(::UndefInitializer, dims::Int...) where {T} return SparseArrayDOK{T}(dims...) end +function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T} + return SparseArrayDOK{T}(dims, zero) +end + # Required `SparseArrayInterface` interface nonzero_structure(a::SparseArrayDOK) = a.data # TODO: Make this a generic function. @@ -67,6 +76,7 @@ function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N} return a[CartesianIndex(I)] end +# TODO: Rename `getindex_sparse` or `sparse_getindex`. function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N}) where {N} if !is_structural_nonzero(a, I) return getindex_zero(a, I) @@ -96,7 +106,15 @@ function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N}) end # similar -# TODO: How does this deal with the converting the zero type? -function Base.similar(a::SparseArrayDOK{T,N,Zero}) where {T,N,Zero} - return SparseArrayDOK{T,N,Zero}(undef, size(a), a.zero) +function Base.similar(a::SparseArrayDOK) + return similar(a, eltype(a)) +end + +function Base.similar(a::SparseArrayDOK, elt::Type) + return similar(a, elt, size(a)) +end + +function Base.similar(a::SparseArrayDOK, elt::Type, dims::Tuple{Vararg{Int}}) + # TODO: Accessor function for `a.zero`. + return SparseArrayDOK{elt}(undef, dims, a.zero) end diff --git a/NDTensors/src/lib/SparseArrayInterface/README.md b/NDTensors/src/lib/SparseArrayInterface/README.md new file mode 100644 index 0000000000..0f434c8746 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/README.md @@ -0,0 +1,59 @@ +# SparseArrayInterface + +Defines a generic interface for sparse arrays in Julia. + +The minimal interface is: +```julia +nonzeros(a::AbstractArray) = ... +nonzero_index_to_index(a::AbstractArray, Inz) = ... +index_to_nonzero_index(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} = ... +Broadcast.BroadcastStyle(arraytype::Type{<:AbstractArray}) = SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() +``` +Once these are defined, along with Julia AbstractArray interface functions like +`Base.size` and `Base.similar`, functions like the following will take advantage of sparsity: +```julia +SparseArrayInterface.nonzero_length # SparseArrays.nnz +SparseArrayInterface.sparse_getindex +SparseArrayInterface.sparse_setindex! +SparseArrayInterface.sparse_map! +SparseArrayInterface.sparse_copy! +SparseArrayInterface.sparse_copyto! +SparseArrayInterface.sparse_permutedims! +``` +which can be used to define the corresponding `Base` functions. + +## TODO +Still need to implement `Base` functions: +```julia +[x] sparse_zero(a::AbstractArray) = similar(a) +[x] sparse_iszero(a::AbstractArray) = iszero(nonzero_length(a)) # Uses `all`, make `sparse_all`? +[x] sparse_one(a::AbstractArray) = ... +[x] sparse_isreal(a::AbstractArray) = ... # Uses `all`, make `sparse_all`? +[x] sparse_isequal(a1::AbstractArray, a2::AbstractArray) = ... +[x] sparse_conj!(a::AbstractArray) = conj!(nonzeros(a)) +[ ] sparse_all(f, a::AbstractArray) = ... +[x] sparse_reshape(a::AbstractArray, dims) = ... +``` +`LinearAlgebra` functions: +```julia +[ ] sparse_mul! +[ ] sparse_lmul! +[ ] sparse_ldiv! +[ ] sparse_rdiv! +[ ] sparse_axpby! +[ ] sparse_axpy! +[ ] sparse_norm +[ ] sparse_dot/sparse_inner +[ ] sparse_adoint! +[ ] sparse_transpose! + +# Using conversion to `SparseMatrixCSC`: +[ ] sparse_qr +[ ] sparse_eigen +[ ] sparse_svd +``` +`TensorAlgebra` functions: +```julia +[ ] add! +[ ] contract! +``` diff --git a/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl new file mode 100644 index 0000000000..afcde0bd48 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl @@ -0,0 +1,11 @@ +module SparseArrayInterface +include("interface.jl") +include("interface_optional.jl") +include("indexing.jl") +include("base.jl") +include("map.jl") +include("copyto.jl") +include("broadcast.jl") +include("wrappers.jl") +include("SparseArrayInterfaceSparseArraysExt.jl") +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterfaceSparseArraysExt.jl b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterfaceSparseArraysExt.jl new file mode 100644 index 0000000000..0f2d5d6b7f --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterfaceSparseArraysExt.jl @@ -0,0 +1,10 @@ +## using SparseArrays: SparseArrays, SparseMatrixCSC +## +## # Julia Base `SparseArrays.AbstractSparseArray` interface +## # SparseMatrixCSC.nnz +## nnz(a::AbstractArray) = nonzero_length(a) +## +## # SparseArrayInterface.SparseMatrixCSC +## function SparseMatrixCSC(a::AbstractMatrix) +## return error("Not implemented") +## end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/backup/basics.jl b/NDTensors/src/lib/SparseArrayInterface/src/backup/basics.jl new file mode 100644 index 0000000000..7c408a2afe --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/backup/basics.jl @@ -0,0 +1,45 @@ +# Also look into: +# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ + +# Minimal interface +# Data structure storing the nonzero values +nonzeros(a::AbstractArray) = a + +# Minimal interface +# Map a `CartesianIndex` to an index into `nonzeros`. +nonzero_index(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} = I + +## # Required `SparseArrayInterface` interface. +## # https://github.com/Jutho/SparseArrayKit.jl interface functions +## nonzero_keys(a::AbstractArray) = error("Not implemented") +## nonzero_values(a::AbstractArray) = error("Not implemented") +## nonzero_pairs(a::AbstractArray) = error("Not implemented") +## +## # A dictionary-like structure +## # TODO: Rename `nonzeros`, `structural_nonzeros`, etc.? +## nonzero_structure(a::AbstractArray) = error("Not implemented") +## +## # Derived `SparseArrayInterface` interface. +## nonzero_length(a::AbstractArray) = length(nonzero_keys(a)) +## is_structural_nonzero(a::AbstractArray, I) = I ∈ nonzero_keys(a) +## +## # Overload if zero value is index dependent or +## # doesn't match element type. +## getindex_nonzero(a::AbstractArray, I) = nonzero_structure(a)[I] +## getindex_zero(a::AbstractArray, I) = zero(eltype(a)) +## function setindex_zero!(a::AbstractArray, value, I) +## # TODO: This may need to be modified. +## nonzero_structure(a)[I] = value +## return a +## end +## function setindex_nonzero!(a::AbstractArray, value, I) +## nonzero_structure(a)[I] = value +## return a +## end +## +## struct Zero end +## (::Zero)(type, I) = zero(type) +## +## default_zero() = Zero() # (eltype, I) -> zero(eltype) +## default_keytype(ndims::Int) = CartesianIndex{ndims} +## default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() diff --git a/NDTensors/src/lib/SparseArrayInterface/src/backup/broadcast.jl b/NDTensors/src/lib/SparseArrayInterface/src/backup/broadcast.jl new file mode 100644 index 0000000000..c539aeb6b3 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/backup/broadcast.jl @@ -0,0 +1,38 @@ +using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted +using ..BroadcastMapConversion: map_function, map_args + +struct SparseArrayDOKStyle{N} <: AbstractArrayStyle{N} end + +function Broadcast.BroadcastStyle(::Type{<:SparseArrayDOK{<:Any,N}}) where {N} + return SparseArrayDOKStyle{N}() +end + +SparseArrayDOKStyle(::Val{N}) where {N} = SparseArrayDOKStyle{N}() +SparseArrayDOKStyle{M}(::Val{N}) where {M,N} = SparseArrayDOKStyle{N}() + +Broadcast.BroadcastStyle(a::SparseArrayDOKStyle, ::DefaultArrayStyle{0}) = a +function Broadcast.BroadcastStyle(::SparseArrayDOKStyle{N}, a::DefaultArrayStyle) where {N} + return BroadcastStyle(DefaultArrayStyle{N}(), a) +end +function Broadcast.BroadcastStyle( + ::SparseArrayDOKStyle{N}, ::Broadcast.Style{Tuple} +) where {N} + return DefaultArrayStyle{N}() +end + +# TODO: Use `allocate_output`, share logic with `map`. +function Base.similar(bc::Broadcasted{<:SparseArrayDOKStyle}, elt::Type) + # TODO: Is this a good definition? Probably should check that + # they have consistent axes. + return similar(first(map_args(bc)), elt) +end + +# Broadcasting implementation +function Base.copyto!( + dest::SparseArrayDOK{<:Any,N}, bc::Broadcasted{SparseArrayDOKStyle{N}} +) where {N} + # convert to map + # flatten and only keep the AbstractArray arguments + map!(map_function(bc), dest, map_args(bc)...) + return dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/base.jl new file mode 100644 index 0000000000..6dd1b9598f --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/base.jl @@ -0,0 +1,101 @@ +# TODO: Use `sparse_mapreduce`? +function sparse_iszero(a::AbstractArray) + iszero(nonzero_length(a)) && return true + return iszero(nonzeros(a)) +end + +# TODO: Use `sparse_mapreduce`? +function sparse_isreal(a::AbstractArray) + iszero(nonzero_length(a)) && return true + return isreal(nonzeros(a)) +end + +function fill_nonzero!(a::AbstractArray, x) + for I in eachindex(a) + a[I] = x + end + return a +end + +function fill_zero!(a::AbstractArray, x) + fill!(nonzeros(a), x) + return a +end + +function sparse_fill!(a::AbstractArray, x) + if !iszero(x) + # TODO: Reverse naming convention? + fill_nonzero!(a, x) + else + fill_zero!(a, x) + end + return a +end + +function sparse_zero(a::AbstractArray) + # Need to overload `similar` for custom types + a = similar(a) + sparse_fill!(a, zero(eltype(a))) + return a +end + +## function sparse_zero(a::AbstractArray, elt::Type) +## return sparse_zero(a, elt, size(a)) +## end +## +## function sparse_zero(a::AbstractArray, dims::Tuple) +## return sparse_zero(a, eltype(a), dims) +## end +## +## function sparse_zero(a::AbstractArray) +## return sparse_zero(a, eltype(a), size(a)) +## end + +# TODO: Make `sparse_one!`? +function sparse_one(a::AbstractMatrix) + m, n = size(a) + @assert m == n + a = zero(a) + # TODO: Use `diagindices` from `DiagonalArrays`? + for i in 1:m + a[i, i] = one(eltype(a)) + end + return a +end + +# TODO: Use `sparse_mapreduce`? +function sparse_isequal(a1::AbstractArray, a2::AbstractArray) + Is = collect(nonzero_indices(a1)) + intersect!(Is, nonzero_indices(a2)) + if !(length(Is) == nonzero_length(a1) == nonzero_length(a2)) + return false + end + for I in Is + a1[I] == a2[I] || return false + end + return true +end + +function sparse_reshape(a::AbstractArray, dims) + @assert length(a) == prod(dims) + a_reshaped = similar(a, dims) + sparse_fill!(a_reshaped, zero(eltype(a))) + linear_inds = LinearIndices(a) + new_cartesian_inds = CartesianIndices(dims) + for I in nonzero_indices(a) + a_reshaped[new_cartesian_inds[linear_inds[I]]] = a[I] + end + return a_reshaped +end + +## function sparse_mapreduce(f, op, as::AbstractArray...) +## # TODO: Make more general. +## @assert iszero(mapreduce(f, op, map(a -> getindex_zero(a, first(eachindex(a))), as)...)) +## # TODO: Create a `promote_nonzero_indices`, for example to help +## # with specialized fast indexing for `DiagonalArrays`. +## Is = union(nonzero_indices.(as)...) +## for I in Is +## a_dest[I] = f(map(a -> a[I], as)...) +## end +## return reduce(op, a_dest) +## end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/broadcast.jl b/NDTensors/src/lib/SparseArrayInterface/src/broadcast.jl new file mode 100644 index 0000000000..d113932d44 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/broadcast.jl @@ -0,0 +1,37 @@ +using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted +using ..BroadcastMapConversion: map_function, map_args + +struct SparseArrayStyle{N} <: AbstractArrayStyle{N} end + +# Define for new sparse array types. +# function Broadcast.BroadcastStyle(arraytype::Type{<:MySparseArray}) +# return SparseArrayStyle{ndims(arraytype)}() +# end + +SparseArrayStyle(::Val{N}) where {N} = SparseArrayStyle{N}() +SparseArrayStyle{M}(::Val{N}) where {M,N} = SparseArrayStyle{N}() + +Broadcast.BroadcastStyle(a::SparseArrayStyle, ::DefaultArrayStyle{0}) = a +function Broadcast.BroadcastStyle(::SparseArrayStyle{N}, a::DefaultArrayStyle) where {N} + return BroadcastStyle(DefaultArrayStyle{N}(), a) +end +function Broadcast.BroadcastStyle(::SparseArrayStyle{N}, ::Broadcast.Style{Tuple}) where {N} + return DefaultArrayStyle{N}() +end + +# TODO: Use `allocate_output`, share logic with `map`. +function Base.similar(bc::Broadcasted{<:SparseArrayStyle}, elt::Type) + # TODO: Is this a good definition? Probably should check that + # they have consistent axes. + return similar(first(map_args(bc)), elt) +end + +# Broadcasting implementation +function Base.copyto!( + dest::AbstractArray{<:Any,N}, bc::Broadcasted{SparseArrayStyle{N}} +) where {N} + # convert to map + # flatten and only keep the AbstractArray arguments + sparse_map!(map_function(bc), dest, map_args(bc)...) + return dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/copyto.jl b/NDTensors/src/lib/SparseArrayInterface/src/copyto.jl new file mode 100644 index 0000000000..218502f3d9 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/copyto.jl @@ -0,0 +1,15 @@ +function sparse_copy!(dest::AbstractArray, src::AbstractArray) + @assert axes(dest) == axes(src) + sparse_map!(identity, dest, src) + return dest +end + +function sparse_copyto!(dest::AbstractArray, src::AbstractArray) + sparse_map!(identity, dest, src) + return dest +end + +function sparse_permutedims!(dest::AbstractArray, src::AbstractArray, perm) + sparse_copyto!(dest, PermutedDimsArray(src, perm)) + return dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl new file mode 100644 index 0000000000..1522eb3d99 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl @@ -0,0 +1,57 @@ +nonzero_length(a::AbstractArray) = length(nonzeros(a)) +function nonzero_indices(a::AbstractArray) + return Iterators.map(Inz -> nonzero_index_to_index(a, Inz), eachindex(nonzeros(a))) +end + +# Derived +function index_to_nonzero_index(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} + return index_to_nonzero_index(a, CartesianIndex(I)) +end + +# Helper type for constructing zero values +struct GetIndexZero end +(::GetIndexZero)(a::AbstractArray, I) = zero(eltype(a)) + +function getindex_nonzero(a::AbstractArray, Inz) + return nonzeros(a)[Inz] +end + +function sparse_getindex(a::AbstractArray, I...) + return sparse_getindex(a, CartesianIndex(to_indices(a, I))) +end + +function sparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} + return sparse_getindex(a, CartesianIndex(I)) +end + +function sparse_getindex(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} + Inz = index_to_nonzero_index(a, I) + if isnothing(Inz) + return getindex_zero(a, I) + end + return getindex_nonzero(a, Inz) +end + +# Update a nonzero value +function setindex_nonzero!(a::AbstractArray, value, Inz) + nonzeros(a)[Inz] = value + return a +end + +function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}) where {N} + sparse_setindex!(a, value, CartesianIndex(I)) + return a +end + +function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::CartesianIndex{N}) where {N} + Inz = index_to_nonzero_index(a, I) + if isnothing(Inz) + if !iszero(value) + # Only try setting if nonzero + setindex_zero!(a, value, I) + end + else + setindex_nonzero!(a, value, Inz) + end + return a +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface.jl new file mode 100644 index 0000000000..1a30d1797e --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface.jl @@ -0,0 +1,17 @@ +# Also look into: +# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ + +# Minimal interface +# Data structure storing the nonzero values +nonzeros(a::AbstractArray) = a + +# Minimal interface +# Map an index in the nonzero data to a CartesianIndex of the +# outer array. +nonzero_index_to_index(a::AbstractArray, Inz) = Inz + +# Minimal interface +# Map a `CartesianIndex` to an index/key into the nonzero data structure +# returned by `nonzeros`. +# Return `nothing` if the index corresponds to a structural zero value. +index_to_nonzero_index(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} = I diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl new file mode 100644 index 0000000000..e2c53cf6eb --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl @@ -0,0 +1,12 @@ +# Optional interface. +# Access a zero value. +function getindex_zero(a::AbstractArray, I) + return zero(eltype(a)) +end + +# Optional interface. +# Insert a new zero value. +# Some types (like `Diagonal`) may not support this. +function setindex_zero!(a::AbstractArray, value, I) + return error("Not implemented") +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/map.jl b/NDTensors/src/lib/SparseArrayInterface/src/map.jl new file mode 100644 index 0000000000..40ad7ed0af --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/map.jl @@ -0,0 +1,12 @@ +# TODO: Rename the preserve zero case to `map_nonzeros!`. +function sparse_map!(f, a_dest::AbstractArray, as::AbstractArray...) + # TODO: Handle nonzero case, fill all values. + @assert iszero(f(map(a -> getindex_zero(a, first(eachindex(a))), as)...)) + # TODO: Create a `promote_nonzero_indices`, for example to help + # with specialized fast indexing for `DiagonalArrays`. + Is = union(nonzero_indices.(as)...) + for I in Is + a_dest[I] = f(map(a -> a[I], as)...) + end + return a_dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl b/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl new file mode 100644 index 0000000000..4ab140de57 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl @@ -0,0 +1,15 @@ +perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P +genperm(v, perm) = map(j -> v[j], perm) +genperm(v::CartesianIndex, perm) = CartesianIndex(map(j -> Tuple(v)[j], perm)) + +nonzeros(a::PermutedDimsArray) = nonzeros(parent(a)) +function nonzero_index_to_index(a::PermutedDimsArray, Inz) + return genperm(nonzero_index_to_index(parent(a), Inz), perm(a)) +end +function index_to_nonzero_index( + a::PermutedDimsArray{<:Any,N}, I::CartesianIndex{N} +) where {N} + return nonzero_index_to_index(parent(a), genperm(I, perm(a))) +end + +# TODO: Add `SubArray`, `ReshapedArray`, `Diagonal`, etc. diff --git a/NDTensors/src/lib/SparseArrayInterface/test/Project.toml b/NDTensors/src/lib/SparseArrayInterface/test/Project.toml new file mode 100644 index 0000000000..a36a39a20a --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/Project.toml @@ -0,0 +1,4 @@ +[deps] +NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl new file mode 100644 index 0000000000..4eaef6082e --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -0,0 +1,401 @@ +@eval module $(gensym()) +using Test: @test, @testset, @test_broken, @test_throws + +@testset "SparseArrayInterface (eltype=$elt)" for elt in + (Float32, ComplexF32, Float64, ComplexF64) + @testset "Array" begin + using NDTensors.SparseArrayInterface: SparseArrayInterface + a = randn(2, 3) + @test SparseArrayInterface.nonzeros(a) == a + @test SparseArrayInterface.index_to_nonzero_index(a, CartesianIndex(1, 2)) == + CartesianIndex(1, 2) + @test SparseArrayInterface.nonzero_index_to_index(a, CartesianIndex(1, 2)) == + CartesianIndex(1, 2) + end + @testset "Custom SparseArray" begin + using LinearAlgebra: norm + using NDTensors.SparseArrayInterface: SparseArrayInterface + struct SparseArray{T,N} <: AbstractArray{T,N} + data::Vector{T} + dims::Tuple{Vararg{Int,N}} + index_to_dataindex::Dict{CartesianIndex{N},Int} + dataindex_to_index::Vector{CartesianIndex{N}} + end + SparseArray{T,N}(dims::Tuple{Vararg{Int,N}}) where {T,N} = SparseArray{T,N}( + T[], dims, Dict{CartesianIndex{N},Int}(), Vector{CartesianIndex{N}}() + ) + 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.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.nonzeros(a::SparseArray) = a.data + function SparseArrayInterface.index_to_nonzero_index( + a::SparseArray{<:Any,N}, I::CartesianIndex{N} + ) where {N} + return get(a.index_to_dataindex, I, nothing) + end + SparseArrayInterface.nonzero_index_to_index(a::SparseArray, I) = a.dataindex_to_index[I] + function SparseArrayInterface.setindex_zero!( + 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 + + # Broadcasting + function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() + end + + # Base + function Base.iszero(a::SparseArray) + return SparseArrayInterface.sparse_iszero(a) + end + function Base.isreal(a::SparseArray) + return SparseArrayInterface.sparse_isreal(a) + end + function Base.zero(a::SparseArray) + return SparseArrayInterface.sparse_zero(a) + end + function Base.one(a::SparseArray) + return SparseArrayInterface.sparse_one(a) + end + function Base.:(==)(a1::SparseArray, a2::SparseArray) + return SparseArrayInterface.sparse_isequal(a1, a2) + end + function Base.reshape(a::SparseArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) + end + + # Map + function Base.map!(f, dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest + end + function Base.copy!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copy!(dest, src) + return dest + end + function Base.copyto!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest + end + function Base.permutedims!(dest::AbstractArray, src::SparseArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest + end + + # TODO: Test `fill!`. + + # Test + a = SparseArray{elt}(2, 3) + @test size(a) == (2, 3) + @test axes(a) == (1:2, 1:3) + @test SparseArrayInterface.nonzeros(a) == elt[] + @test iszero(SparseArrayInterface.nonzero_length(a)) + @test collect(SparseArrayInterface.nonzero_indices(a)) == CartesianIndex{2}[] + @test iszero(a) + @test iszero(norm(a)) + for I in eachindex(a) + @test iszero(a) + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + @test size(a) == (2, 3) + @test axes(a) == (1:2, 1:3) + @test SparseArrayInterface.nonzeros(a) == elt[12] + @test isone(SparseArrayInterface.nonzero_length(a)) + @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test !iszero(a) + @test !iszero(norm(a)) + for I in eachindex(a) + if I == CartesianIndex(1, 2) + @test a[I] == 12 + else + @test iszero(a[I]) + end + end + + a = SparseArray{elt}(2, 2, 2) + a[1, 2, 2] = 122 + a_r = reshape(a, 2, 4) + @test a_r[1, 4] == a[1, 2, 2] == 122 + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + a = zero(a) + @test size(a) == (2, 3) + @test iszero(SparseArrayInterface.nonzero_length(a)) + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = SparseArray{elt}(2, 3) + b[2, 1] = 21 + @test a == a + @test a ≠ b + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + @test isreal(a) + + a = SparseArray{elt}(2, 3) + a[1, 2] = randn(elt) + b = copy(a) + conj!(b) + for I in eachindex(a) + @test conj(a[I]) == b[I] + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = randn(elt) + b = conj(a) + for I in eachindex(a) + @test conj(a[I]) == b[I] + end + + if !(elt <: Real) + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + 12im + @test !isreal(a) + end + + a = SparseArray{elt}(2, 2) + a[1, 2] = 12 + a = one(a) + @test size(a) == (2, 2) + @test isone(a[1, 1]) + @test isone(a[2, 2]) + @test iszero(a[1, 2]) + @test iszero(a[2, 1]) + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + a = zero(a) + @test size(a) == (2, 3) + @test iszero(SparseArrayInterface.nonzero_length(a)) + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + a = copy(a) + @test size(a) == (2, 3) + @test axes(a) == (1:2, 1:3) + @test SparseArrayInterface.nonzeros(a) == elt[12] + @test isone(SparseArrayInterface.nonzero_length(a)) + @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test !iszero(a) + @test !iszero(norm(a)) + for I in eachindex(a) + if I == CartesianIndex(1, 2) + @test a[I] == 12 + else + @test iszero(a[I]) + end + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + a = 2 * a + @test size(a) == (2, 3) + @test axes(a) == (1:2, 1:3) + @test SparseArrayInterface.nonzeros(a) == elt[24] + @test isone(SparseArrayInterface.nonzero_length(a)) + @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test !iszero(a) + @test !iszero(norm(a)) + for I in eachindex(a) + if I == CartesianIndex(1, 2) + @test a[I] == 24 + else + @test iszero(a[I]) + end + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = SparseArray{elt}(2, 3) + b[2, 1] = 21 + c = a + b + @test size(c) == (2, 3) + @test axes(c) == (1:2, 1:3) + @test SparseArrayInterface.nonzeros(c) == elt[12, 21] + @test SparseArrayInterface.nonzero_length(c) == 2 + @test collect(SparseArrayInterface.nonzero_indices(c)) == + [CartesianIndex(1, 2), CartesianIndex(2, 1)] + @test !iszero(c) + @test !iszero(norm(c)) + for I in eachindex(c) + if I == CartesianIndex(1, 2) + @test c[I] == 12 + elseif I == CartesianIndex(2, 1) + @test c[I] == 21 + else + @test iszero(c[I]) + end + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = permutedims(a, (2, 1)) + @test size(b) == (3, 2) + @test axes(b) == (1:3, 1:2) + @test SparseArrayInterface.nonzeros(b) == elt[12] + @test SparseArrayInterface.nonzero_length(b) == 1 + @test collect(SparseArrayInterface.nonzero_indices(b)) == [CartesianIndex(2, 1)] + @test !iszero(b) + @test !iszero(norm(b)) + for I in eachindex(b) + if I == CartesianIndex(2, 1) + @test b[I] == 12 + else + @test iszero(b[I]) + end + end + end + @testset "Custom DiagonalArray" begin + struct DiagonalArray{T,N} <: AbstractArray{T,N} + data::Vector{T} + dims::Tuple{Vararg{Int,N}} + end + DiagonalArray{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int,N}}) where {T,N} = + DiagonalArray{T,N}(Vector{T}(undef, minimum(dims)), dims) + DiagonalArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} = + DiagonalArray{T,N}(undef, dims) + DiagonalArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} = + DiagonalArray{T,length(dims)}(undef, dims) + DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int}) where {T} = + DiagonalArray{T}(undef, dims) + + # AbstractArray interface + Base.size(a::DiagonalArray) = a.dims + function Base.getindex(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_getindex(a, I...) + end + function Base.setindex!(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_setindex!(a, I...) + end + function Base.similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) + return DiagonalArray{elt}(undef, dims) + end + + # Minimal interface + SparseArrayInterface.nonzeros(a::DiagonalArray) = a.data + function SparseArrayInterface.index_to_nonzero_index( + a::DiagonalArray{<:Any,N}, I::CartesianIndex{N} + ) where {N} + !allequal(Tuple(I)) && return nothing + return first(Tuple(I)) + end + function SparseArrayInterface.nonzero_index_to_index(a::DiagonalArray, I) + return CartesianIndex(ntuple(Returns(I), ndims(a))) + end + function SparseArrayInterface.setindex_zero!( + a::DiagonalArray{<:Any,N}, value, I::CartesianIndex{N} + ) where {N} + return throw(ArgumentError("Can't set off-diagonal element of DiagonalArray")) + end + function SparseArrayInterface.sparse_zero(a::DiagonalArray, elt::Type, dims::Tuple) + # TODO: Need to make generic to GPU. + return DiagonalArray(zeros(elt, minimum(dims)), dims) + end + + # Broadcasting + function Broadcast.BroadcastStyle(arraytype::Type{<:DiagonalArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() + end + + # Base + function Base.iszero(a::DiagonalArray) + return SparseArrayInterface.sparse_iszero(a) + end + function Base.isreal(a::DiagonalArray) + return SparseArrayInterface.sparse_isreal(a) + end + function Base.zero(a::DiagonalArray) + return SparseArrayInterface.sparse_zero(a) + end + function Base.one(a::DiagonalArray) + return SparseArrayInterface.sparse_one(a) + end + function Base.:(==)(a1::DiagonalArray, a2::DiagonalArray) + return SparseArrayInterface.sparse_isequal(a1, a2) + end + function Base.reshape(a::DiagonalArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) + end + + # Map + function Base.map!(f, dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest + end + function Base.copy!(dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_copy!(dest, src) + return dest + end + function Base.copyto!(dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest + end + function Base.permutedims!(dest::AbstractArray, src::DiagonalArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest + end + + # TODO: Test `fill!`. + + # Test + a = DiagonalArray{elt}(undef, 2, 3) + @test size(a) == (2, 3) + a[1, 1] = 11 + a[2, 2] = 22 + @test a[1, 1] == 11 + @test a[2, 2] == 22 + @test_throws ArgumentError a[1, 2] = 12 + a[1, 2] = 0 + @test a[1, 1] == 11 + @test a[2, 2] == 22 + + b = similar(a) + @test b isa DiagonalArray + @test size(b) == (2, 3) + + a = DiagonalArray(elt[1, 2, 3], (3, 3)) + @test size(a) == (3, 3) + @test a[1, 1] == 1 + @test a[2, 2] == 2 + @test a[3, 3] == 3 + @test iszero(a[1, 2]) + + a = DiagonalArray(elt[1, 2, 3], (3, 3)) + a = 2 * a + @test size(a) == (3, 3) + @test a[1, 1] == 2 + @test a[2, 2] == 4 + @test a[3, 3] == 6 + @test iszero(a[1, 2]) + + a = DiagonalArray(elt[1, 2, 3], (3, 3)) + a_r = reshape(a, 9) + + a = DiagonalArray(elt[1, 2], (2, 2, 2)) + # TODO: Make `reshape` for `DiagonalArray` return `SparseArray`. + @test_broken reshape(a, 2, 4) + end +end +end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/temp.jl b/NDTensors/src/lib/SparseArrayInterface/test/temp.jl new file mode 100644 index 0000000000..f6265bbfa7 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/temp.jl @@ -0,0 +1,105 @@ +using LinearAlgebra: norm +using NDTensors.SparseArrayInterface: SparseArrayInterface +using Test: @test +struct SparseArray{T,N} <: AbstractArray{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.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.nonzeros(a::SparseArray) = a.data +function SparseArrayInterface.index_to_nonzero_index( + a::SparseArray{<:Any,N}, I::CartesianIndex{N} +) where {N} + return get(a.index_to_dataindex, I, nothing) +end +SparseArrayInterface.nonzero_index_to_index(a::SparseArray, I) = a.dataindex_to_index[I] +function SparseArrayInterface.setindex_zero!( + 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 + +# Base +function Base.zero(a::SparseArray) + return SparseArrayInterface.sparse_zero(a) +end + +# Broadcasting +function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() +end + +# Map +function Base.map!(f, dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest +end +function Base.copy!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copy!(dest, src) + return dest +end +function Base.copyto!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest +end +function Base.permutedims!(dest::AbstractArray, src::SparseArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest +end + +elt = Float64 + +# Test +a = SparseArray{elt}(2, 3) +@test size(a) == (2, 3) +@test axes(a) == (1:2, 1:3) +@test SparseArrayInterface.nonzeros(a) == elt[] +@test iszero(SparseArrayInterface.nonzero_length(a)) +@test collect(SparseArrayInterface.nonzero_indices(a)) == CartesianIndex{2}[] +@test iszero(a) +@test iszero(norm(a)) +for I in eachindex(a) + @test iszero(a) +end + +a[1, 2] = 12 +@test size(a) == (2, 3) +@test axes(a) == (1:2, 1:3) +@test SparseArrayInterface.nonzeros(a) == elt[12] +@test isone(SparseArrayInterface.nonzero_length(a)) +@test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] +@test !iszero(a) +@test !iszero(norm(a)) +for I in eachindex(a) + if I == CartesianIndex(1, 2) + @test a[I] == 12 + else + @test iszero(a[I]) + end +end From b696a7eb6582352529ca7248dd40a1339abe1f42 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 13:04:32 -0500 Subject: [PATCH 05/17] Update interface --- .../src/lib/SparseArrayInterface/src/base.jl | 45 +++---- .../lib/SparseArrayInterface/src/indexing.jl | 124 ++++++++++++++---- .../lib/SparseArrayInterface/src/interface.jl | 15 ++- .../src/interface_optional.jl | 17 ++- .../src/lib/SparseArrayInterface/src/map.jl | 8 +- .../lib/SparseArrayInterface/src/wrappers.jl | 10 +- .../lib/SparseArrayInterface/test/runtests.jl | 104 ++++++++------- 7 files changed, 202 insertions(+), 121 deletions(-) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/base.jl index 6dd1b9598f..2c7c1a3009 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/base.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/base.jl @@ -1,33 +1,35 @@ # TODO: Use `sparse_mapreduce`? function sparse_iszero(a::AbstractArray) - iszero(nonzero_length(a)) && return true - return iszero(nonzeros(a)) + iszero(nstored(a)) && return true + return iszero(storage(a)) end # TODO: Use `sparse_mapreduce`? function sparse_isreal(a::AbstractArray) - iszero(nonzero_length(a)) && return true - return isreal(nonzeros(a)) + iszero(nstored(a)) && return true + return isreal(storage(a)) end -function fill_nonzero!(a::AbstractArray, x) +# Fill all values, stored and unstored +function fill_all!(a::AbstractArray, x) for I in eachindex(a) a[I] = x end return a end -function fill_zero!(a::AbstractArray, x) - fill!(nonzeros(a), x) +# Only fill the stored values +function fill_stored!(a::AbstractArray, x) + fill!(storage(a), x) return a end function sparse_fill!(a::AbstractArray, x) if !iszero(x) # TODO: Reverse naming convention? - fill_nonzero!(a, x) + fill_all!(a, x) else - fill_zero!(a, x) + fill_stored!(a, x) end return a end @@ -35,22 +37,11 @@ end function sparse_zero(a::AbstractArray) # Need to overload `similar` for custom types a = similar(a) + # TODO: Use custom zero value? sparse_fill!(a, zero(eltype(a))) return a end -## function sparse_zero(a::AbstractArray, elt::Type) -## return sparse_zero(a, elt, size(a)) -## end -## -## function sparse_zero(a::AbstractArray, dims::Tuple) -## return sparse_zero(a, eltype(a), dims) -## end -## -## function sparse_zero(a::AbstractArray) -## return sparse_zero(a, eltype(a), size(a)) -## end - # TODO: Make `sparse_one!`? function sparse_one(a::AbstractMatrix) m, n = size(a) @@ -65,9 +56,9 @@ end # TODO: Use `sparse_mapreduce`? function sparse_isequal(a1::AbstractArray, a2::AbstractArray) - Is = collect(nonzero_indices(a1)) - intersect!(Is, nonzero_indices(a2)) - if !(length(Is) == nonzero_length(a1) == nonzero_length(a2)) + Is = collect(stored_indices(a1)) + intersect!(Is, stored_indices(a2)) + if !(length(Is) == nstored(a1) == nstored(a2)) return false end for I in Is @@ -82,7 +73,7 @@ function sparse_reshape(a::AbstractArray, dims) sparse_fill!(a_reshaped, zero(eltype(a))) linear_inds = LinearIndices(a) new_cartesian_inds = CartesianIndices(dims) - for I in nonzero_indices(a) + for I in stored_indices(a) a_reshaped[new_cartesian_inds[linear_inds[I]]] = a[I] end return a_reshaped @@ -91,9 +82,9 @@ end ## function sparse_mapreduce(f, op, as::AbstractArray...) ## # TODO: Make more general. ## @assert iszero(mapreduce(f, op, map(a -> getindex_zero(a, first(eachindex(a))), as)...)) -## # TODO: Create a `promote_nonzero_indices`, for example to help +## # TODO: Create a `promote_stored_indices`, for example to help ## # with specialized fast indexing for `DiagonalArrays`. -## Is = union(nonzero_indices.(as)...) +## Is = union(stored_indices.(as)...) ## for I in Is ## a_dest[I] = f(map(a -> a[I], as)...) ## end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl index 1522eb3d99..22da4d85ca 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl @@ -1,23 +1,63 @@ -nonzero_length(a::AbstractArray) = length(nonzeros(a)) -function nonzero_indices(a::AbstractArray) - return Iterators.map(Inz -> nonzero_index_to_index(a, Inz), eachindex(nonzeros(a))) +# An index into the storage of the sparse array. +struct StorageIndex{I} + i::I +end +index(i::StorageIndex) = i.i + +# Indicate if the index into the sparse array is +# stored or not. +abstract type MaybeStoredIndex{I} end + +# An index into a stored value of the sparse array. +# Stores both the index into the outer array +# as well as into the underlying storage. +struct StoredIndex{Iouter,Istorage} <: MaybeStoredIndex{Iouter} + iouter::Iouter + istorage::StorageIndex{Istorage} +end +index(i::StoredIndex) = i.iouter +StorageIndex(i::StoredIndex) = i.istorage + +nstored(a::AbstractArray) = length(storage(a)) + +struct NotStoredIndex{Iouter} <: MaybeStoredIndex{Iouter} + iouter::Iouter +end +index(i::NotStoredIndex) = i.iouter + +function MaybeStoredIndex(a::AbstractArray, I) + return MaybeStoredIndex(I, index_to_storage_index(a, I)) +end +MaybeStoredIndex(I, I_storage) = StoredIndex(I, StorageIndex(I_storage)) +MaybeStoredIndex(I, I_storage::Nothing) = NotStoredIndex(I) + +function storage_indices(a::AbstractArray) + return eachindex(storage(a)) end # Derived -function index_to_nonzero_index(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} - return index_to_nonzero_index(a, CartesianIndex(I)) +function index_to_storage_index(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} + return index_to_storage_index(a, CartesianIndex(I)) end # Helper type for constructing zero values struct GetIndexZero end (::GetIndexZero)(a::AbstractArray, I) = zero(eltype(a)) -function getindex_nonzero(a::AbstractArray, Inz) - return nonzeros(a)[Inz] +function sparse_getindex( + a::AbstractArray{<:Any,N}, I::NotStoredIndex{CartesianIndex{N}} +) where {N} + return getindex_notstored(a, index(I)) +end + +function sparse_getindex( + a::AbstractArray{<:Any,N}, I::StoredIndex{CartesianIndex{N}} +) where {N} + return sparse_getindex(a, StorageIndex(I)) end -function sparse_getindex(a::AbstractArray, I...) - return sparse_getindex(a, CartesianIndex(to_indices(a, I))) +function sparse_getindex(a::AbstractArray, I::StorageIndex) + return storage(a)[index(I)] end function sparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} @@ -25,16 +65,42 @@ function sparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} end function sparse_getindex(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} - Inz = index_to_nonzero_index(a, I) - if isnothing(Inz) - return getindex_zero(a, I) - end - return getindex_nonzero(a, Inz) + return _sparse_getindex(a, I) +end + +# Ambiguity with linear indexing +function sparse_getindex(a::AbstractArray{<:Any,1}, I::CartesianIndex{1}) + return _sparse_getindex(a, I) +end + +# Implementation of element access +function _sparse_getindex(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} + @boundscheck checkbounds(a, I) + return sparse_getindex(a, MaybeStoredIndex(a, I)) +end + +# Handle trailing indices or linear indexing +function sparse_getindex(a::AbstractArray, I::Vararg{Int}) + return sparse_getindex(a, CartesianIndex(I)) +end + +# Linear indexing +function sparse_getindex(a::AbstractArray, I::CartesianIndex{1}) + return error("Linear indexing not supported yet for sparse arrays") +end + +# Handle trailing indices +function sparse_getindex(a::AbstractArray, I::CartesianIndex) + t = Tuple(I) + length(t) < ndims(a) && error("Not enough indices passed") + I′ = ntuple(i -> t[i], ndims(a)) + @assert all(i -> isone(I[i]), (ndims(a) + 1):length(I)) + return _sparse_getindex(a, I′) end # Update a nonzero value -function setindex_nonzero!(a::AbstractArray, value, Inz) - nonzeros(a)[Inz] = value +function sparse_setindex!(a::AbstractArray, value, I::StorageIndex) + storage(a)[index(I)] = value return a end @@ -43,15 +109,23 @@ function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}) wh return a end -function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::CartesianIndex{N}) where {N} - Inz = index_to_nonzero_index(a, I) - if isnothing(Inz) - if !iszero(value) - # Only try setting if nonzero - setindex_zero!(a, value, I) - end - else - setindex_nonzero!(a, value, Inz) +function sparse_setindex!( + a::AbstractArray{<:Any,N}, value, I::StoredIndex{CartesianIndex{N}} +) where {N} + sparse_setindex!(a, value, StorageIndex(I)) + return a +end + +function sparse_setindex!( + a::AbstractArray{<:Any,N}, value, I::NotStoredIndex{CartesianIndex{N}} +) where {N} + if !iszero(value) + setindex_notstored!(a, value, index(I)) end return a end + +function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::CartesianIndex{N}) where {N} + sparse_setindex!(a, value, MaybeStoredIndex(a, I)) + return a +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface.jl index 1a30d1797e..68257965c8 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/interface.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface.jl @@ -2,16 +2,17 @@ # https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ # Minimal interface -# Data structure storing the nonzero values -nonzeros(a::AbstractArray) = a +# Data structure of the stored (generally nonzero) values +# nonzeros(a::AbstractArray) = a +storage(a::AbstractArray) = a # Minimal interface -# Map an index in the nonzero data to a CartesianIndex of the +# Map an index into the stored data to a CartesianIndex of the # outer array. -nonzero_index_to_index(a::AbstractArray, Inz) = Inz +storage_index_to_index(a::AbstractArray, I) = I # Minimal interface # Map a `CartesianIndex` to an index/key into the nonzero data structure -# returned by `nonzeros`. -# Return `nothing` if the index corresponds to a structural zero value. -index_to_nonzero_index(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} = I +# returned by `storage`. +# Return `nothing` if the index corresponds to a structural zero (unstored) value. +index_to_storage_index(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} = I diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl index e2c53cf6eb..1a881e8e49 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl @@ -1,12 +1,21 @@ # Optional interface. # Access a zero value. -function getindex_zero(a::AbstractArray, I) +function getindex_notstored(a::AbstractArray, I) return zero(eltype(a)) end # Optional interface. -# Insert a new zero value. +# Insert a new value into a location +# where there is not a stored value. # Some types (like `Diagonal`) may not support this. -function setindex_zero!(a::AbstractArray, value, I) - return error("Not implemented") +function setindex_notstored!(a::AbstractArray, value, I) + return throw(ArgumentError("Can't set nonzero values of $(typeof(a)).")) +end + +# Optional interface. +# Iterates over the indices of `a` where there are stored values. +# Can overload for faster iteration when there is more structure, +# for example with DiagonalArrays. +function stored_indices(a::AbstractArray) + return Iterators.map(Inz -> storage_index_to_index(a, Inz), storage_indices(a)) end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/map.jl b/NDTensors/src/lib/SparseArrayInterface/src/map.jl index 40ad7ed0af..dead209bcf 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/map.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/map.jl @@ -1,10 +1,12 @@ # TODO: Rename the preserve zero case to `map_nonzeros!`. function sparse_map!(f, a_dest::AbstractArray, as::AbstractArray...) # TODO: Handle nonzero case, fill all values. - @assert iszero(f(map(a -> getindex_zero(a, first(eachindex(a))), as)...)) - # TODO: Create a `promote_nonzero_indices`, for example to help + @assert iszero( + f(map(a -> sparse_getindex(a, NotStoredIndex(first(eachindex(a)))), as)...) + ) + # TODO: Create a `promote_stored_indices`, for example to help # with specialized fast indexing for `DiagonalArrays`. - Is = union(nonzero_indices.(as)...) + Is = union(stored_indices.(as)...) for I in Is a_dest[I] = f(map(a -> a[I], as)...) end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl b/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl index 4ab140de57..199573fc4e 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/wrappers.jl @@ -2,14 +2,14 @@ perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P genperm(v, perm) = map(j -> v[j], perm) genperm(v::CartesianIndex, perm) = CartesianIndex(map(j -> Tuple(v)[j], perm)) -nonzeros(a::PermutedDimsArray) = nonzeros(parent(a)) -function nonzero_index_to_index(a::PermutedDimsArray, Inz) - return genperm(nonzero_index_to_index(parent(a), Inz), perm(a)) +storage(a::PermutedDimsArray) = storage(parent(a)) +function storage_index_to_index(a::PermutedDimsArray, I) + return genperm(storage_index_to_index(parent(a), I), perm(a)) end -function index_to_nonzero_index( +function index_to_storage_index( a::PermutedDimsArray{<:Any,N}, I::CartesianIndex{N} ) where {N} - return nonzero_index_to_index(parent(a), genperm(I, perm(a))) + return storage_index_to_index(parent(a), genperm(I, perm(a))) end # TODO: Add `SubArray`, `ReshapedArray`, `Diagonal`, etc. diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 4eaef6082e..2201e856d2 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -6,10 +6,10 @@ using Test: @test, @testset, @test_broken, @test_throws @testset "Array" begin using NDTensors.SparseArrayInterface: SparseArrayInterface a = randn(2, 3) - @test SparseArrayInterface.nonzeros(a) == a - @test SparseArrayInterface.index_to_nonzero_index(a, CartesianIndex(1, 2)) == + @test SparseArrayInterface.storage(a) == a + @test SparseArrayInterface.index_to_storage_index(a, CartesianIndex(1, 2)) == CartesianIndex(1, 2) - @test SparseArrayInterface.nonzero_index_to_index(a, CartesianIndex(1, 2)) == + @test SparseArrayInterface.storage_index_to_index(a, CartesianIndex(1, 2)) == CartesianIndex(1, 2) end @testset "Custom SparseArray" begin @@ -21,9 +21,11 @@ using Test: @test, @testset, @test_broken, @test_throws index_to_dataindex::Dict{CartesianIndex{N},Int} dataindex_to_index::Vector{CartesianIndex{N}} end - SparseArray{T,N}(dims::Tuple{Vararg{Int,N}}) where {T,N} = SparseArray{T,N}( - T[], dims, Dict{CartesianIndex{N},Int}(), Vector{CartesianIndex{N}}() - ) + 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) @@ -41,14 +43,14 @@ using Test: @test, @testset, @test_broken, @test_throws end # Minimal interface - SparseArrayInterface.nonzeros(a::SparseArray) = a.data - function SparseArrayInterface.index_to_nonzero_index( + SparseArrayInterface.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.nonzero_index_to_index(a::SparseArray, I) = a.dataindex_to_index[I] - function SparseArrayInterface.setindex_zero!( + 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) @@ -106,9 +108,9 @@ using Test: @test, @testset, @test_broken, @test_throws a = SparseArray{elt}(2, 3) @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) - @test SparseArrayInterface.nonzeros(a) == elt[] - @test iszero(SparseArrayInterface.nonzero_length(a)) - @test collect(SparseArrayInterface.nonzero_indices(a)) == CartesianIndex{2}[] + @test SparseArrayInterface.storage(a) == elt[] + @test iszero(SparseArrayInterface.nstored(a)) + @test collect(SparseArrayInterface.stored_indices(a)) == CartesianIndex{2}[] @test iszero(a) @test iszero(norm(a)) for I in eachindex(a) @@ -119,9 +121,10 @@ using Test: @test, @testset, @test_broken, @test_throws a[1, 2] = 12 @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) - @test SparseArrayInterface.nonzeros(a) == elt[12] - @test isone(SparseArrayInterface.nonzero_length(a)) - @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test a[SparseArrayInterface.StorageIndex(1)] == 12 + @test SparseArrayInterface.storage(a) == elt[12] + @test isone(SparseArrayInterface.nstored(a)) + @test collect(SparseArrayInterface.stored_indices(a)) == [CartesianIndex(1, 2)] @test !iszero(a) @test !iszero(norm(a)) for I in eachindex(a) @@ -141,7 +144,7 @@ using Test: @test, @testset, @test_broken, @test_throws a[1, 2] = 12 a = zero(a) @test size(a) == (2, 3) - @test iszero(SparseArrayInterface.nonzero_length(a)) + @test iszero(SparseArrayInterface.nstored(a)) a = SparseArray{elt}(2, 3) a[1, 2] = 12 @@ -188,16 +191,16 @@ using Test: @test, @testset, @test_broken, @test_throws a[1, 2] = 12 a = zero(a) @test size(a) == (2, 3) - @test iszero(SparseArrayInterface.nonzero_length(a)) + @test iszero(SparseArrayInterface.nstored(a)) a = SparseArray{elt}(2, 3) a[1, 2] = 12 a = copy(a) @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) - @test SparseArrayInterface.nonzeros(a) == elt[12] - @test isone(SparseArrayInterface.nonzero_length(a)) - @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test SparseArrayInterface.storage(a) == elt[12] + @test isone(SparseArrayInterface.nstored(a)) + @test collect(SparseArrayInterface.stored_indices(a)) == [CartesianIndex(1, 2)] @test !iszero(a) @test !iszero(norm(a)) for I in eachindex(a) @@ -213,9 +216,9 @@ using Test: @test, @testset, @test_broken, @test_throws a = 2 * a @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) - @test SparseArrayInterface.nonzeros(a) == elt[24] - @test isone(SparseArrayInterface.nonzero_length(a)) - @test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] + @test SparseArrayInterface.storage(a) == elt[24] + @test isone(SparseArrayInterface.nstored(a)) + @test collect(SparseArrayInterface.stored_indices(a)) == [CartesianIndex(1, 2)] @test !iszero(a) @test !iszero(norm(a)) for I in eachindex(a) @@ -233,9 +236,9 @@ using Test: @test, @testset, @test_broken, @test_throws c = a + b @test size(c) == (2, 3) @test axes(c) == (1:2, 1:3) - @test SparseArrayInterface.nonzeros(c) == elt[12, 21] - @test SparseArrayInterface.nonzero_length(c) == 2 - @test collect(SparseArrayInterface.nonzero_indices(c)) == + @test SparseArrayInterface.storage(c) == elt[12, 21] + @test SparseArrayInterface.nstored(c) == 2 + @test collect(SparseArrayInterface.stored_indices(c)) == [CartesianIndex(1, 2), CartesianIndex(2, 1)] @test !iszero(c) @test !iszero(norm(c)) @@ -254,9 +257,9 @@ using Test: @test, @testset, @test_broken, @test_throws b = permutedims(a, (2, 1)) @test size(b) == (3, 2) @test axes(b) == (1:3, 1:2) - @test SparseArrayInterface.nonzeros(b) == elt[12] - @test SparseArrayInterface.nonzero_length(b) == 1 - @test collect(SparseArrayInterface.nonzero_indices(b)) == [CartesianIndex(2, 1)] + @test SparseArrayInterface.storage(b) == elt[12] + @test SparseArrayInterface.nstored(b) == 1 + @test collect(SparseArrayInterface.stored_indices(b)) == [CartesianIndex(2, 1)] @test !iszero(b) @test !iszero(norm(b)) for I in eachindex(b) @@ -272,14 +275,18 @@ using Test: @test, @testset, @test_broken, @test_throws data::Vector{T} dims::Tuple{Vararg{Int,N}} end - DiagonalArray{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int,N}}) where {T,N} = - DiagonalArray{T,N}(Vector{T}(undef, minimum(dims)), dims) - DiagonalArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} = - DiagonalArray{T,N}(undef, dims) - DiagonalArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} = - DiagonalArray{T,length(dims)}(undef, dims) - DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int}) where {T} = - DiagonalArray{T}(undef, dims) + function DiagonalArray{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int,N}}) where {T,N} + return DiagonalArray{T,N}(Vector{T}(undef, minimum(dims)), dims) + end + function DiagonalArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(undef, dims) + end + function DiagonalArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} + return DiagonalArray{T,length(dims)}(undef, dims) + end + function DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int}) where {T} + return DiagonalArray{T}(undef, dims) + end # AbstractArray interface Base.size(a::DiagonalArray) = a.dims @@ -294,25 +301,16 @@ using Test: @test, @testset, @test_broken, @test_throws end # Minimal interface - SparseArrayInterface.nonzeros(a::DiagonalArray) = a.data - function SparseArrayInterface.index_to_nonzero_index( + SparseArrayInterface.storage(a::DiagonalArray) = a.data + function SparseArrayInterface.index_to_storage_index( a::DiagonalArray{<:Any,N}, I::CartesianIndex{N} ) where {N} !allequal(Tuple(I)) && return nothing return first(Tuple(I)) end - function SparseArrayInterface.nonzero_index_to_index(a::DiagonalArray, I) + function SparseArrayInterface.storage_index_to_index(a::DiagonalArray, I) return CartesianIndex(ntuple(Returns(I), ndims(a))) end - function SparseArrayInterface.setindex_zero!( - a::DiagonalArray{<:Any,N}, value, I::CartesianIndex{N} - ) where {N} - return throw(ArgumentError("Can't set off-diagonal element of DiagonalArray")) - end - function SparseArrayInterface.sparse_zero(a::DiagonalArray, elt::Type, dims::Tuple) - # TODO: Need to make generic to GPU. - return DiagonalArray(zeros(elt, minimum(dims)), dims) - end # Broadcasting function Broadcast.BroadcastStyle(arraytype::Type{<:DiagonalArray}) @@ -367,6 +365,9 @@ using Test: @test, @testset, @test_broken, @test_throws @test a[1, 1] == 11 @test a[2, 2] == 22 @test_throws ArgumentError a[1, 2] = 12 + @test SparseArrayInterface.storage_indices(a) == 1:2 + @test collect(SparseArrayInterface.stored_indices(a)) == + [CartesianIndex(1, 1), CartesianIndex(2, 2)] a[1, 2] = 0 @test a[1, 1] == 11 @test a[2, 2] == 22 @@ -380,6 +381,9 @@ using Test: @test, @testset, @test_broken, @test_throws @test a[1, 1] == 1 @test a[2, 2] == 2 @test a[3, 3] == 3 + @test a[SparseArrayInterface.StorageIndex(1)] == 1 + @test a[SparseArrayInterface.StorageIndex(2)] == 2 + @test a[SparseArrayInterface.StorageIndex(3)] == 3 @test iszero(a[1, 2]) a = DiagonalArray(elt[1, 2, 3], (3, 3)) From e984cbe774e7a519e6267365f5bffe48e34b3243 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 13:06:35 -0500 Subject: [PATCH 06/17] Add another test --- NDTensors/src/lib/SparseArrayInterface/test/runtests.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 2201e856d2..0d3aa6349b 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -200,6 +200,7 @@ using Test: @test, @testset, @test_broken, @test_throws @test axes(a) == (1:2, 1:3) @test SparseArrayInterface.storage(a) == elt[12] @test isone(SparseArrayInterface.nstored(a)) + @test SparseArrayInterface.storage_indices(a) == 1:1 @test collect(SparseArrayInterface.stored_indices(a)) == [CartesianIndex(1, 2)] @test !iszero(a) @test !iszero(norm(a)) From 475ee3124178e5350243e6da70eb68e16082bd58 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 13:35:24 -0500 Subject: [PATCH 07/17] Improved map --- .../src/lib/SparseArrayInterface/README.md | 3 +- .../src/lib/SparseArrayInterface/src/base.jl | 21 +------ .../lib/SparseArrayInterface/src/indexing.jl | 16 ++--- .../src/lib/SparseArrayInterface/src/map.jl | 58 ++++++++++++++++--- .../lib/SparseArrayInterface/test/runtests.jl | 17 +++++- 5 files changed, 71 insertions(+), 44 deletions(-) diff --git a/NDTensors/src/lib/SparseArrayInterface/README.md b/NDTensors/src/lib/SparseArrayInterface/README.md index 0f434c8746..917f5dfdfb 100644 --- a/NDTensors/src/lib/SparseArrayInterface/README.md +++ b/NDTensors/src/lib/SparseArrayInterface/README.md @@ -31,8 +31,9 @@ Still need to implement `Base` functions: [x] sparse_isreal(a::AbstractArray) = ... # Uses `all`, make `sparse_all`? [x] sparse_isequal(a1::AbstractArray, a2::AbstractArray) = ... [x] sparse_conj!(a::AbstractArray) = conj!(nonzeros(a)) -[ ] sparse_all(f, a::AbstractArray) = ... [x] sparse_reshape(a::AbstractArray, dims) = ... +[ ] sparse_all(f, a::AbstractArray) = ... +[ ] sparse_getindex(a::AbstractArray, 1:2, 2:3) = ... # Slicing ``` `LinearAlgebra` functions: ```julia diff --git a/NDTensors/src/lib/SparseArrayInterface/src/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/base.jl index 2c7c1a3009..da01f5f671 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/base.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/base.jl @@ -10,27 +10,8 @@ function sparse_isreal(a::AbstractArray) return isreal(storage(a)) end -# Fill all values, stored and unstored -function fill_all!(a::AbstractArray, x) - for I in eachindex(a) - a[I] = x - end - return a -end - -# Only fill the stored values -function fill_stored!(a::AbstractArray, x) - fill!(storage(a), x) - return a -end - function sparse_fill!(a::AbstractArray, x) - if !iszero(x) - # TODO: Reverse naming convention? - fill_all!(a, x) - else - fill_stored!(a, x) - end + sparse_map!(Returns(x), a, a) return a end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl index 22da4d85ca..7d79e21c13 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl @@ -44,15 +44,11 @@ end struct GetIndexZero end (::GetIndexZero)(a::AbstractArray, I) = zero(eltype(a)) -function sparse_getindex( - a::AbstractArray{<:Any,N}, I::NotStoredIndex{CartesianIndex{N}} -) where {N} +function sparse_getindex(a::AbstractArray, I::NotStoredIndex) return getindex_notstored(a, index(I)) end -function sparse_getindex( - a::AbstractArray{<:Any,N}, I::StoredIndex{CartesianIndex{N}} -) where {N} +function sparse_getindex(a::AbstractArray, I::StoredIndex) return sparse_getindex(a, StorageIndex(I)) end @@ -109,16 +105,12 @@ function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}) wh return a end -function sparse_setindex!( - a::AbstractArray{<:Any,N}, value, I::StoredIndex{CartesianIndex{N}} -) where {N} +function sparse_setindex!(a::AbstractArray, value, I::StoredIndex) sparse_setindex!(a, value, StorageIndex(I)) return a end -function sparse_setindex!( - a::AbstractArray{<:Any,N}, value, I::NotStoredIndex{CartesianIndex{N}} -) where {N} +function sparse_setindex!(a::AbstractArray, value, I::NotStoredIndex) if !iszero(value) setindex_notstored!(a, value, index(I)) end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/map.jl b/NDTensors/src/lib/SparseArrayInterface/src/map.jl index dead209bcf..11ed25329e 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/map.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/map.jl @@ -1,14 +1,54 @@ -# TODO: Rename the preserve zero case to `map_nonzeros!`. -function sparse_map!(f, a_dest::AbstractArray, as::AbstractArray...) - # TODO: Handle nonzero case, fill all values. - @assert iszero( - f(map(a -> sparse_getindex(a, NotStoredIndex(first(eachindex(a)))), as)...) - ) +# Test if the function preserves zero values and therefore +# preserves the sparsity structure. +function preserves_zero(f, a_dest, as...) + return iszero(f(map(a -> sparse_getindex(a, NotStoredIndex(first(eachindex(a)))), as)...)) +end + +# Map a subset of indices +function sparse_map_indices!(f, a_dest::AbstractArray, indices, as::AbstractArray...) + for I in indices + a_dest[I] = f(map(a -> a[I], as)...) + end + return a_dest +end + +# Overload for custom `stored_indices` types. +function promote_indices(I1, I2) + return union(I1, I2) +end + +function promote_indices(I1, I2, Is...) + return promote_indices(promote_indices(I1, I2), Is...) +end + +# Base case +promote_indices(I) = I + +function promote_stored_indices(as::AbstractArray...) + return promote_indices(stored_indices.(as)...) +end + +function sparse_map_stored!(f, a_dest::AbstractArray, as::AbstractArray...) # TODO: Create a `promote_stored_indices`, for example to help # with specialized fast indexing for `DiagonalArrays`. - Is = union(stored_indices.(as)...) - for I in Is - a_dest[I] = f(map(a -> a[I], as)...) + Is = promote_stored_indices(as...) + sparse_map_indices!(f, a_dest, Is, as...) + return a_dest +end + +# Handle nonzero case, fill all values. +function sparse_map_all!(f, a_dest::AbstractArray, as::AbstractArray...) + Is = eachindex(a_dest) + sparse_map_indices!(f, a_dest, Is, as...) + return a_dest +end + +function sparse_map!(f, a_dest::AbstractArray, as::AbstractArray...) + @assert allequal(axes.((a_dest, as...))) + if preserves_zero(f, a_dest, as...) + sparse_map_stored!(f, a_dest, as...) + else + sparse_map_all!(f, a_dest, as...) end return a_dest end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 0d3aa6349b..862123a2b1 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -102,8 +102,6 @@ using Test: @test, @testset, @test_broken, @test_throws return dest end - # TODO: Test `fill!`. - # Test a = SparseArray{elt}(2, 3) @test size(a) == (2, 3) @@ -117,6 +115,21 @@ using Test: @test, @testset, @test_broken, @test_throws @test iszero(a) end + a = SparseArray{elt}(2, 3) + fill!(a, 0) + @test size(a) == (2, 3) + @test iszero(a) + @test iszero(SparseArrayInterface.nstored(a)) + + a = SparseArray{elt}(2, 3) + fill!(a, 2) + @test size(a) == (2, 3) + @test !iszero(a) + @test SparseArrayInterface.nstored(a) == length(a) + for I in eachindex(a) + @test a[I] == 2 + end + a = SparseArray{elt}(2, 3) a[1, 2] = 12 @test size(a) == (2, 3) From 36961f8955e3971ff63ff62e9e921445d90fe348 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 18:04:23 -0500 Subject: [PATCH 08/17] New SparseArrayDOKs using SparseArrayInterface --- .../lib/DiagonalArrays/src/DiagonalArrays.jl | 2 +- .../SparseArrayDOKs/src/SparseArrayDOKs.jl | 7 +- .../src/SparseArrayDOKsSparseArraysExt.jl | 3 +- NDTensors/src/lib/SparseArrayDOKs/src/base.jl | 18 +++ .../lib/SparseArrayDOKs/src/baseinterface.jl | 15 +++ .../src/lib/SparseArrayDOKs/src/broadcast.jl | 40 +------ .../src/lib/SparseArrayDOKs/src/conversion.jl | 19 +++ .../src/lib/SparseArrayDOKs/src/copyto.jl | 12 -- .../src/lib/SparseArrayDOKs/src/defaults.jl | 5 + NDTensors/src/lib/SparseArrayDOKs/src/map.jl | 44 +++++-- .../lib/SparseArrayDOKs/src/sparsearraydok.jl | 63 +--------- .../src/sparsearrayinterface.jl | 49 +++----- .../src/lib/SparseArrayDOKs/test/runtests.jl | 55 ++++++--- .../src/SparseArrayInterface.jl | 4 +- .../src/lib/SparseArrayInterface/src/base.jl | 78 ++++++++----- .../SparseArrayInterface/src/conversion.jl | 5 + .../lib/SparseArrayInterface/src/indexing.jl | 6 +- .../src/interface_optional.jl | 8 ++ .../src/lib/SparseArrayInterface/src/map.jl | 20 +++- .../src/lib/SparseArrayInterface/src/zero.jl | 4 + .../lib/SparseArrayInterface/test/runtests.jl | 108 ++++++++++++++---- .../src/lib/SparseArrayInterface/test/temp.jl | 105 ----------------- 22 files changed, 335 insertions(+), 335 deletions(-) create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/base.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/baseinterface.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/conversion.jl delete mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl create mode 100644 NDTensors/src/lib/SparseArrayDOKs/src/defaults.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/conversion.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/src/zero.jl delete mode 100644 NDTensors/src/lib/SparseArrayInterface/test/temp.jl diff --git a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl index 2df6e33e4d..7ebe04df06 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl @@ -31,7 +31,7 @@ Base.similar(a::DiagonalArray) = set_diag(a, similar(a[DiagIndices()])) Base.size(a::DiagonalArray) = a.dims -# TODO: Rename `nonzer_values`. +# TODO: Rename `storage_values`. diagview(a::DiagonalArray) = a.diag LinearAlgebra.diag(a::DiagonalArray) = copy(diagview(a)) diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl index 3c648f7cba..697f145ff6 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKs.jl @@ -1,8 +1,11 @@ module SparseArrayDOKs -include("sparsearrayinterface.jl") +include("defaults.jl") include("sparsearraydok.jl") +include("sparsearrayinterface.jl") +include("base.jl") include("broadcast.jl") include("map.jl") -include("copyto.jl") +include("baseinterface.jl") +include("conversion.jl") include("SparseArrayDOKsSparseArraysExt.jl") end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl index 93a1104bf0..85d462ca78 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/SparseArrayDOKsSparseArraysExt.jl @@ -1,4 +1,5 @@ using SparseArrays: SparseArrays +using ..SparseArrayInterface: nstored # Julia Base `AbstractSparseArray` interface -SparseArrays.nnz(a::SparseArrayDOK) = nonzero_length(a) +SparseArrays.nnz(a::SparseArrayDOK) = nstored(a) diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/base.jl b/NDTensors/src/lib/SparseArrayDOKs/src/base.jl new file mode 100644 index 0000000000..11cc24d7a0 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/base.jl @@ -0,0 +1,18 @@ +using ..SparseArrayInterface: SparseArrayInterface + +# Base +function Base.:(==)(a1::SparseArrayDOK, a2::SparseArrayDOK) + return SparseArrayInterface.sparse_isequal(a1, a2) +end + +function Base.reshape(a::SparseArrayDOK, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) +end + +function Base.zero(a::SparseArrayDOK) + return SparseArrayInterface.sparse_zero(a) +end + +function Base.one(a::SparseArrayDOK) + return SparseArrayInterface.sparse_one(a) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/baseinterface.jl b/NDTensors/src/lib/SparseArrayDOKs/src/baseinterface.jl new file mode 100644 index 0000000000..bc486df96c --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/baseinterface.jl @@ -0,0 +1,15 @@ +using ..SparseArrayInterface: SparseArrayInterface + +Base.size(a::SparseArrayDOK) = a.dims + +function Base.getindex(a::SparseArrayDOK, I...) + return SparseArrayInterface.sparse_getindex(a, I...) +end + +function Base.setindex!(a::SparseArrayDOK, I...) + return SparseArrayInterface.sparse_setindex!(a, I...) +end + +function Base.similar(a::SparseArrayDOK, elt::Type, dims::Tuple{Vararg{Int}}) + return SparseArrayDOK{elt}(undef, dims) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl b/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl index c539aeb6b3..b7e17b41e0 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/broadcast.jl @@ -1,38 +1,4 @@ -using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted -using ..BroadcastMapConversion: map_function, map_args - -struct SparseArrayDOKStyle{N} <: AbstractArrayStyle{N} end - -function Broadcast.BroadcastStyle(::Type{<:SparseArrayDOK{<:Any,N}}) where {N} - return SparseArrayDOKStyle{N}() -end - -SparseArrayDOKStyle(::Val{N}) where {N} = SparseArrayDOKStyle{N}() -SparseArrayDOKStyle{M}(::Val{N}) where {M,N} = SparseArrayDOKStyle{N}() - -Broadcast.BroadcastStyle(a::SparseArrayDOKStyle, ::DefaultArrayStyle{0}) = a -function Broadcast.BroadcastStyle(::SparseArrayDOKStyle{N}, a::DefaultArrayStyle) where {N} - return BroadcastStyle(DefaultArrayStyle{N}(), a) -end -function Broadcast.BroadcastStyle( - ::SparseArrayDOKStyle{N}, ::Broadcast.Style{Tuple} -) where {N} - return DefaultArrayStyle{N}() -end - -# TODO: Use `allocate_output`, share logic with `map`. -function Base.similar(bc::Broadcasted{<:SparseArrayDOKStyle}, elt::Type) - # TODO: Is this a good definition? Probably should check that - # they have consistent axes. - return similar(first(map_args(bc)), elt) -end - -# Broadcasting implementation -function Base.copyto!( - dest::SparseArrayDOK{<:Any,N}, bc::Broadcasted{SparseArrayDOKStyle{N}} -) where {N} - # convert to map - # flatten and only keep the AbstractArray arguments - map!(map_function(bc), dest, map_args(bc)...) - return dest +# Broadcasting +function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArrayDOK}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/conversion.jl b/NDTensors/src/lib/SparseArrayDOKs/src/conversion.jl new file mode 100644 index 0000000000..b2f50aed31 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/conversion.jl @@ -0,0 +1,19 @@ +function SparseArrayDOK{T,N,Zero}(a::SparseArrayDOK{T,N,Zero}) where {T,N,Zero} + return copy(a) +end + +function Base.convert( + ::Type{SparseArrayDOK{T,N,Zero}}, a::SparseArrayDOK{T,N,Zero} +) where {T,N,Zero} + return a +end + +Base.convert(type::Type{<:SparseArrayDOK}, a::AbstractArray) = type(a) + +SparseArrayDOK(a::AbstractArray) = SparseArrayDOK{eltype(a)}(a) + +SparseArrayDOK{T}(a::AbstractArray) where {T} = SparseArrayDOK{T,ndims(a)}(a) + +function SparseArrayDOK{T,N}(a::AbstractArray) where {T,N} + return SparseArrayInterface.sparse_convert(SparseArrayDOK{T,N}, a) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl b/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl deleted file mode 100644 index 61c3c038b3..0000000000 --- a/NDTensors/src/lib/SparseArrayDOKs/src/copyto.jl +++ /dev/null @@ -1,12 +0,0 @@ -# TODO: Rename `copy_nonzeros!`, move to `sparsearrayinterface.jl`. -function Base.copy!(dest::AbstractArray, src::SparseArrayDOK) - @assert axes(dest) == axes(src) - map!(identity, dest, src) - return dest -end - -# TODO: Rename `copyto_nonzeros!`, move to `sparsearrayinterface.jl`. -function Base.copyto!(dest::AbstractArray, src::SparseArrayDOK) - copy!(dest, src) - return dest -end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/defaults.jl b/NDTensors/src/lib/SparseArrayDOKs/src/defaults.jl new file mode 100644 index 0000000000..044341e8d4 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayDOKs/src/defaults.jl @@ -0,0 +1,5 @@ +using ..SparseArrayInterface: Zero + +default_zero() = Zero() +default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() +default_keytype(ndims::Int) = CartesianIndex{ndims} diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/map.jl b/NDTensors/src/lib/SparseArrayDOKs/src/map.jl index 673fb7fee8..b8b1656ac5 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/map.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/map.jl @@ -1,12 +1,34 @@ -# TODO: Rename the preserve zero case to `map_nonzeros!`, -# move to `sparsearrayinterface.jl`. -function Base.map!(f, a_dest::AbstractArray, as::SparseArrayDOK...) - # TODO: Handle nonzero case, fill all values. - # Also, handle arrays of arrays better, using `getindex_zero(a, first(eachindex(a)))`. - @assert iszero(f(map(a -> getindex_zero(a, first(eachindex(a))), as)...)) - Is = union(nonzero_keys.(as)...) - for I in Is - a_dest[I] = f(map(a -> a[I], as)...) - end - return a_dest +# Map +function Base.map!(f, dest::AbstractArray, src::SparseArrayDOK) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest +end + +function Base.copy!(dest::AbstractArray, src::SparseArrayDOK) + SparseArrayInterface.sparse_copy!(dest, src) + return dest +end + +function Base.copyto!(dest::AbstractArray, src::SparseArrayDOK) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest +end + +function Base.permutedims!(dest::AbstractArray, src::SparseArrayDOK, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest +end + +function Base.mapreduce(f, op, a::SparseArrayDOK; kwargs...) + return SparseArrayInterface.sparse_mapreduce(f, op, a; kwargs...) +end + +# TODO: Why isn't this calling `mapreduce` already? +function Base.iszero(a::SparseArrayDOK) + return SparseArrayInterface.sparse_iszero(a) +end + +# TODO: Why isn't this calling `mapreduce` already? +function Base.isreal(a::SparseArrayDOK) + return SparseArrayInterface.sparse_isreal(a) end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl index a0ded0cc9c..9805e1a5f0 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl @@ -1,6 +1,6 @@ -using Dictionaries: Dictionary, set! +using Dictionaries: Dictionary -# TODO: Define a constructor with a default `zero`. +# TODO: Parametrize by `data`? struct SparseArrayDOK{T,N,Zero} <: AbstractArray{T,N} data::Dictionary{CartesianIndex{N},T} dims::NTuple{N,Int} @@ -59,62 +59,3 @@ end function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) where {T} return SparseArrayDOK{T}(dims, zero) end - -# Required `SparseArrayInterface` interface -nonzero_structure(a::SparseArrayDOK) = a.data -# TODO: Make this a generic function. -nonzero_keys(a::SparseArrayDOK) = keys(nonzero_structure(a)) - -# Optional `SparseArrayInterface` interface -# TODO: Use `SetParameters`. -zero_value(a::SparseArrayDOK, I) = a.zero(eltype(a), I) - -# Accessors -Base.size(a::SparseArrayDOK) = a.dims - -function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N} - return a[CartesianIndex(I)] -end - -# TODO: Rename `getindex_sparse` or `sparse_getindex`. -function Base.getindex(a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N}) where {N} - if !is_structural_nonzero(a, I) - return getindex_zero(a, I) - end - return getindex_nonzero(a, I) -end - -# `SparseArrayInterface` interface -function setindex_zero!(a::SparseArrayDOK, value, I) - # TODO: This is specific to the `Dictionaries.jl` - # interface, make more generic? - set!(nonzero_structure(a), I, value) - return a -end - -function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N} - a[CartesianIndex(I)] = value - return a -end - -function Base.setindex!(a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N}) where {N} - if !is_structural_nonzero(a, I) - setindex_zero!(a, value, I) - end - setindex_nonzero!(a, value, I) - return a -end - -# similar -function Base.similar(a::SparseArrayDOK) - return similar(a, eltype(a)) -end - -function Base.similar(a::SparseArrayDOK, elt::Type) - return similar(a, elt, size(a)) -end - -function Base.similar(a::SparseArrayDOK, elt::Type, dims::Tuple{Vararg{Int}}) - # TODO: Accessor function for `a.zero`. - return SparseArrayDOK{elt}(undef, dims, a.zero) -end diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl index 41f6f8f20b..594582f436 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearrayinterface.jl @@ -1,39 +1,22 @@ -using Dictionaries: Dictionary +using Dictionaries: set! +using ..SparseArrayInterface: SparseArrayInterface -# Also look into: -# https://juliaarrays.github.io/ArrayInterface.jl/stable/sparsearrays/ +SparseArrayInterface.storage(a::SparseArrayDOK) = a.data -# Required `SparseArrayInterface` interface. -# https://github.com/Jutho/SparseArrayKit.jl interface functions -nonzero_keys(a::AbstractArray) = error("Not implemented") -nonzero_values(a::AbstractArray) = error("Not implemented") -nonzero_pairs(a::AbstractArray) = error("Not implemented") - -# A dictionary-like structure -# TODO: Rename `nonzeros`, `structural_nonzeros`, etc.? -nonzero_structure(a::AbstractArray) = error("Not implemented") - -# Derived `SparseArrayInterface` interface. -nonzero_length(a::AbstractArray) = length(nonzero_keys(a)) -is_structural_nonzero(a::AbstractArray, I) = I ∈ nonzero_keys(a) - -# Overload if zero value is index dependent or -# doesn't match element type. -getindex_nonzero(a::AbstractArray, I) = nonzero_structure(a)[I] -getindex_zero(a::AbstractArray, I) = zero(eltype(a)) -function setindex_zero!(a::AbstractArray, value, I) - # TODO: This may need to be modified. - nonzero_structure(a)[I] = value - return a +function SparseArrayInterface.index_to_storage_index( + a::SparseArrayDOK{<:Any,N}, I::CartesianIndex{N} +) where {N} + !isassigned(SparseArrayInterface.storage(a), I) && return nothing + return I end -function setindex_nonzero!(a::AbstractArray, value, I) - nonzero_structure(a)[I] = value + +function SparseArrayInterface.setindex_notstored!( + a::SparseArrayDOK{<:Any,N}, value, I::CartesianIndex{N} +) where {N} + set!(SparseArrayInterface.storage(a), I, value) return a end -struct Zero end -(::Zero)(type, I) = zero(type) - -default_zero() = Zero() # (eltype, I) -> zero(eltype) -default_keytype(ndims::Int) = CartesianIndex{ndims} -default_data(type::Type, ndims::Int) = Dictionary{default_keytype(ndims),type}() +function SparseArrayInterface.empty_storage!(a::SparseArrayDOK) + return empty!(SparseArrayInterface.storage(a)) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl index 621f2e4bde..0ff841d50c 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl @@ -1,7 +1,14 @@ @eval module $(gensym()) + +# TODO: Test: +# zero (PermutedDimsArray) +# Custom zero type +# Slicing + using Test: @test, @testset, @test_broken -using NDTensors.SparseArrayDOKs: SparseArrayDOK, nonzero_keys, nonzero_length -using SparseArrays: nnz +using NDTensors.SparseArrayInterface: storage_indices, nstored +using NDTensors.SparseArrayDOKs: SparseArrayDOK +using SparseArrays: SparseMatrixCSC, nnz @testset "SparseArrayDOK (eltype=$elt)" for elt in (Float32, ComplexF32, Float64, ComplexF64) @testset "Basics" begin @@ -11,14 +18,14 @@ using SparseArrays: nnz @test a == SparseArrayDOK{elt}(undef, (3, 4)) @test iszero(a) @test iszero(nnz(a)) - @test nonzero_length(a) == nnz(a) + @test nstored(a) == nnz(a) @test size(a) == (3, 4) @test eltype(a) == elt for I in eachindex(a) @test iszero(a[I]) @test a[I] isa elt end - @test isempty(nonzero_keys(a)) + @test isempty(storage_indices(a)) x12 = randn(elt) x23 = randn(elt) @@ -31,23 +38,39 @@ using SparseArrays: nnz @test !iszero(b) @test b[1, 2] == x12 @test b[2, 3] == x23 - @test iszero(nonzero_length(a)) - @test_broken nonzero_length(b) == 2 - - # To test: - # reshape - # zero (PermutedDimsArray) - # map[!] - # broadcast - # Custom zero type - # conversion to `SparseMatrixCSC` + @test iszero(nstored(a)) + @test nstored(b) == 2 end @testset "map/broadcast" begin a = SparseArrayDOK{elt}(3, 4) a[1, 1] = 11 a[3, 4] = 34 - @test nonzero_length(a) == 2 - 2 * a + @test nstored(a) == 2 + b = 2 * a + @test nstored(b) == 2 + @test b[1, 1] == 2 * 11 + @test b[3, 4] == 2 * 34 + end + @testset "reshape" begin + a = SparseArrayDOK{elt}(2, 2, 2) + a[1, 2, 2] = 122 + b = reshape(a, 2, 4) + @test b[1, 4] == 122 + end + @testset "SparseMatrixCSC" begin + a = SparseArrayDOK{elt}(2, 2) + a[1, 2] = 12 + for (type, a′) in ((SparseMatrixCSC, a), (SparseArrayDOK, SparseMatrixCSC(a))) + b = type(a′) + @test b isa type{elt} + @test b[1, 2] == 12 + @test isone(nnz(b)) + for I in eachindex(b) + if I ≠ CartesianIndex(1, 2) + @test iszero(b[I]) + end + end + end end end end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl index afcde0bd48..5b76dbe4b4 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/SparseArrayInterface.jl @@ -6,6 +6,8 @@ include("base.jl") include("map.jl") include("copyto.jl") include("broadcast.jl") +include("conversion.jl") include("wrappers.jl") -include("SparseArrayInterfaceSparseArraysExt.jl") +include("zero.jl") +# include("SparseArrayInterfaceSparseArraysExt.jl") end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/base.jl index da01f5f671..ec3d8ffc19 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/base.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/base.jl @@ -1,20 +1,46 @@ -# TODO: Use `sparse_mapreduce`? +function sparse_reduce(op, a::AbstractArray; kwargs...) + return sparse_mapreduce(identity, op, a; kwargs...) +end + +function sparse_all(a::AbstractArray) + return sparse_reduce(&, a; init=true) +end + +function sparse_all(f, a::AbstractArray) + return sparse_mapreduce(f, &, a; init=true) +end + function sparse_iszero(a::AbstractArray) - iszero(nstored(a)) && return true - return iszero(storage(a)) + return sparse_all(iszero, a) end -# TODO: Use `sparse_mapreduce`? function sparse_isreal(a::AbstractArray) - iszero(nstored(a)) && return true - return isreal(storage(a)) + return sparse_all(isreal, a) end +# This is equivalent to: +# +# sparse_map!(Returns(x), a, a) +# +# but we don't use that here since `sparse_fill!` +# is used inside of `sparse_map!`. function sparse_fill!(a::AbstractArray, x) - sparse_map!(Returns(x), a, a) + if iszero(x) + empty_storage!(a) + end + fill!(storage(a), x) return a end +# This could just call `sparse_fill!` +# but it avoids a zero construction and check. +function sparse_zero!(a::AbstractArray) + empty_storage!(a) + fill!(storage(a), zero(eltype(a))) + return a +end + +# TODO: Make `sparse_zero!`? function sparse_zero(a::AbstractArray) # Need to overload `similar` for custom types a = similar(a) @@ -23,19 +49,31 @@ function sparse_zero(a::AbstractArray) return a end -# TODO: Make `sparse_one!`? -function sparse_one(a::AbstractMatrix) +# TODO: Is this a good definition? +function sparse_zero(arraytype::Type{<:AbstractArray}, dims::Tuple{Vararg{Int}}) + a = arraytype(undef, dims) + sparse_fill!(a, zero(eltype(a))) + return a +end + +function sparse_one!(a::AbstractMatrix) + sparse_zero!(a) m, n = size(a) @assert m == n - a = zero(a) - # TODO: Use `diagindices` from `DiagonalArrays`? for i in 1:m a[i, i] = one(eltype(a)) end return a end -# TODO: Use `sparse_mapreduce`? +# TODO: Make `sparse_one!`? +function sparse_one(a::AbstractMatrix) + a = sparse_zero(a) + sparse_one!(a) + return a +end + +# TODO: Use `sparse_mapreduce(==, &, a1, a2)`? function sparse_isequal(a1::AbstractArray, a2::AbstractArray) Is = collect(stored_indices(a1)) intersect!(Is, stored_indices(a2)) @@ -53,21 +91,9 @@ function sparse_reshape(a::AbstractArray, dims) a_reshaped = similar(a, dims) sparse_fill!(a_reshaped, zero(eltype(a))) linear_inds = LinearIndices(a) - new_cartesian_inds = CartesianIndices(dims) + dest_cartesian_inds = CartesianIndices(dims) for I in stored_indices(a) - a_reshaped[new_cartesian_inds[linear_inds[I]]] = a[I] + a_reshaped[dest_cartesian_inds[linear_inds[I]]] = a[I] end return a_reshaped end - -## function sparse_mapreduce(f, op, as::AbstractArray...) -## # TODO: Make more general. -## @assert iszero(mapreduce(f, op, map(a -> getindex_zero(a, first(eachindex(a))), as)...)) -## # TODO: Create a `promote_stored_indices`, for example to help -## # with specialized fast indexing for `DiagonalArrays`. -## Is = union(stored_indices.(as)...) -## for I in Is -## a_dest[I] = f(map(a -> a[I], as)...) -## end -## return reduce(op, a_dest) -## end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/conversion.jl b/NDTensors/src/lib/SparseArrayInterface/src/conversion.jl new file mode 100644 index 0000000000..57f1850ba0 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/conversion.jl @@ -0,0 +1,5 @@ +function sparse_convert(arraytype::Type{<:AbstractArray}, a::AbstractArray) + a_dest = sparse_zero(arraytype, size(a)) + sparse_copyto!(a_dest, a) + return a_dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl index 7d79e21c13..42a50e6264 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl @@ -40,10 +40,6 @@ function index_to_storage_index(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) whe return index_to_storage_index(a, CartesianIndex(I)) end -# Helper type for constructing zero values -struct GetIndexZero end -(::GetIndexZero)(a::AbstractArray, I) = zero(eltype(a)) - function sparse_getindex(a::AbstractArray, I::NotStoredIndex) return getindex_notstored(a, index(I)) end @@ -82,7 +78,7 @@ end # Linear indexing function sparse_getindex(a::AbstractArray, I::CartesianIndex{1}) - return error("Linear indexing not supported yet for sparse arrays") + return sparse_getindex(a, CartesianIndices(a)[I]) end # Handle trailing indices diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl index 1a881e8e49..3110ddc651 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl @@ -19,3 +19,11 @@ end function stored_indices(a::AbstractArray) return Iterators.map(Inz -> storage_index_to_index(a, Inz), storage_indices(a)) end + +# Empty the sparse storage if possible. +# Array types should overload `Base.dataids` to opt-in +# to aliasing detection with `Base.mightalias` +# to avoid emptying an input array in the case of `sparse_map!`. +# `empty_storage!` is used to zero out the output array. +# See also `Base.unalias` and `Base.unaliascopy`. +empty_storage!(a::AbstractArray) = a diff --git a/NDTensors/src/lib/SparseArrayInterface/src/map.jl b/NDTensors/src/lib/SparseArrayInterface/src/map.jl index 11ed25329e..f330768ba3 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/map.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/map.jl @@ -1,6 +1,6 @@ # Test if the function preserves zero values and therefore # preserves the sparsity structure. -function preserves_zero(f, a_dest, as...) +function preserves_zero(f, as...) return iszero(f(map(a -> sparse_getindex(a, NotStoredIndex(first(eachindex(a)))), as)...)) end @@ -29,8 +29,8 @@ function promote_stored_indices(as::AbstractArray...) end function sparse_map_stored!(f, a_dest::AbstractArray, as::AbstractArray...) - # TODO: Create a `promote_stored_indices`, for example to help - # with specialized fast indexing for `DiagonalArrays`. + # Need to zero out the destination. + sparse_zero!(a_dest) Is = promote_stored_indices(as...) sparse_map_indices!(f, a_dest, Is, as...) return a_dest @@ -45,10 +45,22 @@ end function sparse_map!(f, a_dest::AbstractArray, as::AbstractArray...) @assert allequal(axes.((a_dest, as...))) - if preserves_zero(f, a_dest, as...) + if preserves_zero(f, as...) + # Remove aliases to avoid overwriting inputs. + as = map(a -> Base.unalias(a_dest, a), as) sparse_map_stored!(f, a_dest, as...) else sparse_map_all!(f, a_dest, as...) end return a_dest end + +# TODO: Generalize to multiple arguements. +# TODO: Define `sparse_mapreducedim!`. +function sparse_mapreduce(f, op, a::AbstractArray; kwargs...) + output = mapreduce(f, op, storage(a); kwargs...) + # TODO: Use more general `zero` value. + # TODO: Better way to check that zeros don't affect the output? + @assert op(output, f(zero(eltype(a)))) == output + return output +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/zero.jl b/NDTensors/src/lib/SparseArrayInterface/src/zero.jl new file mode 100644 index 0000000000..24757594a4 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/src/zero.jl @@ -0,0 +1,4 @@ +# Represents a zero value and an index +# TODO: Rename `GetIndexZero`? +struct Zero end +(::Zero)(type::Type, I) = zero(type) diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 862123a2b1..9522130ab6 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -59,31 +59,19 @@ using Test: @test, @testset, @test_broken, @test_throws return a end + # Empty the storage, helps with efficiency in `map!` to drop + # zeros. + function SparseArrayInterface.empty_storage!(a::SparseArray) + empty!(a.data) + empty!(a.index_to_dataindex) + return empty!(a.dataindex_to_index) + end + # Broadcasting function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() end - # Base - function Base.iszero(a::SparseArray) - return SparseArrayInterface.sparse_iszero(a) - end - function Base.isreal(a::SparseArray) - return SparseArrayInterface.sparse_isreal(a) - end - function Base.zero(a::SparseArray) - return SparseArrayInterface.sparse_zero(a) - end - function Base.one(a::SparseArray) - return SparseArrayInterface.sparse_one(a) - end - function Base.:(==)(a1::SparseArray, a2::SparseArray) - return SparseArrayInterface.sparse_isequal(a1, a2) - end - function Base.reshape(a::SparseArray, dims::Tuple{Vararg{Int}}) - return SparseArrayInterface.sparse_reshape(a, dims) - end - # Map function Base.map!(f, dest::AbstractArray, src::SparseArray) SparseArrayInterface.sparse_map!(f, dest, src) @@ -102,6 +90,26 @@ using Test: @test, @testset, @test_broken, @test_throws return dest end + # Base + function Base.:(==)(a1::SparseArray, a2::SparseArray) + return SparseArrayInterface.sparse_isequal(a1, a2) + end + function Base.reshape(a::SparseArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) + end + function Base.iszero(a::SparseArray) + return SparseArrayInterface.sparse_iszero(a) + end + function Base.isreal(a::SparseArray) + return SparseArrayInterface.sparse_isreal(a) + end + function Base.zero(a::SparseArray) + return SparseArrayInterface.sparse_zero(a) + end + function Base.one(a::SparseArray) + return SparseArrayInterface.sparse_one(a) + end + # Test a = SparseArray{elt}(2, 3) @test size(a) == (2, 3) @@ -132,6 +140,8 @@ using Test: @test, @testset, @test_broken, @test_throws a = SparseArray{elt}(2, 3) a[1, 2] = 12 + @test a[1, 2] == 12 + @test a[3] == 12 # linear indexing @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) @test a[SparseArrayInterface.StorageIndex(1)] == 12 @@ -283,6 +293,64 @@ using Test: @test, @testset, @test_broken, @test_throws @test iszero(b[I]) end end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = randn(elt, 2, 3) + b .= a + @test a == b + for I in eachindex(a) + @test a[I] == b[I] + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = randn(elt, 2, 3) + b .= 2 .* a + @test 2 * a == b + for I in eachindex(a) + @test 2 * a[I] == b[I] + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = randn(elt, 2, 3) + b .= 2 .+ a + @test 2 .+ a == b + for I in eachindex(a) + @test 2 + a[I] == b[I] + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = randn(elt, 2, 3) + map!(x -> 2x, b, a) + @test 2 * a == b + for I in eachindex(a) + @test 2 * a[I] == b[I] + end + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = zeros(elt, 2, 3) + b[2, 1] = 21 + @test Array(a) == a + @test a + b == Array(a) + b + @test b + a == Array(a) + b + @test b .+ 2 .* a == 2 * Array(a) + b + @test a .+ 2 .* b == Array(a) + 2b + @test a + b isa Matrix{elt} + @test b + a isa Matrix{elt} + @test SparseArrayInterface.nstored(a + b) == length(a) + + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + b = zeros(elt, 2, 3) + b[2, 1] = 21 + a′ = copy(a) + a′ .+= b + @test a′ == a + b + @test SparseArrayInterface.nstored(a′) == 2 end @testset "Custom DiagonalArray" begin struct DiagonalArray{T,N} <: AbstractArray{T,N} diff --git a/NDTensors/src/lib/SparseArrayInterface/test/temp.jl b/NDTensors/src/lib/SparseArrayInterface/test/temp.jl deleted file mode 100644 index f6265bbfa7..0000000000 --- a/NDTensors/src/lib/SparseArrayInterface/test/temp.jl +++ /dev/null @@ -1,105 +0,0 @@ -using LinearAlgebra: norm -using NDTensors.SparseArrayInterface: SparseArrayInterface -using Test: @test -struct SparseArray{T,N} <: AbstractArray{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.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.nonzeros(a::SparseArray) = a.data -function SparseArrayInterface.index_to_nonzero_index( - a::SparseArray{<:Any,N}, I::CartesianIndex{N} -) where {N} - return get(a.index_to_dataindex, I, nothing) -end -SparseArrayInterface.nonzero_index_to_index(a::SparseArray, I) = a.dataindex_to_index[I] -function SparseArrayInterface.setindex_zero!( - 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 - -# Base -function Base.zero(a::SparseArray) - return SparseArrayInterface.sparse_zero(a) -end - -# Broadcasting -function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) - return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() -end - -# Map -function Base.map!(f, dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_map!(f, dest, src) - return dest -end -function Base.copy!(dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_copy!(dest, src) - return dest -end -function Base.copyto!(dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_copyto!(dest, src) - return dest -end -function Base.permutedims!(dest::AbstractArray, src::SparseArray, perm) - SparseArrayInterface.sparse_permutedims!(dest, src, perm) - return dest -end - -elt = Float64 - -# Test -a = SparseArray{elt}(2, 3) -@test size(a) == (2, 3) -@test axes(a) == (1:2, 1:3) -@test SparseArrayInterface.nonzeros(a) == elt[] -@test iszero(SparseArrayInterface.nonzero_length(a)) -@test collect(SparseArrayInterface.nonzero_indices(a)) == CartesianIndex{2}[] -@test iszero(a) -@test iszero(norm(a)) -for I in eachindex(a) - @test iszero(a) -end - -a[1, 2] = 12 -@test size(a) == (2, 3) -@test axes(a) == (1:2, 1:3) -@test SparseArrayInterface.nonzeros(a) == elt[12] -@test isone(SparseArrayInterface.nonzero_length(a)) -@test collect(SparseArrayInterface.nonzero_indices(a)) == [CartesianIndex(1, 2)] -@test !iszero(a) -@test !iszero(norm(a)) -for I in eachindex(a) - if I == CartesianIndex(1, 2) - @test a[I] == 12 - else - @test iszero(a[I]) - end -end From 25f82291ac1b59ecf3d77c111039e4a2f8455cc8 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 19:08:58 -0500 Subject: [PATCH 09/17] Fix namespace issues --- NDTensors/src/abstractarray/fill.jl | 3 +++ NDTensors/src/arraystorage/arraystorage/tensor/svd.jl | 2 ++ 2 files changed, 5 insertions(+) diff --git a/NDTensors/src/abstractarray/fill.jl b/NDTensors/src/abstractarray/fill.jl index dc4d10c1b7..b3be9331c7 100644 --- a/NDTensors/src/abstractarray/fill.jl +++ b/NDTensors/src/abstractarray/fill.jl @@ -1,3 +1,6 @@ +using .SetParameters: DefaultParameters, set_unspecified_parameters +using .Unwrap: unwrap_type + function generic_randn( arraytype::Type{<:AbstractArray}, dim::Integer=0; rng=Random.default_rng() ) diff --git a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl index 69f446ac53..e7ccec1513 100644 --- a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl +++ b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl @@ -1,3 +1,5 @@ +using .DiagonalArrays: DiagonalMatrix + backup_svd_alg(::Algorithm"divide_and_conquer") = Algorithm"qr_iteration"() backup_svd_alg(::Algorithm"qr_iteration") = Algorithm"recursive"() From 3b3c50ddc8ae7d594cd73070c2cecb7c8f3abdfe Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 19:24:14 -0500 Subject: [PATCH 10/17] One more namespace issue --- NDTensors/src/arraystorage/arraystorage/tensor/svd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl index e7ccec1513..2f97f41978 100644 --- a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl +++ b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl @@ -1,4 +1,4 @@ -using .DiagonalArrays: DiagonalMatrix +using .DiagonalArrays: DiagIndices, DiagonalMatrix backup_svd_alg(::Algorithm"divide_and_conquer") = Algorithm"qr_iteration"() backup_svd_alg(::Algorithm"qr_iteration") = Algorithm"recursive"() From a20e8198623d17f2a5effb8e6cf57b7f0ee9585f Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 28 Nov 2023 20:02:55 -0500 Subject: [PATCH 11/17] More namespace issues --- NDTensors/src/array/set_types.jl | 2 +- NDTensors/test/Project.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/NDTensors/src/array/set_types.jl b/NDTensors/src/array/set_types.jl index 739c562dfc..81b3bd88df 100644 --- a/NDTensors/src/array/set_types.jl +++ b/NDTensors/src/array/set_types.jl @@ -1,4 +1,4 @@ -using .SetParameters: Position, set_parameters +using .SetParameters: Position, get_parameter, set_parameters """ TODO: Use `Accessors.jl` notation: diff --git a/NDTensors/test/Project.toml b/NDTensors/test/Project.toml index 5301386cba..f47f1f7f21 100644 --- a/NDTensors/test/Project.toml +++ b/NDTensors/test/Project.toml @@ -13,6 +13,7 @@ Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" From d571bd9766e382f1e3d9a0dd7c5857441c99999f Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 29 Nov 2023 07:46:53 -0500 Subject: [PATCH 12/17] Julia 1.6 backwards compatibility --- .../src/lib/DiagonalArrays/src/DiagonalArrays.jl | 11 +++++++++-- NDTensors/src/lib/DiagonalArrays/src/diagview.jl | 2 ++ NDTensors/src/lib/SparseArrayInterface/src/map.jl | 2 ++ .../src/lib/SparseArrayInterface/test/Project.toml | 1 + .../src/lib/SparseArrayInterface/test/runtests.jl | 1 + 5 files changed, 15 insertions(+), 2 deletions(-) diff --git a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl index 7ebe04df06..077ae41fa7 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl @@ -1,7 +1,12 @@ module DiagonalArrays +# include("diagonalarrayinterface.jl") +# include("diagonalarray.jl") +# include("sparsearrayinterface.jl") +# include("DiagonalArraysLinearAlgebraExt.jl") +# include("diagonalmatrix.jl") +# include("diagonalvector.jl") -using Compat # allequal -using LinearAlgebra +using LinearAlgebra: LinearAlgebra export DiagonalArray, DiagonalMatrix, DiagonalVector, DiagIndex, DiagIndices, densearray @@ -33,6 +38,8 @@ Base.size(a::DiagonalArray) = a.dims # TODO: Rename `storage_values`. diagview(a::DiagonalArray) = a.diag + +# DiagonalArraysLinearAlgebraExt.jl LinearAlgebra.diag(a::DiagonalArray) = copy(diagview(a)) function DiagonalArray{T,N}( diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagview.jl b/NDTensors/src/lib/DiagonalArrays/src/diagview.jl index b3e2547362..cde0b7b561 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/diagview.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/diagview.jl @@ -1,3 +1,5 @@ +using Compat: allequal + # Convert to an offset along the diagonal. # Otherwise, return `nothing`. function diagindex(a::AbstractArray{T,N}, I::CartesianIndex{N}) where {T,N} diff --git a/NDTensors/src/lib/SparseArrayInterface/src/map.jl b/NDTensors/src/lib/SparseArrayInterface/src/map.jl index f330768ba3..6b50167130 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/map.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/map.jl @@ -1,3 +1,5 @@ +using Compat: allequal + # Test if the function preserves zero values and therefore # preserves the sparsity structure. function preserves_zero(f, as...) diff --git a/NDTensors/src/lib/SparseArrayInterface/test/Project.toml b/NDTensors/src/lib/SparseArrayInterface/test/Project.toml index a36a39a20a..4dd438d7ae 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/Project.toml +++ b/NDTensors/src/lib/SparseArrayInterface/test/Project.toml @@ -1,4 +1,5 @@ [deps] +Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 9522130ab6..0e2de3ef0a 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -1,4 +1,5 @@ @eval module $(gensym()) +using Compat: Returns, allequal using Test: @test, @testset, @test_broken, @test_throws @testset "SparseArrayInterface (eltype=$elt)" for elt in From 410df32cb841bd86d3b7b224e490a2c7924b85f8 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 29 Nov 2023 13:24:04 -0500 Subject: [PATCH 13/17] Use SparseArrayInterface in DiagonalArrays --- .../arraystorage/arraystorage/tensor/svd.jl | 2 +- .../lib/DiagonalArrays/examples/Project.toml | 3 + .../src/lib/DiagonalArrays/examples/README.jl | 27 ++- .../lib/DiagonalArrays/src/DiagonalArrays.jl | 218 +----------------- .../src/lib/DiagonalArrays/src/defaults.jl | 3 + .../src/lib/DiagonalArrays/src/diagindex.jl | 14 ++ .../src/lib/DiagonalArrays/src/diagindices.jl | 14 ++ .../lib/DiagonalArrays/src/diaginterface.jl | 52 +++++ .../lib/DiagonalArrays/src/diagonalarray.jl | 67 ++++++ .../src/diagonalarraydiaginterface.jl | 23 ++ .../lib/DiagonalArrays/src/diagonalmatrix.jl | 5 + .../lib/DiagonalArrays/src/diagonalvector.jl | 5 + .../src/lib/DiagonalArrays/src/diagview.jl | 84 ------- .../src/sparsearrayinterface.jl | 69 ++++++ .../src/lib/SparseArrayInterface/src/base.jl | 38 ++- .../lib/SparseArrayInterface/src/indexing.jl | 55 ++++- .../src/interface_optional.jl | 5 + .../DiagonalArrays.jl | 93 ++++++++ .../SparseArrayInterfaceTestUtils.jl | 4 + .../SparseArrays.jl | 98 ++++++++ .../lib/SparseArrayInterface/test/runtests.jl | 212 +++-------------- 21 files changed, 594 insertions(+), 497 deletions(-) create mode 100644 NDTensors/src/lib/DiagonalArrays/examples/Project.toml create mode 100644 NDTensors/src/lib/DiagonalArrays/src/defaults.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagindex.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagindices.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diaginterface.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagonalarray.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagonalarraydiaginterface.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagonalmatrix.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagonalvector.jl delete mode 100644 NDTensors/src/lib/DiagonalArrays/src/diagview.jl create mode 100644 NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl create mode 100644 NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl diff --git a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl index 2f97f41978..947f5e82a7 100644 --- a/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl +++ b/NDTensors/src/arraystorage/arraystorage/tensor/svd.jl @@ -113,7 +113,7 @@ function svd( # Make the new indices to go onto U and V # TODO: Put in a separate function, such as # `rewrap_inds` or something like that. - dS = length(S[DiagIndices()]) + dS = length(S[DiagIndices(:)]) indstype = typeof(inds(T)) u = eltype(indstype)(dS) v = eltype(indstype)(dS) diff --git a/NDTensors/src/lib/DiagonalArrays/examples/Project.toml b/NDTensors/src/lib/DiagonalArrays/examples/Project.toml new file mode 100644 index 0000000000..ef491a529c --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/examples/Project.toml @@ -0,0 +1,3 @@ +[deps] +NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/NDTensors/src/lib/DiagonalArrays/examples/README.jl b/NDTensors/src/lib/DiagonalArrays/examples/README.jl index fb463d9ed9..26733d3ba7 100644 --- a/NDTensors/src/lib/DiagonalArrays/examples/README.jl +++ b/NDTensors/src/lib/DiagonalArrays/examples/README.jl @@ -2,7 +2,7 @@ # # A Julia `DiagonalArray` type. -using NDTensors.DiagonalArrays: DiagonalArray, DiagIndex, DiagIndices +using NDTensors.DiagonalArrays: DiagonalArray, DiagIndex, DiagIndices, isdiagindex using Test function main() @@ -14,7 +14,11 @@ function main() d[2, 2, 2] = 22 @test d[2, 2, 2] == 22 - @test length(d[DiagIndices()]) == 3 + d_r = reshape(d, 3, 20) + @test size(d_r) == (3, 20) + @test all(I -> d_r[I] == d[I], LinearIndices(d)) + + @test length(d[DiagIndices(:)]) == 3 @test Array(d) == d @test d[DiagIndex(2)] == d[2, 2, 2] @@ -23,15 +27,24 @@ function main() a = randn(3, 4, 5) new_diag = randn(3) - a[DiagIndices()] = new_diag - d[DiagIndices()] = a[DiagIndices()] + a[DiagIndices(:)] = new_diag + d[DiagIndices(:)] = a[DiagIndices(:)] - @test a[DiagIndices()] == new_diag - @test d[DiagIndices()] == new_diag + @test a[DiagIndices(:)] == new_diag + @test d[DiagIndices(:)] == new_diag permuted_d = permutedims(d, (3, 2, 1)) @test permuted_d isa DiagonalArray - @test permuted_d == d + @test permuted_d[DiagIndices(:)] == d[DiagIndices(:)] + @test size(d) == (3, 4, 5) + @test size(permuted_d) == (5, 4, 3) + for I in eachindex(d) + if !isdiagindex(d, I) + @test iszero(d[I]) + else + @test !iszero(d[I]) + end + end mapped_d = map(x -> 2x, d) @test mapped_d isa DiagonalArray diff --git a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl index 077ae41fa7..60e259a17e 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl @@ -1,212 +1,10 @@ module DiagonalArrays -# include("diagonalarrayinterface.jl") -# include("diagonalarray.jl") -# include("sparsearrayinterface.jl") -# include("DiagonalArraysLinearAlgebraExt.jl") -# include("diagonalmatrix.jl") -# include("diagonalvector.jl") - -using LinearAlgebra: LinearAlgebra - -export DiagonalArray, DiagonalMatrix, DiagonalVector, DiagIndex, DiagIndices, densearray - -include("diagview.jl") - -# TODO: Make use of `Zero` type in `SparseArrayInterface`. -struct DefaultZero end - -function (::DefaultZero)(eltype::Type, I::CartesianIndex) - return zero(eltype) -end - -# TODO: Rename `diag` to `nonzero_values`, to make more -# generic? -struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractArray{T,N} - diag::Diag - dims::NTuple{N,Int} - zero::Zero -end - -# TODO: Use `Accessors.jl`. -# TODO: Make more generic, call `set_nonzero_values`. -set_diag(a::DiagonalArray, diag) = DiagonalArray(diag, size(a), a.zero) - -Base.copy(a::DiagonalArray) = set_diag(a, copy(a[DiagIndices()])) -Base.similar(a::DiagonalArray) = set_diag(a, similar(a[DiagIndices()])) - -Base.size(a::DiagonalArray) = a.dims - -# TODO: Rename `storage_values`. -diagview(a::DiagonalArray) = a.diag - -# DiagonalArraysLinearAlgebraExt.jl -LinearAlgebra.diag(a::DiagonalArray) = copy(diagview(a)) - -function DiagonalArray{T,N}( - diag::AbstractVector{T}, d::Tuple{Vararg{Int,N}}, zero=DefaultZero() -) where {T,N} - return DiagonalArray{T,N,typeof(diag),typeof(zero)}(diag, d, zero) -end - -function DiagonalArray{T,N}( - diag::AbstractVector, d::Tuple{Vararg{Int,N}}, zero=DefaultZero() -) where {T,N} - return DiagonalArray{T,N}(T.(diag), d, zero) -end - -function DiagonalArray{T,N}(diag::AbstractVector, d::Vararg{Int,N}) where {T,N} - return DiagonalArray{T,N}(diag, d) -end - -function DiagonalArray{T}( - diag::AbstractVector, d::Tuple{Vararg{Int,N}}, zero=DefaultZero() -) where {T,N} - return DiagonalArray{T,N}(diag, d, zero) -end - -function DiagonalArray{T}(diag::AbstractVector, d::Vararg{Int,N}) where {T,N} - return DiagonalArray{T,N}(diag, d) -end - -function DiagonalArray(diag::AbstractVector{T}, d::Tuple{Vararg{Int,N}}) where {T,N} - return DiagonalArray{T,N}(diag, d) -end - -function DiagonalArray(diag::AbstractVector{T}, d::Vararg{Int,N}) where {T,N} - return DiagonalArray{T,N}(diag, d) -end - -default_size(diag::AbstractVector, n) = ntuple(Returns(length(diag)), n) - -# Infer size from diagonal -function DiagonalArray{T,N}(diag::AbstractVector) where {T,N} - return DiagonalArray{T,N}(diag, default_size(diag, N)) -end - -function DiagonalArray{<:Any,N}(diag::AbstractVector{T}) where {T,N} - return DiagonalArray{T,N}(diag) -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) -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) -end - -function DiagonalArray{T}(::UndefInitializer, d::Vararg{Int,N}) where {T,N} - return DiagonalArray{T,N}(undef, d) -end - -# TODO: Use `sparse_getindex` from `SparseArrayInterface`. -function Base.getindex(a::DiagonalArray{T,N}, I::CartesianIndex{N}) where {T,N} - i = diagindex(a, I) - isnothing(i) && return a.zero(T, I) - return a[DiagIndex(i)] -end - -function Base.getindex(a::DiagonalArray{T,N}, I::Vararg{Int,N}) where {T,N} - return a[CartesianIndex(I)] -end - -# TODO: Use `sparse_setindex!` from `SparseArrayInterface`. -function Base.setindex!(a::DiagonalArray{T,N}, v, I::CartesianIndex{N}) where {T,N} - i = diagindex(a, I) - isnothing(i) && return error("Can't set off-diagonal element of DiagonalArray") - a[DiagIndex(i)] = v - return a -end - -function Base.setindex!(a::DiagonalArray{T,N}, v, I::Vararg{Int,N}) where {T,N} - a[CartesianIndex(I)] = v - return a -end - -# Make dense. -function densearray(a::DiagonalArray) - # TODO: Check this works on GPU. - # TODO: Make use of `a.zero`? - d = similar(diagview(a), size(a)) - # TODO: Use `getindex_zero(a)` or `zero_value(a)`. - fill!(d, zero(eltype(a))) - diagcopyto!(d, a) - return d -end - -# TODO: Make generic in `SparseArrayInterface`, use a `SparseIndices(:)` -# iterator? -function Base.permutedims!(a_dest::AbstractArray, a_src::DiagonalArray, perm) - @assert ndims(a_src) == ndims(a_dest) == length(perm) - a_dest[DiagIndices()] = a_src[DiagIndices()] - return a_dest -end - -# TODO: Should this copy? `LinearAlgebra.Diagonal` does not copy -# with `permutedims`. -# TODO: Use `copy(a)` or `permutedims!!`? May be better for immutable diagonals. -function Base.permutedims(a::DiagonalArray, perm) - a_dest = similar(a) - permutedims!(a_dest, a, perm) - return a_dest -end - -# This function is used by Julia broadcasting for sparse arrays -# to decide how to allocated the output: -# LinearAlgebra.fzeropreserving(Broadcast.Broadcasted(f, (a,))) -# If it preserves zeros value, it keeps it structured, if not -# it allocates an Array. -# TODO: Introduce an `IsSparse` trait? This would be generic -# for any type that implements `map_nonzeros!`. -function Base.map(f, a::DiagonalArray) - # TODO: Test `iszero(f(x))`. - return map_nonzeros(f, a) -end - -# TODO: Make this generic in `SparseArrayInterface`. -function map_diag(f, a::DiagonalArray) - return set_diag(a, map(f, a[DiagIndices()])) -end - -# API from `SparseArray`/`BlockSparseArray` -function map_nonzeros(f, a::DiagonalArray) - return map_diag(f, a) -end - -# TODO: Introduce an `IsSparse` trait? This would be generic -# for any type that implements `map_nonzeros!`. -function Base.map!(f, a_dest::AbstractArray, a_src::DiagonalArray) - # TODO: Test `iszero(f(x))`. - map_nonzeros!(f, a_dest, a_src) - return a_dest -end - -function map_diag!(f, a_dest::AbstractArray, a_src::DiagonalArray) - map!(f, a_dest[DiagIndices()], a_src[DiagIndices()]) - return a_dest -end - -# API from `SparseArray`/`BlockSparseArray` -function map_nonzeros!(f, a_dest::AbstractArray, a_src::DiagonalArray) - map_diag!(f, a_dest, a_src) - return a_dest -end - -const DiagonalMatrix{T,Diag,Zero} = DiagonalArray{T,2,Diag,Zero} - -function DiagonalMatrix(diag::AbstractVector) - return DiagonalArray{<:Any,2}(diag) -end - -const DiagonalVector{T,Diag,Zero} = DiagonalArray{T,1,Diag,Zero} - -function DiagonalVector(diag::AbstractVector) - return DiagonalArray{<:Any,1}(diag) -end - +include("diaginterface.jl") +include("diagindex.jl") +include("diagindices.jl") +include("diagonalarray.jl") +include("sparsearrayinterface.jl") +include("diagonalarraydiaginterface.jl") +include("diagonalmatrix.jl") +include("diagonalvector.jl") end diff --git a/NDTensors/src/lib/DiagonalArrays/src/defaults.jl b/NDTensors/src/lib/DiagonalArrays/src/defaults.jl new file mode 100644 index 0000000000..39e9395eb6 --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/defaults.jl @@ -0,0 +1,3 @@ +using Compat: Returns + +default_size(diag::AbstractVector, n) = ntuple(Returns(length(diag)), n) diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagindex.jl b/NDTensors/src/lib/DiagonalArrays/src/diagindex.jl new file mode 100644 index 0000000000..e177e80697 --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagindex.jl @@ -0,0 +1,14 @@ +# Represents a linear offset along the diagonal +struct DiagIndex{I} + i::I +end +index(i::DiagIndex) = i.i + +function Base.getindex(a::AbstractArray, i::DiagIndex) + return getdiagindex(a, index(i)) +end + +function Base.setindex!(a::AbstractArray, value, i::DiagIndex) + setdiagindex!(a, value, index(i)) + return a +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagindices.jl b/NDTensors/src/lib/DiagonalArrays/src/diagindices.jl new file mode 100644 index 0000000000..08590d148e --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagindices.jl @@ -0,0 +1,14 @@ +# Represents a set of linear offsets along the diagonal +struct DiagIndices{I} + i::I +end +indices(i::DiagIndices) = i.i + +function Base.getindex(a::AbstractArray, I::DiagIndices) + return getdiagindices(a, indices(I)) +end + +function Base.setindex!(a::AbstractArray, value, i::DiagIndices) + setdiagindices!(a, value, indices(i)) + return a +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diaginterface.jl b/NDTensors/src/lib/DiagonalArrays/src/diaginterface.jl new file mode 100644 index 0000000000..de7c8ad51c --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diaginterface.jl @@ -0,0 +1,52 @@ +using Compat: allequal + +function isdiagindex(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} + @boundscheck checkbounds(a, I) + return allequal(Tuple(I)) +end + +function diagstride(a::AbstractArray) + s = 1 + p = 1 + for i in 1:(ndims(a) - 1) + p *= size(a, i) + s += p + end + return s +end + +function diagindices(a::AbstractArray) + diaglength = minimum(size(a)) + maxdiag = LinearIndices(a)[CartesianIndex(ntuple(Returns(diaglength), ndims(a)))] + return 1:diagstride(a):maxdiag +end + +function diagindices(a::AbstractArray{<:Any,0}) + return Base.OneTo(1) +end + +function diagview(a::AbstractArray) + return @view a[diagindices(a)] +end + +function getdiagindex(a::AbstractArray, i::Integer) + return diagview(a)[i] +end + +function setdiagindex!(a::AbstractArray, v, i::Integer) + diagview(a)[i] = v + return a +end + +function getdiagindices(a::AbstractArray, I) + return @view diagview(a)[I] +end + +function getdiagindices(a::AbstractArray, I::Colon) + return diagview(a) +end + +function setdiagindices!(a::AbstractArray, v, i::Colon) + diagview(a) .= v + return a +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagonalarray.jl b/NDTensors/src/lib/DiagonalArrays/src/diagonalarray.jl new file mode 100644 index 0000000000..13dbcc59f4 --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagonalarray.jl @@ -0,0 +1,67 @@ +using ..SparseArrayInterface: Zero + +struct DiagonalArray{T,N,Diag<:AbstractVector{T},Zero} <: AbstractArray{T,N} + diag::Diag + dims::NTuple{N,Int} + zero::Zero +end + +function DiagonalArray{T,N}( + diag::AbstractVector{T}, d::Tuple{Vararg{Int,N}}, zero=Zero() +) where {T,N} + return DiagonalArray{T,N,typeof(diag),typeof(zero)}(diag, d, zero) +end + +function DiagonalArray{T,N}( + diag::AbstractVector, d::Tuple{Vararg{Int,N}}, zero=Zero() +) where {T,N} + return DiagonalArray{T,N}(T.(diag), d, zero) +end + +function DiagonalArray{T,N}(diag::AbstractVector, d::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(diag, d) +end + +function DiagonalArray{T}( + diag::AbstractVector, d::Tuple{Vararg{Int,N}}, zero=Zero() +) where {T,N} + return DiagonalArray{T,N}(diag, d, zero) +end + +function DiagonalArray{T}(diag::AbstractVector, d::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(diag, d) +end + +function DiagonalArray(diag::AbstractVector{T}, d::Tuple{Vararg{Int,N}}) where {T,N} + return DiagonalArray{T,N}(diag, d) +end + +function DiagonalArray(diag::AbstractVector{T}, d::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(diag, d) +end + +# Infer size from diagonal +function DiagonalArray{T,N}(diag::AbstractVector) where {T,N} + return DiagonalArray{T,N}(diag, default_size(diag, N)) +end + +function DiagonalArray{<:Any,N}(diag::AbstractVector{T}) where {T,N} + return DiagonalArray{T,N}(diag) +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) +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) +end + +function DiagonalArray{T}(::UndefInitializer, d::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(undef, d) +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagonalarraydiaginterface.jl b/NDTensors/src/lib/DiagonalArrays/src/diagonalarraydiaginterface.jl new file mode 100644 index 0000000000..3b6c17bfae --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagonalarraydiaginterface.jl @@ -0,0 +1,23 @@ +using ..SparseArrayInterface: SparseArrayInterface, StorageIndex, StorageIndices + +SparseArrayInterface.StorageIndex(i::DiagIndex) = StorageIndex(index(i)) + +function Base.getindex(a::DiagonalArray, i::DiagIndex) + return a[StorageIndex(i)] +end + +function Base.setindex!(a::DiagonalArray, value, i::DiagIndex) + a[StorageIndex(i)] = value + return a +end + +SparseArrayInterface.StorageIndices(i::DiagIndices) = StorageIndices(indices(i)) + +function Base.getindex(a::DiagonalArray, i::DiagIndices) + return a[StorageIndices(i)] +end + +function Base.setindex!(a::DiagonalArray, value, i::DiagIndices) + a[StorageIndices(i)] = value + return a +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagonalmatrix.jl b/NDTensors/src/lib/DiagonalArrays/src/diagonalmatrix.jl new file mode 100644 index 0000000000..873410db78 --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagonalmatrix.jl @@ -0,0 +1,5 @@ +const DiagonalMatrix{T,Diag,Zero} = DiagonalArray{T,2,Diag,Zero} + +function DiagonalMatrix(diag::AbstractVector) + return DiagonalArray{<:Any,2}(diag) +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagonalvector.jl b/NDTensors/src/lib/DiagonalArrays/src/diagonalvector.jl new file mode 100644 index 0000000000..40e35e409d --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/diagonalvector.jl @@ -0,0 +1,5 @@ +const DiagonalVector{T,Diag,Zero} = DiagonalArray{T,1,Diag,Zero} + +function DiagonalVector(diag::AbstractVector) + return DiagonalArray{<:Any,1}(diag) +end diff --git a/NDTensors/src/lib/DiagonalArrays/src/diagview.jl b/NDTensors/src/lib/DiagonalArrays/src/diagview.jl deleted file mode 100644 index cde0b7b561..0000000000 --- a/NDTensors/src/lib/DiagonalArrays/src/diagview.jl +++ /dev/null @@ -1,84 +0,0 @@ -using Compat: allequal - -# Convert to an offset along the diagonal. -# Otherwise, return `nothing`. -function diagindex(a::AbstractArray{T,N}, I::CartesianIndex{N}) where {T,N} - !allequal(Tuple(I)) && return nothing - return first(Tuple(I)) -end - -function diagindex(a::AbstractArray{T,N}, I::Vararg{Int,N}) where {T,N} - return diagindex(a, CartesianIndex(I)) -end - -function getdiagindex(a::AbstractArray, i::Integer) - return diagview(a)[i] -end - -function setdiagindex!(a::AbstractArray, v, i::Integer) - diagview(a)[i] = v - return a -end - -struct DiagIndex - I::Int -end - -function Base.getindex(a::AbstractArray, i::DiagIndex) - return getdiagindex(a, i.I) -end - -function Base.setindex!(a::AbstractArray, v, i::DiagIndex) - setdiagindex!(a, v, i.I) - return a -end - -function setdiag!(a::AbstractArray, v) - copyto!(diagview(a), v) - return a -end - -function diaglength(a::AbstractArray) - # length(diagview(a)) - return minimum(size(a)) -end - -function diagstride(a::AbstractArray) - s = 1 - p = 1 - for i in 1:(ndims(a) - 1) - p *= size(a, i) - s += p - end - return s -end - -function diagindices(a::AbstractArray) - diaglength = minimum(size(a)) - maxdiag = LinearIndices(a)[CartesianIndex(ntuple(Returns(diaglength), ndims(a)))] - return 1:diagstride(a):maxdiag -end - -function diagindices(a::AbstractArray{<:Any,0}) - return Base.OneTo(1) -end - -function diagview(a::AbstractArray) - return @view a[diagindices(a)] -end - -function diagcopyto!(dest::AbstractArray, src::AbstractArray) - copyto!(diagview(dest), diagview(src)) - return dest -end - -struct DiagIndices end - -function Base.getindex(a::AbstractArray, ::DiagIndices) - return diagview(a) -end - -function Base.setindex!(a::AbstractArray, v, ::DiagIndices) - setdiag!(a, v) - return a -end diff --git a/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl b/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl new file mode 100644 index 0000000000..dc390c6e30 --- /dev/null +++ b/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl @@ -0,0 +1,69 @@ +using Compat: Returns, allequal +using ..SparseArrayInterface: SparseArrayInterface +# TODO: Put into `DiagonalArraysSparseArrayDOKsExt`? +using ..SparseArrayDOKs: SparseArrayDOK + +# Minimal interface +SparseArrayInterface.storage(a::DiagonalArray) = a.diag + +function SparseArrayInterface.index_to_storage_index( + a::DiagonalArray{<:Any,N}, I::CartesianIndex{N} +) where {N} + !allequal(Tuple(I)) && return nothing + return first(Tuple(I)) +end + +function SparseArrayInterface.storage_index_to_index(a::DiagonalArray, I) + return CartesianIndex(ntuple(Returns(I), ndims(a))) +end + +# 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) +end + +# 1-dimensional case can be `DiagonalArray`. +function SparseArrayInterface.sparse_similar(a::DiagonalArray, elt::Type, dims::Tuple{Int}) + return similar(a, elt, dims) +end + +# AbstractArray interface +Base.size(a::DiagonalArray) = a.dims + +function Base.getindex(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_getindex(a, I...) +end + +function Base.setindex!(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_setindex!(a, I...) +end + +function Base.similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) + return DiagonalArray{elt}(undef, dims) +end + +# AbstractArray functionality +# broadcast +function Broadcast.BroadcastStyle(arraytype::Type{<:DiagonalArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() +end + +# map +function Base.map!(f, dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest +end + +# permutedims +function Base.permutedims!(dest::AbstractArray, src::DiagonalArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest +end + +# reshape +function Base.reshape(a::DiagonalArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/base.jl index ec3d8ffc19..cc43a3a224 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/base.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/base.jl @@ -1,3 +1,22 @@ +# This is used when a sparse output structure not matching +# the input structure is needed, for example when reshaping +# a DiagonalArray. Overload: +# +# sparse_similar(a::AbstractArray, elt::Type, dims::Tuple{Vararg{Int}}) +# +# as needed. +function sparse_similar(a::AbstractArray, elt::Type) + return similar(a, elt, size(a)) +end + +function sparse_similar(a::AbstractArray, dims::Tuple{Vararg{Int}}) + return sparse_similar(a, eltype(a), dims) +end + +function sparse_similar(a::AbstractArray) + return sparse_similar(a, eltype(a), size(a)) +end + function sparse_reduce(op, a::AbstractArray; kwargs...) return sparse_mapreduce(identity, op, a; kwargs...) end @@ -86,14 +105,19 @@ function sparse_isequal(a1::AbstractArray, a2::AbstractArray) return true end -function sparse_reshape(a::AbstractArray, dims) - @assert length(a) == prod(dims) - a_reshaped = similar(a, dims) - sparse_fill!(a_reshaped, zero(eltype(a))) - linear_inds = LinearIndices(a) +function sparse_reshape!(a_dest::AbstractArray, a_src::AbstractArray, dims) + @assert length(a_src) == prod(dims) + sparse_fill!(a_dest, zero(eltype(a_src))) + linear_inds = LinearIndices(a_src) dest_cartesian_inds = CartesianIndices(dims) - for I in stored_indices(a) - a_reshaped[dest_cartesian_inds[linear_inds[I]]] = a[I] + for I in stored_indices(a_src) + a_dest[dest_cartesian_inds[linear_inds[I]]] = a_src[I] end + return a_dest +end + +function sparse_reshape(a::AbstractArray, dims) + a_reshaped = sparse_similar(a, dims) + sparse_reshape!(a_reshaped, a, dims) return a_reshaped end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl index 42a50e6264..0e94b39c38 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/indexing.jl @@ -87,7 +87,7 @@ function sparse_getindex(a::AbstractArray, I::CartesianIndex) length(t) < ndims(a) && error("Not enough indices passed") I′ = ntuple(i -> t[i], ndims(a)) @assert all(i -> isone(I[i]), (ndims(a) + 1):length(I)) - return _sparse_getindex(a, I′) + return _sparse_getindex(a, CartesianIndex(I′)) end # Update a nonzero value @@ -96,11 +96,40 @@ function sparse_setindex!(a::AbstractArray, value, I::StorageIndex) return a end -function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}) where {N} +# Implementation of element access +function _sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::CartesianIndex{N}) where {N} + @boundscheck checkbounds(a, I) + sparse_setindex!(a, value, MaybeStoredIndex(a, I)) + return a +end + +# Ambiguity with linear indexing +function sparse_setindex!(a::AbstractArray{<:Any,1}, value, I::CartesianIndex{1}) + _sparse_setindex!(a, value, I) + return a +end + +# Handle trailing indices or linear indexing +function sparse_setindex!(a::AbstractArray, value, I::Vararg{Int}) sparse_setindex!(a, value, CartesianIndex(I)) return a end +# Linear indexing +function sparse_setindex!(a::AbstractArray, value, I::CartesianIndex{1}) + sparse_setindex!(a, value, CartesianIndices(a)[I]) + return a +end + +# Handle trailing indices +function sparse_setindex!(a::AbstractArray, value, I::CartesianIndex) + t = Tuple(I) + length(t) < ndims(a) && error("Not enough indices passed") + I′ = ntuple(i -> t[i], ndims(a)) + @assert all(i -> isone(I[i]), (ndims(a) + 1):length(I)) + return _sparse_setindex!(a, value, CartesianIndex(I′)) +end + function sparse_setindex!(a::AbstractArray, value, I::StoredIndex) sparse_setindex!(a, value, StorageIndex(I)) return a @@ -113,7 +142,25 @@ function sparse_setindex!(a::AbstractArray, value, I::NotStoredIndex) return a end -function sparse_setindex!(a::AbstractArray{<:Any,N}, value, I::CartesianIndex{N}) where {N} - sparse_setindex!(a, value, MaybeStoredIndex(a, I)) +# A set of indices into the storage of the sparse array. +struct StorageIndices{I} + i::I +end +indices(i::StorageIndices) = i.i + +function sparse_getindex(a::AbstractArray, I::StorageIndices{Colon}) + return storage(a) +end + +function sparse_getindex(a::AbstractArray, I::StorageIndices) + return error("Not implemented") +end + +function sparse_setindex!(a::AbstractArray, value, I::StorageIndices{Colon}) + storage(a) .= value return a end + +function sparse_setindex!(a::AbstractArray, value, I::StorageIndices) + return error("Not implemented") +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl index 3110ddc651..d1f8f0bb2e 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/interface_optional.jl @@ -27,3 +27,8 @@ end # `empty_storage!` is used to zero out the output array. # See also `Base.unalias` and `Base.unaliascopy`. empty_storage!(a::AbstractArray) = a + +# Overload +function sparse_similar(a::AbstractArray, elt::Type, dims::Tuple{Vararg{Int}}) + return similar(a, elt, dims) +end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl new file mode 100644 index 0000000000..2ea8d4690c --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl @@ -0,0 +1,93 @@ +module DiagonalArrays +using NDTensors.SparseArrayInterface: SparseArrayInterface + +struct DiagonalArray{T,N} <: AbstractArray{T,N} + data::Vector{T} + dims::Tuple{Vararg{Int,N}} +end +function DiagonalArray{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int,N}}) where {T,N} + return DiagonalArray{T,N}(Vector{T}(undef, minimum(dims)), dims) +end +function DiagonalArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} + return DiagonalArray{T,N}(undef, dims) +end +function DiagonalArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} + return DiagonalArray{T,length(dims)}(undef, dims) +end +function DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int}) where {T} + return DiagonalArray{T}(undef, dims) +end + +# AbstractArray interface +Base.size(a::DiagonalArray) = a.dims +function Base.getindex(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_getindex(a, I...) +end +function Base.setindex!(a::DiagonalArray, I...) + return SparseArrayInterface.sparse_setindex!(a, I...) +end +function Base.similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) + return DiagonalArray{elt}(undef, dims) +end + +# Minimal interface +SparseArrayInterface.storage(a::DiagonalArray) = a.data +function SparseArrayInterface.index_to_storage_index( + a::DiagonalArray{<:Any,N}, I::CartesianIndex{N} +) where {N} + !allequal(Tuple(I)) && return nothing + return first(Tuple(I)) +end +function SparseArrayInterface.storage_index_to_index(a::DiagonalArray, I) + return CartesianIndex(ntuple(Returns(I), ndims(a))) +end +function SparseArrayInterface.sparse_similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) + return Array{elt}(undef, dims) +end +function SparseArrayInterface.sparse_similar(a::DiagonalArray, elt::Type, dims::Tuple{Int}) + return similar(a, elt, dims) +end + +# Broadcasting +function Broadcast.BroadcastStyle(arraytype::Type{<:DiagonalArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() +end + +# Base +function Base.iszero(a::DiagonalArray) + return SparseArrayInterface.sparse_iszero(a) +end +function Base.isreal(a::DiagonalArray) + return SparseArrayInterface.sparse_isreal(a) +end +function Base.zero(a::DiagonalArray) + return SparseArrayInterface.sparse_zero(a) +end +function Base.one(a::DiagonalArray) + return SparseArrayInterface.sparse_one(a) +end +function Base.:(==)(a1::DiagonalArray, a2::DiagonalArray) + return SparseArrayInterface.sparse_isequal(a1, a2) +end +function Base.reshape(a::DiagonalArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) +end + +# Map +function Base.map!(f, dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest +end +function Base.copy!(dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_copy!(dest, src) + return dest +end +function Base.copyto!(dest::AbstractArray, src::DiagonalArray) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest +end +function Base.permutedims!(dest::AbstractArray, src::DiagonalArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest +end +end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl new file mode 100644 index 0000000000..01d00f8dc0 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl @@ -0,0 +1,4 @@ +module SparseArrayInterfaceTestUtils +include("SparseArrays.jl") +include("DiagonalArrays.jl") +end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl new file mode 100644 index 0000000000..8a0c244f95 --- /dev/null +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl @@ -0,0 +1,98 @@ +module SparseArrays +using NDTensors.SparseArrayInterface: SparseArrayInterface + +struct SparseArray{T,N} <: AbstractArray{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.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.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.empty_storage!(a::SparseArray) + empty!(a.data) + empty!(a.index_to_dataindex) + return empty!(a.dataindex_to_index) +end + +# Broadcasting +function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) + return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() +end + +# Map +function Base.map!(f, dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_map!(f, dest, src) + return dest +end +function Base.copy!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copy!(dest, src) + return dest +end +function Base.copyto!(dest::AbstractArray, src::SparseArray) + SparseArrayInterface.sparse_copyto!(dest, src) + return dest +end +function Base.permutedims!(dest::AbstractArray, src::SparseArray, perm) + SparseArrayInterface.sparse_permutedims!(dest, src, perm) + return dest +end + +# Base +function Base.:(==)(a1::SparseArray, a2::SparseArray) + return SparseArrayInterface.sparse_isequal(a1, a2) +end +function Base.reshape(a::SparseArray, dims::Tuple{Vararg{Int}}) + return SparseArrayInterface.sparse_reshape(a, dims) +end +function Base.iszero(a::SparseArray) + return SparseArrayInterface.sparse_iszero(a) +end +function Base.isreal(a::SparseArray) + return SparseArrayInterface.sparse_isreal(a) +end +function Base.zero(a::SparseArray) + return SparseArrayInterface.sparse_zero(a) +end +function Base.one(a::SparseArray) + return SparseArrayInterface.sparse_one(a) +end +end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl index 0e2de3ef0a..f00698db26 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/runtests.jl @@ -1,5 +1,9 @@ @eval module $(gensym()) using Compat: Returns, allequal +using LinearAlgebra: norm +include("SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl") +using .SparseArrayInterfaceTestUtils.DiagonalArrays: DiagonalArray +using .SparseArrayInterfaceTestUtils.SparseArrays: SparseArray using Test: @test, @testset, @test_broken, @test_throws @testset "SparseArrayInterface (eltype=$elt)" for elt in @@ -14,104 +18,6 @@ using Test: @test, @testset, @test_broken, @test_throws CartesianIndex(1, 2) end @testset "Custom SparseArray" begin - using LinearAlgebra: norm - using NDTensors.SparseArrayInterface: SparseArrayInterface - struct SparseArray{T,N} <: AbstractArray{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.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.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.empty_storage!(a::SparseArray) - empty!(a.data) - empty!(a.index_to_dataindex) - return empty!(a.dataindex_to_index) - end - - # Broadcasting - function Broadcast.BroadcastStyle(arraytype::Type{<:SparseArray}) - return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() - end - - # Map - function Base.map!(f, dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_map!(f, dest, src) - return dest - end - function Base.copy!(dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_copy!(dest, src) - return dest - end - function Base.copyto!(dest::AbstractArray, src::SparseArray) - SparseArrayInterface.sparse_copyto!(dest, src) - return dest - end - function Base.permutedims!(dest::AbstractArray, src::SparseArray, perm) - SparseArrayInterface.sparse_permutedims!(dest, src, perm) - return dest - end - - # Base - function Base.:(==)(a1::SparseArray, a2::SparseArray) - return SparseArrayInterface.sparse_isequal(a1, a2) - end - function Base.reshape(a::SparseArray, dims::Tuple{Vararg{Int}}) - return SparseArrayInterface.sparse_reshape(a, dims) - end - function Base.iszero(a::SparseArray) - return SparseArrayInterface.sparse_iszero(a) - end - function Base.isreal(a::SparseArray) - return SparseArrayInterface.sparse_isreal(a) - end - function Base.zero(a::SparseArray) - return SparseArrayInterface.sparse_zero(a) - end - function Base.one(a::SparseArray) - return SparseArrayInterface.sparse_one(a) - end - - # Test a = SparseArray{elt}(2, 3) @test size(a) == (2, 3) @test axes(a) == (1:2, 1:3) @@ -159,6 +65,17 @@ using Test: @test, @testset, @test_broken, @test_throws end end + a = SparseArray{elt}(2, 3) + a[1, 2] = 12 + a = map(x -> 2x, a) + for I in eachindex(a) + if I == CartesianIndex(1, 2) + @test a[I] == 2 * 12 + else + @test iszero(a[I]) + end + end + a = SparseArray{elt}(2, 2, 2) a[1, 2, 2] = 122 a_r = reshape(a, 2, 4) @@ -354,90 +271,6 @@ using Test: @test, @testset, @test_broken, @test_throws @test SparseArrayInterface.nstored(a′) == 2 end @testset "Custom DiagonalArray" begin - struct DiagonalArray{T,N} <: AbstractArray{T,N} - data::Vector{T} - dims::Tuple{Vararg{Int,N}} - end - function DiagonalArray{T,N}(::UndefInitializer, dims::Tuple{Vararg{Int,N}}) where {T,N} - return DiagonalArray{T,N}(Vector{T}(undef, minimum(dims)), dims) - end - function DiagonalArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} - return DiagonalArray{T,N}(undef, dims) - end - function DiagonalArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} - return DiagonalArray{T,length(dims)}(undef, dims) - end - function DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int}) where {T} - return DiagonalArray{T}(undef, dims) - end - - # AbstractArray interface - Base.size(a::DiagonalArray) = a.dims - function Base.getindex(a::DiagonalArray, I...) - return SparseArrayInterface.sparse_getindex(a, I...) - end - function Base.setindex!(a::DiagonalArray, I...) - return SparseArrayInterface.sparse_setindex!(a, I...) - end - function Base.similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) - return DiagonalArray{elt}(undef, dims) - end - - # Minimal interface - SparseArrayInterface.storage(a::DiagonalArray) = a.data - function SparseArrayInterface.index_to_storage_index( - a::DiagonalArray{<:Any,N}, I::CartesianIndex{N} - ) where {N} - !allequal(Tuple(I)) && return nothing - return first(Tuple(I)) - end - function SparseArrayInterface.storage_index_to_index(a::DiagonalArray, I) - return CartesianIndex(ntuple(Returns(I), ndims(a))) - end - - # Broadcasting - function Broadcast.BroadcastStyle(arraytype::Type{<:DiagonalArray}) - return SparseArrayInterface.SparseArrayStyle{ndims(arraytype)}() - end - - # Base - function Base.iszero(a::DiagonalArray) - return SparseArrayInterface.sparse_iszero(a) - end - function Base.isreal(a::DiagonalArray) - return SparseArrayInterface.sparse_isreal(a) - end - function Base.zero(a::DiagonalArray) - return SparseArrayInterface.sparse_zero(a) - end - function Base.one(a::DiagonalArray) - return SparseArrayInterface.sparse_one(a) - end - function Base.:(==)(a1::DiagonalArray, a2::DiagonalArray) - return SparseArrayInterface.sparse_isequal(a1, a2) - end - function Base.reshape(a::DiagonalArray, dims::Tuple{Vararg{Int}}) - return SparseArrayInterface.sparse_reshape(a, dims) - end - - # Map - function Base.map!(f, dest::AbstractArray, src::DiagonalArray) - SparseArrayInterface.sparse_map!(f, dest, src) - return dest - end - function Base.copy!(dest::AbstractArray, src::DiagonalArray) - SparseArrayInterface.sparse_copy!(dest, src) - return dest - end - function Base.copyto!(dest::AbstractArray, src::DiagonalArray) - SparseArrayInterface.sparse_copyto!(dest, src) - return dest - end - function Base.permutedims!(dest::AbstractArray, src::DiagonalArray, perm) - SparseArrayInterface.sparse_permutedims!(dest, src, perm) - return dest - end - # TODO: Test `fill!`. # Test @@ -479,10 +312,21 @@ using Test: @test, @testset, @test_broken, @test_throws a = DiagonalArray(elt[1, 2, 3], (3, 3)) a_r = reshape(a, 9) + @test a_r isa DiagonalArray{elt,1} + for I in LinearIndices(a) + @test a[I] == a_r[I] + end + # This needs `Base.reshape` with a custom destination + # calling `SparseArrayInterface.sparse_reshape!` + # in order to specify an appropriate output + # type to work. a = DiagonalArray(elt[1, 2], (2, 2, 2)) - # TODO: Make `reshape` for `DiagonalArray` return `SparseArray`. - @test_broken reshape(a, 2, 4) + a_r = reshape(a, 2, 4) + @test a_r isa Matrix{elt} + for I in LinearIndices(a) + @test a[I] == a_r[I] + end end end end From 41c9816b691ffefc29b8ac25d375c2729450fdf7 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 29 Nov 2023 13:25:14 -0500 Subject: [PATCH 14/17] Format --- .../src/lib/DiagonalArrays/src/sparsearrayinterface.jl | 6 ++++-- .../test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl | 4 +++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl b/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl index dc390c6e30..bbaed1be99 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/sparsearrayinterface.jl @@ -21,7 +21,9 @@ end # 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}}) +function SparseArrayInterface.sparse_similar( + a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}} +) return SparseArrayDOK{elt}(undef, dims) end @@ -55,7 +57,7 @@ end function Base.map!(f, dest::AbstractArray, src::DiagonalArray) SparseArrayInterface.sparse_map!(f, dest, src) return dest -end +end # permutedims function Base.permutedims!(dest::AbstractArray, src::DiagonalArray, perm) diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl index 2ea8d4690c..9285498935 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/DiagonalArrays.jl @@ -41,7 +41,9 @@ end function SparseArrayInterface.storage_index_to_index(a::DiagonalArray, I) return CartesianIndex(ntuple(Returns(I), ndims(a))) end -function SparseArrayInterface.sparse_similar(a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}}) +function SparseArrayInterface.sparse_similar( + a::DiagonalArray, elt::Type, dims::Tuple{Vararg{Int}} +) return Array{elt}(undef, dims) end function SparseArrayInterface.sparse_similar(a::DiagonalArray, elt::Type, dims::Tuple{Int}) From 420692aa4df39836706d49adec7fbaefb5aa4ba9 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 29 Nov 2023 13:26:46 -0500 Subject: [PATCH 15/17] Fix loading issue --- NDTensors/src/NDTensors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index 6d2284387b..43139bbd56 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -28,8 +28,8 @@ for lib in [ :RankFactorization, :TensorAlgebra, :SparseArrayInterface, - :DiagonalArrays, :SparseArrayDOKs, + :DiagonalArrays, :BlockSparseArrays, :NamedDimsArrays, :SmallVectors, From c6dcefe799b0309628d73b2c8b4f8a388b81dec4 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 29 Nov 2023 13:34:17 -0500 Subject: [PATCH 16/17] Missing include, improve README --- NDTensors/src/lib/DiagonalArrays/README.md | 41 +++++++++++++++---- .../lib/DiagonalArrays/examples/Project.toml | 1 + .../src/lib/DiagonalArrays/examples/README.jl | 16 ++++++-- .../lib/DiagonalArrays/src/DiagonalArrays.jl | 1 + 4 files changed, 47 insertions(+), 12 deletions(-) diff --git a/NDTensors/src/lib/DiagonalArrays/README.md b/NDTensors/src/lib/DiagonalArrays/README.md index ca541f70cb..07cc9e40c8 100644 --- a/NDTensors/src/lib/DiagonalArrays/README.md +++ b/NDTensors/src/lib/DiagonalArrays/README.md @@ -3,19 +3,33 @@ A Julia `DiagonalArray` type. ````julia -using NDTensors.DiagonalArrays: DiagonalArray, DiagIndex, DiagIndices +using NDTensors.DiagonalArrays: DiagonalArray, DiagonalMatrix, DiagIndex, DiagIndices, isdiagindex using Test function main() - d = DiagonalArray([1.0, 2, 3], 3, 4, 5) + d = DiagonalMatrix([1.0, 2.0, 3.0]) + @test eltype(d) == Float64 + @test size(d) == (3, 3) + @test d[1, 1] == 1 + @test d[2, 2] == 2 + @test d[3, 3] == 3 + @test d[1, 2] == 0 + + d = DiagonalArray([1.0, 2.0, 3.0], 3, 4, 5) + @test eltype(d) == Float64 @test d[1, 1, 1] == 1 @test d[2, 2, 2] == 2 + @test d[3, 3, 3] == 3 @test d[1, 2, 1] == 0 d[2, 2, 2] = 22 @test d[2, 2, 2] == 22 - @test length(d[DiagIndices()]) == 3 + d_r = reshape(d, 3, 20) + @test size(d_r) == (3, 20) + @test all(I -> d_r[I] == d[I], LinearIndices(d)) + + @test length(d[DiagIndices(:)]) == 3 @test Array(d) == d @test d[DiagIndex(2)] == d[2, 2, 2] @@ -24,15 +38,24 @@ function main() a = randn(3, 4, 5) new_diag = randn(3) - a[DiagIndices()] = new_diag - d[DiagIndices()] = a[DiagIndices()] + a[DiagIndices(:)] = new_diag + d[DiagIndices(:)] = a[DiagIndices(:)] - @test a[DiagIndices()] == new_diag - @test d[DiagIndices()] == new_diag + @test a[DiagIndices(:)] == new_diag + @test d[DiagIndices(:)] == new_diag permuted_d = permutedims(d, (3, 2, 1)) @test permuted_d isa DiagonalArray - @test permuted_d == d + @test permuted_d[DiagIndices(:)] == d[DiagIndices(:)] + @test size(d) == (3, 4, 5) + @test size(permuted_d) == (5, 4, 3) + for I in eachindex(d) + if !isdiagindex(d, I) + @test iszero(d[I]) + else + @test !iszero(d[I]) + end + end mapped_d = map(x -> 2x, d) @test mapped_d isa DiagonalArray @@ -48,7 +71,7 @@ You can generate this README with: ```julia using Literate using NDTensors.DiagonalArrays -dir = joinpath(pkgdir(DiagonalArrays), "src", "DiagonalArrays") +dir = joinpath(pkgdir(DiagonalArrays), "src", "lib", "DiagonalArrays") Literate.markdown(joinpath(dir, "examples", "README.jl"), dir; flavor=Literate.CommonMarkFlavor()) ``` diff --git a/NDTensors/src/lib/DiagonalArrays/examples/Project.toml b/NDTensors/src/lib/DiagonalArrays/examples/Project.toml index ef491a529c..b46e6ac7ac 100644 --- a/NDTensors/src/lib/DiagonalArrays/examples/Project.toml +++ b/NDTensors/src/lib/DiagonalArrays/examples/Project.toml @@ -1,3 +1,4 @@ [deps] +Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/NDTensors/src/lib/DiagonalArrays/examples/README.jl b/NDTensors/src/lib/DiagonalArrays/examples/README.jl index 26733d3ba7..c61f7b35c1 100644 --- a/NDTensors/src/lib/DiagonalArrays/examples/README.jl +++ b/NDTensors/src/lib/DiagonalArrays/examples/README.jl @@ -2,13 +2,23 @@ # # A Julia `DiagonalArray` type. -using NDTensors.DiagonalArrays: DiagonalArray, DiagIndex, DiagIndices, isdiagindex +using NDTensors.DiagonalArrays: DiagonalArray, DiagonalMatrix, DiagIndex, DiagIndices, isdiagindex using Test function main() - d = DiagonalArray([1.0, 2, 3], 3, 4, 5) + d = DiagonalMatrix([1.0, 2.0, 3.0]) + @test eltype(d) == Float64 + @test size(d) == (3, 3) + @test d[1, 1] == 1 + @test d[2, 2] == 2 + @test d[3, 3] == 3 + @test d[1, 2] == 0 + + d = DiagonalArray([1.0, 2.0, 3.0], 3, 4, 5) + @test eltype(d) == Float64 @test d[1, 1, 1] == 1 @test d[2, 2, 2] == 2 + @test d[3, 3, 3] == 3 @test d[1, 2, 1] == 0 d[2, 2, 2] = 22 @@ -60,7 +70,7 @@ You can generate this README with: ```julia using Literate using NDTensors.DiagonalArrays -dir = joinpath(pkgdir(DiagonalArrays), "src", "DiagonalArrays") +dir = joinpath(pkgdir(DiagonalArrays), "src", "lib", "DiagonalArrays") Literate.markdown(joinpath(dir, "examples", "README.jl"), dir; flavor=Literate.CommonMarkFlavor()) ``` =# diff --git a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl index 60e259a17e..76335ef77a 100644 --- a/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl +++ b/NDTensors/src/lib/DiagonalArrays/src/DiagonalArrays.jl @@ -1,4 +1,5 @@ module DiagonalArrays +include("defaults.jl") include("diaginterface.jl") include("diagindex.jl") include("diagindices.jl") From c5b8ff4b47c5295ba2196dbed26a9022b66b711e Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Wed, 29 Nov 2023 13:36:44 -0500 Subject: [PATCH 17/17] Format Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- NDTensors/src/lib/DiagonalArrays/examples/README.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/NDTensors/src/lib/DiagonalArrays/examples/README.jl b/NDTensors/src/lib/DiagonalArrays/examples/README.jl index c61f7b35c1..204ac937e5 100644 --- a/NDTensors/src/lib/DiagonalArrays/examples/README.jl +++ b/NDTensors/src/lib/DiagonalArrays/examples/README.jl @@ -2,7 +2,8 @@ # # A Julia `DiagonalArray` type. -using NDTensors.DiagonalArrays: DiagonalArray, DiagonalMatrix, DiagIndex, DiagIndices, isdiagindex +using NDTensors.DiagonalArrays: + DiagonalArray, DiagonalMatrix, DiagIndex, DiagIndices, isdiagindex using Test function main()