[NDTensors] Get more Array storage functionality working

NDTensors/src/NDTensors.jl

@@ -132,11 +132,29 @@ include("empty/adapt.jl")
 
 #####################################
 # Array Tensor (experimental)
-# TODO: Move to `Experimental` module.
+# 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/storage/combiner.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/svd.jl")
+include("arraystorage/arraystorage/tensor/combiner.jl")
+
+# DiagonalArray storage
+include("arraystorage/diagonalarray/tensor/contract.jl")
+
+# BlockSparseArray storage
+include("arraystorage/blocksparsearray/storage/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

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

@@ -0,0 +1 @@
+promote_rule(::Type{<:Combiner}, arraytype::Type{<:MatrixOrArrayStorage}) = arraytype

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

@@ -0,0 +1,2 @@
+conj(as::AliasStyle, A::AbstractArray) = conj(A)
+conj(as::AllowAlias, A::Array{<:Real}) = A

NDTensors/src/arraytensor/array.jl → NDTensors/src/arraystorage/arraystorage/storage/contract.jl

@@ -1,6 +1,3 @@
-# Combiner
-promote_rule(::Type{<:Combiner}, arraytype::Type{<:MatrixOrArrayStorage}) = arraytype
-
 # Generic AbstractArray code
 function contract(
   array1::MatrixOrArrayStorage,
@@ -43,26 +40,21 @@ function contraction_output(
 end
 
 # Required interface for specific AbstractArray types
+# TODO: Define `default_α` and `default_β`.
+# TODO: Define this as a `ttgt` or `matricize` backend.
 function contract!(
-  arrayR::MatrixOrArrayStorage,
-  labelsR,
+  array_dest::MatrixOrArrayStorage,
+  labels_dest,
   array1::MatrixOrArrayStorage,
   labels1,
   array2::MatrixOrArrayStorage,
   labels2,
+  α=one(eltype(array_dest)),
+  β=zero(eltype(array_dest));
 )
-  props = ContractionProperties(labels1, labels2, labelsR)
-  compute_contraction_properties!(props, array1, array2, arrayR)
+  props = ContractionProperties(labels1, labels2, labels_dest)
+  compute_contraction_properties!(props, array1, array2, array_dest)
   # TODO: Change this to just `contract!`, or maybe `contract_ttgt!`?
-  _contract!(arrayR, array1, array2, props)
-  return arrayR
-end
-
-function permutedims!(
-  output_array::MatrixOrArrayStorage, array::MatrixOrArrayStorage, perm, f::Function
-)
-  output_array = permutedims!!(
-    leaf_parenttype(output_array), output_array, leaf_parenttype(array), array, perm, f
-  )
-  return output_array
+  _contract!(array_dest, array1, array2, props, α, β)
+  return array_dest
 end

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

@@ -0,0 +1,8 @@
+function permutedims!(
+  output_array::MatrixOrArrayStorage, array::MatrixOrArrayStorage, perm, f::Function
+)
+  output_array = permutedims!!(
+    leaf_parenttype(output_array), output_array, leaf_parenttype(array), array, perm, f
+  )
+  return output_array
+end

NDTensors/src/arraystorage/arraystorage/tensor/arraystorage.jl

@@ -0,0 +1,27 @@
+const ArrayStorageTensor{T,N,S,I} = Tensor{T,N,S,I} where {S<:ArrayStorage{T,N}}
+const MatrixStorageTensor{T,S,I} = Tensor{T,2,S,I} where {S<:MatrixStorage{T}}
+const MatrixOrArrayStorageTensor{T,S,I} =
+  Tensor{T,N,S,I} where {N,S<:MatrixOrArrayStorage{T}}
+
+function Tensor(storage::MatrixOrArrayStorageTensor, inds::Tuple)
+  return Tensor(NeverAlias(), storage, inds)
+end
+
+function Tensor(as::AliasStyle, storage::MatrixOrArrayStorage, inds::Tuple)
+  return Tensor{eltype(storage),length(inds),typeof(storage),typeof(inds)}(
+    as, storage, inds
+  )
+end
+
+array(tensor::MatrixOrArrayStorageTensor) = storage(tensor)
+
+# Linear algebra (matrix algebra)
+function Base.adjoint(tens::MatrixStorageTensor)
+  return tensor(adjoint(storage(tens)), reverse(inds(tens)))
+end
+
+# Conversion from a tensor with TensorStorage storage
+# to AbstractArray storage,
+function to_arraytensor(x::DenseTensor)
+  return tensor(reshape(data(x), size(x)), inds(x))
+end

NDTensors/src/arraystorage/arraystorage/tensor/combiner.jl

@@ -0,0 +1,80 @@
+function contraction_output(
+  tensor1::MatrixOrArrayStorageTensor, tensor2::CombinerTensor, indsR
+)
+  tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR)
+  return NDTensors.similar(tensortypeR, indsR)
+end
+
+function contract!!(
+  output_tensor::ArrayStorageTensor,
+  output_tensor_labels,
+  combiner_tensor::CombinerTensor,
+  combiner_tensor_labels,
+  tensor::ArrayStorageTensor,
+  tensor_labels,
+)
+  if ndims(combiner_tensor) ≤ 1
+    # Empty combiner, acts as multiplying by 1
+    output_tensor = permutedims!!(
+      output_tensor, tensor, getperm(output_tensor_labels, tensor_labels)
+    )
+    return output_tensor
+  end
+  if is_index_replacement(tensor, tensor_labels, combiner_tensor, combiner_tensor_labels)
+    ui = setdiff(combiner_tensor_labels, tensor_labels)[]
+    newind = inds(combiner_tensor)[findfirst(==(ui), combiner_tensor_labels)]
+    cpos1, cpos2 = intersect_positions(combiner_tensor_labels, tensor_labels)
+    output_tensor_storage = copy(storage(tensor))
+    output_tensor_inds = setindex(inds(tensor), newind, cpos2)
+    return NDTensors.tensor(output_tensor_storage, output_tensor_inds)
+  end
+  is_combining_contraction = is_combining(
+    tensor, tensor_labels, combiner_tensor, combiner_tensor_labels
+  )
+  if is_combining_contraction
+    Alabels, Blabels = tensor_labels, combiner_tensor_labels
+    final_labels = contract_labels(Blabels, Alabels)
+    final_labels_n = contract_labels(combiner_tensor_labels, tensor_labels)
+    output_tensor_inds = inds(output_tensor)
+    if final_labels != final_labels_n
+      perm = getperm(final_labels_n, final_labels)
+      output_tensor_inds = permute(inds(output_tensor), perm)
+      output_tensor_labels = permute(output_tensor_labels, perm)
+    end
+    cpos1, output_tensor_cpos = intersect_positions(
+      combiner_tensor_labels, output_tensor_labels
+    )
+    labels_comb = deleteat(combiner_tensor_labels, cpos1)
+    output_tensor_vl = [output_tensor_labels...]
+    for (ii, li) in enumerate(labels_comb)
+      insert!(output_tensor_vl, output_tensor_cpos + ii, li)
+    end
+    deleteat!(output_tensor_vl, output_tensor_cpos)
+    labels_perm = tuple(output_tensor_vl...)
+    perm = getperm(labels_perm, tensor_labels)
+    # TODO: Add a `reshape` for `ArrayStorageTensor`.
+    ## tensorp = reshape(output_tensor, NDTensors.permute(inds(tensor), perm))
+    tensorp_inds = permute(inds(tensor), perm)
+    tensorp = NDTensors.tensor(
+      reshape(storage(output_tensor), dims(tensorp_inds)), tensorp_inds
+    )
+    permutedims!(tensorp, tensor, perm)
+    # TODO: Add a `reshape` for `ArrayStorageTensor`.
+    ## reshape(tensorp, output_tensor_inds)
+    return NDTensors.tensor(
+      reshape(storage(tensorp), dims(output_tensor_inds)), output_tensor_inds
+    )
+  else # Uncombining
+    cpos1, cpos2 = intersect_positions(combiner_tensor_labels, tensor_labels)
+    output_tensor_storage = copy(storage(tensor))
+    indsC = deleteat(inds(combiner_tensor), cpos1)
+    output_tensor_inds = insertat(inds(tensor), indsC, cpos2)
+    # TODO: Add a `reshape` for `ArrayStorageTensor`.
+    return NDTensors.tensor(
+      reshape(output_tensor_storage, dims(output_tensor_inds)), output_tensor_inds
+    )
+  end
+  return invalid_combiner_contraction_error(
+    tensor, tensor_labels, combiner_tensor, combiner_tensor_labels
+  )
+end

NDTensors/src/arraystorage/arraystorage/tensor/contract.jl

@@ -0,0 +1,22 @@
+# TODO: Just call `contraction_output(storage(tensor1), storage(tensor2), indsR)`
+function contraction_output(
+  tensor1::MatrixOrArrayStorageTensor, tensor2::MatrixOrArrayStorageTensor, indsR
+)
+  tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR)
+  return NDTensors.similar(tensortypeR, indsR)
+end
+
+# TODO: Define `default_α` and `default_β`.
+function contract!(
+  tensor_dest::MatrixOrArrayStorageTensor,
+  labels_dest,
+  tensor1::MatrixOrArrayStorageTensor,
+  labels1,
+  tensor2::MatrixOrArrayStorageTensor,
+  labels2,
+  α=one(eltype(tensor_dest)),
+  β=zero(eltype(tensor_dest));
+)
+  contract!(storage(tensor_dest), labels_dest, storage(tensor1), labels1, storage(tensor2), labels2, α, β)
+  return tensor_dest
+end

NDTensors/src/arraystorage/arraystorage/tensor/indexing.jl

@@ -0,0 +1,8 @@
+function getindex(tensor::MatrixOrArrayStorageTensor, I::Integer...)
+  return storage(tensor)[I...]
+end
+
+function setindex!(tensor::MatrixOrArrayStorageTensor, v, I::Integer...)
+  storage(tensor)[I...] = v
+  return tensor
+end

NDTensors/src/arraystorage/arraystorage/tensor/mul.jl

@@ -0,0 +1,6 @@
+function LinearAlgebra.mul!(
+  C::MatrixStorageTensor, A::MatrixStorageTensor, B::MatrixStorageTensor
+)
+  mul!(storage(C), storage(A), storage(B))
+  return C
+end

NDTensors/src/arraystorage/arraystorage/tensor/permutedims.jl

@@ -0,0 +1,9 @@
+function permutedims!(
+  output_tensor::MatrixOrArrayStorageTensor,
+  tensor::MatrixOrArrayStorageTensor,
+  perm,
+  f::Function,
+)
+  permutedims!(storage(output_tensor), storage(tensor), perm, f)
+  return output_tensor
+end

NDTensors/src/arraystorage/arraystorage/tensor/qr.jl

@@ -0,0 +1,10 @@
+function qr(A::ArrayStorageTensor)
+  Q, R = qr(storage(A))
+  Q = convert(typeof(R), Q)
+  i, j = inds(A)
+  q = size(A, 1) < size(A, 2) ? i : j
+  q = sim(q)
+  Qₜ = tensor(Q, (i, q))
+  Rₜ = tensor(R, (dag(q), j))
+  return Qₜ, Rₜ
+end