[NDTensors] Add DiagonalArrays submodule

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}

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

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).*
+

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., 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())
+# ```

NDTensors/src/DiagonalArrays/src/DiagonalArrays.jl

@@ -0,0 +1,105 @@
+module DiagonalArrays
+
+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

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

NDTensors/src/DiagonalArrays/test/runtests.jl

@@ -0,0 +1,12 @@
+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

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")
@@ -132,11 +134,32 @@ include("empty/adapt.jl")
 
 #####################################
 # Array Tensor (experimental)
-# TODO: Move to `Experimental` module.
 #
-include("arraytensor/arraytensor.jl")
-include("arraytensor/array.jl")
-include("arraytensor/blocksparsearray.jl")
+include("arraystorage/arraystorage/storage/arraystorage.jl")
+include("arraystorage/arraystorage/storage/conj.jl")
+include("arraystorage/arraystorage/storage/permutedims.jl")
+include("arraystorage/arraystorage/storage/contract.jl")
+
+include("arraystorage/arraystorage/tensor/arraystorage.jl")
+include("arraystorage/arraystorage/tensor/zeros.jl")
+include("arraystorage/arraystorage/tensor/indexing.jl")
+include("arraystorage/arraystorage/tensor/permutedims.jl")
+include("arraystorage/arraystorage/tensor/mul.jl")
+include("arraystorage/arraystorage/tensor/contract.jl")
+include("arraystorage/arraystorage/tensor/qr.jl")
+include("arraystorage/arraystorage/tensor/eigen.jl")
+include("arraystorage/arraystorage/tensor/svd.jl")
+
+# DiagonalArray storage
+include("arraystorage/diagonalarray/tensor/contract.jl")
+
+# BlockSparseArray storage
+include("arraystorage/blocksparsearray/storage/contract.jl")
+
+# Combiner storage
+include("arraystorage/combiner/storage/contract.jl")
+
+include("arraystorage/combiner/tensor/contract.jl")
 
 #####################################
 # Deprecations

NDTensors/src/abstractarray/fill.jl

@@ -8,10 +8,10 @@ function generic_randn(
   return randn!(rng, data)
 end
 
-function generic_zeros(arraytype::Type{<:AbstractArray}, dim::Integer=0)
+function generic_zeros(arraytype::Type{<:AbstractArray}, dims...)
   arraytype_specified = set_unspecified_parameters(
     leaf_parenttype(arraytype), DefaultParameters()
   )
   ElT = eltype(arraytype_specified)
-  return fill!(similar(arraytype_specified, dim), zero(ElT))
+  return fill!(similar(arraytype_specified, dims...), zero(ElT))
 end

NDTensors/src/abstractarray/similar.jl

@@ -77,6 +77,11 @@ function similar(array::AbstractArray, eltype::Type, dims::Tuple)
   return NDTensors.similar(similartype(typeof(array), eltype), dims)
 end
 
+# NDTensors.similar
+function similar(array::AbstractArray, eltype::Type, dims::Int)
+  return NDTensors.similar(similartype(typeof(array), eltype), dims)
+end
+
 # NDTensors.similar
 similar(array::AbstractArray, dims::Tuple) = NDTensors.similar(typeof(array), dims)
 

NDTensors/src/arraystorage/arraystorage/storage/arraystorage.jl

@@ -0,0 +1,30 @@
+# Used for dispatch to distinguish from Tensors wrapping TensorStorage.
+# Remove once TensorStorage is removed.
+const ArrayStorage{T,N} = Union{
+  Array{T,N},
+  ReshapedArray{T,N},
+  SubArray{T,N},
+  PermutedDimsArray{T,N},
+  StridedView{T,N},
+  BlockSparseArray{T,N},
+}
+
+const MatrixStorage{T} = Union{
+  ArrayStorage{T,2},
+  Transpose{T},
+  Adjoint{T},
+  Symmetric{T},
+  Hermitian{T},
+  UpperTriangular{T},
+  LowerTriangular{T},
+  UnitUpperTriangular{T},
+  UnitLowerTriangular{T},
+  Diagonal{T},
+}
+
+const MatrixOrArrayStorage{T} = Union{MatrixStorage{T},ArrayStorage{T}}
+
+# TODO: Delete once `Dense` is removed.
+function to_arraystorage(x::DenseTensor)
+  return tensor(reshape(data(x), size(x)), inds(x))
+end