Update gauge_walk functionality and fix bug in VidalITensorNetwork (#222)

Author: JoeyT1994

src/ITensorNetworks.jl

@@ -26,11 +26,11 @@ include("specialitensornetworks.jl")
 include("boundarymps.jl")
 include("partitioneditensornetwork.jl")
 include("edge_sequences.jl")
+include("caches/abstractbeliefpropagationcache.jl")
+include("caches/beliefpropagationcache.jl")
 include("formnetworks/abstractformnetwork.jl")
 include("formnetworks/bilinearformnetwork.jl")
 include("formnetworks/quadraticformnetwork.jl")
-include("caches/abstractbeliefpropagationcache.jl")
-include("caches/beliefpropagationcache.jl")
 include("contraction_tree_to_graph.jl")
 include("gauging.jl")
 include("utils.jl")

src/abstractitensornetwork.jl

@@ -623,18 +623,42 @@ function gauge_walk(
   return gauge_walk(alg, tn, edgetype(tn).(edges); kwargs...)
 end
 
+function tree_gauge(alg::Algorithm, ψ::AbstractITensorNetwork, region)
+  return tree_gauge(alg, ψ, [region])
+end
+
+#Get the path that moves the gauge from a to b on a tree
+#TODO: Move to NamedGraphs
+function edge_sequence_between_regions(g::AbstractGraph, region_a::Vector, region_b::Vector)
+  issetequal(region_a, region_b) && return edgetype(g)[]
+  st = steiner_tree(g, union(region_a, region_b))
+  path = post_order_dfs_edges(st, first(region_b))
+  path = filter(e -> !((src(e) ∈ region_b) && (dst(e) ∈ region_b)), path)
+  return path
+end
+
+# Gauge a ITensorNetwork from cur_region towards new_region, treating
+# the network as a tree spanned by a spanning tree.
+function tree_gauge(
+  alg::Algorithm,
+  ψ::AbstractITensorNetwork,
+  cur_region::Vector,
+  new_region::Vector;
+  kwargs...,
+)
+  es = edge_sequence_between_regions(ψ, cur_region, new_region)
+  ψ = gauge_walk(alg, ψ, es; kwargs...)
+  return ψ
+end
+
 # Gauge a ITensorNetwork towards a region, treating
 # the network as a tree spanned by a spanning tree.
 function tree_gauge(alg::Algorithm, ψ::AbstractITensorNetwork, region::Vector)
-  region_center =
-    length(region) != 1 ? first(center(steiner_tree(ψ, region))) : only(region)
-  path = post_order_dfs_edges(bfs_tree(ψ, region_center), region_center)
-  path = filter(e -> !((src(e) ∈ region) && (dst(e) ∈ region)), path)
-  return gauge_walk(alg, ψ, path)
+  return tree_gauge(alg, ψ, collect(vertices(ψ)), region)
 end
 
-function tree_gauge(alg::Algorithm, ψ::AbstractITensorNetwork, region)
-  return tree_gauge(alg, ψ, [region])
+function tree_orthogonalize(ψ::AbstractITensorNetwork, cur_region, new_region; kwargs...)
+  return tree_gauge(Algorithm("orthogonalize"), ψ, cur_region, new_region; kwargs...)
 end
 
 function tree_orthogonalize(ψ::AbstractITensorNetwork, region; kwargs...)

src/apply.jl

@@ -378,7 +378,7 @@ function ITensors.apply(o, ψ::VidalITensorNetwork; normalize=false, apply_kwarg
 
   else
     updated_ψ = apply(o, updated_ψ; normalize)
-    return VidalITensorNetwork(ψ, updated_bond_tensors)
+    return VidalITensorNetwork(updated_ψ, updated_bond_tensors)
   end
 end
 

src/caches/abstractbeliefpropagationcache.jl

@@ -40,9 +40,6 @@ default_messages(ptn::PartitionedGraph) = Dictionary()
   return default_bp_maxiter(undirected_graph(underlying_graph(g)))
 end
 default_partitioned_vertices(ψ::AbstractITensorNetwork) = group(v -> v, vertices(ψ))
-function default_partitioned_vertices(f::AbstractFormNetwork)
-  return group(v -> original_state_vertex(f, v), vertices(f))
-end
 
 partitioned_tensornetwork(bpc::AbstractBeliefPropagationCache) = not_implemented()
 messages(bpc::AbstractBeliefPropagationCache) = not_implemented()

src/formnetworks/abstractformnetwork.jl

@@ -80,3 +80,7 @@ operator_vertex(f::AbstractFormNetwork, v) = operator_vertex_map(f)(v)
 bra_vertex(f::AbstractFormNetwork, v) = bra_vertex_map(f)(v)
 ket_vertex(f::AbstractFormNetwork, v) = ket_vertex_map(f)(v)
 original_state_vertex(f::AbstractFormNetwork, v) = inv_vertex_map(f)(v)
+
+function default_partitioned_vertices(f::AbstractFormNetwork)
+  return group(v -> original_state_vertex(f, v), vertices(f))
+end

src/treetensornetworks/abstracttreetensornetwork.jl

@@ -36,13 +36,7 @@ function set_ortho_region(tn::AbstractTTN, new_region)
 end
 
 function gauge(alg::Algorithm, ttn::AbstractTTN, region::Vector; kwargs...)
-  issetequal(region, ortho_region(ttn)) && return ttn
-  st = steiner_tree(ttn, union(region, ortho_region(ttn)))
-  path = post_order_dfs_edges(st, first(region))
-  path = filter(e -> !((src(e) ∈ region) && (dst(e) ∈ region)), path)
-  if !isempty(path)
-    ttn = typeof(ttn)(gauge_walk(alg, ITensorNetwork(ttn), path; kwargs...))
-  end
+  ttn = tree_gauge(alg, ttn, collect(ortho_region(ttn)), region; kwargs...)
   return set_ortho_region(ttn, region)
 end