Improvements

Author: JoeyT1994

src/caches/abstractbeliefpropagationcache.jl

@@ -283,3 +283,69 @@ function update(
 )
   return update(Algorithm(alg), bpc; kwargs...)
 end
+
+function rescale_message(bp_cache::AbstractBeliefPropagationCache, partitionpair)
+  return rescale_messages(bp_cache, typeof(partitionpair)[partitionpair])
+end
+
+function rescale_messages(bp_cache::AbstractBeliefPropagationCache)
+  return rescale_messages(bp_cache, partitionpairs(bp_cache))
+end
+
+function rescale_partitions(
+  bpc::AbstractBeliefPropagationCache,
+  partitions::Vector;
+  vs_to_rescale::Vector=collect(vertices(tensornetwork(bpc))),
+)
+  tn = tensornetwork(bpc)
+  for pv in partitions
+    pv_vs = filter(v -> v ∈ vs_to_rescale, vertices(bpc, pv))
+
+    isempty(pv_vs) && continue
+
+    vn = region_scalar(bpc, pv)
+    if isreal(vn)
+      tn[first(pv_vs)] *= sign(vn)
+      vn *= sign(vn)
+    end
+
+    vn = vn^(1 / length(pv_vs))
+    for v in pv_vs
+      tn[v] /= vn
+    end
+  end
+
+  return bpc
+end
+
+function rescale_partitions(bpc::AbstractBeliefPropagationCache; kwargs...)
+  return rescale_partitions(bpc, collect(partitions(bpc)); kwargs...)
+end
+
+function rescale_partition(bpc::AbstractBeliefPropagationCache, partition; kwargs...)
+  return rescale_partitions(bpc, typeof(partition)[partition]; kwargs...)
+end
+
+function rescale(bpc::AbstractBeliefPropagationCache; kwargs...)
+  bpc = rescale_messages(bpc)
+  bpc = rescale_partitions(bpc; kwargs...)
+  return bpc
+end
+
+function logscalar(bpc::AbstractBeliefPropagationCache)
+  numerator_terms, denominator_terms = scalar_factors_quotient(bpc)
+  numerator_terms =
+    any(t -> real(t) < 0, numerator_terms) ? complex.(numerator_terms) : numerator_terms
+  denominator_terms = if any(t -> real(t) < 0, denominator_terms)
+    complex.(denominator_terms)
+  else
+    denominator_terms
+  end
+
+  any(iszero, denominator_terms) && return -Inf
+  return sum(log.(numerator_terms)) - sum(log.((denominator_terms)))
+end
+
+function ITensors.scalar(bpc::AbstractBeliefPropagationCache)
+  return exp(logscalar(bpc))
+end

src/caches/beliefpropagationcache.jl

@@ -107,28 +107,20 @@ function region_scalar(bp_cache::BeliefPropagationCache, pe::PartitionEdge)
   return contract(ts; sequence)[]
 end
 
-function normalize_messages(bp_cache::BeliefPropagationCache, pes::Vector{<:PartitionEdge})
+function rescale_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)])
     me, mer = normalize(me), normalize(mer)
     n = dot(me, mer)
-    if isreal(n) && n < 0
-      set!(mts, pe, ITensor[(sgn(n) / sqrt(abs(n))) * me])
-      set!(mts, reverse(pe), ITensor[(1 / sqrt(abs(n))) * mer])
-    else
-      set!(mts, pe, ITensor[(1 / sqrt(n)) * me])
-      set!(mts, reverse(pe), ITensor[(1 / sqrt(n)) * mer])
+    if isreal(n)
+      me *= sign(n)
+      n *= sign(n)
     end
+
+    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/contract.jl

