Improvements

Author: JoeyT1994

src/abstractitensornetwork.jl

@@ -935,24 +935,29 @@ 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)
+""" Scale each tensor of the network via a function vertex -> Number"""
+function scale!(
+  weight_function::Function, tn::AbstractITensorNetwork; verts=collect(vertices(tn))
+)
+  for v in verts
+    setindex_preserve_graph!(tn, weight_function(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)
+""" Scale each tensor of the network by a scale factor for each vertex in the keys of the dictionary"""
+function scale!(tn::AbstractITensorNetwork, vertices_weights::Dictionary)
+  return scale!(v -> vertices_weights[v], tn; verts=keys(vertices_weights))
+end
+
+function scale(weight_function::Function, tn; kwargs...)
+  tn = copy(tn)
+  return scale!(weight_function, tn; kwargs...)
 end
 
-function scale(tn, args...)
+function scale(tn, vertices_weights::Dictionary; kwargs...)
   tn = copy(tn)
-  return scale!(tn, args...)
+  return scale!(tn, vertices_weights; kwargs...)
 end
 
 Base.:+(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork) = add(tn1, tn2)

src/caches/abstractbeliefpropagationcache.jl

@@ -305,16 +305,16 @@ end
 function rescale_partitions(
   bpc::AbstractBeliefPropagationCache,
   partitions::Vector;
-  verts_to_rescale::Vector=vertices(bpc, partitions),
+  verts::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))
+  norms = map(v -> inv(norm(tn[v])), verts)
+  scale!(bpc, Dictionary(verts, norms))
 
   vertices_weights = Dictionary()
   for pv in partitions
-    pv_vs = filter(v -> v ∈ verts_to_rescale, vertices(bpc, pv))
+    pv_vs = filter(v -> v ∈ verts, vertices(bpc, pv))
     isempty(pv_vs) && continue
 
     vn = region_scalar(bpc, pv)

src/normalize.jl

@@ -1,4 +1,4 @@
-using LinearAlgebra
+using LinearAlgebra: normalize
 
 function rescale(tn::AbstractITensorNetwork; alg="exact", kwargs...)
   return rescale(Algorithm(alg), tn; kwargs...)