Skip to content

Commit e406418

Browse files
authored
Merge fefae9d into 10f7eee
2 parents 10f7eee + fefae9d commit e406418

File tree

8 files changed

+244
-8
lines changed

8 files changed

+244
-8
lines changed

NDTensors/src/NDTensors.jl

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,13 @@ include("empty/EmptyTensor.jl")
109109
include("empty/tensoralgebra/contract.jl")
110110
include("empty/adapt.jl")
111111

112+
#####################################
113+
# Array Tensor (experimental)
114+
# TODO: Move to `Experimental` module.
115+
#
116+
include("arraytensor/arraytensor.jl")
117+
include("arraytensor/array.jl")
118+
112119
#####################################
113120
# Deprecations
114121
#

NDTensors/src/arraytensor/array.jl

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Combiner
2+
promote_rule(::Type{<:Combiner}, arraytype::Type{<:MatrixOrArrayStorage}) = arraytype
3+
4+
# Generic AbstractArray code
5+
function contract(
6+
array1::MatrixOrArrayStorage,
7+
labels1,
8+
array2::MatrixOrArrayStorage,
9+
labels2,
10+
labelsR=contract_labels(labels1, labels2),
11+
)
12+
output_array = contraction_output(array1, labels1, array2, labels2, labelsR)
13+
contract!(output_array, labelsR, array1, labels1, array2, labels2)
14+
return output_array
15+
end
16+
17+
function contraction_output(
18+
array1::MatrixOrArrayStorage, array2::MatrixOrArrayStorage, indsR
19+
)
20+
arraytypeR = contraction_output_type(typeof(array1), typeof(array2), indsR)
21+
return NDTensors.similar(arraytypeR, indsR)
22+
end
23+
24+
function contraction_output_type(
25+
arraytype1::Type{<:MatrixOrArrayStorage}, arraytype2::Type{<:MatrixOrArrayStorage}, inds
26+
)
27+
return similartype(promote_type(arraytype1, arraytype2), inds)
28+
end
29+
30+
function contraction_output(
31+
array1::MatrixOrArrayStorage,
32+
labelsarray1,
33+
array2::MatrixOrArrayStorage,
34+
labelsarray2,
35+
labelsoutput_array,
36+
)
37+
# TODO: Maybe use `axes` here to be more generic, for example for BlockArrays?
38+
indsoutput_array = contract_inds(
39+
size(array1), labelsarray1, size(array2), labelsarray2, labelsoutput_array
40+
)
41+
output_array = contraction_output(array1, array2, indsoutput_array)
42+
return output_array
43+
end
44+
45+
# Required interface for specific AbstractArray types
46+
function contract!(
47+
arrayR::MatrixOrArrayStorage,
48+
labelsR,
49+
array1::MatrixOrArrayStorage,
50+
labels1,
51+
array2::MatrixOrArrayStorage,
52+
labels2,
53+
)
54+
props = ContractionProperties(labels1, labels2, labelsR)
55+
compute_contraction_properties!(props, array1, array2, arrayR)
56+
# TODO: Change this to just `contract!`, or maybe `contract_ttgt!`?
57+
_contract!(arrayR, array1, array2, props)
58+
return arrayR
59+
end
60+
61+
function permutedims!(
62+
output_array::MatrixOrArrayStorage, array::MatrixOrArrayStorage, perm, f::Function
63+
)
64+
@strided output_array .= f.(output_array, permutedims(array, perm))
65+
return output_array
66+
end
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
# Used for dispatch to distinguish from Tensors wrapping TensorStorage.
2+
# Remove once TensorStorage is removed.
3+
const ArrayStorage{T,N} = Union{
4+
Array{T,N},ReshapedArray{T,N},SubArray{T,N},PermutedDimsArray{T,N},StridedView{T,N}
5+
}
6+
const MatrixStorage{T} = Union{
7+
ArrayStorage{T,2},
8+
Transpose{T},
9+
Adjoint{T},
10+
Symmetric{T},
11+
Hermitian{T},
12+
UpperTriangular{T},
13+
LowerTriangular{T},
14+
UnitUpperTriangular{T},
15+
UnitLowerTriangular{T},
16+
Diagonal{T},
17+
}
18+
const MatrixOrArrayStorage{T} = Union{MatrixStorage{T},ArrayStorage{T}}
19+
20+
const ArrayStorageTensor{T,N,S,I} = Tensor{T,N,S,I} where {S<:ArrayStorage{T,N}}
21+
const MatrixStorageTensor{T,S,I} = Tensor{T,2,S,I} where {S<:MatrixStorage{T}}
22+
const MatrixOrArrayStorageTensor{T,S,I} =
23+
Tensor{T,N,S,I} where {N,S<:MatrixOrArrayStorage{T}}
24+
25+
function Tensor(storage::MatrixOrArrayStorageTensor, inds::Tuple)
26+
return Tensor(NeverAlias(), storage, inds)
27+
end
28+
29+
function Tensor(as::AliasStyle, storage::MatrixOrArrayStorage, inds::Tuple)
30+
return Tensor{eltype(storage),length(inds),typeof(storage),typeof(inds)}(
31+
as, storage, inds
32+
)
33+
end
34+
35+
function getindex(tensor::MatrixOrArrayStorageTensor, I::Integer...)
36+
return storage(tensor)[I...]
37+
end
38+
39+
function setindex!(tensor::MatrixOrArrayStorageTensor, v, I::Integer...)
40+
storage(tensor)[I...] = v
41+
return tensor
42+
end
43+
44+
function contraction_output(
45+
tensor1::MatrixOrArrayStorageTensor, tensor2::MatrixOrArrayStorageTensor, indsR
46+
)
47+
tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR)
48+
return NDTensors.similar(tensortypeR, indsR)
49+
end
50+
51+
function contract!(
52+
tensorR::MatrixOrArrayStorageTensor,
53+
labelsR,
54+
tensor1::MatrixOrArrayStorageTensor,
55+
labels1,
56+
tensor2::MatrixOrArrayStorageTensor,
57+
labels2,
58+
)
59+
contract!(storage(tensorR), labelsR, storage(tensor1), labels1, storage(tensor2), labels2)
60+
return tensorR
61+
end
62+
63+
function permutedims!(
64+
output_tensor::MatrixOrArrayStorageTensor,
65+
tensor::MatrixOrArrayStorageTensor,
66+
perm,
67+
f::Function,
68+
)
69+
permutedims!(storage(output_tensor), storage(tensor), perm, f)
70+
return output_tensor
71+
end
72+
73+
# Linear algebra (matrix algebra)
74+
function Base.adjoint(tens::MatrixStorageTensor)
75+
return tensor(adjoint(storage(tens)), reverse(inds(tens)))
76+
end
77+
78+
function LinearAlgebra.mul!(
79+
C::MatrixStorageTensor, A::MatrixStorageTensor, B::MatrixStorageTensor
80+
)
81+
mul!(storage(C), storage(A), storage(B))
82+
return C
83+
end
84+
85+
function LinearAlgebra.svd(tens::MatrixStorageTensor)
86+
F = svd(storage(tens))
87+
U, S, V = F.U, F.S, F.Vt
88+
i, j = inds(tens)
89+
# TODO: Make this more general with a `similar_ind` function,
90+
# so the dimension can be determined from the length of `S`.
91+
min_ij = dim(i) dim(j) ? i : j
92+
α = sim(min_ij) # similar_ind(i, space(S))
93+
β = sim(min_ij) # similar_ind(i, space(S))
94+
Utensor = tensor(U, (i, α))
95+
# TODO: Remove conversion to `Diagonal` to make more general, or make a generic `Diagonal` concept that works for `BlockSparseArray`.
96+
# Used for now to avoid introducing wrapper types.
97+
Stensor = tensor(Diagonal(S), (α, β))
98+
Vtensor = tensor(V, (β, j))
99+
return Utensor, Stensor, Vtensor, Spectrum(nothing, 0.0)
100+
end
101+
102+
array(tensor::MatrixOrArrayStorageTensor) = storage(tensor)
103+
104+
# Combiner
105+
function contraction_output(
106+
tensor1::MatrixOrArrayStorageTensor, tensor2::CombinerTensor, indsR
107+
)
108+
tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR)
109+
return NDTensors.similar(tensortypeR, indsR)
110+
end

