Belief Propagation Version 3

examples/belief_propagation/bpexample.jl

@@ -6,12 +6,12 @@ using Random
 using SplitApplyCombine
 
 using ITensorNetworks:
-  compute_message_tensors, calculate_contraction, contract_inner, nested_graph_leaf_vertices
+  belief_propagation, approx_network_region, contract_inner, nested_graph_leaf_vertices
 
 function main()
   n = 4
-  dims = (n, n)
-  g = named_grid(dims)
+  g_dims = (n, n)
+  g = named_grid(g_dims)
   s = siteinds("S=1/2", g)
   chi = 2
 
@@ -31,11 +31,12 @@ function main()
   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)])[]
+  mts = belief_propagation(ψψ; vertex_groups=vertex_groups)
+  numerator_network = approx_network_region(
+    ψψ, mts, [(v, 1)]; verts_tn=ITensorNetwork([apply(op("Sz", s[v]), ψ[v])])
+  )
+  denominator_network = approx_network_region(ψψ, mts, [(v, 1)])
+  sz_bp = contract(numerator_network)[] / contract(denominator_network)[]
 
   println(
     "Simple Belief Propagation Gives Sz on Site " * string(v) * " as " * string(sz_bp)
@@ -46,23 +47,62 @@ function main()
   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)])[]
+  mts = belief_propagation(ψψ; vertex_groups=vertex_groups)
+  numerator_network = approx_network_region(
+    ψψ, mts, [(v, 1)]; verts_tn=ITensorNetwork([apply(op("Sz", s[v]), ψ[v])])
+  )
+  denominator_network = approx_network_region(ψψ, mts, [(v, 1)])
+  sz_bp = contract(numerator_network)[] / contract(denominator_network)[]
 
   println(
-    "General Belief Propagation (2-site subgraphs) Gives Sz on Site " *
+    "General Belief Propagation (4-site subgraphs) Gives Sz on Site " *
     string(v) *
     " as " *
     string(sz_bp),
   )
 
-  #Now do it exactly
+  #Now do General Belief Propagation with Matrix Product State Message Tensors Measure Sz on Site v
+  ψψ = flatten_networks(ψ, dag(ψ); combine_linkinds=false, map_bra_linkinds=prime)
   Oψ = copy(ψ)
   Oψ[v] = apply(op("Sz", s[v]), ψ[v])
