Clean up belief propagation

README.md

@@ -161,7 +161,7 @@ with vertex data:
  3 │ Index[(dim=2|id=683|"S=1/2,Site,n=3")]
 
 and edge data:
-0-element Dictionaries.Dictionary{NamedGraphs.NamedEdge{Int64}, Vector{Index}}
+0-element Dictionaries.Dictionary{NamedEdge{Int64}, Vector{Index}}
 
 julia> tn1 = ITensorNetwork(s; link_space=2)
 ITensorNetwork{Int64} with 3 vertices:

examples/belief_propagation/bpexample.jl

@@ -2,47 +2,69 @@ using Compat
 using ITensors
 using Metis
 using ITensorNetworks
+using Random
+using SplitApplyCombine
 
-using ITensorNetworks: construct_initial_mts, update_all_mts, get_single_site_expec
-
-n = 4
-system_dims = (n, n)
-g = named_grid(system_dims)
-# g = named_comb_tree(system_dims)
-s = siteinds("S=1/2", g)
-chi = 2
-
-ψ = randomITensorNetwork(s; link_space=chi)
-
-ψψ = norm_sqr_network(ψ; flatten=true, map_bra_linkinds=prime)
-combiners = linkinds_combiners(ψψ)
-ψψ = combine_linkinds(ψψ, combiners)
-
-# Apply Sz to site v
-v = one.(system_dims)
-Oψ = copy(ψ)
-Oψ[v] = apply(op("Sz", s[v]), ψ[v])
-ψOψ = inner_network(ψ, Oψ; flatten=true, map_bra_linkinds=prime)
-ψOψ = combine_linkinds(ψOψ, combiners)
-
-# Get the value of sz on v via exact contraction
-contract_seq = contraction_sequence(ψψ)
-actual_sz = contract(ψOψ; sequence=contract_seq)[] / contract(ψψ; sequence=contract_seq)[]
-
-println("Actual value of Sz on site " * string(v) * " is " * string(actual_sz))
-
-niters = 20
-
-nsites = 1
-println("\nFirst " * string(nsites) * " sites form a subgraph")
-mts = construct_initial_mts(ψψ, nsites; init=(I...) -> @compat allequal(I) ? 1 : 0)
-@show get_single_site_expec(ψψ, mts, ψOψ, v)
-mts = update_all_mts(ψψ, mts, niters)
-@show get_single_site_expec(ψψ, mts, ψOψ, v)
-
-nsites = 4
-println("\nNow " * string(nsites) * " sites form a subgraph")
-mts = construct_initial_mts(ψψ, nsites; init=(I...) -> @compat allequal(I) ? 1 : 0)
-@show get_single_site_expec(ψψ, mts, ψOψ, v)
-mts = update_all_mts(ψψ, mts, niters)
-@show get_single_site_expec(ψψ, mts, ψOψ, v)
+using ITensorNetworks:
+  compute_message_tensors, calculate_contraction, contract_inner, nested_graph_leaf_vertices
+
+function main()
+  n = 4
+  dims = (n, n)
+  g = named_grid(dims)
+  s = siteinds("S=1/2", g)
+  chi = 2
+
+  Random.seed!(5467)
+
+  #bra
+  ψ = randomITensorNetwork(s; link_space=chi)
+  #bra-ket (but not actually contracted)
+  ψψ = ψ ⊗ prime(dag(ψ); sites=[])
+
+  #Site to take expectation value on
+  v = (1, 1)
+
+  #Now do Simple Belief Propagation to Measure Sz on Site v
+  nsites = 1
+
+  vertex_groups = nested_graph_leaf_vertices(
+    partition(partition(ψψ, group(v -> v[1], vertices(ψψ))); nvertices_per_partition=nsites)
+  )
+  mts = compute_message_tensors(ψψ; vertex_groups=vertex_groups)
+  sz_bp =
+    calculate_contraction(
+      ψψ, mts, [(v, 1)]; verts_tensors=ITensor[apply(op("Sz", s[v]), ψ[v])]
+    )[] / calculate_contraction(ψψ, mts, [(v, 1)])[]
+
+  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
+  vertex_groups = nested_graph_leaf_vertices(
+    partition(partition(ψψ, group(v -> v[1], vertices(ψψ))); nvertices_per_partition=nsites)
+  )
+  mts = compute_message_tensors(ψψ; vertex_groups=vertex_groups)
+  sz_bp =
+    calculate_contraction(
+      ψψ, mts, [(v, 1)]; verts_tensors=ITensor[apply(op("Sz", s[v]), ψ[v])]
+    )[] / calculate_contraction(ψψ, mts, [(v, 1)])[]
+
+  println(
+    "General Belief Propagation (2-site subgraphs) Gives Sz on Site " *
+    string(v) *
+    " as " *
+    string(sz_bp),
+  )
+
+  #Now do it exactly
+  Oψ = copy(ψ)
+  Oψ[v] = apply(op("Sz", s[v]), ψ[v])
+  sz_exact = contract_inner(Oψ, ψ) / contract_inner(ψ, ψ)
+
+  return println("The exact value of Sz on Site " * string(v) * " is " * string(sz_exact))
+end
+
+main()

src/abstractitensornetwork.jl

@@ -593,6 +593,22 @@ function combine_linkinds(tn::AbstractITensorNetwork, combiners=linkinds_combine
   return combined_tn
 end
 
+function split_index(
+  tn::AbstractITensorNetwork,
+  edges_to_split;
+  src_ind_map::Function=identity,
+  dst_ind_map::Function=prime,
+)
+  tn = copy(tn)
+  for e in edges_to_split
+    inds = commoninds(tn[src(e)], tn[dst(e)])
+    tn[src(e)] = replaceinds(tn[src(e)], inds, src_ind_map(inds))
+    tn[dst(e)] = replaceinds(tn[dst(e)], inds, dst_ind_map(inds))
+  end
+
+  return tn
+end
+
 function flatten_networks(
   tn1::AbstractITensorNetwork,
   tn2::AbstractITensorNetwork;