Optimizations, norm_network, faster dag, gauging example (#94)

examples/gauging/gauging_itns.jl

@@ -0,0 +1,116 @@
+using Compat
+using ITensors
+using Metis
+using ITensorNetworks
+using Random
+using SplitApplyCombine
+using ProfileView
+
+using ITensorNetworks:
+  message_tensors,
+  nested_graph_leaf_vertices,
+  belief_propagation_iteration,
+  belief_propagation,
+  find_subgraph,
+  vidal_gauge,
+  symmetric_gauge,
+  vidal_itn_canonicalness,
+  vidal_to_symmetric_gauge,
+  initialize_bond_tensors,
+  vidal_itn_isometries,
+  norm_network
+
+using NamedGraphs
+using NamedGraphs: add_edges!, rem_vertex!, hexagonal_lattice_graph
+using Graphs
+
+"""Eager Gauging"""
+function eager_gauging(ψ::ITensorNetwork, bond_tensors::DataGraph, mts::DataGraph)
+  isometries = vidal_itn_isometries(ψ, bond_tensors)
+
+  ψ = copy(ψ)
+  mts = copy(mts)
+  for e in edges(ψ)
+    s1, s2 = find_subgraph((src(e), 1), mts), find_subgraph((dst(e), 1), mts)
+    normalize!(isometries[e])
+    normalize!(isometries[reverse(e)])
+    mts[s1 => s2], mts[s2 => s1] = ITensorNetwork(isometries[e]),
+    ITensorNetwork(isometries[reverse(e)])
+  end
+
+  ψ, bond_tensors = vidal_gauge(ψ, mts, bond_tensors)
+
+  return ψ, bond_tensors, mts
+end
+
+"""Bring an ITN into the Vidal gauge, various methods possible. Result is timed"""
+function benchmark_state_gauging(
+  ψ::ITensorNetwork; mode="BeliefPropagation", no_iterations=50
+)
+  s = siteinds(ψ)
+
+  C = zeros((no_iterations))
+  times_iters = zeros((no_iterations))
+  times_gauging = zeros((no_iterations))
+
+  ψψ = norm_network(ψ)
+  ψinit = copy(ψ)
+  vertex_groups = nested_graph_leaf_vertices(partition(ψψ, group(v -> v[1], vertices(ψψ))))
+  mts = message_tensors(partition(ψψ; subgraph_vertices=vertex_groups))
+  bond_tensors = initialize_bond_tensors(ψ)
+  for e in edges(mts)
+    mts[e] = ITensorNetwork(dense(delta(inds(ITensors.contract(ITensor(mts[e]))))))
+  end
+
+  for i in 1:no_iterations
+    println("On Iteration " * string(i))
+
+    if mode == "BeliefPropagation"
+      times_iters[i] = @elapsed mts, _ = belief_propagation_iteration(
+        ψψ, mts; contract_kwargs=(; alg="exact")
+      )
+      times_gauging[i] = @elapsed ψ, bond_tensors = vidal_gauge(ψinit, mts)
+    elseif mode == "Eager"
+      times_iters[i] = @elapsed ψ, bond_tensors, mts = eager_gauging(ψ, bond_tensors, mts)
+    else
+      times_iters[i] = @elapsed begin
+        for e in edges(ψ)
+          ψ, bond_tensors = apply(e, ψ, bond_tensors; normalize=true, cutoff=1e-16)
+        end
+      end
+    end
+
+    C[i] = vidal_itn_canonicalness(ψ, bond_tensors)
+  end
+
+  simulation_times = cumsum(times_iters) + times_gauging
+
+  return simulation_times, C
+end
+
+L, χ = 10, 10
+g = named_grid((L, L))
+s = siteinds("S=1/2", g)
+ψ = randomITensorNetwork(s; link_space=χ)
+no_iterations = 30
+
+BPG_simulation_times, BPG_Cs = benchmark_state_gauging(ψ; no_iterations)
+Eager_simulation_times, Eager_Cs = benchmark_state_gauging(ψ; mode="Eager", no_iterations)
+SU_simulation_times, SU_Cs = benchmark_state_gauging(ψ; mode="SU", no_iterations)
+
+epsilon = 1e-6
+println(
+  "Time for BPG to reach C < epsilon was " *
+  string(BPG_simulation_times[findfirst(x -> x < 0, BPG_Cs .- epsilon)]) *
+  " seconds",
+)
+println(
+  "Time for Eager to reach C < epsilon was " *
+  string(Eager_simulation_times[findfirst(x -> x < 0, Eager_Cs .- epsilon)]) *
+  " seconds",
+)
+println(
+  "Time for SU to reach C < epsilon was " *
+  string(SU_simulation_times[findfirst(x -> x < 0, SU_Cs .- epsilon)]) *
+  " seconds",
+)

src/abstractitensornetwork.jl

@@ -330,7 +330,14 @@ end
 
 adjoint(tn::Union{IndsNetwork,AbstractITensorNetwork}) = prime(tn)
 
-dag(tn::AbstractITensorNetwork) = map_vertex_data(dag, tn)
+#dag(tn::AbstractITensorNetwork) = map_vertex_data(dag, tn)
+function dag(tn::AbstractITensorNetwork)
+  tndag = copy(tn)
+  for v in vertices(tndag)
+    setindex_preserve_graph!(tndag, dag(tndag[v]), v)
+  end
+  return tndag
+end
 
 # TODO: should this make sure that internal indices
 # don't clash?
@@ -646,6 +653,23 @@ function flatten_networks(
   return flatten_networks(flatten_networks(tn1, tn2; kwargs...), tn3, tn_tail...; kwargs...)
 end
 
+#Ideally this will dispatch to inner_network but this is a temporary fast version for now
+function norm_network(tn::AbstractITensorNetwork)
+  tnbra = rename_vertices(v -> (v, 1), data_graph(tn))
+  tndag = copy(tn)
+  for v in vertices(tndag)
+    setindex_preserve_graph!(tndag, dag(tndag[v]), v)
+  end
+  tnket = rename_vertices(v -> (v, 2), data_graph(prime(tndag; sites=[])))
+  tntn = ITensorNetwork(union(tnbra, tnket))
+  for v in vertices(tn)
+    if !isempty(commoninds(tntn[(v, 1)], tntn[(v, 2)]))
+      add_edge!(tntn, (v, 1) => (v, 2))
+    end
+  end
+  return tntn
+end
+
 # TODO: Use or replace with `flatten_networks`
 function inner_network(
   tn1::AbstractITensorNetwork,

src/apply.jl

@@ -19,7 +19,7 @@ function ITensors.apply(
     if normalize
       oψᵥ ./= norm(oψᵥ)
     end
-    ψ[v⃗[1]] = oψᵥ
+    setindex_preserve_graph!(ψ, oψᵥ, v⃗[1])
   elseif length(v⃗) == 2
     e = v⃗[1] => v⃗[2]
     if !has_edge(ψ, e)
@@ -72,8 +72,8 @@ function ITensors.apply(
       ψᵥ₂ ./= norm(ψᵥ₂)
     end
 
-    ψ[v⃗[1]] = ψᵥ₁
-    ψ[v⃗[2]] = ψᵥ₂
+    setindex_preserve_graph!(ψ, ψᵥ₁, v⃗[1])
+    setindex_preserve_graph!(ψ, ψᵥ₂, v⃗[2])
 
   elseif length(v⃗) < 1
     error("Gate being applied does not share indices with tensor network.")
@@ -123,15 +123,18 @@ end
 _gate_vertices(o::ITensor, ψ) = neighbor_vertices(ψ, o)
 _gate_vertices(o::AbstractEdge, ψ) = [src(o), dst(o)]
 
-function _contract_factorized_gate(o::ITensor, ψv1, ψv2)
-  G1, G2 = factorize(o, Index[commonind(ψv1, o), commonind(ψv1, o)']; cutoff=1e-16)
-  ψv1 = noprime(ψv1 * G1)
-  ψv2 = noprime(ψv2 * G2)
-  return ψv1, ψv2
+function _contract_gate(o::ITensor, ψv1, ψv2)
+  Qᵥ₁, Rᵥ₁ = qr(ψv1, setdiff(uniqueinds(ψv1, ψv2), commoninds(ψv1, o)))
+  Qᵥ₂, Rᵥ₂ = qr(ψv2, setdiff(uniqueinds(ψv2, ψv1), commoninds(ψv2, o)))
+  theta = noprime(Rᵥ₁ * Rᵥ₂ * o)
+  return Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta
 end
 
-function _contract_factorized_gate(o::AbstractEdge, ψv1, ψv2)
-  return ψv1, ψv2
+function _contract_gate(o::AbstractEdge, ψv1, ψv2)
+  Qᵥ₁, Rᵥ₁ = qr(ψv1, uniqueinds(ψv1, ψv2))
+  Qᵥ₂, Rᵥ₂ = qr(ψv2, uniqueinds(ψv2, ψv1))
+  theta = Rᵥ₁ * Rᵥ₂
+  return Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta
 end
 
 #In the future we will try to unify this into apply() above but currently leave it mostly as a separate function
@@ -149,7 +152,7 @@ function ITensors.apply(
   v⃗ = _gate_vertices(o, ψ)
   if length(v⃗) == 2
     e = NamedEdge(v⃗[1] => v⃗[2])
-    ψv1, ψv2 = copy(ψ[src(e)]), copy(ψ[dst(e)])
+    ψv1, ψv2 = ψ[src(e)], ψ[dst(e)]
     e_ind = commonind(ψv1, ψv2)
 
     for vn in neighbors(ψ, src(e))
@@ -164,14 +167,9 @@ function ITensors.apply(
       end
     end
 
-    ψv1, ψv2 = _contract_factorized_gate(o, ψv1, ψv2)
-
     ψv2 = noprime(ψv2 * bond_tensors[e])
 
-    Qᵥ₁, Rᵥ₁ = factorize(ψv1, uniqueinds(ψv1, ψv2); cutoff=1e-16)
-    Qᵥ₂, Rᵥ₂ = factorize(ψv2, uniqueinds(ψv2, ψv1); cutoff=1e-16)
-
-    theta = Rᵥ₁ * Rᵥ₂
+    Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta = _contract_gate(o, ψv1, ψv2)
 
     U, S, V = ITensors.svd(
       theta,
@@ -201,12 +199,13 @@ function ITensors.apply(
     end
 
     if normalize
-      normalize!(ψv1)
-      normalize!(ψv2)
+      ψv1 /= norm(ψv1)
+      ψv2 /= norm(ψv2)
       normalize!(bond_tensors[e])
     end
 
-    ψ[src(e)], ψ[dst(e)] = ψv1, ψv2
+    setindex_preserve_graph!(ψ, ψv1, src(e))
+    setindex_preserve_graph!(ψ, ψv2, dst(e))
 
     return ψ, bond_tensors
 

src/gauging.jl

@@ -23,6 +23,7 @@ function vidal_gauge(
 
   for e in edges(ψ_vidal)
     vsrc, vdst = src(e), dst(e)
+    ψvsrc, ψvdst = ψ_vidal[vsrc], ψ_vidal[vdst]
 
     s1, s2 = find_subgraph((vsrc, 1), mts), find_subgraph((vdst, 1), mts)
     edge_ind = commoninds(ψ_vidal[vsrc], ψ_vidal[vdst])
@@ -44,8 +45,7 @@ function vidal_gauge(
     inv_rootX = X_U * inv_rootX_D * prime(dag(X_U))
     inv_rootY = Y_U * inv_rootY_D * prime(dag(Y_U))
 
-    ψ_vidal[vsrc] = noprime(ψ_vidal[vsrc] * inv_rootX)
-    ψ_vidal[vdst] = noprime(ψ_vidal[vdst] * inv_rootY)
+    ψvsrc, ψvdst = noprime(ψvsrc * inv_rootX), noprime(ψvdst * inv_rootY)
 
     Ce = rootX * prime(bond_tensors[e])
     replaceinds!(Ce, edge_ind'', edge_ind')
@@ -55,9 +55,12 @@ function vidal_gauge(
 
     new_edge_ind = Index[Index(dim(commoninds(S, U)), tags(first(edge_ind)))]
 
-    ψ_vidal[vsrc] = replaceinds(ψ_vidal[vsrc] * U, commoninds(S, U), new_edge_ind)
-    ψ_vidal[vdst] = replaceinds(ψ_vidal[vdst], edge_ind, edge_ind_sim)
-    ψ_vidal[vdst] = replaceinds(ψ_vidal[vdst] * V, commoninds(V, S), new_edge_ind)
+    ψvsrc = replaceinds(ψvsrc * U, commoninds(S, U), new_edge_ind)
+    ψvdst = replaceinds(ψvdst, edge_ind, edge_ind_sim)
+    ψvdst = replaceinds(ψvdst * V, commoninds(V, S), new_edge_ind)
+
+    setindex_preserve_graph!(ψ_vidal, ψvsrc, vsrc)
+    setindex_preserve_graph!(ψ_vidal, ψvdst, vdst)
 
     S = replaceinds(
       S,
@@ -93,7 +96,7 @@ function vidal_gauge(
   target_canonicalness::Union{Nothing,Float64}=nothing,
   svd_kwargs...,
 )
-  ψψ = ψ ⊗ prime(dag(ψ); sites=[])
+  ψψ = norm_network(ψ)
   Z = partition(ψψ; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψ)))))
   mts = message_tensors(Z)
 
@@ -108,17 +111,18 @@ end
 """Transform from an ITensor in the Vidal Gauge (bond tensors) to the Symmetric Gauge (message tensors)"""
 function vidal_to_symmetric_gauge(ψ::ITensorNetwork, bond_tensors::DataGraph)
   ψsymm = copy(ψ)
-  ψψsymm = ψsymm ⊗ prime(dag(ψsymm); sites=[])
+  ψψsymm = norm_network(ψsymm)
   Z = partition(
     ψψsymm; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψsymm))))
   )
   ψsymm_mts = message_tensors_skeleton(Z)
 
   for e in edges(ψsymm)
-    s1, s2 = find_subgraph((src(e), 1), ψsymm_mts), find_subgraph((dst(e), 1), ψsymm_mts)
+    vsrc, vdst = src(e), dst(e)
+    s1, s2 = find_subgraph((vsrc, 1), ψsymm_mts), find_subgraph((vdst, 1), ψsymm_mts)
     root_S = sqrt_diag(bond_tensors[e])
-    ψsymm[src(e)] = noprime(ψsymm[src(e)] * root_S)
-    ψsymm[dst(e)] = noprime(ψsymm[dst(e)] * root_S)
+    setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vsrc]), vsrc)
+    setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vdst]), vdst)
 
     ψsymm_mts[s1 => s2], ψsymm_mts[s2 => s1] = ITensorNetwork(bond_tensors[e]),
     ITensorNetwork(bond_tensors[e])
@@ -157,11 +161,12 @@ function symmetric_to_vidal_gauge(
   ψ_vidal = copy(ψ)
 
   for e in edges(ψ)
-    s1, s2 = find_subgraph((src(e), 1), mts), find_subgraph((dst(e), 1), mts)
+    vsrc, vdst = src(e), dst(e)
+    s1, s2 = find_subgraph((vsrc, 1), mts), find_subgraph((vdst, 1), mts)
     bond_tensors[e] = ITensor(mts[s1 => s2])
     invroot_S = invsqrt_diag(map_diag(x -> x + regularization, bond_tensors[e]))
-    ψ_vidal[src(e)] = noprime(invroot_S * ψ_vidal[src(e)])
-    ψ_vidal[dst(e)] = noprime(invroot_S * ψ_vidal[dst(e)])
+    setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vsrc]), vsrc)
+    setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vdst]), vdst)
   end
 
   return ψ_vidal, bond_tensors

test/test_apply.jl

@@ -6,7 +6,8 @@ using ITensorNetworks:
   message_tensors,
   nested_graph_leaf_vertices,
   vidal_gauge,
-  vidal_to_symmetric_gauge
+  vidal_to_symmetric_gauge,
+  norm_network
 using Test
 using Compat
 using ITensors
@@ -26,7 +27,7 @@ using SplitApplyCombine
   ψ = randomITensorNetwork(s; link_space=χ)
   v1, v2 = (2, 2), (1, 2)
 
-  ψψ = ψ ⊗ prime(dag(ψ); sites=[])
+  ψψ = norm_network(ψ)
 
   #Simple Belief Propagation Grouping
   vertex_groupsSBP = nested_graph_leaf_vertices(