Updated to work with latest builds of NamedGraphs and DataGraphs

examples/apply/apply_bp/apply_bp.jl

@@ -61,13 +61,13 @@ function simple_update_bp(
       v⃗ = neighbor_vertices(ψ, o)
       for e in edges
         @assert order(only(mts[e])) == 2
-        @assert order(only(mts[PartitionEdge(reverse(NamedGraphs.parent(e)))])) == 2
+        @assert order(only(mts[reverse(e)])) == 2
       end
 
       @assert length(v⃗) == 2
       v1, v2 = v⃗
 
-      pe = NamedGraphs.partition_edge(pψψ, NamedEdge((v1, 1) => (v2, 1)))
+      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,
@@ -91,7 +91,7 @@ function simple_update_bp(
       ψψ = norm_network(ψ)
       pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ)))
       mts[pe] = dense.(obs.singular_values)
-      mts[PartitionEdge(reverse(NamedGraphs.parent(pe)))] = dense.(obs.singular_values)
+      mts[reverse(pe)] = dense.(obs.singular_values)
     end
     if regauge
       println("regauge")

examples/boundary.jl

@@ -9,7 +9,7 @@ tn = ITensorNetwork(named_grid((6, 3)); link_space=4)
 @visualize tn
 
 ptn = PartitionedGraph(
-  tn, NamedGraphs.partition_vertices(underlying_graph(tn); nvertices_per_partition=2)
+  tn, partitioned_vertices(underlying_graph(tn); nvertices_per_partition=2)
 )
 sub_vs_1, sub_vs_2 = vertices(ptn, PartitionVertex(1)), vertices(ptn, PartitionVertex(2))
 

examples/distances.jl

@@ -14,4 +14,4 @@ t = dijkstra_tree(ψ, only(center(ψ)))
 @show a_star(ψ, (2, 1), (2, 5))
 @show mincut_partitions(ψ)
 @show mincut_partitions(ψ, (1, 1), (3, 5))
-@show NamedGraphs.partition_vertices(underlying_graph(ψ); npartitions=2)
+@show partitioned_vertices(underlying_graph(ψ); npartitions=2)

examples/gauging/gauging_itns.jl

