Centre around scaling function

Author: JoeyT1994

src/abstractitensornetwork.jl

@@ -935,6 +935,26 @@ function add(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork)
   return tn12
 end
 
+""" Scale each tensor of the network by a scale factor on each vertex"""
+function scale!(tn::AbstractITensorNetwork, vertices_weights::Dictionary)
+  for v in keys(vertices_weights)
+    setindex_preserve_graph!(tn, vertices_weights[v] * tn[v], v)
+  end
+  return tn
+end
+
+""" Scale each tensor of the network via a function (vertex, ITensor) -> Number"""
+function scale!(tn::AbstractITensorNetwork, weight_function::Function)
+  vs = collect(vertices(tn))
+  vertices_weights = Dictionary(vs, [weight_function(v, tn[v]) for v in vs])
+  return scale!(tn, vertices_weights)
+end
+
+function scale(tn, args...)
+  tn = copy(tn)
+  return scale!(tn, args...)
+end
+
 Base.:+(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork) = add(tn1, tn2)
 
 ITensors.hasqns(tn::AbstractITensorNetwork) = any(v -> hasqns(tn[v]), vertices(tn))

src/caches/abstractbeliefpropagationcache.jl

@@ -1,4 +1,4 @@
-using Graphs: IsDirected
+using Graphs: Graphs, IsDirected
 using SplitApplyCombine: group
 using LinearAlgebra: diag, dot
 using ITensors: dir
@@ -88,6 +88,10 @@ function tensornetwork(bpc::AbstractBeliefPropagationCache)
   return unpartitioned_graph(partitioned_tensornetwork(bpc))
 end
 
+function setindex_preserve_graph!(bpc::AbstractBeliefPropagationCache, args...)
+  return setindex_preserve_graph!(tensornetwork(bpc), args...)
+end
+
 function factors(bpc::AbstractBeliefPropagationCache, verts::Vector)
   return ITensor[tensornetwork(bpc)[v] for v in verts]
 end
@@ -284,6 +288,10 @@ function update(
   return update(Algorithm(alg), bpc; kwargs...)
 end
 
+function scale!(bp_cache::AbstractBeliefPropagationCache, args...)
+  return scale!(tensornetwork(bp_cache), args...)
+end
+
 function rescale_messages(
   bp_cache::AbstractBeliefPropagationCache, partitionedge::PartitionEdge
 )
@@ -297,48 +305,45 @@ end
 function rescale_partitions(
   bpc::AbstractBeliefPropagationCache,
   partitions::Vector;
-  verts_to_rescale::Vector=collect(vertices(tensornetwork(bpc))),
+  verts_to_rescale::Vector=vertices(bpc, partitions),
 )
+  bpc = copy(bpc)
   tn = tensornetwork(bpc)
+  norms = map(v -> inv(norm(tn[v])), verts_to_rescale)
+  scale!(bpc, Dictionary(verts_to_rescale, norms))
+
+  vertices_weights = Dictionary()
   for pv in partitions
     pv_vs = filter(v -> v ∈ verts_to_rescale, vertices(bpc, pv))
-
     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)
-      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
-      t = tn[v]
-      setindex_preserve_graph!(tn, t / vn, v)
+    s = isreal(vn) ? sign(vn) : 1.0
+    vn = s * inv(vn^(1 / length(pv_vs)))
+    set!(vertices_weights, first(pv_vs), s*vn)
+    for v in pv_vs[2:length(pv_vs)]
+      set!(vertices_weights, v, vn)
     end
   end
 
+  scale!(bpc, vertices_weights)
+
   return bpc
 end
 
-function rescale_partitions(bpc::AbstractBeliefPropagationCache; kwargs...)
-  return rescale_partitions(bpc, collect(partitions(bpc)); kwargs...)
+function rescale_partitions(bpc::AbstractBeliefPropagationCache, args...; kwargs...)
+  return rescale_partitions(bpc, collect(partitions(bpc)), args...; kwargs...)
 end
 
-function rescale_partition(bpc::AbstractBeliefPropagationCache, partition; kwargs...)
-  return rescale_partitions(bpc, [partition]; kwargs...)
+function rescale_partition(
+  bpc::AbstractBeliefPropagationCache, partition, args...; kwargs...
+)
+  return rescale_partitions(bpc, [partition], args...; kwargs...)
 end
 
-function rescale(bpc::AbstractBeliefPropagationCache; kwargs...)
+function rescale(bpc::AbstractBeliefPropagationCache, args...; kwargs...)
   bpc = rescale_messages(bpc)
-  bpc = rescale_partitions(bpc; kwargs...)
+  bpc = rescale_partitions(bpc, args...; kwargs...)
   return bpc
 end
 

src/caches/beliefpropagationcache.jl

@@ -9,7 +9,9 @@ using NamedGraphs.PartitionedGraphs:
   boundary_partitionedges,
   partitionvertices,
   partitionedges,
-  unpartitioned_graph
+  partitioned_vertices,
+  unpartitioned_graph,
+  which_partition
 using SimpleTraits: SimpleTraits, Not, @traitfn
 using NDTensors: NDTensors
 

src/normalize.jl

@@ -4,25 +4,18 @@ function rescale(tn::AbstractITensorNetwork; alg="exact", kwargs...)
   return rescale(Algorithm(alg), tn; kwargs...)
 end
 
-function rescale(
-  alg::Algorithm"exact",
-  tn::AbstractITensorNetwork;
-  verts_to_rescale=collect(vertices(tn)),
-  kwargs...,
-)
+function rescale(alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...)
   logn = logscalar(alg, tn; kwargs...)
   vs = collect(vertices(tn))
-  c = inv(exp(logn / length(verts_to_rescale)))
-  tn = copy(tn)
-  for v in verts_to_rescale
-    tn[v] *= c
-  end
-  return tn
+  c = inv(exp(logn / length(vs)))
+  vertices_weights = Dictionary(vs, [c for v in vs])
+  return scale(tn, vertices_weights)
 end
 
 function rescale(
   alg::Algorithm,
-  tn::AbstractITensorNetwork;
+  tn::AbstractITensorNetwork,
+  args...;
   (cache!)=nothing,
   cache_construction_kwargs=default_cache_construction_kwargs(alg, tn),
   update_cache=isnothing(cache!),
@@ -37,7 +30,7 @@ function rescale(
     cache![] = update(cache![]; cache_update_kwargs...)
   end
 
-  cache![] = rescale(cache![]; kwargs...)
+  cache![] = rescale(cache![], args...; kwargs...)
 
   return tensornetwork(cache![])
 end
@@ -49,12 +42,11 @@ end
 function LinearAlgebra.normalize(
   alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...
 )
-  norm_tn = inner_network(tn, tn)
+  logn = logscalar(alg, inner_network(tn, tn); kwargs...)
   vs = collect(vertices(tn))
-  verts_to_rescale = vcat(
-    [ket_vertex(norm_tn, v) for v in vs], [bra_vertex(norm_tn, v) for v in vs]
-  )
-  return ket_network(rescale(alg, norm_tn; verts_to_rescale, kwargs...))
+  c = inv(exp(logn / (2*length(vs))))
+  vertices_weights = Dictionary(vs, [c for v in vs])
+  return scale(tn, vertices_weights)
 end
 
 function LinearAlgebra.normalize(