Refactor entropy function to also use spectrum directly (#377)

* Refactor entropy function to also use spectrum directly

* specialize to SectorVector and add test

* add infinite test

* update docstring

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: AFeuerpfeil

src/algorithms/toolbox.jl

@@ -1,15 +1,24 @@
 """
     entropy(state, [site::Int])
+    entropy(spectrum::SectorVector)
 
-Calculate the von Neumann entanglement entropy of a given MPS. If an integer `site` is
-given, the entropy is across the entanglement cut to the right of site `site`. Otherwise, a
-vector of entropies is returned, one for each site.
+Calculate the von Neumann entanglement entropy. The entropy can be computed from either an
+MPS state or directly from an entanglement spectrum as obtained from
+[`entanglement_spectrum`](@ref).
+
+When called on an MPS with an integer `site`, the entropy is computed across the
+entanglement cut to the right of site `site`. For `InfiniteMPS`, omitting `site` returns a
+vector of entropies, one for each site. For `FiniteMPS` and `WindowMPS`, `site` is
+required.
 """
 entropy(state::InfiniteMPS) = map(Base.Fix1(entropy, state), 1:length(state))
 function entropy(state::Union{FiniteMPS, WindowMPS, InfiniteMPS}, loc::Int)
-    S = zero(real(scalartype(state)))
-    tol = eps(typeof(S))
-    for (c, b) in pairs(entanglement_spectrum(state, loc))
+    return entropy(entanglement_spectrum(state, loc))
+end
+function entropy(spectrum::TensorKit.SectorVector{T}) where {T}
+    S = zero(T)
+    tol = eps(T)
+    for (c, b) in pairs(spectrum)
         s = zero(S)
         for x in b
             x < tol && break

test/states/finitemps.jl

@@ -127,3 +127,34 @@ end
     @test (ismissing(mps1.Cs[end]) && ismissing(mps2.Cs[end])) ||
         mps1.Cs[end] !== mps2.Cs[end]
 end
+
+@testset "FiniteMPS entropy ($(sectortype(D)), $elt)" for (D, d, elt) in [
+        (ℙ^10, ℙ^2, ComplexF64),
+        (
+            Rep[U₁](-1 => 3, 0 => 3, 1 => 3),
+            Rep[U₁](-1 => 1, 0 => 1, 1 => 1),
+            ComplexF64,
+        ),
+    ]
+    L = 6
+    ψ = FiniteMPS(rand, elt, L, d, D)
+
+    # entropy is non-negative at all sites
+    for site in 1:L
+        @test real(entropy(ψ, site)) >= 0
+    end
+
+    # entropy is consistent with entanglement_spectrum
+    for site in 1:L
+        @test entropy(ψ, site) ≈ entropy(entanglement_spectrum(ψ, site))
+    end
+
+    # entropy is zero at the right boundary (trivial bond)
+    @test entropy(ψ, L) ≈ 0 atol = 1.0e-10
+
+    # product state has zero entropy everywhere
+    ψ_product = FiniteMPS(rand, elt, L, d, oneunit(D))
+    for site in 1:L
+        @test entropy(ψ_product, site) ≈ 0 atol = 1.0e-10
+    end
+end

test/states/infinitemps.jl

@@ -70,3 +70,25 @@ end
     @test mps1.AC[1] !== mps2.AC[1]
     @test mps1.C[1] !== mps2.C[1]
 end
+
+@testset "InfiniteMPS entropy ($(sectortype(D)), $elt)" for (D, d, elt) in
+    [(ℙ^10, ℙ^2, ComplexF64), (Rep[U₁](1 => 3), Rep[U₁](0 => 1), ComplexF64)]
+    ψ = InfiniteMPS([d, d], [D, D])
+
+    # entropy(ψ) returns one value per site, all non-negative
+    Ss = entropy(ψ)
+    @test length(Ss) == length(ψ)
+    @test all(isreal, Ss)
+    @test all(>=(0), Ss)
+
+    # entropy(ψ, site) is non-negative and consistent with entropy(ψ)
+    for site in 1:length(ψ)
+        @test entropy(ψ, site) ≈ Ss[site]
+        @test entropy(ψ, site) ≈ entropy(entanglement_spectrum(ψ, site))
+    end
+
+    # product state has zero entropy everywhere
+    ψ_product = InfiniteMPS([d, d], [oneunit(D), oneunit(D)])
+    Ss_product = entropy(ψ_product)
+    @test all(S -> isapprox(S, 0; atol = 1.0e-10), Ss_product)
+end