diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index b980a52f..3da4ca67 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: @@ -12,10 +12,10 @@ using NamedGraphs.PartitionedGraphs: partitionedges, unpartitioned_graph using SimpleTraits: SimpleTraits, Not, @traitfn +using NDTensors: NDTensors -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() -default_message_norm(m::ITensor) = norm(m) function default_message_update(contract_list::Vector{ITensor}; kwargs...) sequence = optimal_contraction_sequence(contract_list) updated_messages = contract(contract_list; sequence, kwargs...) @@ -33,17 +33,16 @@ 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}; message_norm=default_message_norm -) +#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) - norm_lhs, norm_rhs = message_norm(lhs), message_norm(rhs) - return 0.5 * norm((denseblocks(lhs) / norm_lhs) - (denseblocks(rhs) / norm_rhs)) + f = abs2(dot(lhs / norm(lhs), rhs / norm(rhs))) + return 1 - f end struct BeliefPropagationCache{PTN,MTS,DM} @@ -99,8 +98,10 @@ for f in [ end end +NDTensors.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 728b06d6..a151e4d1 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 @@ -34,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, s; link_space=χ) - ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - bpc = BeliefPropagationCache(ψψ) - bpc = update(bpc; maxiter=50, 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 - 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 -> imag(eig) ≈ 0, eigs) - @test all(eig -> real(eig) >= -eps(eltype(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 * ψ[(1, 1)] + @test iszero(scalar(ψ; alg="bp")) + end end 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)