Refactored

Author: shinaoka

src/contraction.jl

@@ -62,7 +62,7 @@ _localdims(obj::TensorTrain{<:Any,4}, n::Int)::Tuple{Int,Int} =
 _localdims(obj::Contraction{<:Any}, n::Int)::Tuple{Int,Int} =
     (size(obj.mpo[1][n], 2), size(obj.mpo[2][n], 3))
 
-_getindex(x, indices) = ntuple(i->x[indices[i]], length(indices))
+_getindex(x, indices) = ntuple(i -> x[indices[i]], length(indices))
 
 function _contract(
     a::AbstractArray{T1,N1},

test/test_contraction.jl

@@ -22,7 +22,7 @@ end
     @test vec(reshape(permutedims(a, (2, 1, 3)), 3, :) * reshape(permutedims(b, (1, 3, 2)), :, 5)) ≈ vec(ab)
 end
 
-@testset "MPO-MPO naive contraction" begin
+@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]
@@ -39,29 +39,13 @@ end
         for n = 1:N
     ])
 
-    ab = contract(a, b; algorithm="naive")
-
-    @test _tomat(ab) ≈ _tomat(a) * _tomat(b)
-end
-
-@testset "MPO-MPO contraction" for f in [x -> x, x -> 2 * x]
-    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
-    ])
-
-    ab = contract(a, b; f = f, algorithm="TCI")
-    @test sitedims(ab) == [[localdims1[i], localdims3[i]] for i = 1:N]
-    @test _tomat(ab) ≈ f.(_tomat(a) * _tomat(b))
+    if f === nothing || algorithm != "naive"
+        ab = contract(a, b; f=f, algorithm=algorithm)
+        @test sitedims(ab) == [[localdims1[i], localdims3[i]] for i = 1:N]
+        if f === nothing
+            @test _tomat(ab) ≈ _tomat(a) * _tomat(b)
+        else
+            @test _tomat(ab) ≈ f.(_tomat(a) * _tomat(b))
+        end
+    end
 end