Added More Flexible Partitioning and Subgraph Creation

src/beliefpropagation.jl

@@ -1,5 +1,5 @@
 function construct_initial_mts(tn::ITensorNetwork, nvertices_per_subgraph::Integer; subgraph_kwargs=(;), kwargs...)
-  return construct_initial_mts(tn, subgraphs(tn, nvertices_per_subgraph; subgraph_kwargs...); kwargs...)
+  return construct_initial_mts(tn, create_subgraphs(tn, nvertices_per_subgraph; subgraph_kwargs...); kwargs...)
 end
 
 function construct_initial_mts(tn::ITensorNetwork, subgraphs::DataGraph; init)

src/subgraphs.jl

@@ -25,69 +25,72 @@ function find_subgraph(vertex, subgraphs::DataGraph)
   return findfirst_on_vertices(subgraph_vertices -> vertex ∈ subgraph_vertices, subgraphs)
 end
 
+
+"""
+Take an abstract graph and a grouping of the vertices (specified by a dict: group_name -> vertices in group) and create a datagraph of the subgraphs.
+Edges connect subgraphs which contain vertices connected in g. Edge data is the number of edges in g between subgraphs
+"""
+function create_subgraphs(g::AbstractGraph, vertex_groups::Dictionary)
+  dg_subgraph = DataGraph{vertextype(NamedGraph(keys(vertex_groups))), Vector{Any}, Int64}(NamedGraph(keys(vertex_groups)), vertex_groups)
+  for e in edges(g)
+    s1 = findfirst_on_vertices(vertex_groups -> src(e) ∈ vertex_groups, dg_subgraph)
+    s2 = findfirst_on_vertices(vertex_groups -> dst(e) ∈ vertex_groups, dg_subgraph)
+    if (s1 != s2)
+      if(!has_edge(dg_subgraph, s1, s2))
+          add_edge!(dg_subgraph, s1, s2)
+          dg_subgraph[s1=>s2] = 1
+      else
+          dg_subgraph[s1=>s2] += 1
+      end
+    end
+
+  end
+  return dg_subgraph
+end
+
 """
-Given a graph g, form 'nv(g) ÷ nvertices_per_partition' subgraphs. Try to keep all subgraphs the same size and minimise edges cut between them
-Returns a datagraph where each vertex contains the list of vertices involved in that subgraph. The edges state which subgraphs are connected
-KaHyPar needs to be installed for this function to work
+Graph partitioning functions. One of Metis or KaHyPar needs to be installed for these to work
 """
-function subgraphs(g::AbstractGraph, nvertices_per_partition::Integer; kwargs...)
+
+"""create a partitioning of the vertices of g, nvertices per partition
+Returns a dictionary of partition_no -> vertices in partition
+"""
+function partition_vertices(g::AbstractGraph, nvertices_per_partition::Integer; kwargs...)
   nvertices_per_partition = min(nv(g), nvertices_per_partition)
   npartitions = nv(g) ÷ nvertices_per_partition
   vertex_to_partition = partition(g, npartitions; kwargs...)
-  partition_vertices = groupfind(vertex_to_partition)
-  if length(partition_vertices) ≠ npartitions
+  partitioned_vertices = groupfind(vertex_to_partition)
+  if length(partitioned_vertices) ≠ npartitions
     @warn "Requested $nvertices_per_partition vertices per partition for a graph with $(nv(g)) vertices. Attempting to partition the graph into $npartitions partitions but instead it was partitioned into $(length(partition_vertices)) partitions. Partitions are $partition_vertices."
   end
-  if !issetequal(keys(partition_vertices), 1:length(partition_vertices))
-    @warn "Vertex partioning is $partition_vertices, which may not be what you were hoping for. If not, try a different graph partitioning algorithm or backend."
+  if !issetequal(keys(partitioned_vertices), 1:length(partitioned_vertices))
+    @warn "Vertex partioning is $partitioned_vertices, which may not be what you were hoping for. If not, try a different graph partitioning algorithm or backend."
   end
-  dg_subgraphs = DataGraph(NamedGraph(keys(partition_vertices)), partition_vertices)
-  for e in edges(g)
-    s1 = findfirst_on_vertices(subgraph_vertices -> src(e) ∈ subgraph_vertices, dg_subgraphs)
-    s2 = findfirst_on_vertices(subgraph_vertices -> dst(e) ∈ subgraph_vertices, dg_subgraphs)
-    if (!has_edge(dg_subgraphs, s1, s2) && s1 != s2)
-      add_edge!(dg_subgraphs, s1, s2)
-    end
-  end
-  return dg_subgraphs
+  return partitioned_vertices
+end
+
+"""create a partitioning of the vertices of g, npartitions
+Returns a dictionary of partition_no -> vertices in partition
+"""
+function partition_vertices(g::AbstractGraph; npartitions=nv(g), kwargs...)
+  nvertices_per_partition = floor(Int, nv(g)/npartitions)
+  return partition_vertices(g, nvertices_per_partition; kwargs...)
 end
 
-## """
-## Given a graph g on a d-dimensional grid of size Ls[1] x Ls[2] x ..., form subgraphs of size ls[1] x ls[2] in a regular fashion
-## Return the subgraphs (1... npartitions) with their contained vertices. Also return a dictionary of the subgraphs connected to each sgraph
-## """
-## function subgraphs_grid(g::NamedGraph, Ls::Vector{Int64}, ls::Vector{Int64})
-##   lengths = Ls ./ ls
-##   ps = Dict{Tuple,Int64}()
-##   for v in vertices(g)
-##     pos = []
-##     count = 1
-##     for i in 1:length(v)
-##       push!(pos, ceil(Int64, v[i] / ls[i]))
-##       if (pos[i] == 1)
-##         p = 1
-##       else
-##         p = prod(lengths[1:(i - 1)])
-##       end
-##       count = Int(count + (pos[i] - 1) * p)
-##     end
-##     ps[v] = count
-##   end
-## 
-##   nsubgraphs = Int(prod(lengths))
-## 
-##   dg_subgraphs = DataGraph{Vector{Tuple},Any}(NamedGraph([(i,) for i in 1:nsubgraphs]))
-##   for s in 1:nsubgraphs
-##     dg_subgraphs[(s,)] = [v for v in vertices(g) if ps[v] == s]
-##   end
-## 
-##   for e in edges(g)
-##     v1, v2 = src(e), dst(e)
-##     s1, s2 = find_subgraph(v1, dg_subgraphs), find_subgraph(v2, dg_subgraphs)
-##     if (!has_edge(dg_subgraphs, s1, s2) && s1 != s2)
-##       add_edge!(dg_subgraphs, s1, s2)
-##     end
-##   end
-## 
-##   return dg_subgraphs
-## end
+"""
+Form subgraphs (returns a datagraph where vertices of g are grouped into larger vertices of the datagraph. 
+Aims to have nvertices_per_subgraph grouped into each of the datagraph vertices)
+"""
+function create_subgraphs(g::AbstractGraph, nvertices_per_subgraph::Integer; kwargs...)
+  vertex_groups = partition_vertices(g, nvertices_per_subgraph; kwargs...)
+  return create_subgraphs(g, vertex_groups)
+end
+
+"""
+Form subgraphs (returns a datagraph where vertices of g are grouped into larger vertices of the datagraph. 
+Aims to have number of vertices in the datagraph be nsubgraphs)
+"""
+function create_subgraphs(g::AbstractGraph; nsubgraphs = 1, kwargs...)
+  vertex_groups = partition_vertices(g; nsubgraphs = nsubgraphs, kwargs...)
+  return create_subgraphs(g, vertex_groups)
+end