Unit test of canonization in MPO

src/MPS.jl

@@ -511,12 +511,12 @@ function canonize!(ψ::AbstractMPO; normalize=false)
     Λ = Tensor[]
 
     # right-to-left QR sweep, get right-canonical tensors
-    for i in nsites(ψ):-1:2
+    for i in nlanes(ψ):-1:2
         canonize_site!(ψ, Lane(i); direction=:left, method=:qr)
     end
 
     # left-to-right SVD sweep, get left-canonical tensors and singular values without reversing
-    for i in 1:(nsites(ψ) - 1)
+    for i in 1:(nlanes(ψ) - 1)
         canonize_site!(ψ, Lane(i); direction=:right, method=:svd)
 
         # extract the singular values and contract them with the next tensor
@@ -527,7 +527,7 @@ function canonize!(ψ::AbstractMPO; normalize=false)
         push!(Λ, Λᵢ)
     end
 
-    for i in 2:nsites(ψ) # tensors at i in "A" form, need to contract (Λᵢ)⁻¹ with A to get Γᵢ
+    for i in 2:nlanes(ψ) # tensors at i in "A" form, need to contract (Λᵢ)⁻¹ with A to get Γᵢ
         Λᵢ = Λ[i - 1] # singular values start between site 1 and 2
         A = tensors(ψ; at=Lane(i))
         Γᵢ = contract(A, Tensor(diag(pinv(Diagonal(parent(Λᵢ)); atol=1e-64)), inds(Λᵢ)); dims=())

test/MPO_test.jl

@@ -74,4 +74,86 @@
         @test isapprox(norm(H), 1.0)
         @test maximum(last, size(H)) <= χ
     end
+
+    @testset "canonize!" begin
+        using Tenet: isleftcanonical, isrightcanonical
+
+        ψ = MPO([rand(4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4)])
+        canonized = canonize(ψ)
+
+        @test form(canonized) isa Canonical
+
+        @test length(tensors(canonized)) == 9 # 5 tensors + 4 singular values vectors
+        @test isapprox(contract(transform(TensorNetwork(canonized), Tenet.HyperFlatten())), contract(ψ))
+        @test isapprox(norm(ψ), norm(canonized))
+
+        # Extract the singular values between each adjacent pair of sites in the canonized chain
+        Λ = [tensors(canonized; between=(Lane(i), Lane(i + 1))) for i in 1:4]
+
+        norm_psi = norm(ψ)
+        @test all(λ -> sqrt(sum(abs2, λ)) ≈ norm_psi, Λ)
+
+        for i in 1:5
+            canonized = canonize(ψ)
+
+            if i == 1
+                @test isleftcanonical(canonized, Lane(i))
+            elseif i == 5 # in the limits of the chain, we get the norm of the state
+                normalize!(tensors(canonized; bond=(Lane(i - 1), Lane(i))))
+                contract!(canonized; between=(Lane(i - 1), Lane(i)), direction=:right)
+                @test isleftcanonical(canonized, Lane(i))
+            else
+                contract!(canonized; between=(Lane(i - 1), Lane(i)), direction=:right)
+                @test isleftcanonical(canonized, Lane(i))
+            end
+        end
+
+        for i in 1:5
+            canonized = canonize(ψ)
+
+            if i == 1 # in the limits of the chain, we get the norm of the state
+                normalize!(tensors(canonized; bond=(Lane(i), Lane(i + 1))))
+                contract!(canonized; between=(Lane(i), Lane(i + 1)), direction=:left)
+                @test isrightcanonical(canonized, Lane(i))
+            elseif i == 5
+                @test isrightcanonical(canonized, Lane(i))
+            else
+                contract!(canonized; between=(Lane(i), Lane(i + 1)), direction=:left)
+                @test isrightcanonical(canonized, Lane(i))
+            end
+        end
+    end
+
+    @testset "mixed_canonize!" begin
+        ψ = MPO([rand(4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4, 4), rand(4, 4, 4)])
+
+        @testset "single Site" begin
+            canonized = mixed_canonize(ψ, lane"3")
+            @test Tenet.check_form(canonized)
+
+            @test form(canonized) isa MixedCanonical
+            @test form(canonized).orthog_center == lane"3"
+
+            @test isisometry(canonized, lane"1"; dir=:right)
+            @test isisometry(canonized, lane"2"; dir=:right)
+            @test isisometry(canonized, lane"4"; dir=:left)
+            @test isisometry(canonized, lane"5"; dir=:left)
+
+            @test contract(canonized) ≈ contract(ψ)
+        end
+
+        @testset "multiple Sites" begin
+            canonized = mixed_canonize(ψ, [lane"2", lane"3"])
+
+            @test Tenet.check_form(canonized)
+            @test form(canonized) isa MixedCanonical
+            @test form(canonized).orthog_center == [lane"2", lane"3"]
+
+            @test isisometry(canonized, lane"1"; dir=:right)
+            @test isisometry(canonized, lane"4"; dir=:left)
+            @test isisometry(canonized, lane"5"; dir=:left)
+
+            @test contract(canonized) ≈ contract(ψ)
+        end
+    end
 end

test/MPS_test.jl

@@ -249,7 +249,9 @@ using LinearAlgebra
 
         # Extract the singular values between each adjacent pair of sites in the canonized chain
         Λ = [tensors(canonized; between=(Lane(i), Lane(i + 1))) for i in 1:4]
-        @test map(λ -> sum(abs2, λ), Λ) ≈ ones(length(Λ)) * norm(canonized)^2
+
+        norm_psi = norm(ψ)
+        @test all(λ -> sqrt(sum(abs2, λ)) ≈ norm_psi, Λ)
 
         for i in 1:5
             canonized = canonize(ψ)