Merge pull request #13 from tensor4all/8-tto-tto-contraction-with-zip-up-algorithm

Implement contract_zipup using SVD or LU

Author: shinaoka

src/contraction.jl

@@ -399,6 +399,56 @@ function contract_TCI(
     )
 end
 
+"""
+See SVD version:
+https://tensornetwork.org/mps/algorithms/zip_up_mpo/
+"""
+function contract_zipup(
+    A::TensorTrain{ValueType,4},
+    B::TensorTrain{ValueType,4};
+    tolerance::Float64=1e-12,
+    method::Symbol=:SVD, # :SVD, :LU
+    maxbonddim::Int=typemax(Int)
+) where {ValueType}
+    if length(A) != length(B)
+        throw(ArgumentError("Cannot contract tensor trains with different length."))
+    end
+    R::Array{ValueType,3} = ones(ValueType, 1, 1, 1)
+
+    sitetensors = Vector{Array{ValueType,4}}(undef, length(A))
+    for n in 1:length(A)
+        # R:     (link_ab, link_an, link_bn)
+        # A[n]:  (link_an, s_n, s_n', link_anp1)
+        RA = _contract(R, A[n], (2,), (1,))
+
+        # RA[n]: (link_ab, link_bn, s_n, s_n' link_anp1)
+        # B[n]:  (link_bn, s_n', s_n'', link_bnp1)
+        # C:     (link_ab, s_n, link_anp1, s_n'', link_bnp1)
+        #  =>    (link_ab, s_n, s_n'', link_anp1, link_bnp1)
+        C = permutedims(_contract(RA, B[n], (2, 4), (1, 2)), (1, 2, 4, 3, 5))
+        if n == length(A)
+            sitetensors[n] = reshape(C, size(C)[1:3]..., 1)
+            break
+        end
+
+        # Cmat:  (link_ab * s_n * s_n'', link_anp1 * link_bnp1)
+
+        #lu = rrlu(Cmat; reltol, abstol, leftorthogonal=true)
+        left, right, newbonddim = _factorize(
+            reshape(C, prod(size(C)[1:3]), prod(size(C)[4:5])),
+            method; tolerance, maxbonddim
+        )
+
+        # U:     (link_ab, s_n, s_n'', link_ab_new)
+        sitetensors[n] = reshape(left, size(C)[1:3]..., newbonddim)
+
+        # R:     (link_ab_new, link_an, link_bn)
+        R = reshape(right, newbonddim, size(C)[4:5]...)
+    end
+
+    return TensorTrain{ValueType,4}(sitetensors)
+end
+
 """
     function contract(
         A::TensorTrain{ValueType,4},
@@ -415,13 +465,15 @@ Contract two tensor trains `A` and `B`.
 Currently, two implementations are available:
  1. `algorithm=:TCI` constructs a new TCI that fits the contraction of `A` and `B`.
  2. `algorithm=:naive` uses a naive tensor contraction and subsequent SVD recompression of the tensor train.
+ 2. `algorithm=:zipup` uses a naive tensor contraction with on-the-fly LU decomposition.
 
 Arguments:
 - `A` and `B` are the tensor trains to be contracted.
 - `algorithm` chooses the algorithm used to evaluate the contraction.
 - `tolerance` is the tolerance of the TCI or SVD recompression.
 - `maxbonddim` sets the maximum bond dimension of the resulting tensor train.
 - `f` is a function to be applied elementwise to the result. This option is only available with `algorithm=:TCI`.
+- `method` chooses the method used for the factorization in the `algorithm=:zipup` case (`:SVD` or `:LU`).
 - `kwargs...` are forwarded to [`crossinterpolate2`](@ref) if `algorithm=:TCI`.
 """
 function contract(
@@ -440,6 +492,11 @@ function contract(
             error("Naive contraction implementation cannot contract matrix product with a function. Use algorithm=:TCI instead.")
         end
         return contract_naive(A, B; tolerance=tolerance, maxbonddim=maxbonddim)
+    elseif algorithm === :zipup
+        if f !== nothing
+            error("Zipup contraction implementation cannot contract matrix product with a function. Use algorithm=:TCI instead.")
+        end
+        return contract_zipup(A, B; tolerance, maxbonddim)
     else
         throw(ArgumentError("Unknown algorithm $algorithm."))
     end

test/test_contraction.jl

@@ -15,7 +15,7 @@ function _tomat(tto::TensorTrain{T,4}) where {T}
     return mat
 end
 
-function _tovec(tt::TensorTrain{T, 3}) where {T}
+function _tovec(tt::TensorTrain{T,3}) where {T}
     sitedims = TCI.sitedims(tt)
     localdims1 = [s[1] for s in sitedims]
     return evaluate.(Ref(tt), CartesianIndices(Tuple(localdims1))[:])
@@ -28,7 +28,45 @@ end
     @test vec(reshape(permutedims(a, (2, 1, 3)), 3, :) * reshape(permutedims(b, (1, 3, 2)), :, 5)) ≈ vec(ab)
 end
 
+function _gen_testdata_TTO_TTO()
+    N = 4
+    bonddims_a = [1, 2, 3, 2, 1]
+    bonddims_b = [1, 2, 3, 2, 1]
+    localdims1 = [2, 2, 2, 2]
+    localdims2 = [3, 3, 3, 3]
+    localdims3 = [2, 2, 2, 2]
+
+    a = TensorTrain{ComplexF64,4}([
+        rand(ComplexF64, bonddims_a[n], localdims1[n], localdims2[n], bonddims_a[n+1])
+        for n = 1:N
+    ])
+    b = TensorTrain{ComplexF64,4}([
+        rand(ComplexF64, bonddims_b[n], localdims2[n], localdims3[n], bonddims_b[n+1])
+        for n = 1:N
+    ])
+    return N, a, b, localdims1, localdims2, localdims3
+end
+
+function _gen_testdata_TTO_TTS()
+    N = 4
+    bonddims_a = [1, 2, 3, 2, 1]
+    bonddims_b = [1, 2, 3, 2, 1]
+    localdims1 = [3, 3, 3, 3]
+    localdims2 = [3, 3, 3, 3]
+
+    a = TensorTrain{ComplexF64,4}([
+        rand(ComplexF64, bonddims_a[n], localdims1[n], localdims2[n], bonddims_a[n+1])
+        for n = 1:N
+    ])
+    b = TensorTrain{ComplexF64,3}([
+        rand(ComplexF64, bonddims_b[n], localdims2[n], bonddims_b[n+1])
+        for n = 1:N
+    ])
+    return N, a, b, localdims1, localdims2
+end
+
 @testset "MPO-MPO contraction" for f in [nothing, x -> 2 * x], algorithm in [:TCI, :naive]
+    #==
     N = 4
     bonddims_a = [1, 2, 3, 2, 1]
     bonddims_b = [1, 2, 3, 2, 1]
@@ -44,6 +82,8 @@ end
         rand(ComplexF64, bonddims_b[n], localdims2[n], localdims3[n], bonddims_b[n+1])
         for n = 1:N
     ])
