diff --git a/NDTensors/src/BlockSparseArrays/src/blocksparsearray.jl b/NDTensors/src/BlockSparseArrays/src/blocksparsearray.jl index 572ed929e2..b14e28ac34 100644 --- a/NDTensors/src/BlockSparseArrays/src/blocksparsearray.jl +++ b/NDTensors/src/BlockSparseArrays/src/blocksparsearray.jl @@ -65,7 +65,6 @@ function BlockArrays.viewblock(block_arr::BlockSparseArray, block) # TODO: Make this `Zeros`? ## zero = zeros(eltype(block_arr), block_size) return block_arr.blocks[blks...] # Fails because zero isn't defined - ## return get_nonzero(block_arr.blocks, blks, zero) end function Base.getindex(block_arr::BlockSparseArray{T,N}, bi::BlockIndex{N}) where {T,N} diff --git a/NDTensors/src/BlockSparseArrays/src/sparsearray.jl b/NDTensors/src/BlockSparseArrays/src/sparsearray.jl index a3f5b6dcbd..e140377c7f 100644 --- a/NDTensors/src/BlockSparseArrays/src/sparsearray.jl +++ b/NDTensors/src/BlockSparseArrays/src/sparsearray.jl @@ -1,6 +1,7 @@ +# TODO: Define a constructor with a default `zero`. struct SparseArray{T,N,Zero} <: AbstractArray{T,N} data::Dictionary{CartesianIndex{N},T} - dims::NTuple{N,Int64} + dims::NTuple{N,Int} zero::Zero end @@ -20,13 +21,3 @@ end function Base.getindex(a::SparseArray{T,N}, I::Vararg{Int,N}) where {T,N} return getindex(a, CartesianIndex(I)) end - -## # `getindex` but uses a default if the value is -## # structurally zero. -## function get_nonzero(a::SparseArray{T,N}, I::CartesianIndex{N}, zero) where {T,N} -## @boundscheck checkbounds(a, I) -## return get(a.data, I, zero) -## end -## function get_nonzero(a::SparseArray{T,N}, I::NTuple{N,Int}, zero) where {T,N} -## return get_nonzero(a, CartesianIndex(I), zero) -## end diff --git a/NDTensors/src/DiagonalArrays/README.md b/NDTensors/src/DiagonalArrays/README.md new file mode 100644 index 0000000000..33996cf52d --- /dev/null +++ b/NDTensors/src/DiagonalArrays/README.md @@ -0,0 +1,49 @@ +# DiagonalArrays.jl + +A Julia `DiagonalArray` type. + +````julia +using NDTensors.DiagonalArrays: + DiagonalArray, + densearray, + diagview, + diaglength, + getdiagindex, + setdiagindex!, + setdiag!, + diagcopyto! + +d = DiagonalArray([1., 2, 3], 3, 4, 5) +@show d[1, 1, 1] == 1 +@show d[2, 2, 2] == 2 +@show d[1, 2, 1] == 0 + +d[2, 2, 2] = 22 +@show d[2, 2, 2] == 22 + +@show diaglength(d) == 3 +@show densearray(d) == d +@show getdiagindex(d, 2) == d[2, 2, 2] + +setdiagindex!(d, 222, 2) +@show d[2, 2, 2] == 222 + +a = randn(3, 4, 5) +new_diag = randn(3) +setdiag!(a, new_diag) +diagcopyto!(d, a) + +@show diagview(a) == new_diag +@show diagview(d) == new_diag +```` + +You can generate this README with: +```julia +using Literate +Literate.markdown("examples/README.jl", "."; flavor=Literate.CommonMarkFlavor()) +``` + +--- + +*This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* + diff --git a/NDTensors/src/DiagonalArrays/examples/README.jl b/NDTensors/src/DiagonalArrays/examples/README.jl new file mode 100644 index 0000000000..ede04d6c2e --- /dev/null +++ b/NDTensors/src/DiagonalArrays/examples/README.jl @@ -0,0 +1,42 @@ +# # DiagonalArrays.jl +# +# A Julia `DiagonalArray` type. + +using NDTensors.DiagonalArrays: + DiagonalArray, + densearray, + diagview, + diaglength, + getdiagindex, + setdiagindex!, + setdiag!, + diagcopyto! + +d = DiagonalArray([1.0, 2, 3], 3, 4, 5) +@show d[1, 1, 1] == 1 +@show d[2, 2, 2] == 2 +@show d[1, 2, 1] == 0 + +d[2, 2, 2] = 22 +@show d[2, 2, 2] == 22 + +@show diaglength(d) == 3 +@show densearray(d) == d +@show getdiagindex(d, 2) == d[2, 2, 2] + +setdiagindex!(d, 222, 2) +@show d[2, 2, 2] == 222 + +a = randn(3, 4, 5) +new_diag = randn(3) +setdiag!(a, new_diag) +diagcopyto!(d, a) + +@show diagview(a) == new_diag +@show diagview(d) == new_diag + +# You can generate this README with: +# ```julia +# using Literate +# Literate.markdown("examples/README.jl", "."; flavor=Literate.CommonMarkFlavor()) +# ``` diff --git a/NDTensors/src/DiagonalArrays/src/DiagonalArrays.jl b/NDTensors/src/DiagonalArrays/src/DiagonalArrays.jl new file mode 100644 index 0000000000..45319ad8ce --- /dev/null +++ b/NDTensors/src/DiagonalArrays/src/DiagonalArrays.jl @@ -0,0 +1,110 @@ +module DiagonalArrays + +using Compat # allequal +using LinearAlgebra + +export DiagonalArray + +include("diagview.jl") + +struct DefaultZero end + +function (::DefaultZero)(eltype::Type, I::CartesianIndex) + return zero(eltype) +end + +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=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 + +# 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 + +Base.size(a::DiagonalArray) = a.dims + +diagview(a::DiagonalArray) = a.diag +LinearAlgebra.diag(a::DiagonalArray) = copy(diagview(a)) + +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 getdiagindex(a, i) +end + +function Base.getindex(a::DiagonalArray{T,N}, I::Vararg{Int,N}) where {T,N} + return getindex(a, CartesianIndex(I)) +end + +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") + setdiagindex!(a, v, i) + 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)) + fill!(d, zero(eltype(a))) + diagcopyto!(d, a) + return d +end + +end diff --git a/NDTensors/src/DiagonalArrays/src/diagview.jl b/NDTensors/src/DiagonalArrays/src/diagview.jl new file mode 100644 index 0000000000..21f7ac8a06 --- /dev/null +++ b/NDTensors/src/DiagonalArrays/src/diagview.jl @@ -0,0 +1,54 @@ +# 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 + +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 diagview(A::AbstractArray) + return @view A[diagindices(A)] +end + +function diagcopyto!(dest::AbstractArray, src::AbstractArray) + copyto!(diagview(dest), diagview(src)) + return dest +end diff --git a/NDTensors/src/DiagonalArrays/test/runtests.jl b/NDTensors/src/DiagonalArrays/test/runtests.jl new file mode 100644 index 0000000000..45d8b55e5c --- /dev/null +++ b/NDTensors/src/DiagonalArrays/test/runtests.jl @@ -0,0 +1,10 @@ +using Test +using NDTensors.DiagonalArrays + +@testset "Test NDTensors.DiagonalArrays" begin + @testset "README" begin + @test include( + joinpath(pkgdir(DiagonalArrays), "src", "DiagonalArrays", "examples", "README.jl") + ) isa Any + end +end diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index 446b3eedb2..296bfa4f99 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -20,6 +20,8 @@ using TupleTools include("SetParameters/src/SetParameters.jl") using .SetParameters +include("DiagonalArrays/src/DiagonalArrays.jl") +using .DiagonalArrays include("BlockSparseArrays/src/BlockSparseArrays.jl") using .BlockSparseArrays include("SmallVectors/src/SmallVectors.jl") diff --git a/NDTensors/test/DiagonalArrays.jl b/NDTensors/test/DiagonalArrays.jl new file mode 100644 index 0000000000..fa9b5a0c2b --- /dev/null +++ b/NDTensors/test/DiagonalArrays.jl @@ -0,0 +1,4 @@ +using Test +using NDTensors + +include(joinpath(pkgdir(NDTensors), "src", "DiagonalArrays", "test", "runtests.jl")) diff --git a/NDTensors/test/runtests.jl b/NDTensors/test/runtests.jl index 4e934a4c0d..9e5bba62f5 100644 --- a/NDTensors/test/runtests.jl +++ b/NDTensors/test/runtests.jl @@ -20,6 +20,7 @@ end @safetestset "NDTensors" begin @testset "$filename" for filename in [ "BlockSparseArrays.jl", + "DiagonalArrays.jl", "SetParameters.jl", "SmallVectors.jl", "SortedSets.jl",