Enhance simple_update! for MPS in the Canonical form

src/Ansatz.jl

@@ -321,10 +321,10 @@ end
 Truncate the dimension of the virtual `bond` of a [`Canonical`](@ref) Tensor Network by keeping the `maxdim` largest
 **Schmidt coefficients** or those larger than `threshold`, and then recanonizes the Tensor Network if `recanonize` is `true`.
 """
-function truncate!(::Canonical, tn::AbstractAnsatz, bond; threshold, maxdim, recanonize=true)
+function truncate!(::Canonical, tn::AbstractAnsatz, bond; threshold, maxdim, recanonize=false)
     truncate!(NonCanonical(), tn, bond; threshold, maxdim, compute_local_svd=false)
 
-    recanonize && canonize!(tn)
+    recanonize && canonize!(tn; normalize=true)
 
     return tn
 end
@@ -387,11 +387,11 @@ function simple_update!(ψ::AbstractAnsatz, gate; threshold=nothing, maxdim=noth
 
     if nlanes(gate) == 1
         return simple_update_1site!(ψ, gate)
+    elseif nlanes(gate) == 2
+        return simple_update_2site!(form(ψ), ψ, gate; threshold, maxdim, kwargs...)
+    else
+        throw(ArgumentError("Only 1-site and 2-site gates are currently supported"))
     end
-
-    @assert has_edge(ψ, lanes(gate)...) "Gate must act on neighboring sites"
-
-    return simple_update!(form(ψ), ψ, gate; threshold, maxdim, kwargs...)
 end
 
 # TODO a lot of problems with merging... maybe we shouldn't merge manually
@@ -419,9 +419,16 @@ function simple_update_1site!(ψ::AbstractAnsatz, gate)
     return contract!(ψ, contracting_index)
 end
 
+function simple_update_2site!(
+    ::MixedCanonical, ψ::AbstractAnsatz, gate; threshold=nothing, maxdim=nothing, renormalize=false
+)
+    return simple_update_2site!(NonCanonical(), ψ, gate; threshold, maxdim, renormalize)
+end
+
 # TODO remove `renormalize` argument?
-function simple_update!(::NonCanonical, ψ::AbstractAnsatz, gate; threshold=nothing, maxdim=nothing, renormalize=false)
-    @assert nlanes(gate) == 2 "Only 2-site gates are supported currently"
+function simple_update_2site!(
+    ::NonCanonical, ψ::AbstractAnsatz, gate; threshold=nothing, maxdim=nothing, renormalize=false
+)
     @assert has_edge(ψ, lanes(gate)...) "Gate must act on neighboring sites"
 
     # shallow copy to avoid problems if errors in mid execution
@@ -455,16 +462,38 @@ function simple_update!(::NonCanonical, ψ::AbstractAnsatz, gate; threshold=noth
 
     # truncate virtual index
     if any(!isnothing, (threshold, maxdim))
-        truncate!(ψ, bond; threshold, maxdim)
-        renormalize && normalize!(ψ, bond[1])
+        truncate!(ψ, collect(bond); threshold, maxdim)
+        renormalize && normalize!(ψ, bond)
     end
 
     return ψ
 end
 
 # TODO remove `renormalize` argument?
-# TODO optimize correctly -> avoid recanonization + use lateral Λs
-function simple_update!(::Canonical, ψ::AbstractAnsatz, gate; threshold, maxdim, renormalize=false)
-    simple_update!(NonCanonical(), ψ, gate; threshold, maxdim, renormalize)
-    return canonize!(ψ)
+function simple_update_2site!(::Canonical, ψ::AbstractAnsatz, gate; threshold, maxdim, renormalize=false)
+    # Contract the exterior Λ tensors
+    sitel, siter = extrema(lanes(gate))
+    (0 < id(sitel) < nsites(ψ) || 0 < id(siter) < nsites(ψ)) ||
+        throw(ArgumentError("The sites in the bond must be between 1 and $(nsites(ψ))"))
+
+    Λᵢ₋₁ = id(sitel) == 1 ? nothing : tensors(ψ; between=(Site(id(sitel) - 1), sitel))
+    Λᵢ₊₁ = id(sitel) == nsites(ψ) - 1 ? nothing : tensors(ψ; between=(siter, Site(id(siter) + 1)))
+
+    !isnothing(Λᵢ₋₁) && contract!(ψ; between=(Site(id(sitel) - 1), sitel), direction=:right, delete_Λ=false)
+    !isnothing(Λᵢ₊₁) && contract!(ψ; between=(siter, Site(id(siter) + 1)), direction=:left, delete_Λ=false)
+
+    simple_update_2site!(NonCanonical(), ψ, gate; threshold, maxdim, renormalize)
+
+    # contract the updated tensors with the inverse of Λᵢ and Λᵢ₊₂, to get the new Γ tensors
+    U, Vt = tensors(ψ; at=sitel), tensors(ψ; at=siter)
+    Γᵢ₋₁ =
+        isnothing(Λᵢ₋₁) ? U : contract(U, Tensor(diag(pinv(Diagonal(parent(Λᵢ₋₁)); atol=1e-32)), inds(Λᵢ₋₁)); dims=())
+    Γᵢ =
+        isnothing(Λᵢ₊₁) ? Vt : contract(Tensor(diag(pinv(Diagonal(parent(Λᵢ₊₁)); atol=1e-32)), inds(Λᵢ₊₁)), Vt; dims=())
+
+    # Update the tensors in the tensor network
+    replace!(ψ, tensors(ψ; at=sitel) => Γᵢ₋₁)
+    replace!(ψ, tensors(ψ; at=siter) => Γᵢ)
+
+    return ψ
 end

src/MPS.jl

@@ -481,7 +481,7 @@ function canonize_site!(ψ::MPS, site::Site; direction::Symbol, method=:qr)
     return ψ
 end
 
-function canonize!(ψ::AbstractMPO)
+function canonize!(ψ::AbstractMPO; normalize=false)
     Λ = Tensor[]
 
     # right-to-left QR sweep, get right-canonical tensors
@@ -495,6 +495,7 @@ function canonize!(ψ::AbstractMPO)
 
         # extract the singular values and contract them with the next tensor
         Λᵢ = pop!(ψ, tensors(ψ; between=(Site(i), Site(i + 1))))
+        normalize && (Λᵢ ./= norm(Λᵢ))
         Aᵢ₊₁ = tensors(ψ; at=Site(i + 1))
         replace!(ψ, Aᵢ₊₁ => contract(Aᵢ₊₁, Λᵢ; dims=()))
         push!(Λ, Λᵢ)
@@ -540,20 +541,35 @@ function mixed_canonize!(tn::AbstractMPO, orthog_center)
     return tn
 end
 
+LinearAlgebra.normalize(ψ::AbstractMPO, site) = normalize!(copy(ψ), site)
+
 LinearAlgebra.normalize!(ψ::AbstractMPO; kwargs...) = normalize!(form(ψ), ψ; kwargs...)
+LinearAlgebra.normalize!(ψ::AbstractMPO, site) = normalize!(form(ψ), ψ; at=site)
 
 function LinearAlgebra.normalize!(::NonCanonical, ψ::AbstractMPO; at=Site(nsites(ψ) ÷ 2))
     tensor = tensors(ψ; at)
     tensor ./= norm(ψ)
     return ψ
 end
 
-LinearAlgebra.normalize!(ψ::AbstractMPO, site::Site) = normalize!(mixed_canonize!(ψ, site); at=site)
-
 function LinearAlgebra.normalize!(config::MixedCanonical, ψ::AbstractMPO; at=config.orthog_center)
     mixed_canonize!(ψ, at)
     normalize!(tensors(ψ; at), 2)
     return ψ
 end
 
-# TODO function LinearAlgebra.normalize!(::Canonical, ψ::AbstractMPO) end
+function LinearAlgebra.normalize!(config::Canonical, ψ::AbstractMPO; at=nothing)
+    if isnothing(at) # Normalize all λ tensors
+        normalizer = (norm(ψ))^(1 / (nsites(ψ) - 1))
+
+        for i in 1:(nsites(ψ) - 1)
+            λ = tensors(ψ; between=(Site(i), Site(i + 1)))
+            replace!(ψ, λ => λ ./ normalizer)
+        end
+    else
+        λ = tensors(ψ; between=at)
+        replace!(ψ, λ => λ ./ norm(ψ))
+    end
+
+    return ψ
+end

src/Quantum.jl

@@ -408,6 +408,8 @@ function Base.merge!(a::AbstractQuantum, b::AbstractQuantum; reset=true)
     return a
 end
 
+LinearAlgebra.normalize(ψ::AbstractQuantum; kwargs...) = normalize!(copy(ψ); kwargs...)
+
 function LinearAlgebra.norm(ψ::AbstractQuantum, p::Real=2; kwargs...)
     p == 2 || throw(ArgumentError("only L2-norm is implemented yet"))
     return LinearAlgebra.norm2(ψ; kwargs...)

test/MPS_test.jl

@@ -115,8 +115,9 @@ using LinearAlgebra
             ψ = rand(MPS; n=5, maxdim=16)
             canonize!(ψ)
 
-            truncated = truncate(ψ, [site"2", site"3"]; maxdim=2)
+            truncated = truncate(ψ, [site"2", site"3"]; maxdim=2, recanonize=true)
             @test size(truncated, inds(truncated; bond=[site"2", site"3"])) == 2
+            @test Tenet.check_form(truncated)
         end
 
         @testset "MixedCanonical" begin
@@ -144,11 +145,42 @@ using LinearAlgebra
     end
 
     @testset "normalize!" begin
-        using LinearAlgebra: normalize!
+        using LinearAlgebra: normalize, normalize!
 
-        ψ = MPS([rand(4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4)])
-        normalize!(ψ, Site(3))
-        @test isapprox(norm(ψ), 1.0)
+        @testset "NonCanonical" begin
+            ψ = MPS([rand(4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4)])
+
+            normalized = normalize(ψ)
+            @test norm(normalized) ≈ 1.0
+
+            normalize!(ψ, Site(3))
+            @test norm(ψ) ≈ 1.0
+        end
+
+        @testset "MixedCanonical" begin
+            ψ = rand(MPS; n=5, maxdim=16)
+
+            # Perturb the state to make it non-normalized
+            t = tensors(ψ; at=site"3")
+            replace!(ψ, t => Tensor(rand(size(t)...), inds(t)))
+
+            normalized = normalize(ψ)
+            @test norm(normalized) ≈ 1.0
+
+            normalize!(ψ, Site(3))
+            @test norm(ψ) ≈ 1.0
+        end
+
+        @testset "Canonical" begin
+            ψ = MPS([rand(4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4, 4), rand(4, 4)])
+            canonize!(ψ)
+
+            normalized = normalize(ψ)
+            @test norm(normalized) ≈ 1.0
+
+            normalize!(ψ, (Site(3), Site(4)))
+            @test norm(ψ) ≈ 1.0
+        end
     end
 
     @testset "canonize_site!" begin
@@ -306,11 +338,16 @@ using LinearAlgebra
             end
 
             @testset "Canonical" begin
-                ψ = deepcopy(ψ)
+                ψ = rand(MPS; n=5, maxdim=20)
+                ϕ = deepcopy(ψ)
                 canonize!(ψ)
-                evolved = evolve!(deepcopy(ψ), gate; threshold=1e-14)
-                @test isapprox(contract(evolved), contract(ψ))
-                @test issetequal(size.(tensors(evolved)), [(2, 2), (2,), (2, 2, 2), (2,), (2, 2, 2), (2,), (2, 2)])
+                evolved = evolve!(deepcopy(ψ), gate)
+
+                @test Tenet.check_form(evolved)
+                @test isapprox(contract(evolved), contract(ϕ)) # Identity gate should not change the state
+
+                # Ensure that the original MixedCanonical state evolves into the same state as the canonicalized one
+                @test contract(evolve!(ϕ, gate; threshold=1e-14)) ≈ contract(ψ)
             end
         end
     end