Partitioned ITensorNetwork

examples/apply/apply_bp/apply_bp.jl

@@ -4,10 +4,8 @@ using ITensorNetworks:
   belief_propagation,
   get_environment,
   contract_inner,
-  find_subgraph,
   message_tensors,
   neighbor_vertices,
-  nested_graph_leaf_vertices,
   symmetric_gauge,
   vidal_gauge,
   vidal_to_symmetric_gauge,
@@ -24,14 +22,13 @@ using LinearAlgebra
 using SplitApplyCombine
 using OMEinsumContractionOrders
 
-function expect_bp(opname, v, ψ, mts)
+function expect_bp(opname, v, ψ, pψψ, mts)
   s = siteinds(ψ)
-  ψψ = norm_network(ψ)
-  numerator_network = approx_network_region(
-    ψψ, mts, [(v, 1)]; verts_tn=ITensorNetwork(ITensor[apply(op(opname, s[v]), ψ[v])])
+  numerator_tensors = approx_network_region(
+    pψψ, mts, [(v, 1)]; verts_tensors=ITensor[apply(op(opname, s[v]), ψ[v])]
   )
-  denominator_network = approx_network_region(ψψ, mts, [(v, 1)])
-  return contract(numerator_network)[] / contract(denominator_network)[]
+  denominator_tensors = approx_network_region(pψψ, mts, [(v, 1)])
+  return contract(numerator_tensors)[] / contract(denominator_tensors)[]
 end
 
 function vertex_array(ψ, v, v⃗ⱼ)
@@ -52,26 +49,26 @@ function simple_update_bp(
 )
   println("Simple update, BP")
   ψψ = norm_network(ψ)
-  mts = message_tensors(partition(ψψ, group(v -> v[1], vertices(ψψ))))
+  pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ)))
   mts = belief_propagation(
-    ψψ, mts; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
+    pψψ; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
   )
+  edges = PartitionEdge.(NamedGraphs.edges(partitioned_graph(pψψ)))
   for layer in eachindex(os)
     @show layer
     o⃗ = os[layer]
     for o in o⃗
       v⃗ = neighbor_vertices(ψ, o)
-      for e in edges(mts)
+      for e in edges
         @assert order(only(mts[e])) == 2
         @assert order(only(mts[reverse(e)])) == 2
       end
 
       @assert length(v⃗) == 2
       v1, v2 = v⃗
 
-      s1 = find_subgraph((v1, 1), mts)
-      s2 = find_subgraph((v2, 1), mts)
-      envs = get_environment(ψψ, mts, [(v1, 1), (v1, 2), (v2, 1), (v2, 2)])
+      pe = partitionedge(pψψ, NamedEdge((v1, 1) => (v2, 1)))
+      envs = get_environment(pψψ, mts, [(v1, 1), (v1, 2), (v2, 1), (v2, 2)])
       obs = observer()
       # TODO: Make a version of `apply` that accepts message tensors,
       # and computes the environment and does the message tensor update of the bond internally.
@@ -92,29 +89,28 @@ function simple_update_bp(
 
       # Update message tensor
       ψψ = norm_network(ψ)
-      mts[s1] = ITensorNetwork(dictionary([(v1, 1) => ψψ[v1, 1], (v1, 2) => ψψ[v1, 2]]))
-      mts[s2] = ITensorNetwork(dictionary([(v2, 1) => ψψ[v2, 1], (v2, 2) => ψψ[v2, 2]]))
-      mts[s1 => s2] = ITensorNetwork(obs.singular_values)
-      mts[s2 => s1] = ITensorNetwork(obs.singular_values)
+      pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ)))
+      mts[pe] = dense.(obs.singular_values)
+      mts[reverse(pe)] = dense.(obs.singular_values)
     end
     if regauge
       println("regauge")
       mts = belief_propagation(
-        ψψ, mts; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
+        pψψ, mts; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
       )
     end
   end
-  return ψ, mts
+  return ψ, pψψ, mts
 end
 
 function simple_update_vidal(os, ψ::ITensorNetwork; maxdim, regauge=false)
   println("Simple update, Vidal gauge")
   ψψ = norm_network(ψ)
