From 262f0de20ffe0c8df57716811d91f463ad1d1a88 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Thu, 7 Mar 2024 16:19:54 -0600
Subject: [PATCH] Make isometry code more modular

---
 src/gauging.jl       | 42 +++++++++++++++++++++++-------------------
 test/test_gauging.jl |  2 +-
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/gauging.jl b/src/gauging.jl
index 16e05831..2898673b 100644
--- a/src/gauging.jl
+++ b/src/gauging.jl
@@ -26,7 +26,7 @@ end
 default_cache_update_kwargs(cache) = (; maxiter=20, tol=1e-5)
 
 function ITensorNetwork(
-  ψ_vidal::VidalITensorNetwork; (bp_cache!)=nothing, update_gauge=false, update_kwargs...
+  ψ_vidal::VidalITensorNetwork; (cache!)=nothing, update_gauge=false, update_kwargs...
 )
   if update_gauge
     ψ_vidal = update(ψ_vidal; update_kwargs...)
@@ -41,7 +41,7 @@ function ITensorNetwork(
     setindex_preserve_graph!(ψ, noprime(root_S * ψ[vdst]), vdst)
   end
 
-  if !isnothing(bp_cache!)
+  if !isnothing(cache!)
     bp_cache = default_norm_cache(ψ)
     mts = messages(bp_cache)
 
@@ -53,7 +53,7 @@ function ITensorNetwork(
     end
 
     bp_cache = set_messages(bp_cache, mts)
-    bp_cache![] = bp_cache
+    cache![] = bp_cache
   end
 
   return ψ
@@ -139,25 +139,29 @@ function update(ψ::VidalITensorNetwork; kwargs...)
   return VidalITensorNetwork(ITensorNetwork(ψ; update_gauge=false); kwargs...)
 end
 
-"""Function to measure the 'isometries' of a state in the Vidal Gauge"""
-function vidal_gauge_isometries(
-  ψ::VidalITensorNetwork; edges=vcat(NamedGraphs.edges(ψ), reverse.(NamedGraphs.edges(ψ)))
-)
-  isometries = Dict()
-
-  for e in edges
-    vsrc, vdst = src(e), dst(e)
-    ψ_vsrc = copy(ψ[vsrc])
-    for vn in setdiff(neighbors(ψ, vsrc), [vdst])
-      ψ_vsrc = noprime(ψ_vsrc * bond_tensor(ψ, vn => vsrc))
-    end
+"""Function to construct the 'isometry' of a state in the Vidal Gauge on the given edge"""
+function vidal_gauge_isometry(ψ::VidalITensorNetwork, edge)
+  vsrc, vdst = src(edge), dst(edge)
+  ψ_vsrc = copy(ψ[vsrc])
 
-    ψ_vsrcdag = dag(ψ_vsrc)
-    replaceind!(ψ_vsrcdag, commonind(ψ_vsrc, ψ[vdst]), commonind(ψ_vsrc, ψ[vdst])')
-    isometries[e] = ψ_vsrcdag * ψ_vsrc
+  for vn in setdiff(neighbors(ψ, vsrc), [vdst])
+    ψ_vsrc = noprime(ψ_vsrc * bond_tensor(ψ, vn => vsrc))
   end
 
-  return isometries
+  ψ_vsrcdag = dag(ψ_vsrc)
+  replaceind!(ψ_vsrcdag, commonind(ψ_vsrc, ψ[vdst]), commonind(ψ_vsrc, ψ[vdst])')
+
+  return ψ_vsrcdag * ψ_vsrc
+end
+
+function vidal_gauge_isometries(ψ::VidalITensorNetwork, edges::Vector)
+  return Dict([e => vidal_gauge_isometry(ψ, e) for e in edges])
+end
+
+function vidal_gauge_isometries(ψ::VidalITensorNetwork)
+  return vidal_gauge_isometries(
+    ψ, vcat(NamedGraphs.edges(ψ), reverse.(NamedGraphs.edges(ψ)))
+  )
 end
 
 """Function to measure the 'distance' of a state from the Vidal Gauge"""
diff --git a/test/test_gauging.jl b/test/test_gauging.jl
index d2b17710..f0a7d10b 100644
--- a/test/test_gauging.jl
+++ b/test/test_gauging.jl
@@ -24,7 +24,7 @@ using SplitApplyCombine
 
   # Move to symmetric gauge
   cache_ref = Ref{BeliefPropagationCache}()
-  ψ_symm = ITensorNetwork(ψ_vidal; (bp_cache!)=cache_ref)
+  ψ_symm = ITensorNetwork(ψ_vidal; (cache!)=cache_ref)
   bp_cache = cache_ref[]
 
   # Test we just did a gauge transform and didn't change the overall network