@@ -31,12 +31,10 @@ function eager_gauging(
   mts = copy(mts)
   for e in edges(ψ)
     vsrc, vdst = src(e), dst(e)
-    pe = NamedGraphs.partition_edge(pψψ, NamedEdge((vsrc, 1) => (vdst, 1)))
+    pe = which_partitionedge(pψψ, NamedEdge((vsrc, 1) => (vdst, 1)))
     normalize!(isometries[e])
     normalize!(isometries[reverse(e)])
-    mts[pe], mts[PartitionEdge(reverse(NamedGraphs.parent(pe)))] = ITensorNetwork(
-      isometries[e]
-    ),
+    mts[pe], mts[reverse(pe)] = ITensorNetwork(isometries[e]),
     ITensorNetwork(isometries[reverse(e)])
   end
 

src/ITensorNetworks.jl

@@ -53,7 +53,8 @@ using NamedGraphs:
   parent_graph,
   vertex_to_parent_vertex,
   parent_vertices_to_vertices,
-  not_implemented
+  not_implemented,
+  parent
 
 include("imports.jl")
 

src/approx_itensornetwork/approx_itensornetwork.jl

@@ -69,7 +69,7 @@ function approx_itensornetwork(
   contraction_sequence_alg="optimal",
   contraction_sequence_kwargs=(;),
 )
-  par = partition(tn, inds_btree; alg="mincut_recursive_bisection")
+  par = _partition(tn, inds_btree; alg="mincut_recursive_bisection")
   output_tn, log_root_norm = approx_itensornetwork(
     alg,
     par;

src/approx_itensornetwork/binary_tree_partition.jl

@@ -59,7 +59,7 @@ Note: for a given binary tree with n indices, the output partition will contain
 Note: name of vertices in the output partition are the same as the name of vertices
   in `inds_btree`.
 """
-function partition(
+function _partition(
   ::Algorithm"mincut_recursive_bisection", tn::ITensorNetwork, inds_btree::DataGraph
 )
   @assert _is_rooted_directed_binary_tree(inds_btree)
@@ -115,7 +115,7 @@ function partition(
   end
   tn_deltas = ITensorNetwork(vcat(output_deltas_vector...))
   tn_scalars = ITensorNetwork(vcat(scalars_vector...))
-  par = partition(disjoint_union(out_tn, tn_deltas, tn_scalars), subgraph_vs)
+  par = _partition(disjoint_union(out_tn, tn_deltas, tn_scalars), subgraph_vs)
   @assert is_tree(par)
   name_map = Dict()
   for (i, v) in enumerate(pre_order_dfs_vertices(inds_btree, root))
@@ -124,6 +124,6 @@ function partition(
   return rename_vertices(par, name_map)
 end
 
-function partition(tn::ITensorNetwork, inds_btree::DataGraph; alg::String)
-  return partition(Algorithm(alg), tn, inds_btree)
+function _partition(tn::ITensorNetwork, inds_btree::DataGraph; alg::String)
+  return _partition(Algorithm(alg), tn, inds_btree)
 end

src/approx_itensornetwork/partition.jl

@@ -1,49 +1,26 @@
-"""
-    partition_vertices(g::AbstractGraph, subgraph_vertices::Vector)
-
-Given a graph (`g`) and groups of vertices defining subgraphs of that
-graph (`subgraph_vertices`), return a DataGraph storing the subgraph
-vertices on the vertices of the graph and with edges denoting
-which subgraphs of the original graph have edges connecting them, along with
-edge data storing the original edges that were connecting the subgraphs.
-"""
-function partition_vertices(g::AbstractGraph, subgraph_vertices)
-  partitioned_vertices = DataGraph(
-    NamedGraph(eachindex(subgraph_vertices)), Dictionary(subgraph_vertices)
+function _partition(g::AbstractGraph, subgraph_vertices)
+  partitioned_graph = DataGraph(
+    NamedGraph(eachindex(subgraph_vertices)),
+    map(vs -> subgraph(g, vs), Dictionary(subgraph_vertices)),
   )
   for e in edges(g)
-    s1 = findfirst_on_vertices(
-      subgraph_vertices -> src(e) ∈ subgraph_vertices, partitioned_vertices
-    )
-    s2 = findfirst_on_vertices(
-      subgraph_vertices -> dst(e) ∈ subgraph_vertices, partitioned_vertices
-    )
-    if (!has_edge(partitioned_vertices, s1, s2) && s1 ≠ s2)
-      add_edge!(partitioned_vertices, s1, s2)
-      partitioned_vertices[s1 => s2] = Vector{edgetype(g)}()
+    s1 = findfirst_on_vertices(subgraph -> src(e) ∈ vertices(subgraph), partitioned_graph)
+    s2 = findfirst_on_vertices(subgraph -> dst(e) ∈ vertices(subgraph), partitioned_graph)
+    if (!has_edge(partitioned_graph, s1, s2) && s1 ≠ s2)
+      add_edge!(partitioned_graph, s1, s2)
+      partitioned_graph[s1 => s2] = Dictionary(
+        [:edges, :edge_data],
+        [Vector{edgetype(g)}(), Dictionary{edgetype(g),edge_data_type(g)}()],
+      )
     end
-    if has_edge(partitioned_vertices, s1, s2)
-      push!(partitioned_vertices[s1 => s2], e)
+    if has_edge(partitioned_graph, s1, s2)
+      push!(partitioned_graph[s1 => s2][:edges], e)
+      if isassigned(g, e)
+        set!(partitioned_graph[s1 => s2][:edge_data], e, g[e])
+      end
     end
   end
-  return partitioned_vertices
-end
-
-"""
-    partition_vertices(g::AbstractGraph; npartitions, nvertices_per_partition, kwargs...)
-
-Given a graph `g`, partition the vertices of `g` into 'npartitions' partitions
-or into partitions with `nvertices_per_partition` vertices per partition.
-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.
-A graph partitioning backend such as Metis or KaHyPar needs to be installed for this function to work.
-"""
-function partition_vertices(
-  g::AbstractGraph; npartitions=nothing, nvertices_per_partition=nothing, kwargs...
-)
-  return partition_vertices(
-    g, subgraph_vertices(g; npartitions, nvertices_per_partition, kwargs...)
-  )
+  return partitioned_graph
 end
 
 """
@@ -74,109 +51,34 @@ function findfirst_on_edges(f::Function, graph::AbstractDataGraph)
   return findfirst(f, edge_data(graph))
 end
 
-function subgraphs(g::AbstractSimpleGraph, subgraph_vertices)
-  return subgraphs(NamedGraph(g), subgraph_vertices)
-end
-
-"""
-    subgraphs(g::AbstractGraph, subgraph_vertices)
-
-Return a collection of subgraphs of `g` defined by the subgraph
-vertices `subgraph_vertices`.
-"""
-function subgraphs(g::AbstractGraph, subgraph_vertices)
-  return map(vs -> subgraph(g, vs), subgraph_vertices)
-end
-
-"""
-    subgraphs(g::AbstractGraph; npartitions::Integer, kwargs...)
-
-Given a graph `g`, partition `g` into `npartitions` partitions
-or into partitions with `nvertices_per_partition` vertices per partition,
-returning a list of subgraphs.
-Try to keep all subgraphs the same size and minimise edges cut between them.
-A graph partitioning backend such as Metis or KaHyPar needs to be installed for this function to work.
-"""
-function subgraphs(
-  g::AbstractGraph; npartitions=nothing, nvertices_per_partition=nothing, kwargs...
-)
-  return subgraphs(g, subgraph_vertices(g; npartitions, nvertices_per_partition, kwargs...))
-end
-
-function partition(g::AbstractSimpleGraph, subgraph_vertices)
-  return partition(NamedGraph(g), subgraph_vertices)
-end
-
-function partition(g::AbstractGraph, subgraph_vertices)
-  partitioned_graph = DataGraph(
-    NamedGraph(eachindex(subgraph_vertices)), subgraphs(g, Dictionary(subgraph_vertices))
-  )
-  for e in edges(g)
-    s1 = findfirst_on_vertices(subgraph -> src(e) ∈ vertices(subgraph), partitioned_graph)
-    s2 = findfirst_on_vertices(subgraph -> dst(e) ∈ vertices(subgraph), partitioned_graph)
-    if (!has_edge(partitioned_graph, s1, s2) && s1 ≠ s2)
-      add_edge!(partitioned_graph, s1, s2)
-      partitioned_graph[s1 => s2] = Dictionary(
-        [:edges, :edge_data],
-        [Vector{edgetype(g)}(), Dictionary{edgetype(g),edge_data_type(g)}()],
-      )
-    end
-    if has_edge(partitioned_graph, s1, s2)
-      push!(partitioned_graph[s1 => s2][:edges], e)
-      if isassigned(g, e)
-        set!(partitioned_graph[s1 => s2][:edge_data], e, g[e])
-      end
-    end
-  end
-  return partitioned_graph
-end
-
-"""
-    partition(g::AbstractGraph; npartitions::Integer, kwargs...)
-    partition(g::AbstractGraph, subgraph_vertices)
-
-Given a graph `g`, partition `g` into `npartitions` partitions
-or into partitions with `nvertices_per_partition` vertices per partition.
-The partitioning tries to keep all subgraphs the same size and minimize
-edges cut between them.
-
-Alternatively, specify a desired partitioning with a collection of sugraph
-vertices.
-
-Returns a data graph where each vertex contains the corresponding subgraph as vertex data.
-The edges indicates which subgraphs are connected, and the edge data stores a dictionary
-with two fields. The field `:edges` stores a list of the edges of the original graph
-that were connecting the two subgraphs, and `:edge_data` stores a dictionary
-mapping edges of the original graph to the data living on the edges of the original
-graph, if it existed.
-
-Therefore, one should be able to extract that data and recreate the original
-graph from the results of `partition`.
-
-A graph partitioning backend such as Metis or KaHyPar needs to be installed for this function to work
-if the subgraph vertices aren't specified explicitly.
-"""
-function partition(
-  g::AbstractGraph;
-  npartitions=nothing,
-  nvertices_per_partition=nothing,
-  subgraph_vertices=nothing,
-  kwargs...,
-)
-  if count(isnothing, (npartitions, nvertices_per_partition, subgraph_vertices)) != 2
-    error(
-      "Error: Cannot give multiple/ no partitioning options. Please specify exactly one."
-    )
-  end
-
-  if isnothing(subgraph_vertices)
-    subgraph_vertices = NamedGraphs.partition_vertices(
-      g; npartitions, nvertices_per_partition, kwargs...
-    )
-  end
-
-  return partition(g, subgraph_vertices)
-end
+# function subgraphs(g::AbstractSimpleGraph, subgraph_vertices)
+#   return subgraphs(NamedGraph(g), subgraph_vertices)
+# end
+
+# """
+#     subgraphs(g::AbstractGraph, subgraph_vertices)
+
+# Return a collection of subgraphs of `g` defined by the subgraph
+# vertices `subgraph_vertices`.
+# """
+# function subgraphs(g::AbstractGraph, subgraph_vertices)
+#   return map(vs -> subgraph(g, vs), subgraph_vertices)
+# end
+
+# """
+#     subgraphs(g::AbstractGraph; npartitions::Integer, kwargs...)
+
+# Given a graph `g`, partition `g` into `npartitions` partitions
+# or into partitions with `nvertices_per_partition` vertices per partition,
+# returning a list of subgraphs.
+# Try to keep all subgraphs the same size and minimise edges cut between them.
+# A graph partitioning backend such as Metis or KaHyPar needs to be installed for this function to work.
+# """
+# function subgraphs(
+#   g::AbstractGraph; npartitions=nothing, nvertices_per_partition=nothing, kwargs...
+# )
+#   return subgraphs(g, subgraph_vertices(g; npartitions, nvertices_per_partition, kwargs...))
+# end
 
 """
   TODO: do we want to make it a public function?

src/beliefpropagation/beliefpropagation.jl

@@ -2,12 +2,12 @@ function message_tensors(
   ptn::PartitionedGraph; itensor_constructor=inds_e -> ITensor[dense(delta(inds_e))]
 )
   mts = Dict()
-  for e in PartitionEdge.(edges(partitioned_graph(ptn)))
+  for e in partitionedges(ptn)
     src_e_itn = unpartitioned_graph(subgraph(ptn, [src(e)]))
     dst_e_itn = unpartitioned_graph(subgraph(ptn, [dst(e)]))
     inds_e = commoninds(src_e_itn, dst_e_itn)
     mts[e] = itensor_constructor(inds_e)
-    mts[PartitionEdge(reverse(NamedGraphs.parent(e)))] = dag.(mts[e])
+    mts[reverse(e)] = dag.(mts[e])
   end
   return mts
 end
@@ -21,12 +21,11 @@ function update_message_tensor(
   mts;
   contract_kwargs=(; alg="density_matrix", output_structure=path_graph_structure, maxdim=1),
 )
-  incoming_messages = [
-    mts[PartitionEdge(e_in)] for e_in in setdiff(
-      boundary_edges(partitioned_graph(ptn), [NamedGraphs.parent(src(edge))]; dir=:in),
-      [reverse(NamedGraphs.parent(edge))],
-    )
-  ]
+  pedges = setdiff(
+    partitionedges(ptn, boundary_edges(ptn, vertices(ptn, src(edge)); dir=:in)),
+    [reverse(edge)],
+  )
+  incoming_messages = [mts[e_in] for e_in in pedges]
   incoming_messages = reduce(vcat, incoming_messages; init=ITensor[])
 
   contract_list = ITensor[
@@ -35,18 +34,14 @@ function update_message_tensor(
   ]
 
   if contract_kwargs.alg != "exact"
-    mt, _ = contract(ITensorNetwork(contract_list); contract_kwargs...)
+    mt = first(contract(ITensorNetwork(contract_list); contract_kwargs...))
   else
     mt = contract(
       contract_list; sequence=contraction_sequence(contract_list; alg="optimal")
     )
   end
 
-  if isa(mt, ITensor)
-    mt = ITensor[mt]
-  elseif isa(mt, ITensorNetwork)
-    mt = ITensor(mt)
-  end
+  mt = isa(mt, ITensor) ? ITensor[mt] : ITensor(mt)
   normalize!.(mt)
 
   return mt
@@ -152,19 +147,18 @@ Given a subet of vertices of a given Tensor Network and the Message Tensors for
 Specifically, the contraction of the environment tensors and tn[vertices] will be a scalar.
 """
 function get_environment(ptn::PartitionedGraph, mts, verts::Vector; dir=:in)
-  partition_vertices = unique([NamedGraphs.which_partition(ptn, v) for v in verts])
+  partition_verts = partitionvertices(ptn, verts)
+  central_verts = vertices(ptn, partition_verts)
 
   if dir == :out
     return get_environment(ptn, mts, setdiff(vertices(ptn), verts))
   end
 
-  env_tensors = [
-    mts[PartitionEdge(e)] for e in
-    boundary_edges(partitioned_graph(ptn), NamedGraphs.parent.(partition_vertices); dir=:in)
-  ]
+  pedges = partitionedges(ptn, boundary_edges(ptn, central_verts; dir=:in))
+  env_tensors = [mts[e] for e in pedges]
   env_tensors = reduce(vcat, env_tensors; init=ITensor[])
   central_tensors = ITensor[
-    (unpartitioned_graph(ptn))[v] for v in setdiff(vertices(ptn, partition_vertices), verts)
+    (unpartitioned_graph(ptn))[v] for v in setdiff(central_verts, verts)
   ]
 
   return vcat(env_tensors, central_tensors)