-  mts = message_tensors(partition(ψψ, group(v -> v[1], vertices(ψψ))))
+  pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ)))
   mts = belief_propagation(
-    ψψ, mts; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
+    pψψ; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
   )
-  ψ, bond_tensors = vidal_gauge(ψ, mts)
+  ψ, bond_tensors = vidal_gauge(ψ, pψψ, mts)
   for layer in eachindex(os)
     @show layer
     o⃗ = os[layer]
@@ -124,12 +120,15 @@ function simple_update_vidal(os, ψ::ITensorNetwork; maxdim, regauge=false)
     end
     if regauge
       println("regauge")
-      ψ_symmetric, mts = vidal_to_symmetric_gauge(ψ, bond_tensors)
-      ψψ = norm_network(ψ_symmetric)
+      ψ_symmetric, pψψ_symmetric, mts = vidal_to_symmetric_gauge(ψ, bond_tensors)
       mts = belief_propagation(
-        ψψ, mts; contract_kwargs=(; alg="exact"), niters=50, target_precision=1e-5
+        pψψ_symmetric,
+        mts;
+        contract_kwargs=(; alg="exact"),
+        niters=50,
+        target_precision=1e-5,
       )
-      ψ, bond_tensors = vidal_gauge(ψ_symmetric, mts)
+      ψ, bond_tensors = vidal_gauge(ψ_symmetric, pψψ_symmetric, mts)
     end
   end
   return ψ, bond_tensors
@@ -157,14 +156,14 @@ function main(;
   ]
 
   # BP SU
-  ψ_bp, mts = simple_update_bp(
+  ψ_bp, pψψ_bp, mts_bp = simple_update_bp(
     os, ψ; maxdim=χ, variational_optimization_only, regauge, reduced
   )
   # ψ_bp, mts = vidal_to_symmetric_gauge(vidal_gauge(ψ_bp, mts)...)
 
   # Vidal SU
   ψ_vidal, bond_tensors = simple_update_vidal(os, ψ; maxdim=χ, regauge)
-  ψ_vidal, mts_vidal = vidal_to_symmetric_gauge(ψ_vidal, bond_tensors)
+  ψ_vidal, pψψ_vidal, mts_vidal = vidal_to_symmetric_gauge(ψ_vidal, bond_tensors)
 
-  return ψ_bp, mts, ψ_vidal, mts_vidal
+  return ψ_bp, pψψ_bp, mts_bp, ψ_vidal, pψψ_vidal, mts_vidal
 end

examples/apply/apply_bp/apply_bp_run.jl

@@ -10,7 +10,7 @@ graph = named_grid
 
 dims = (6, 6)
 
-ψ_bp, mts_bp, ψ_vidal, mts_vidal = main(;
+ψ_bp, pψψ_bp, mts_bp, ψ_vidal, pψψ_vidal, mts_vidal = main(;
   seed=1234,
   opname,
   graph,
@@ -24,31 +24,33 @@ dims = (6, 6)
 
 v = dims .÷ 2
 
-sz_bp = @show expect_bp("Sz", v, ψ_bp, mts_bp)
-sz_vidal = @show expect_bp("Sz", v, ψ_vidal, mts_vidal)
+sz_bp = @show expect_bp("Sz", v, ψ_bp, pψψ_bp, mts_bp)
+sz_vidal = @show expect_bp("Sz", v, ψ_vidal, pψψ_vidal, mts_vidal)
 @show abs(sz_bp - sz_vidal) / abs(sz_vidal)
 
 # Run BP again
 mts_bp = belief_propagation(
-  norm_network(ψ_bp),
+  pψψ_bp,
   mts_bp;
   contract_kwargs=(; alg="exact"),
   niters=50,
-  target_precision=1e-5,
+  target_precision=1e-7,
+  verbose=true,
 )
 mts_vidal = belief_propagation(
-  norm_network(ψ_vidal),
+  pψψ_vidal,
   mts_vidal;
   contract_kwargs=(; alg="exact"),
   niters=50,
-  target_precision=1e-5,
+  target_precision=1e-7,
+  verbose=true,
 )
 
-sz_bp = @show expect_bp("Sz", v, ψ_bp, mts_bp)
-sz_vidal = @show expect_bp("Sz", v, ψ_vidal, mts_vidal)
+sz_bp = @show expect_bp("Sz", v, ψ_bp, pψψ_bp, mts_bp)
+sz_vidal = @show expect_bp("Sz", v, ψ_vidal, pψψ_vidal, mts_vidal)
 @show abs(sz_bp - sz_vidal) / abs(sz_vidal)
 
-ψ_symmetric, _ = symmetric_gauge(ψ_bp)
+ψ_symmetric, _, _ = symmetric_gauge(ψ_bp)
 
 v⃗ⱼ = [v .+ (1, 0), v .- (1, 0), v .+ (0, 1), v .- (0, 1)]
 ψ_bp_v = vertex_array(ψ_bp, v, v⃗ⱼ)

examples/belief_propagation/bpexample.jl

@@ -7,11 +7,7 @@ using SplitApplyCombine
 using NamedGraphs
 
 using ITensorNetworks:
-  belief_propagation,
-  approx_network_region,
-  contract_inner,
-  message_tensors,
-  nested_graph_leaf_vertices
+  belief_propagation, approx_network_region, contract_inner, message_tensors
 
 function main()
   n = 4
@@ -31,41 +27,33 @@ function main()
   v = (1, 1)
 
   #Now do Simple Belief Propagation to Measure Sz on Site v
-  mts = message_tensors(
-    ψψ; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψ))))
+  pψψ = PartitionedGraph(ψψ, collect(values(group(v -> v[1], vertices(ψψ)))))
+  mts = belief_propagation(
+    pψψ; contract_kwargs=(; alg="exact"), verbose=true, niters=10, target_precision=1e-3
   )
