Merge branch '42-conversion-from-3-leg-tt-to-4-leg-tt-for-ttos' into 'main'

Implement conversion between TTs with different shapes

Closes #42

See merge request tensors4fields/TensorCrossInterpolation.jl!83

Author: shinaoka

src/contraction.jl

@@ -299,15 +299,15 @@ function (obj::Contraction{T})(
 end
 
 
-function _contractsitetensors(a::Array{T, 4}, b::Array{T, 4})::Array{T, 4} where {T}
+function _contractsitetensors(a::Array{T,4}, b::Array{T,4})::Array{T,4} where {T}
     # indices: (link_a, s1, s2, link_a') * (link_b, s2, s3, link_b')
-    ab::Array{T, 6} = _contract(a, b, (3,), (2,))
+    ab::Array{T,6} = _contract(a, b, (3,), (2,))
     # => indices: (link_a, s1, link_a', link_b, s3, link_b')
     abpermuted = permutedims(ab, (1, 4, 2, 5, 3, 6))
     # => indices: (link_a, link_b, s1, s3, link_a', link_b')
     return reshape(abpermuted,
         size(a, 1) * size(b, 1),  # link_a * link_b
-        size(a, 2),  size(b, 3),  # s1, s3
+        size(a, 2), size(b, 3),  # s1, s3
         size(a, 4) * size(b, 4)   # link_a' * link_b'
     )
 end
@@ -328,7 +328,7 @@ function contract_naive(
     end
 
     a, b = obj.mpo
-    tt = TensorTrain{T, 4}(_contractsitetensors.(sitetensors(a), sitetensors(b)))
+    tt = TensorTrain{T,4}(_contractsitetensors.(sitetensors(a), sitetensors(b)))
     if tolerance > 0 || maxbonddim < typemax(Int)
         compress!(tt, :SVD; tolerance, maxbonddim)
     end
@@ -445,27 +445,20 @@ function contract(
     end
 end
 
-function _promoteMPStoMPO(
-    tt::AbstractTensorTrain{V},
-    unfusedlocalshape::Union{AbstractVector{Int}, Tuple}
-)::TensorTrain{V, 4} where {V}
-    return TensorTrain{V, 4}(_reshape_splitsites.(sitetensors(tt), Ref(unfusedlocalshape)))
-end
-
 function contract(
-    A::Union{TensorCI1{V}, TensorCI2{V}, TensorTrain{V, 3}},
-    B::TensorTrain{V, 4};
+    A::Union{TensorCI1{V},TensorCI2{V},TensorTrain{V,3}},
+    B::TensorTrain{V,4};
     kwargs...
-)::TensorTrain{V, 3} where {V}
-    tt = contract(_promoteMPStoMPO(A, (1, sitedim(A, 1)...)), B; kwargs...)
-    return TensorTrain{V, 3}([T for (T, shape) in _reshape_fusesites.(sitetensors(tt))])
+)::TensorTrain{V,3} where {V}
+    tt = contract(TensorTrain{4}(A, [(1, s...) for s in sitedims(A)]), B; kwargs...)
+    return TensorTrain{3}(tt, prod.(sitedims(tt)))
 end
 
 function contract(
-    A::TensorTrain{V, 4},
-    B::Union{TensorCI1{V}, TensorCI2{V}, TensorTrain{V, 3}};
+    A::TensorTrain{V,4},
+    B::Union{TensorCI1{V},TensorCI2{V},TensorTrain{V,3}};
     kwargs...
-)::TensorTrain{V, 3} where {V}
-    tt = contract(A, _promoteMPStoMPO(B, (sitedim(B, 1)..., 1)); kwargs...)
-    return TensorTrain{V, 3}([T for (T, shape) in _reshape_fusesites.(sitetensors(tt))])
+)::TensorTrain{V,3} where {V}
+    tt = contract(A, TensorTrain{4}(B, [(s..., 1) for s in sitedims(B)]); kwargs...)
+    return TensorTrain{3}(tt, prod.(sitedims(tt)))
 end

src/tensortrain.jl

@@ -52,13 +52,38 @@ function TensorTrain(tci::AbstractTensorTrain{V})::TensorTrain{V,3} where {V}
     return TensorTrain{V,3}(sitetensors(tci))
 end
 
+"""
+    function TensorTrain{N}(tci::AbstractTensorTrain{V}) where {V,N}
+
+Convert a tensor-train-like object into a tensor train.
+
+Arguments:
+- `tt::AbstractTensorTrain{V}`: a tensor-train-like object.
+- `localdims`: a vector of local dimensions for each tensor in the tensor train. A each element
+  of `localdims` should be an array-like object of `N-2` integers.
+"""
+function TensorTrain{V,N}(tt::AbstractTensorTrain{V}, localdims)::TensorTrain{V,N} where {V,N}
+    for d in localdims
+        length(d) == N - 2 || error("Each element of localdims be a list of N-2 integers.")
+    end
+    for n in 1:length(tt)
+        prod(size(tt[n])[2:end-1]) == prod(localdims[n]) || error("The local dimensions at n=$n must match the tensor sizes.")
+    end
+    return TensorTrain{V,N}(
+        [reshape(t, size(t, 1), localdims[n]..., size(t)[end]) for (n, t) in enumerate(sitetensors(tt))])
+end
+
+function TensorTrain{N}(tt::AbstractTensorTrain{V}, localdims)::TensorTrain{V,N} where {V,N}
+    return TensorTrain{V,N}(tt, localdims)
+end
+
 function tensortrain(tci)
     return TensorTrain(tci)
 end
 
 function _factorize(
     A::Matrix{V}, method::Symbol; tolerance::Float64, maxbonddim::Int
-)::Tuple{Matrix{V}, Matrix{V}, Int} where {V}
+)::Tuple{Matrix{V},Matrix{V},Int} where {V}
     if method === :LU
         factorization = rrlu(A, abstol=tolerance, maxrank=maxbonddim)
         return left(factorization), right(factorization), npivots(factorization)