-  sz_exact = contract_inner(Oψ, ψ) / contract_inner(ψ, ψ)
+  ψOψ = flatten_networks(ψ, dag(Oψ); combine_linkinds=false, map_bra_linkinds=prime)
+
+  combiners = linkinds_combiners(ψψ)
+  ψψ = combine_linkinds(ψψ, combiners)
+  ψOψ = combine_linkinds(ψOψ, combiners)
+
+  vertex_groups = nested_graph_leaf_vertices(partition(ψψ, group(v -> v[1], vertices(ψψ))))
+  maxdim = 8
+
+  mts = belief_propagation(
+    ψψ;
+    vertex_groups=vertex_groups,
+    contract_kwargs=(;
+      alg="density_matrix",
+      output_structure=path_graph_structure,
+      maxdim,
+      contraction_sequence_alg="optimal",
+    ),
+    init_contract_kwargs=(;
+      alg="density_matrix",
+      output_structure=path_graph_structure,
+      cutoff=1e-16,
+      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)[]
+
+  println(
+    "General Belief Propagation with Column Partitioning and MPS Message Tensors (Max dim 8) Gives Sz on Site " *
+    string(v) *
+    " as " *
+    string(sz_bp),
+  )
+
+  #Now do it exactly
+  sz_exact = contract(ψOψ)[] / contract(ψψ)[]
 
   return println("The exact value of Sz on Site " * string(v) * " is " * string(sz_exact))
 end

src/apply.jl

@@ -46,6 +46,7 @@ function ITensors.apply(
       )
     end
 
+    envs = Vector{ITensor}(envs)
     if !isempty(envs)
       extended_envs = vcat(envs, Qᵥ₁, prime(dag(Qᵥ₁)), Qᵥ₂, prime(dag(Qᵥ₂)))
       Rᵥ₁, Rᵥ₂ = optimise_p_q(

src/beliefpropagation.jl

@@ -1,41 +1,67 @@
-function construct_initial_mts(
+function message_tensors(
   tn::ITensorNetwork, nvertices_per_partition::Integer; partition_kwargs=(;), kwargs...
 )
-  return construct_initial_mts(
-    tn, partition(tn; nvertices_per_partition, partition_kwargs...); kwargs...
+  return message_tensors(
+    partition(tn; nvertices_per_partition, partition_kwargs...); kwargs...
   )
 end
 
-function construct_initial_mts(
-  tn::ITensorNetwork, subgraphs::DataGraph; init=(I...) -> @compat allequal(I) ? 1 : 0
+function message_tensors(
+  subgraphs::DataGraph; contract_kwargs=(;), init=(I...) -> allequal(I) ? 1 : 0
 )
-  # TODO: This is dropping the vertex data for some reason.
-  # mts = DataGraph{vertextype(subgraphs),vertex_data_type(subgraphs),ITensor}(subgraphs)
-  mts = DataGraph{vertextype(subgraphs),vertex_data_type(subgraphs),ITensor}(
+  mts = DataGraph{vertextype(subgraphs),vertex_data_type(subgraphs),ITensorNetwork}(
     directed_graph(underlying_graph(subgraphs))
   )
   for v in vertices(mts)
     mts[v] = subgraphs[v]
   end
   for subgraph in vertices(subgraphs)
-    tns_to_contract = ITensor[]
     for subgraph_neighbor in neighbors(subgraphs, subgraph)
-      edge_inds = Index[]
-      for vertex in vertices(subgraphs[subgraph])
-        psiv = tn[vertex]
-        for e in [edgetype(tn)(vertex => neighbor) for neighbor in neighbors(tn, vertex)]
-          if (find_subgraph(dst(e), subgraphs) == subgraph_neighbor)
-            append!(edge_inds, commoninds(tn, e))
-          end
-        end
-      end
-      mt = normalize!(
-        itensor(
-          [init(Tuple(I)...) for I in CartesianIndices(tuple(dim.(edge_inds)...))],
-          edge_inds,
+      boundary_vertices_s = setdiff(
+        vertices(subgraphs[subgraph]),
+        unique(
+          flatten([
+            neighbors(subgraphs[subgraph_neighbor], vn) for
+            vn in vertices(subgraphs[subgraph_neighbor])
+          ]),
+        ),
+      )
+      boundary_vertices_sn = setdiff(
+        vertices(subgraphs[subgraph_neighbor]),
+        unique(
+          flatten([
+            neighbors(subgraphs[subgraph], vn) for vn in vertices(subgraphs[subgraph])
+          ]),
         ),
       )
-      mts[subgraph => subgraph_neighbor] = mt
+      boundary_tensors = ITensor[]
+      for v in boundary_vertices_s
+        inds_boundary = flatten(
+          unique([inds(subgraphs[subgraph_neighbor][vn]) for vn in boundary_vertices_sn])
+        )
+        inds_internal = flatten(
+          unique([inds(subgraphs[subgraph][v]) for v in setdiff(boundary_vertices_s, [v])])
+        )
+        new_inds = intersect(
+          inds(subgraphs[subgraph][v]), vcat(inds_boundary, inds_internal)
+        )
+        push!(
+          boundary_tensors,
+          itensor(
+            [init(Tuple(I)...) for I in CartesianIndices(tuple(dim.(new_inds)...))],
+            new_inds,
+          ),
+        )
+      end
+
+      mt = ITensorNetwork(Dictionaries.Dictionary(boundary_vertices_s, boundary_tensors))
+
+      contract_output = contract(mt; contract_kwargs...)
+      mts[subgraph => subgraph_neighbor] = if typeof(contract_output) == ITensor
+        ITensorNetwork(contract_output)
+      else
+        first(contract_output)
+      end
     end
   end
   return mts
@@ -44,60 +70,85 @@ end
 """
 DO a single update of a message tensor using the current subgraph and the incoming mts
 """
-function update_mt(
+function update_message_tensor(
   tn::ITensorNetwork,
   subgraph_vertices::Vector,
-  mts::Vector{ITensor};
-  contraction_sequence::Function=tn -> contraction_sequence(tn; alg="optimal"),
+  mts::Vector{ITensorNetwork};
+  contract_kwargs=(;),
 )
-  contract_list = [mts; [tn[v] for v in subgraph_vertices]]
+  contract_list = ITensorNetwork[mts; ITensorNetwork([tn[v] for v in subgraph_vertices])]
 
-  new_mt = if isone(length(contract_list))
+  tn = if isone(length(contract_list))
     copy(only(contract_list))
   else
-    contract(contract_list; sequence=contraction_sequence(contract_list))
+    reduce(⊗, contract_list)
   end
-  return normalize!(new_mt)
+
+  contract_output = contract(tn; contract_kwargs...)
+  itn = if typeof(contract_output) == ITensor
+    ITensorNetwork(contract_output)
+  else
+    first(contract_output)
+  end
+  normalize!.(vertex_data(itn))
+
+  return itn
 end
 
-function update_mt(
-  tn::ITensorNetwork, subgraph::ITensorNetwork, mts::Vector{ITensor}; kwargs...
+function update_message_tensor(
+  tn::ITensorNetwork, subgraph::ITensorNetwork, mts::Vector{ITensorNetwork}; kwargs...
 )
-  return update_mt(tn, vertices(subgraph), mts; kwargs...)
+  return update_message_tensor(tn, vertices(subgraph), mts; kwargs...)
 end
 
 """
 Do an update of all message tensors for a given ITensornetwork and its partition into sub graphs
 """
-function update_all_mts(
-  tn::ITensorNetwork,
-  mts::DataGraph;
-  contraction_sequence::Function=tn -> contraction_sequence(tn; alg="optimal"),
+function belief_propagation_iteration(
+  tn::ITensorNetwork, mts::DataGraph; contract_kwargs=(;)
 )
-  update_mts = copy(mts)
+  new_mts = copy(mts)
   for e in edges(mts)
-    environment_tensors = ITensor[
+    environment_tensornetworks = ITensorNetwork[
       mts[e_in] for e_in in setdiff(boundary_edges(mts, src(e); dir=:in), [reverse(e)])
     ]
-    update_mts[src(e) => dst(e)] = update_mt(
-      tn, mts[src(e)], environment_tensors; contraction_sequence
+
+    new_mts[src(e) => dst(e)] = update_message_tensor(
+      tn, mts[src(e)], environment_tensornetworks; contract_kwargs
     )
   end
-  return update_mts
+  return new_mts
 end
 
-function update_all_mts(
-  tn::ITensorNetwork,
-  mts::DataGraph,
-  niters::Int;
-  contraction_sequence::Function=tn -> contraction_sequence(tn; alg="optimal"),
+function belief_propagation(
+  tn::ITensorNetwork, mts::DataGraph, niters::Int; contract_kwargs=(;)
 )
   for i in 1:niters
-    mts = update_all_mts(tn, mts; contraction_sequence)
+    mts = belief_propagation_iteration(tn, mts; contract_kwargs)
   end
   return mts
 end
 
+"""
+Simulaneously initialise and update message tensors of a tensornetwork
+"""
+function belief_propagation(
+  tn::ITensorNetwork;
+  niters=10,
+  nvertices_per_partition=nothing,
+  npartitions=nothing,
+  vertex_groups=nothing,
+  contract_kwargs=(;),
+  init_contract_kwargs=(;),
+  init_kwargs...,
+)
+  Z = partition(tn; nvertices_per_partition, npartitions, subgraph_vertices=vertex_groups)
+
+  mts = message_tensors(Z; contract_kwargs=init_contract_kwargs, init_kwargs...)
+  mts = belief_propagation(tn, mts, niters; contract_kwargs)
+  return mts
+end
+
 """
 Given a subet of vertices of a given Tensor Network and the Message Tensors for that network, return a Dictionary with the involved subgraphs as keys and the vector of tensors associated with that subgraph as values
 Specifically, the contraction of the environment tensors and tn[vertices] will be a scalar.
@@ -109,43 +160,24 @@ function get_environment(tn::ITensorNetwork, mts::DataGraph, verts::Vector; dir=
     return get_environment(tn, mts, setdiff(vertices(tn), verts))
   end
 
-  env_tensors = ITensor[mts[e] for e in boundary_edges(mts, subgraphs; dir=:in)]
-  return vcat(
-    env_tensors,
-    ITensor[tn[v] for v in setdiff(flatten([vertices(mts[s]) for s in subgraphs]), verts)],
-  )
+  env_tns = ITensorNetwork[mts[e] for e in boundary_edges(mts, subgraphs; dir=:in)]
+  central_tn = ITensorNetwork([
+    tn[v] for v in setdiff(flatten([vertices(mts[s]) for s in subgraphs]), verts)
+  ])
+  return ITensorNetwork(vcat(env_tns, ITensorNetwork[central_tn]))
 end
 
 """
 Calculate the contraction of a tensor network centred on the vertices verts. Using message tensors.
 Defaults to using tn[verts] as the local network but can be overriden
 """
-function calculate_contraction(
+function approx_network_region(
   tn::ITensorNetwork,
   mts::DataGraph,
   verts::Vector;
-  verts_tensors=ITensor[tn[v] for v in verts],
-  contraction_sequence::Function=tn -> contraction_sequence(tn; alg="optimal"),
+  verts_tn=ITensorNetwork([tn[v] for v in verts]),
 )
-  environment_tensors = get_environment(tn, mts, verts)
-  tensors_to_contract = vcat(environment_tensors, verts_tensors)
-  return contract(tensors_to_contract; sequence=contraction_sequence(tensors_to_contract))
-end
+  environment_tn = get_environment(tn, mts, verts)
 
-"""
-Simulaneously initialise and update message tensors of a tensornetwork
-"""
-function compute_message_tensors(
-  tn::ITensorNetwork;
-  niters=10,
-  nvertices_per_partition=nothing,
-  npartitions=nothing,
-  vertex_groups=nothing,
-  kwargs...,
-)
-  Z = partition(tn; nvertices_per_partition, npartitions, subgraph_vertices=vertex_groups)
-
-  mts = construct_initial_mts(tn, Z; kwargs...)
-  mts = update_all_mts(tn, mts, niters)
-  return mts
+  return environment_tn ⊗ verts_tn
 end