Improved gauging interface, TNO construction, bug fixes, apply functi…

…on for a Vidal TN (#88)

src/ITensorNetworks.jl

@@ -64,7 +64,6 @@ function iterate(::AbstractDataGraph)
 end
 
 include("observers.jl")
-include("utils.jl")
 include("visualize.jl")
 include("graphs.jl")
 include("itensors.jl")
@@ -94,6 +93,8 @@ include("boundarymps.jl")
 include("beliefpropagation.jl")
 include("contraction_tree_to_graph.jl")
 include("gauging.jl")
+include("utils.jl")
+include("tensornetworkoperators.jl")
 include(joinpath("ITensorsExt", "itensorutils.jl"))
 include(joinpath("Graphs", "abstractgraph.jl"))
 include(joinpath("Graphs", "abstractdatagraph.jl"))

src/ITensorsExt/itensorutils.jl

@@ -19,3 +19,32 @@ inv_diag(it::ITensor) = map_diag(inv, it)
 invsqrt_diag(it::ITensor) = map_diag(inv ∘ sqrt, it)
 pinv_diag(it::ITensor) = map_diag(pinv, it)
 pinvsqrt_diag(it::ITensor) = map_diag(pinv ∘ sqrt, it)
+
+"""Given a vector of ITensors, separate them into groups of commuting itensors (i.e. itensors in the same group do not share any common indices)"""
+function group_commuting_itensors(its::Vector{ITensor})
+  remaining_its = copy(its)
+  it_groups = Vector{ITensor}[]
+
+  while !isempty(remaining_its)
+    cur_group = ITensor[]
+    cur_indices = Index[]
+    inds_to_remove = []
+    for i in 1:length(remaining_its)
+      it = remaining_its[i]
+      it_inds = inds(it)
+
+      if all([i ∉ cur_indices for i in it_inds])
+        push!(cur_group, it)
+        push!(cur_indices, it_inds...)
+        push!(inds_to_remove, i)
+      end
+    end
+    remaining_its = ITensor[
+      remaining_its[i] for
+      i in setdiff([i for i in 1:length(remaining_its)], inds_to_remove)
+    ]
+    push!(it_groups, cur_group)
+  end
+
+  return it_groups
+end

src/abstractitensornetwork.jl

@@ -590,8 +590,10 @@ end
 function combine_linkinds(tn::AbstractITensorNetwork, combiners=linkinds_combiners(tn))
   combined_tn = copy(tn)
   for e in edges(tn)
-    combined_tn[src(e)] = combined_tn[src(e)] * combiners[e]
-    combined_tn[dst(e)] = combined_tn[dst(e)] * combiners[reverse(e)]
+    if !isempty(linkinds(tn, e))
+      combined_tn[src(e)] = combined_tn[src(e)] * combiners[e]
+      combined_tn[dst(e)] = combined_tn[dst(e)] * combiners[reverse(e)]
+    end
   end
   return combined_tn
 end

src/apply.jl

@@ -1,14 +1,13 @@
 function ITensors.apply(
   o::ITensor,
   ψ::AbstractITensorNetwork;
-  cutoff=nothing,
-  maxdim=nothing,
   normalize=false,
   ortho=false,
   envs=ITensor[],
   nfullupdatesweeps=10,
   print_fidelity_loss=true,
   envisposdef=false,
+  apply_kwargs...,
 )
   ψ = copy(ψ)
   v⃗ = neighbor_vertices(ψ, o)
@@ -55,13 +54,13 @@ function ITensors.apply(
         extended_envs,
         o;
         nfullupdatesweeps,
-        maxdim,
         print_fidelity_loss,
         envisposdef,
+        apply_kwargs...,
       )
     else
       Rᵥ₁, Rᵥ₂ = factorize(
-        apply(o, Rᵥ₁ * Rᵥ₂), inds(Rᵥ₁); cutoff, maxdim, tags=ITensorNetworks.edge_tag(e)
+        apply(o, Rᵥ₁ * Rᵥ₂), inds(Rᵥ₁); tags=ITensorNetworks.edge_tag(e), apply_kwargs...
       )
     end
 
@@ -87,52 +86,134 @@ end
 function ITensors.apply(
   o⃗::Vector{ITensor},
   ψ::AbstractITensorNetwork;
-  cutoff,
-  maxdim=typemax(Int),
   normalize=false,
   ortho=false,
-  kwargs...,
+  apply_kwargs...,
 )
   o⃗ψ = ψ
   for oᵢ in o⃗
-    o⃗ψ = apply(oᵢ, o⃗ψ; cutoff, maxdim, normalize, ortho)
+    o⃗ψ = apply(oᵢ, o⃗ψ; normalize, ortho, apply_kwargs...)
   end
   return o⃗ψ
 end
 
 function ITensors.apply(
-  o⃗::Scaled,
-  ψ::AbstractITensorNetwork;
-  cutoff,
-  maxdim,
-  normalize=false,
-  ortho=false,
-  kwargs...,
+  o⃗::Scaled, ψ::AbstractITensorNetwork; normalize=false, ortho=false, apply_kwargs...
 )
   return maybe_real(Ops.coefficient(o⃗)) *
-         apply(Ops.argument(o⃗), ψ; cutoff, maxdim, normalize, ortho, kwargs...)
+         apply(Ops.argument(o⃗), ψ; cutoff, maxdim, normalize, ortho, apply_kwargs...)
 end
 
 function ITensors.apply(
-  o⃗::Prod,
-  ψ::AbstractITensorNetwork;
-  cutoff,
-  maxdim,
-  normalize=false,
-  ortho=false,
-  kwargs...,
+  o⃗::Prod, ψ::AbstractITensorNetwork; normalize=false, ortho=false, apply_kwargs...
 )
   o⃗ψ = ψ
   for oᵢ in o⃗
-    o⃗ψ = apply(oᵢ, o⃗ψ; cutoff, maxdim, normalize, ortho, kwargs...)
+    o⃗ψ = apply(oᵢ, o⃗ψ; normalize, ortho, apply_kwargs...)
   end
   return o⃗ψ
 end
 
 function ITensors.apply(
-  o::Op, ψ::AbstractITensorNetwork; cutoff, maxdim, normalize=false, ortho=false, kwargs...
+  o::Op, ψ::AbstractITensorNetwork; normalize=false, ortho=false, apply_kwargs...
 )
-  return apply(ITensor(o, siteinds(ψ)), ψ; cutoff, maxdim, normalize, ortho, kwargs...)
+  return apply(ITensor(o, siteinds(ψ)), ψ; normalize, ortho, apply_kwargs...)
+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
+end
+
+function _contract_factorized_gate(o::AbstractEdge, ψv1, ψv2)
+  return ψv1, ψv2
+end
+
+#In the future we will try to unify this into apply() above but currently leave it mostly as a separate function
+"""Apply() function for an ITN in the Vidal Gauge. Hence the bond tensors are required.
+Gate does not necessarily need to be passed. Can supply an edge to do an identity update instead. Uses Simple Update procedure assuming gate is two-site"""
+function ITensors.apply(
+  o::Union{ITensor,NamedEdge},
+  ψ::AbstractITensorNetwork,
+  bond_tensors::DataGraph;
+  normalize=false,
+  apply_kwargs...,
+)
+  ψ = copy(ψ)
+  bond_tensors = copy(bond_tensors)
+  v⃗ = _gate_vertices(o, ψ)
+  if length(v⃗) == 2
+    e = NamedEdge(v⃗[1] => v⃗[2])
+    ψv1, ψv2 = copy(ψ[src(e)]), copy(ψ[dst(e)])
+    e_ind = commonind(ψv1, ψv2)
+
+    for vn in neighbors(ψ, src(e))
+      if (vn != dst(e))
+        ψv1 = noprime(ψv1 * bond_tensors[vn => src(e)])
+      end
+    end
+
+    for vn in neighbors(ψ, dst(e))
+      if (vn != src(e))
+        ψv2 = noprime(ψv2 * bond_tensors[vn => dst(e)])
+      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ᵥ₂
+
+    U, S, V = ITensors.svd(
+      theta,
+      uniqueinds(Rᵥ₁, Rᵥ₂);
+      lefttags=ITensorNetworks.edge_tag(e),
+      righttags=ITensorNetworks.edge_tag(e),
+      apply_kwargs...,
+    )
+
+    ind_to_replace = commonind(V, S)
+    ind_to_replace_with = commonind(U, S)
+    replaceind!(S, ind_to_replace, ind_to_replace_with')
+    replaceind!(V, ind_to_replace, ind_to_replace_with)
+
+    ψv1, bond_tensors[e], ψv2 = U * Qᵥ₁, S, V * Qᵥ₂
+
+    for vn in neighbors(ψ, src(e))
+      if (vn != dst(e))
+        ψv1 = noprime(ψv1 * inv_diag(bond_tensors[vn => src(e)]))
+      end
+    end
+
+    for vn in neighbors(ψ, dst(e))
+      if (vn != src(e))
+        ψv2 = noprime(ψv2 * inv_diag(bond_tensors[vn => dst(e)]))
+      end
+    end
+
+    if normalize
+      normalize!(ψv1)
+      normalize!(ψv2)
+      normalize!(bond_tensors[e])
+    end
+
+    ψ[src(e)], ψ[dst(e)] = ψv1, ψv2
+
+    return ψ, bond_tensors
+
+  else
+    ψ = ITensors.apply(o, ψ; normalize)
+    return ψ, bond_tensors
+  end
 end
 
 ### Full Update Routines ###
@@ -193,11 +274,13 @@ function optimise_p_q(
   envs::Vector{ITensor},
   o::ITensor;
   nfullupdatesweeps=10,
-  maxdim=nothing,
   print_fidelity_loss=false,
   envisposdef=true,
+  apply_kwargs...,
 )
-  p_cur, q_cur = factorize(apply(o, p * q), inds(p); maxdim, tags=tags(commonind(p, q)))
+  p_cur, q_cur = factorize(
+    apply(o, p * q), inds(p); tags=tags(commonind(p, q)), apply_kwargs...
+  )
 
   fstart = print_fidelity_loss ? fidelity(envs, p_cur, q_cur, p, q, o) : 0
 

src/beliefpropagation.jl

@@ -10,15 +10,21 @@ function message_tensors(
   )
 end
 
-function message_tensors(
-  subgraphs::DataGraph; itensor_constructor=inds_e -> dense(delta(inds_e))
-)
+function message_tensors_skeleton(subgraphs::DataGraph)
   mts = DataGraph{vertextype(subgraphs),vertex_data_type(subgraphs),ITensorNetwork}(
     directed_graph(underlying_graph(subgraphs))
   )
   for v in vertices(mts)
     mts[v] = subgraphs[v]
   end
+
+  return mts
+end
+
+function message_tensors(
+  subgraphs::DataGraph; itensor_constructor=inds_e -> dense(delta(inds_e))
+)
+  mts = message_tensors_skeleton(subgraphs)
   for e in edges(subgraphs)
     inds_e = commoninds(subgraphs[src(e)], subgraphs[dst(e)])
     mts[e] = ITensorNetwork(map(itensor_constructor, inds_e))
@@ -36,15 +42,21 @@ function update_message_tensor(
   mts::Vector{ITensorNetwork};
   contract_kwargs=(; alg="density_matrix", output_structure=path_graph_structure, maxdim=1),
 )
-  contract_list = ITensorNetwork[mts; ITensorNetwork([tn[v] for v in subgraph_vertices])]
+  mts_itensors = reduce(vcat, ITensor.(mts); init=ITensor[])
 
+  contract_list = ITensor[mts_itensors; ITensor[tn[v] for v in subgraph_vertices]]
   tn = if isone(length(contract_list))
     copy(only(contract_list))
   else
-    reduce(⊗, contract_list)
+    ITensorNetwork(contract_list)
+  end
+
+  if contract_kwargs.alg != "exact"
+    contract_output = contract(tn; contract_kwargs...)
+  else
+    contract_output = contract(tn; sequence=contraction_sequence(tn; alg="optimal"))
   end
 
-  contract_output = contract(tn; contract_kwargs...)
   itn = if typeof(contract_output) == ITensor
     ITensorNetwork(contract_output)
   else
@@ -68,28 +80,49 @@ function belief_propagation_iteration(
   tn::ITensorNetwork,
   mts::DataGraph;
   contract_kwargs=(; alg="density_matrix", output_structure=path_graph_structure, maxdim=1),
+  compute_norm=false,
 )
   new_mts = copy(mts)
-  for e in edges(mts)
+  c = 0
+  es = edges(mts)
+  for e in es
     environment_tensornetworks = ITensorNetwork[
-      mts[e_in] for e_in in setdiff(boundary_edges(mts, src(e); dir=:in), [reverse(e)])
+      mts[e_in] for e_in in setdiff(boundary_edges(mts, [src(e)]; dir=:in), [reverse(e)])
     ]
 
     new_mts[src(e) => dst(e)] = update_message_tensor(
       tn, mts[src(e)], environment_tensornetworks; contract_kwargs
     )
+
+    if compute_norm
+      LHS, RHS = ITensors.contract(ITensor(mts[src(e) => dst(e)])),
+      ITensors.contract(ITensor(new_mts[src(e) => dst(e)]))
+      LHS /= sum(diag(LHS))
+      RHS /= sum(diag(RHS))
+      c += 0.5 * norm(LHS - RHS)
+    end
   end
-  return new_mts
+  return new_mts, c / (length(es))
 end
 
 function belief_propagation(
   tn::ITensorNetwork,
   mts::DataGraph;
   contract_kwargs=(; alg="density_matrix", output_structure=path_graph_structure, maxdim=1),
   niters=20,
+  target_precision::Union{Float64,Nothing}=nothing,
 )
+  compute_norm = target_precision == nothing ? false : true
   for i in 1:niters
-    mts = belief_propagation_iteration(tn, mts; contract_kwargs)
+    mts, c = belief_propagation_iteration(tn, mts; contract_kwargs, compute_norm)
+    if compute_norm && c <= target_precision
+      println(
+        "Belief Propagation finished. Reached a canonicalness of " *
+        string(c) *
+        " after $i iterations. ",
+      )
+      break
+    end
   end
   return mts
 end
@@ -101,12 +134,10 @@ function belief_propagation(
   npartitions=nothing,
   subgraph_vertices=nothing,
   niters=20,
+  target_precision::Union{Float64,Nothing}=nothing,
 )
   mts = message_tensors(tn; nvertices_per_partition, npartitions, subgraph_vertices)
-  for i in 1:niters
-    mts = belief_propagation_iteration(tn, mts; contract_kwargs)
-  end
-  return mts
+  return belief_propagation(tn, mts; contract_kwargs, niters, target_precision)
 end
 
 """