@@ -92,11 +117,11 @@ end
 Compress the tensor train `tt` using `LU`, `CI` or `SVD` decompositions.
 """
 function compress!(
-    tt::TensorTrain{V, N},
+    tt::TensorTrain{V,N},
     method::Symbol=:LU;
     tolerance::Float64=1e-12,
     maxbonddim::Int=typemax(Int)
-) where {V, N}
+) where {V,N}
     for ell in 1:length(tt)-1
         shapel = size(tt.sitetensors[ell])
         left, right, newbonddim = _factorize(
@@ -125,48 +150,48 @@ function compress!(
 end
 
 
-function multiply!(tt::TensorTrain{V, N}, a) where {V, N}
+function multiply!(tt::TensorTrain{V,N}, a) where {V,N}
     tt.sitetensors[end] .= tt.sitetensors[end] .* a
     nothing
 end
 
-function multiply!(a, tt::TensorTrain{V, N}) where {V, N}
+function multiply!(a, tt::TensorTrain{V,N}) where {V,N}
     tt.sitetensors[end] .= a .* tt.sitetensors[end]
     nothing
 end
 
-function multiply(tt::TensorTrain{V, N}, a)::TensorTrain{V, N} where {V, N}
+function multiply(tt::TensorTrain{V,N}, a)::TensorTrain{V,N} where {V,N}
     tt2 = deepcopy(tt)
     multiply!(tt2, a)
     return tt2
 end
 
-function multiply(a, tt::TensorTrain{V, N})::TensorTrain{V, N} where {V, N}
+function multiply(a, tt::TensorTrain{V,N})::TensorTrain{V,N} where {V,N}
     tt2 = deepcopy(tt)
     multiply!(a, tt2)
     return tt2
 end
 
-function Base.:*(tt::TensorTrain{V, N}, a)::TensorTrain{V, N} where {V, N}
+function Base.:*(tt::TensorTrain{V,N}, a)::TensorTrain{V,N} where {V,N}
     return multiply(tt, a)
 end
 
-function Base.:*(a, tt::TensorTrain{V, N})::TensorTrain{V, N} where {V, N}
+function Base.:*(a, tt::TensorTrain{V,N})::TensorTrain{V,N} where {V,N}
     return multiply(a, tt)
 end
 
-function divide!(tt::TensorTrain{V, N}, a) where {V, N}
+function divide!(tt::TensorTrain{V,N}, a) where {V,N}
     tt.sitetensors[end] .= tt.sitetensors[end] ./ a
     nothing
 end
 
-function divide(tt::TensorTrain{V, N}, a) where {V, N}
+function divide(tt::TensorTrain{V,N}, a) where {V,N}
     tt2 = deepcopy(tt)
     divide!(tt2, a)
     return tt2
 end
 
-function Base.:/(tt::TensorTrain{V, N}, a) where {V, N}
+function Base.:/(tt::TensorTrain{V,N}, a) where {V,N}
     return divide(tt, a)
 end
 
@@ -201,7 +226,7 @@ function to_tensors(obj::TensorTrainFit{ValueType}, x::Vector{ValueType}) where
     ]
 end
 
-function _evaluate(tt::Vector{Array{V, 3}}, indexset) where {V}
+function _evaluate(tt::Vector{Array{V,3}}, indexset) where {V}
     only(prod(T[:, i, :] for (T, i) in zip(tt, indexset)))
 end
 

test/test_tensortrain.jl

@@ -35,6 +35,23 @@ using Optim
     end
 end
 
+
+@testset "TT shape conversion" for T in [Float64, ComplexF64]
+    linkdims = [1, 2, 3, 1]
+    L = length(linkdims) - 1
+    localdims = fill(4, L)
+    tts = TCI.TensorTrain{T,3}([randn(T, linkdims[n], localdims[n], linkdims[n+1]) for n in 1:L])
+    tto = TCI.TensorTrain{4}(tts, fill([2,2], L))
+    tts_reconst = TCI.TensorTrain{3}(tto, localdims)
+
+    for n in 1:L
+        @test all(tts[n] .== tts_reconst[n])
+    end 
+
+    @test_throws ErrorException TCI.TensorTrain{4}(tts, fill([2,3], L)) # Wrong shape
+    @test_throws ErrorException TCI.TensorTrain{4}(tts, fill([1,2,3], L)) # Wrong shape
+end
+
 @testset "batchevaluate" begin
     N = 4
     #bonddims = fill(3, N + 1)