NDTensors/src/tensor/tensor.jl

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
1-
21
"""
32
Tensor{StoreT,IndsT}
43
54
A plain old tensor (with order independent
65
interface and no assumption of labels)
76
"""
8-
struct Tensor{ElT,N,StoreT<:TensorStorage,IndsT} <: AbstractArray{ElT,N}
7+
struct Tensor{ElT,N,StoreT,IndsT} <: AbstractArray{ElT,N}
98
storage::StoreT
109
inds::IndsT
1110

@@ -21,8 +20,8 @@ struct Tensor{ElT,N,StoreT<:TensorStorage,IndsT} <: AbstractArray{ElT,N}
2120
and tensor(store::TensorStorage, inds) constructors.
2221
"""
2322
function Tensor{ElT,N,StoreT,IndsT}(
24-
::AllowAlias, storage::TensorStorage, inds::Tuple
25-
) where {ElT,N,StoreT<:TensorStorage,IndsT}
23+
::AllowAlias, storage, inds::Tuple
24+
) where {ElT,N,StoreT,IndsT}
2625
@assert ElT == eltype(StoreT)
2726
@assert length(inds) == N
2827
return new{ElT,N,StoreT,IndsT}(storage, inds)
@@ -74,7 +73,7 @@ end
7473
# already (like a Vector). In the future this may be lifted
7574
# to allow for very large tensor orders in which case Tuple
7675
# operations may become too slow.
77-
function Tensor(as::AliasStyle, storage::TensorStorage, inds)
76+
function Tensor(as::AliasStyle, storage, inds)
7877
return Tensor(as, storage, Tuple(inds))
7978
end
8079

NDTensors/src/tensoralgebra/generic_tensor_operations.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ end
1818
function permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)
1919
Base.checkdims_perm(output_tensor, tensor, perm)
2020
error(
21-
"`perutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.",
21+
"`permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.",
2222
)
2323
return output_tensor
2424
end
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
using NDTensors
2+
using LinearAlgebra
3+
using Test
4+
5+
using NDTensors: storage, storagetype
6+
7+
@testset "Tensor wrapping Array" begin
8+
is1 = (2, 3)
9+
D1 = randn(is1)
10+
11+
is2 = (3, 4)
12+
D2 = randn(is2)
13+
14+
T1 = tensor(D1, is1)
15+
T2 = tensor(D2, is2)
16+
17+
@test T1[1, 1] == D1[1, 1]
18+
19+
x = rand()
20+
T1[1, 1] = x
21+
22+
@test T1[1, 1] == x
23+
@test array(T1) == D1
24+
@test storagetype(T1) <: Matrix{Float64}
25+
@test storage(T1) == D1
26+
@test eltype(T1) == eltype(D1)
27+
@test inds(T1) == is1
28+
29+
R = T1 * T2
30+
@test storagetype(R) <: Matrix{Float64}
31+
@test Array(R) Array(T1) * Array(T2)
32+
33+
T1r = randn!(similar(T1))
34+
@test Array(T1r + T1) Array(T1r) + Array(T1)
35+
@test Array(permutedims(T1, (2, 1))) permutedims(Array(T1), (2, 1))
36+
37+
U, S, V = svd(T1)
38+
@test U * S * V T1
39+
40+
T12 = contract(T1, (1, -1), T2, (-1, 2))
41+
@test T12 T1 * T2
42+
43+
D12 = contract(D1, (1, -1), D2, (-1, 2))
44+
@test D12 Array(T12)
45+
46+
## IT1 = itensor(T1)
47+
## IT2 = itensor(T2)
48+
## IR = IT1 * IT2
49+
## IX = IT1 + IT1
50+
51+
## IU, IS, IV = svd(IT1, i)
52+
## @test IU * IS * IV ≈ IT1
53+
end

NDTensors/test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ end
2828
"emptynumber.jl",
2929
"emptystorage.jl",
3030
"combiner.jl",
31+
"arraytensor/arraytensor.jl",
3132
]
3233
println("Running $filename")
3334
include(filename)

src/itensor.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ NDTensors.Dense{Float64,Array{Float64,1}}
8282
## Accessor Functions, Index Functions and Operations
8383
mutable struct ITensor
8484
tensor::Tensor
85-
function ITensor(::AllowAlias, T::Tensor{<:Any,<:Any,<:TensorStorage,<:Tuple})
85+
function ITensor(::AllowAlias, T::Tensor{<:Any,<:Any,<:Any,<:Tuple})
8686
@debug_check begin
8787
is = inds(T)
8888
if !allunique(is)
@@ -761,7 +761,7 @@ end
761761
762762
Return a view of the TensorStorage of the ITensor.
763763
"""
764-
storage(T::ITensor)::TensorStorage = storage(tensor(T))
764+
storage(T::ITensor) = storage(tensor(T))
765765

766766
storagetype(x::ITensor) = storagetype(tensor(x))
767767

0 commit comments

Comments
 (0)