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,11 @@ 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)
-function default_message_update(contract_list::Vector{ITensor}; kwargs...)
+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]
@@ -157,6 +157,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)))
@@ -309,3 +318,23 @@ 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
+    n = region_scalar(bp_cache, pe)
+    me, mer = only(mts[pe]), only(mts[reverse(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