From 9e954dd97006b3dee893c637bc94fca5848ebbdd Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 09:24:35 -0400 Subject: [PATCH 1/9] Improved BP Convergence MEasure --- src/caches/beliefpropagationcache.jl | 15 ++++++++------- test/test_belief_propagation.jl | 7 ++++--- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index b980a52f..b5584a0a 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -1,6 +1,6 @@ using Graphs: IsDirected using SplitApplyCombine: group -using LinearAlgebra: diag +using LinearAlgebra: diag, dot using ITensors: dir using ITensorMPS: ITensorMPS using NamedGraphs.PartitionedGraphs: @@ -15,7 +15,7 @@ using SimpleTraits: SimpleTraits, Not, @traitfn default_message(inds_e) = ITensor[denseblocks(delta(i)) for i in inds_e] default_messages(ptn::PartitionedGraph) = Dictionary() -default_message_norm(m::ITensor) = norm(m) +euclidean_dist(x, y) = sqrt(abs(norm(x)^2 + norm(y)^2 - 2 * real(dot(x, y)))) function default_message_update(contract_list::Vector{ITensor}; kwargs...) sequence = optimal_contraction_sequence(contract_list) updated_messages = contract(contract_list; sequence, kwargs...) @@ -38,12 +38,13 @@ function default_cache_construction_kwargs(alg::Algorithm"bp", ψ::AbstractITens return (; partitioned_vertices=default_partitioned_vertices(ψ)) end -function message_diff( - message_a::Vector{ITensor}, message_b::Vector{ITensor}; message_norm=default_message_norm -) +function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) lhs, rhs = contract(message_a), contract(message_b) - norm_lhs, norm_rhs = message_norm(lhs), message_norm(rhs) - return 0.5 * norm((denseblocks(lhs) / norm_lhs) - (denseblocks(rhs) / norm_rhs)) + @assert issetequal(inds(lhs), inds(rhs)) + c = combiner(inds(lhs)) + lhs *= c + rhs *= c + return euclidean_dist(lhs / norm(lhs), rhs / norm(rhs)) end struct BeliefPropagationCache{PTN,MTS,DM} diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index 728b06d6..3bfa25c7 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -23,7 +23,8 @@ using ITensorNetworks: tensornetwork, update, update_factor, - update_message + update_message, + message_diff using ITensors: ITensors, ITensor, combiner, dag, inds, inner, op, prime, random_itensor using ITensorNetworks.ModelNetworks: ModelNetworks using ITensors.NDTensors: array @@ -42,8 +43,8 @@ using Test: @test, @testset rng = StableRNG(1234) ψ = random_tensornetwork(rng, s; link_space=χ) ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - bpc = BeliefPropagationCache(ψψ) - bpc = update(bpc; maxiter=50, tol=1e-10) + bpc = BeliefPropagationCache(ψψ, group(v -> first(v), vertices(ψψ))) + bpc = update(bpc; maxiter=20, tol=1e-8) #Test messages are converged for pe in partitionedges(partitioned_tensornetwork(bpc)) @test update_message(bpc, pe) ≈ message(bpc, pe) atol = 1e-8 From cf8e4e9808a11b6e605da9aa1c9d0d462926b6af Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 12:15:33 -0400 Subject: [PATCH 2/9] Switch to ComplexF64 for BP Examples --- test/test_belief_propagation.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index 3bfa25c7..367a4c24 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -41,13 +41,13 @@ using Test: @test, @testset s = siteinds("S=1/2", g) χ = 2 rng = StableRNG(1234) - ψ = random_tensornetwork(rng, s; link_space=χ) + ψ = random_tensornetwork(rng, ComplexF64, s; link_space=χ) ψψ = ψ ⊗ prime(dag(ψ); sites=[]) bpc = BeliefPropagationCache(ψψ, group(v -> first(v), vertices(ψψ))) - bpc = update(bpc; maxiter=20, tol=1e-8) + bpc = update(bpc; maxiter=25, tol=1e-10) #Test messages are converged for pe in partitionedges(partitioned_tensornetwork(bpc)) - @test update_message(bpc, pe) ≈ message(bpc, pe) atol = 1e-8 + @test message_diff(update_message(bpc, pe), message(bpc, pe)) < 1e-8 end #Test updating the underlying tensornetwork in the cache v = first(vertices(ψψ)) @@ -70,8 +70,8 @@ using Test: @test, @testset eigs = eigvals(rdm) @test size(rdm) == (2^length(vs), 2^length(vs)) - @test all(eig -> imag(eig) ≈ 0, eigs) - @test all(eig -> real(eig) >= -eps(eltype(eig)), eigs) + @test all(eig -> abs(imag(eig)) <= 1e-16, eigs) + @test all(eig -> real(eig) >= -eps(eltype(real(eig))), eigs) #Test edge case of network which evalutes to 0 χ = 2 From 6eb3ceda2faaf3bf9d6ad2f37a8bd7b579f2e5ff Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 12:17:30 -0400 Subject: [PATCH 3/9] Switch test --- src/caches/beliefpropagationcache.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index b5584a0a..c1f7f778 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -15,7 +15,6 @@ using SimpleTraits: SimpleTraits, Not, @traitfn default_message(inds_e) = ITensor[denseblocks(delta(i)) for i in inds_e] default_messages(ptn::PartitionedGraph) = Dictionary() -euclidean_dist(x, y) = sqrt(abs(norm(x)^2 + norm(y)^2 - 2 * real(dot(x, y)))) function default_message_update(contract_list::Vector{ITensor}; kwargs...) sequence = optimal_contraction_sequence(contract_list) updated_messages = contract(contract_list; sequence, kwargs...) @@ -44,7 +43,9 @@ function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) c = combiner(inds(lhs)) lhs *= c rhs *= c - return euclidean_dist(lhs / norm(lhs), rhs / norm(rhs)) + lhs /= norm(lhs) + rhs /= norm(rhs) + return 1 - sqrt(abs(dot(lhs, rhs))) end struct BeliefPropagationCache{PTN,MTS,DM} From b23f5a61861a634c72a45fb6ee932a63d31dbddd Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 12:24:50 -0400 Subject: [PATCH 4/9] Better fidelity def --- src/caches/beliefpropagationcache.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index c1f7f778..1f465fd9 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -45,7 +45,8 @@ function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) rhs *= c lhs /= norm(lhs) rhs /= norm(rhs) - return 1 - sqrt(abs(dot(lhs, rhs))) + f = abs(dot(lhs, rhs))^2 + return 1 - f end struct BeliefPropagationCache{PTN,MTS,DM} From 8ad449319d0473bde7ab867fd867e6e35468944e Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 12:58:20 -0400 Subject: [PATCH 5/9] Fix gauging tests. Lower tol for default --- src/caches/beliefpropagationcache.jl | 10 ++-------- test/test_gauging.jl | 6 ++---- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 1f465fd9..3b77a0b1 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -32,20 +32,14 @@ default_partitioned_vertices(ψ::AbstractITensorNetwork) = group(v -> v, vertice function default_partitioned_vertices(f::AbstractFormNetwork) return group(v -> original_state_vertex(f, v), vertices(f)) end -default_cache_update_kwargs(cache) = (; maxiter=20, tol=1e-5) +default_cache_update_kwargs(cache) = (; maxiter=25, tol=1e-8) function default_cache_construction_kwargs(alg::Algorithm"bp", ψ::AbstractITensorNetwork) return (; partitioned_vertices=default_partitioned_vertices(ψ)) end function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) lhs, rhs = contract(message_a), contract(message_b) - @assert issetequal(inds(lhs), inds(rhs)) - c = combiner(inds(lhs)) - lhs *= c - rhs *= c - lhs /= norm(lhs) - rhs /= norm(rhs) - f = abs(dot(lhs, rhs))^2 + f = abs(dot(lhs / norm(lhs), rhs / norm(rhs)))^2 return 1 - f end diff --git a/test/test_gauging.jl b/test/test_gauging.jl index 9eb4ebbf..67b1b5dd 100644 --- a/test/test_gauging.jl +++ b/test/test_gauging.jl @@ -27,9 +27,7 @@ using Test: @test, @testset ψ = random_tensornetwork(rng, s; link_space=χ) # Move directly to vidal gauge - ψ_vidal = VidalITensorNetwork( - ψ; cache_update_kwargs=(; maxiter=20, tol=1e-12, verbose=true) - ) + ψ_vidal = VidalITensorNetwork(ψ; cache_update_kwargs=(; maxiter=30, verbose=true)) @test gauge_error(ψ_vidal) < 1e-8 # Move to symmetric gauge @@ -38,7 +36,7 @@ using Test: @test, @testset bp_cache = cache_ref[] # Test we just did a gauge transform and didn't change the overall network - @test inner(ψ_symm, ψ) / sqrt(inner(ψ_symm, ψ_symm) * inner(ψ, ψ)) ≈ 1.0 + @test inner(ψ_symm, ψ) / sqrt(inner(ψ_symm, ψ_symm) * inner(ψ, ψ)) ≈ 1.0 atol = 1e-8 #Test all message tensors are approximately diagonal even when we keep running BP bp_cache = update(bp_cache; maxiter=10) From 173e19bd6979d83067b8716c50805ceae08171cf Mon Sep 17 00:00:00 2001 From: Joseph Tindall <51231103+JoeyT1994@users.noreply.github.com> Date: Tue, 25 Jun 2024 13:55:30 -0400 Subject: [PATCH 6/9] Update src/caches/beliefpropagationcache.jl Co-authored-by: Matt Fishman --- src/caches/beliefpropagationcache.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 3b77a0b1..79ac1249 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -39,7 +39,7 @@ end function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) lhs, rhs = contract(message_a), contract(message_b) - f = abs(dot(lhs / norm(lhs), rhs / norm(rhs)))^2 + f = abs2(dot(lhs / norm(lhs), rhs / norm(rhs))) return 1 - f end From b4c6532ba8736154aa54c09caefe985c2d8c7900 Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 16:14:25 -0400 Subject: [PATCH 7/9] Eltype support --- src/caches/beliefpropagationcache.jl | 7 ++- test/test_belief_propagation.jl | 84 +++++++++++++++------------- 2 files changed, 50 insertions(+), 41 deletions(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 79ac1249..6cc0ef42 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -13,7 +13,7 @@ using NamedGraphs.PartitionedGraphs: unpartitioned_graph using SimpleTraits: SimpleTraits, Not, @traitfn -default_message(inds_e) = ITensor[denseblocks(delta(i)) for i in inds_e] +default_message(elt, inds_e) = ITensor[denseblocks(delta(elt, i)) for i in inds_e] default_messages(ptn::PartitionedGraph) = Dictionary() function default_message_update(contract_list::Vector{ITensor}; kwargs...) sequence = optimal_contraction_sequence(contract_list) @@ -37,6 +37,7 @@ function default_cache_construction_kwargs(alg::Algorithm"bp", ψ::AbstractITens return (; partitioned_vertices=default_partitioned_vertices(ψ)) end +#TODO: Take `dot` without precontracting the messages to allow scaling to more complex messages function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) lhs, rhs = contract(message_a), contract(message_b) f = abs2(dot(lhs / norm(lhs), rhs / norm(rhs))) @@ -96,8 +97,10 @@ for f in [ end end +ITensorNetworks.scalartype(bp_cache) = scalartype(tensornetwork(bp_cache)) + function default_message(bp_cache::BeliefPropagationCache, edge::PartitionEdge) - return default_message(bp_cache)(linkinds(bp_cache, edge)) + return default_message(bp_cache)(scalartype(bp_cache), linkinds(bp_cache, edge)) end function message(bp_cache::BeliefPropagationCache, edge::PartitionEdge) diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index 367a4c24..d17565a6 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -35,50 +35,56 @@ using NamedGraphs.PartitionedGraphs: PartitionVertex, partitionedges using SplitApplyCombine: group using StableRNGs: StableRNG using Test: @test, @testset -@testset "belief_propagation" begin - ITensors.disable_warn_order() - g = named_grid((3, 3)) - s = siteinds("S=1/2", g) - χ = 2 - rng = StableRNG(1234) - ψ = random_tensornetwork(rng, ComplexF64, s; link_space=χ) - ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - bpc = BeliefPropagationCache(ψψ, group(v -> first(v), vertices(ψψ))) - bpc = update(bpc; maxiter=25, tol=1e-10) - #Test messages are converged - for pe in partitionedges(partitioned_tensornetwork(bpc)) - @test message_diff(update_message(bpc, pe), message(bpc, pe)) < 1e-8 - end - #Test updating the underlying tensornetwork in the cache - v = first(vertices(ψψ)) - rng = StableRNG(1234) - new_tensor = random_itensor(rng, inds(ψψ[v])) - bpc_updated = update_factor(bpc, v, new_tensor) - ψψ_updated = tensornetwork(bpc_updated) - @test ψψ_updated[v] == new_tensor - #Test forming a two-site RDM. Check it has the correct size, trace 1 and is PSD - vs = [(2, 2), (2, 3)] +@testset "belief_propagation (eltype=$elt)" for elt in ( + Float32, Float64, Complex{Float32}, Complex{Float64} +) + begin + ITensors.disable_warn_order() + g = named_grid((3, 3)) + s = siteinds("S=1/2", g) + χ = 2 + rng = StableRNG(1234) + ψ = random_tensornetwork(rng, elt, s; link_space=χ) + ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + bpc = BeliefPropagationCache(ψψ, group(v -> first(v), vertices(ψψ))) + bpc = update(bpc; maxiter=25, tol=eps(real(elt))) + #Test messages are converged + for pe in partitionedges(partitioned_tensornetwork(bpc)) + @test message_diff(update_message(bpc, pe), message(bpc, pe)) < 10 * eps(real(elt)) + @test eltype(only(message(bpc, pe))) == elt + end + #Test updating the underlying tensornetwork in the cache + v = first(vertices(ψψ)) + rng = StableRNG(1234) + new_tensor = random_itensor(rng, inds(ψψ[v])) + bpc_updated = update_factor(bpc, v, new_tensor) + ψψ_updated = tensornetwork(bpc_updated) + @test ψψ_updated[v] == new_tensor + + #Test forming a two-site RDM. Check it has the correct size, trace 1 and is PSD + vs = [(2, 2), (2, 3)] - ψψsplit = split_index(ψψ, NamedEdge.([(v, 1) => (v, 2) for v in vs])) - env_tensors = environment(bpc, [(v, 2) for v in vs]) - rdm = contract(vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]])) + ψψsplit = split_index(ψψ, NamedEdge.([(v, 1) => (v, 2) for v in vs])) + env_tensors = environment(bpc, [(v, 2) for v in vs]) + rdm = contract(vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]])) - rdm = array((rdm * combiner(inds(rdm; plev=0)...)) * combiner(inds(rdm; plev=1)...)) - rdm /= tr(rdm) + rdm = array((rdm * combiner(inds(rdm; plev=0)...)) * combiner(inds(rdm; plev=1)...)) + rdm /= tr(rdm) - eigs = eigvals(rdm) - @test size(rdm) == (2^length(vs), 2^length(vs)) + eigs = eigvals(rdm) + @test size(rdm) == (2^length(vs), 2^length(vs)) - @test all(eig -> abs(imag(eig)) <= 1e-16, eigs) - @test all(eig -> real(eig) >= -eps(eltype(real(eig))), eigs) + @test all(eig -> abs(imag(eig)) <= eps(real(elt)), eigs) + @test all(eig -> real(eig) >= -eps(real(elt)), eigs) - #Test edge case of network which evalutes to 0 - χ = 2 - g = named_grid((3, 1)) - rng = StableRNG(1234) - ψ = random_tensornetwork(rng, ComplexF64, g; link_space=χ) - ψ[(1, 1)] = 0.0 * ψ[(1, 1)] - @test iszero(scalar(ψ; alg="bp")) + #Test edge case of network which evalutes to 0 + χ = 2 + g = named_grid((3, 1)) + rng = StableRNG(1234) + ψ = random_tensornetwork(rng, elt, g; link_space=χ) + ψ[(1, 1)] = 0.0 * ψ[(1, 1)] + @test iszero(scalar(ψ; alg="bp")) + end end end From 754e46ccd7f9668aa112f1a95154322d89555013 Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 16:17:59 -0400 Subject: [PATCH 8/9] Further test change --- test/test_belief_propagation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index d17565a6..a151e4d1 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -83,7 +83,7 @@ using Test: @test, @testset g = named_grid((3, 1)) rng = StableRNG(1234) ψ = random_tensornetwork(rng, elt, g; link_space=χ) - ψ[(1, 1)] = 0.0 * ψ[(1, 1)] + ψ[(1, 1)] = 0 * ψ[(1, 1)] @test iszero(scalar(ψ; alg="bp")) end end From 6d275f125d613c11568ed013ef4ec9700bcefe97 Mon Sep 17 00:00:00 2001 From: Joseph Tindall Date: Tue, 25 Jun 2024 16:30:46 -0400 Subject: [PATCH 9/9] Import scalartype from NDTensors --- src/caches/beliefpropagationcache.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 6cc0ef42..3da4ca67 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -12,6 +12,7 @@ using NamedGraphs.PartitionedGraphs: partitionedges, unpartitioned_graph using SimpleTraits: SimpleTraits, Not, @traitfn +using NDTensors: NDTensors default_message(elt, inds_e) = ITensor[denseblocks(delta(elt, i)) for i in inds_e] default_messages(ptn::PartitionedGraph) = Dictionary() @@ -97,7 +98,7 @@ for f in [ end end -ITensorNetworks.scalartype(bp_cache) = scalartype(tensornetwork(bp_cache)) +NDTensors.scalartype(bp_cache) = scalartype(tensornetwork(bp_cache)) function default_message(bp_cache::BeliefPropagationCache, edge::PartitionEdge) return default_message(bp_cache)(scalartype(bp_cache), linkinds(bp_cache, edge))