@@ -63,17 +63,7 @@ function logscalar(
     cache![] = update(cache![]; cache_update_kwargs...)
   end
 
-  numerator_terms, denominator_terms = scalar_factors_quotient(cache![])
-  numerator_terms =
-    any(t -> real(t) < 0, numerator_terms) ? complex.(numerator_terms) : numerator_terms
-  denominator_terms = if any(t -> real(t) < 0, denominator_terms)
-    complex.(denominator_terms)
-  else
-    denominator_terms
-  end
-
-  any(iszero, denominator_terms) && return -Inf
-  return sum(log.(numerator_terms)) - sum(log.((denominator_terms)))
+  return logscalar(cache![])
 end
 
 function ITensors.scalar(alg::Algorithm, tn::AbstractITensorNetwork; kwargs...)

src/normalize.jl

@@ -5,25 +5,28 @@ function rescale(tn::AbstractITensorNetwork; alg="exact", kwargs...)
 end
 
 function rescale(
-  alg::Algorithm"exact", tn::AbstractITensorNetwork, vs=collect(vertices(tn)); kwargs...
+  alg::Algorithm"exact",
+  tn::AbstractITensorNetwork;
+  vs_to_rescale=collect(vertices(tn)),
+  kwargs...,
 )
   logn = logscalar(alg, tn; kwargs...)
-  c = 1.0 / (exp(logn / length(vs)))
+  c = 1.0 / (exp(logn / length(vs_to_rescale)))
   tn = copy(tn)
-  for v in vs
+  for v in vs_to_rescale
     tn[v] *= c
   end
   return tn
 end
 
 function rescale(
   alg::Algorithm,
-  tn::AbstractITensorNetwork,
-  vs=collect(vertices(tn));
+  tn::AbstractITensorNetwork;
+  vs_to_rescale=collect(vertices(tn)),
   (cache!)=nothing,
   cache_construction_kwargs=default_cache_construction_kwargs(alg, tn),
   update_cache=isnothing(cache!),
-  cache_update_kwargs=default_cache_update_kwargs(cache!),
+  cache_update_kwargs=default_cache_update_kwargs(alg),
 )
   if isnothing(cache!)
     cache! = Ref(cache(alg, tn; cache_construction_kwargs...))
@@ -33,29 +36,9 @@ function rescale(
     cache![] = update(cache![]; cache_update_kwargs...)
   end
 
-  tn = copy(tn)
-  cache![] = normalize_messages(cache![])
-  vertices_states = Dictionary()
-  for pv in partitionvertices(cache![])
-    pv_vs = filter(v -> v ∈ vs, vertices(cache![], pv))
-
-    isempty(pv_vs) && continue
+  cache![] = rescale(cache![]; vs_to_rescale)
 
-    vn = region_scalar(cache![], pv)
-    if isreal(vn) && vn < 0
-      tn[first(pv_vs)] *= -1
-      vn = abs(vn)
-    end
-
-    vn = vn^(1 / length(pv_vs))
-    for v in pv_vs
-      tn[v] /= vn
-      set!(vertices_states, v, tn[v])
-    end
-  end
-
-  cache![] = update_factors(cache![], vertices_states)
-  return tn
+  return tensornetwork(cache![])
 end
 
 function LinearAlgebra.normalize(tn::AbstractITensorNetwork; alg="exact", kwargs...)
@@ -65,9 +48,8 @@ end
 function LinearAlgebra.normalize(
   alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...
 )
-  norm_tn = QuadraticFormNetwork(tn)
-  vs = filter(v -> v ∉ operator_vertices(norm_tn), collect(vertices(norm_tn)))
-  return ket_network(rescale(alg, norm_tn, vs; kwargs...))
+  norm_tn = inner_network(tn, tn)
+  return ket_network(rescale(alg, norm_tn; kwargs...))
 end
 
 function LinearAlgebra.normalize(
@@ -77,15 +59,19 @@ function LinearAlgebra.normalize(
   cache_construction_function=tn ->
     cache(alg, tn; default_cache_construction_kwargs(alg, tn)...),
   update_cache=isnothing(cache!),
-  cache_update_kwargs=default_cache_update_kwargs(cache!),
+  cache_update_kwargs=default_cache_update_kwargs(alg),
+  cache_construction_kwargs=(;),
 )
-  norm_tn = QuadraticFormNetwork(tn)
+  norm_tn = inner_network(tn, tn)
   if isnothing(cache!)
-    cache! = Ref(cache_construction_function(norm_tn))
+    cache! = Ref(cache(alg, norm_tn; cache_construction_kwargs...))
   end
 
-  vs = filter(v -> v ∉ operator_vertices(norm_tn), collect(vertices(norm_tn)))
-  norm_tn = rescale(alg, norm_tn, vs; cache!, update_cache, cache_update_kwargs)
+  vs = collect(vertices(tn))
+  vs_to_rescale = vcat(
+    [ket_vertex(norm_tn, v) for v in vs], [bra_vertex(norm_tn, v) for v in vs]
+  )
+  norm_tn = rescale(alg, norm_tn; vs_to_rescale, cache!, update_cache, cache_update_kwargs)
 
   return ket_network(norm_tn)
 end

test/test_normalize.jl

@@ -5,9 +5,10 @@ using ITensorNetworks:
   edge_scalars,
   norm_sqr_network,
   random_tensornetwork,
+  siteinds,
   vertex_scalars,
   rescale
-using ITensors: dag, inner, siteinds, scalar
+using ITensors: dag, inner, scalar
 using Graphs: SimpleGraph, uniform_tree
 using LinearAlgebra: normalize
 using NamedGraphs: NamedGraph