Rescale local tensors first

Author: JoeyT1994

src/caches/abstractbeliefpropagationcache.jl

@@ -65,6 +65,9 @@ end
 function region_scalar(bpc::AbstractBeliefPropagationCache, pe::PartitionEdge; kwargs...)
   return not_implemented()
 end
+function message_overlap(bpc::AbstractBeliefPropagationCache, partitionedge; kwargs...)
+  return not_implemented()
+end
 partitions(bpc::AbstractBeliefPropagationCache) = not_implemented()
 PartitionedGraphs.partitionedges(bpc::AbstractBeliefPropagationCache) = not_implemented()
 
@@ -285,7 +288,7 @@ function update(
 end
 
 function rescale_message(bp_cache::AbstractBeliefPropagationCache, partitionedge)
-  return rescale_messages(bp_cache, typeof(partitionedge)[partitionedge])
+  return rescale_messages(bp_cache, [partitionedge])
 end
 
 function rescale_messages(bp_cache::AbstractBeliefPropagationCache)
@@ -303,15 +306,23 @@ function rescale_partitions(
 
     isempty(pv_vs) && continue
 
+    for v in pv_vs
+      t = tn[v]
+      setindex_preserve_graph!(tn, t / norm(t), v)
+    end
+
     vn = region_scalar(bpc, pv)
     if isreal(vn)
-      tn[first(pv_vs)] *= sign(vn)
+      v = first(pv_vs)
+      t = tn[v]
+      setindex_preserve_graph!(tn, t * sign(vn), v)
       vn *= sign(vn)
     end
 
     vn = vn^(1 / length(pv_vs))
     for v in pv_vs
-      tn[v] /= vn
+      t = tn[v]
+      setindex_preserve_graph!(tn, t / vn, v)
     end
   end
 
@@ -323,7 +334,7 @@ function rescale_partitions(bpc::AbstractBeliefPropagationCache; kwargs...)
 end
 
 function rescale_partition(bpc::AbstractBeliefPropagationCache, partition; kwargs...)
-  return rescale_partitions(bpc, typeof(partition)[partition]; kwargs...)
+  return rescale_partitions(bpc, [partition]; kwargs...)
 end
 
 function rescale(bpc::AbstractBeliefPropagationCache; kwargs...)

src/caches/beliefpropagationcache.jl

@@ -109,20 +109,27 @@ function region_scalar(bp_cache::BeliefPropagationCache, pe::PartitionEdge)
   return contract(ts; sequence)[]
 end
 
+function message_overlap(bpc::BeliefPropagationCache, partitionedge; kwargs...)
+  me, mer = only(message(bpc, partitionedge)), only(message(bpc, reverse(partitionedge)))
+  return dot(me, mer)
+end
+
 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)
+    me, mer = normalize.(mts[pe]), normalize.(mts[reverse(pe)])
+    set!(mts, pe, me)
+    set!(mts, reverse(pe), mer)
+    n = message_overlap(bp_cache, pe)
     if isreal(n)
-      me *= sign(n)
+      me[1] *= sign(n)
       n *= sign(n)
     end
 
-    set!(mts, pe, ITensor[(1 / sqrt(n)) * me])
-    set!(mts, reverse(pe), ITensor[(1 / sqrt(n)) * mer])
+    sf = (1 / sqrt(n)) ^ (1 / length(me))
+    set!(mts, pe, sf .* me)
+    set!(mts, reverse(pe), sf .* mer)
   end
   return bp_cache
 end

test/test_normalize.jl

@@ -14,6 +14,7 @@ using LinearAlgebra: normalize
 using NamedGraphs: NamedGraph
 using NamedGraphs.NamedGraphGenerators: named_grid, named_comb_tree
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "Normalize" begin