+    ==#
+    N, a, b, localdims1, localdims2, localdims3 = _gen_testdata_TTO_TTO()
 
     if f !== nothing && algorithm === :naive
         @test_throws ErrorException contract(a, b; f=f, algorithm=algorithm)
@@ -59,6 +99,7 @@ end
 end
 
 @testset "MPO-MPS contraction" for f in [nothing, x -> 2 * x], algorithm in [:TCI, :naive]
+    #==
     N = 4
     bonddims_a = [1, 2, 3, 2, 1]
     bonddims_b = [1, 2, 3, 2, 1]
@@ -73,6 +114,8 @@ end
         rand(ComplexF64, bonddims_b[n], localdims2[n], bonddims_b[n+1])
         for n = 1:N
     ])
+    ==#
+    N, a, b, localdims1, localdims2 = _gen_testdata_TTO_TTS()
 
     if f !== nothing && algorithm === :naive
         @test_throws ErrorException contract(a, b; f=f, algorithm=algorithm)
@@ -90,3 +133,16 @@ end
         end
     end
 end
+
+
+@testset "MPO-MPO contraction (zipup)" for method in [:SVD, :LU]
+    N, a, b, localdims1, localdims2, localdims3 = _gen_testdata_TTO_TTO()
+    ab = contract(a, b; algorithm=:zipup, method=method)
+    @test _tomat(ab) ≈ _tomat(a) * _tomat(b)
+end
+
+@testset "MPO-MPS contraction (zipup)" for method in [:SVD, :LU]
+    N, a, b, localdims1, localdims2 = _gen_testdata_TTO_TTS()
+    ab = contract(a, b; algorithm=:zipup, method=method)
+    @test _tovec(ab) ≈ _tomat(a) * _tovec(b)
+end