Normalize

src/ITensorNetworks.jl

@@ -61,6 +61,7 @@ include("solvers/linsolve.jl")
 include("solvers/sweep_plans/sweep_plans.jl")
 include("apply.jl")
 include("inner.jl")
+include("normalize.jl")
 include("expect.jl")
 include("environment.jl")
 include("exports.jl")

src/caches/beliefpropagationcache.jl

@@ -16,11 +16,12 @@ using NDTensors: NDTensors
 
 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...)
+default_message_norm(m::ITensor) = norm(m)
+function default_message_update(contract_list::Vector{ITensor}; normalize=true, kwargs...)
  sequence = optimal_contraction_sequence(contract_list)
  updated_messages = contract(contract_list; sequence, kwargs...)
  message_norm = norm(updated_messages)
- if !iszero(message_norm)
+ if normalize && !iszero(message_norm)
  updated_messages /= message_norm
  end
  return ITensor[updated_messages]
@@ -158,6 +159,15 @@ function environment(bp_cache::BeliefPropagationCache, verts::Vector)
  return vcat(messages, central_tensors)
 end
 
+function factors(bp_cache::BeliefPropagationCache, vertices)
+ tn = tensornetwork(bp_cache)
+ return map(vertex -> tn[vertex], vertices)
+end
+
+function factor(bp_cache::BeliefPropagationCache, vertex)
+ return only(factors(bp_cache, [vertex]))
+end
+
 function factor(bp_cache::BeliefPropagationCache, vertex::PartitionVertex)
  ptn = partitioned_tensornetwork(bp_cache)
  return collect(eachtensor(subgraph(ptn, vertex)))
@@ -310,3 +320,25 @@ end
 function scalar_factors_quotient(bp_cache::BeliefPropagationCache)
  return vertex_scalars(bp_cache), edge_scalars(bp_cache)
 end
+
+function normalize_messages(bp_cache::BeliefPropagationCache, pes::Vector{<:PartitionEdge})
+ bp_cache = copy(bp_cache)
+ mts = messages(bp_cache)
+ for pe in pes
+ me, mer = only(mts[pe]), only(mts[reverse(pe)])
+ set!(mts, pe, ITensor[me / norm(me)])
+ set!(mts, reverse(pe), ITensor[mer / norm(mer)])
+ n = region_scalar(bp_cache, pe)
+ set!(mts, pe, ITensor[(1 / sqrt(n)) * me])
+ set!(mts, reverse(pe), ITensor[(1 / sqrt(n)) * mer])
+ end
+ return bp_cache
+end
+
+function normalize_message(bp_cache::BeliefPropagationCache, pe::PartitionEdge)
+ return normalize_messages(bp_cache, PartitionEdge[pe])
+end
+
+function normalize_messages(bp_cache::BeliefPropagationCache)
+ return normalize_messages(bp_cache, partitionedges(partitioned_tensornetwork(bp_cache)))
+end

src/normalize.jl

@@ -0,0 +1,56 @@
+using LinearAlgebra
+
+function LinearAlgebra.normalize(tn::AbstractITensorNetwork; alg="exact", kwargs...)
+ return normalize(Algorithm(alg), tn; kwargs...)
+end
+
+function LinearAlgebra.normalize(alg::Algorithm"exact", tn::AbstractITensorNetwork)
+ norm_tn = norm_sqr_network(tn)
+ log_norm = logscalar(alg, norm_tn)
+ tn = copy(tn)
+ L = length(vertices(tn))
+ c = exp(log_norm / L)
+ for v in vertices(tn)
+ tn[v] = tn[v] / sqrt(c)
+ end
+ return tn
+end
+
+function LinearAlgebra.normalize(
+ alg::Algorithm"bp",
+ tn::AbstractITensorNetwork;
+ (cache!)=nothing,
+ update_cache=isnothing(cache!),
+ cache_update_kwargs=default_cache_update_kwargs(cache!),
+)
+ if isnothing(cache!)
+ cache! = Ref(BeliefPropagationCache(QuadraticFormNetwork(tn)))
+ end
+
+ if update_cache
+ cache![] = update(cache![]; cache_update_kwargs...)
+ end
+
+ tn = copy(tn)
+ cache![] = normalize_messages(cache![])
+ norm_tn = tensornetwork(cache![])
+
+ vertices_states = Dictionary()
+ for v in vertices(tn)
+ v_ket, v_bra = ket_vertex(norm_tn, v), bra_vertex(norm_tn, v)
+ pv = only(partitionvertices(cache![], [v_ket]))
+ vn = region_scalar(cache![], pv)
+ state = tn[v] / sqrt(vn)
+ state_dag = copy(dag(state))
+ state_dag = replaceinds(
+ state_dag, inds(state_dag), dual_index_map(norm_tn).(inds(state_dag))
+ )
+ set!(vertices_states, v_ket, state)
+ set!(vertices_states, v_bra, state_dag)
+ tn[v] = state
+ end
+
+ cache![] = update_factors(cache![], vertices_states)
+
+ return tn
+end

test/test_normalize.jl

@@ -0,0 +1,52 @@
+@eval module $(gensym())
+using ITensorNetworks:
+ BeliefPropagationCache,
+ QuadraticFormNetwork,
+ edge_scalars,
+ norm_sqr_network,
+ random_tensornetwork,
+ vertex_scalars
+using ITensors: dag, inner, siteinds, scalar
+using Graphs: SimpleGraph, uniform_tree
+using LinearAlgebra: normalize
+using NamedGraphs: NamedGraph
+using NamedGraphs.NamedGraphGenerators: named_grid
+using StableRNGs: StableRNG
+using Test: @test, @testset
+@testset "Normalize" begin
+
+ #First lets do a tree 
+ L = 6
+ χ = 2
+ rng = StableRNG(1234)
+
+ g = NamedGraph(SimpleGraph(uniform_tree(L)))
+ s = siteinds("S=1/2", g)
+ x = random_tensornetwork(rng, s; link_space=χ)
+
+ ψ = normalize(x; alg="exact")
+ @test scalar(norm_sqr_network(ψ); alg="exact") ≈ 1.0
+
+ ψ = normalize(x; alg="bp")
+ @test scalar(norm_sqr_network(ψ); alg="exact") ≈ 1.0
+
+ #Now a loopy graph
+ Lx, Ly = 3, 2
+ χ = 2
+ rng = StableRNG(1234)
+
+ g = named_grid((Lx, Ly))
+ s = siteinds("S=1/2", g)
+ x = random_tensornetwork(rng, s; link_space=χ)
+
+ ψ = normalize(x; alg="exact")
+ @test scalar(norm_sqr_network(ψ); alg="exact") ≈ 1.0
+
+ ψIψ_bpc = Ref(BeliefPropagationCache(QuadraticFormNetwork(x)))
+ ψ = normalize(x; alg="bp", (cache!)=ψIψ_bpc, update_cache=true)
+ ψIψ_bpc = ψIψ_bpc[]
+ @test all(x -> x ≈ 1.0, edge_scalars(ψIψ_bpc))
+ @test all(x -> x ≈ 1.0, vertex_scalars(ψIψ_bpc))
+ @test scalar(QuadraticFormNetwork(ψ); alg="bp") ≈ 1.0
+end
+end