-
-  mts = belief_propagation(ψψ, mts; contract_kwargs=(; alg="exact"), niters=20)
-
-  numerator_network = approx_network_region(
-    ψψ, mts, [(v, 1)]; verts_tn=ITensorNetwork([apply(op("Sz", s[v]), ψ[v])])
+  numerator_tensors = approx_network_region(
+    pψψ, mts, [(v, 1)]; verts_tensors=[apply(op("Sz", s[v]), ψ[v])]
   )
-  denominator_network = approx_network_region(ψψ, mts, [(v, 1)])
-  sz_bp = contract(numerator_network)[] / contract(denominator_network)[]
+  denominator_tensors = approx_network_region(pψψ, mts, [(v, 1)])
+  sz_bp = contract(numerator_tensors)[] / contract(denominator_tensors)[]
 
   println(
     "Simple Belief Propagation Gives Sz on Site " * string(v) * " as " * string(sz_bp)
   )
 
-  #Now do General Belief Propagation to Measure Sz on Site v
-  nsites = 4
-  Zp = partition(
-    partition(ψψ, group(v -> v[1], vertices(ψψ))); nvertices_per_partition=nsites
+  #Now do Column-wise General Belief Propagation to Measure Sz on Site v
+  pψψ = PartitionedGraph(ψψ, collect(values(group(v -> v[1][1], vertices(ψψ)))))
+  mts = belief_propagation(
+    pψψ; contract_kwargs=(; alg="exact"), verbose=true, niters=10, target_precision=1e-3
   )
-  Zpp = partition(ψψ; subgraph_vertices=nested_graph_leaf_vertices(Zp))
-  mts = message_tensors(Zpp)
-  mts = belief_propagation(ψψ, mts; contract_kwargs=(; alg="exact"), niters=20)
-  numerator_network = approx_network_region(
-    ψψ, mts, [(v, 1)]; verts_tn=ITensorNetwork([apply(op("Sz", s[v]), ψ[v])])
+  numerator_tensors = approx_network_region(
+    pψψ, mts, [(v, 1)]; verts_tensors=[apply(op("Sz", s[v]), ψ[v])]
   )
-  denominator_network = approx_network_region(ψψ, mts, [(v, 1)])
-  sz_bp = contract(numerator_network)[] / contract(denominator_network)[]
+  denominator_tensors = approx_network_region(pψψ, mts, [(v, 1)])
+  sz_gen_bp = contract(numerator_tensors)[] / contract(denominator_tensors)[]
 
   println(
-    "General Belief Propagation (4-site subgraphs) Gives Sz on Site " *
-    string(v) *
-    " as " *
-    string(sz_bp),
+    "General Belief Propagation Gives Sz on Site " * string(v) * " as " * string(sz_gen_bp)
   )
 
   #Now do General Belief Propagation with Matrix Product State Message Tensors Measure Sz on Site v
@@ -78,30 +66,26 @@ function main()
   ψψ = combine_linkinds(ψψ, combiners)
   ψOψ = combine_linkinds(ψOψ, combiners)
 
-  Z = partition(ψψ, group(v -> v[1], vertices(ψψ)))
-  maxdim = 8
-  mts = message_tensors(Z)
-
+  pψψ = PartitionedGraph(ψψ, collect(values(group(v -> v[1], vertices(ψψ)))))
   mts = belief_propagation(
-    ψψ,
-    mts;
+    pψψ;
+    itensor_constructor=inds_e -> ITensor[dense(delta(i)) for i in inds_e],
     contract_kwargs=(;
       alg="density_matrix",
       output_structure=path_graph_structure,
-      maxdim,
+      maxdim=8,
       contraction_sequence_alg="optimal",
     ),
   )
-
-  numerator_network = approx_network_region(ψψ, mts, [v]; verts_tn=ITensorNetwork(ψOψ[v]))
-  denominator_network = approx_network_region(ψψ, mts, [v])
-  sz_bp = contract(numerator_network)[] / contract(denominator_network)[]
+  numerator_tensors = approx_network_region(pψψ, mts, [v]; verts_tensors=[ψOψ[v]])
+  denominator_tensors = approx_network_region(pψψ, mts, [v])
+  sz_MPS_bp = contract(numerator_tensors)[] / contract(denominator_tensors)[]
 
   println(
-    "General Belief Propagation with Column Partitioning and MPS Message Tensors (Max dim 8) Gives Sz on Site " *
+    "Column-Wise MPS Belief Propagation Gives Sz on Site " *
     string(v) *
     " as " *
-    string(sz_bp),
+    string(sz_gen_bp),
   )
 
   #Now do it exactly

examples/belief_propagation/bpsequences.jl

@@ -8,12 +8,7 @@ using Graphs
 using NamedGraphs
 
 using ITensorNetworks:
-  belief_propagation,
-  approx_network_region,
-  contract_inner,
-  message_tensors,
-  nested_graph_leaf_vertices,
-  edge_sequence
+  belief_propagation, approx_network_region, contract_inner, message_tensors, edge_sequence
 
 function main()
   g_labels = [
@@ -39,36 +34,40 @@ function main()
     ψψ = ψ ⊗ prime(dag(ψ); sites=[])
 
     #Initial message tensors for BP
-    mts_init = message_tensors(
-      ψψ; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψ))))
-    )
+    pψψ = PartitionedGraph(ψψ, collect(values(group(v -> v[1], vertices(ψψ)))))
+    mts_init = message_tensors(pψψ)
 
     println("\nFirst testing out a $g_label. Random network with bond dim $χ")
 
     #Now test out various sequences
     print("Parallel updates (sequence is irrelevant): ")
     belief_propagation(
-      ψψ,
+      pψψ,
       mts_init;
       contract_kwargs=(; alg="exact"),
       target_precision=1e-10,
       niters=100,
-      edges=edge_sequence(mts_init; alg="parallel"),
+      edges=[
+        PartitionEdge.(e) for e in edge_sequence(partitioned_graph(pψψ); alg="parallel")
+      ],
       verbose=true,
     )
     print("Sequential updates (sequence is default edge list of the message tensors): ")
     belief_propagation(
-      ψψ,
+      pψψ,
       mts_init;
       contract_kwargs=(; alg="exact"),
       target_precision=1e-10,
       niters=100,
-      edges=[e for e in edges(mts_init)],
+      edges=PartitionEdge.([
+        e for
+        e in vcat(edges(partitioned_graph(pψψ)), reverse.(edges(partitioned_graph(pψψ))))
+      ]),
       verbose=true,
     )
     print("Sequential updates (sequence is our custom sequence finder): ")
     belief_propagation(
-      ψψ,
+      pψψ,
       mts_init;
       contract_kwargs=(; alg="exact"),
       target_precision=1e-10,