Add DataGraphsPartitionedGraphsExt glue for TensorNetwork type

Also includes some fixes to the way `TensorNetwork` types are
constructed based on index structure.

Author: Jack Dunham

src/tensornetwork.jl

@@ -1,10 +1,21 @@
 using Combinatorics: combinations
 using DataGraphs: DataGraphs, AbstractDataGraph, DataGraph
 using Dictionaries: AbstractDictionary, Indices, dictionary
-using Graphs: AbstractSimpleGraph
+using Graphs: AbstractSimpleGraph, rem_vertex!, rem_edge!
 using NamedDimsArrays: AbstractNamedDimsArray, dimnames
 using NamedGraphs: NamedGraphs, NamedEdge, NamedGraph, vertextype
-using NamedGraphs.GraphsExtensions: add_edges!, arrange_edge, arranged_edges, vertextype
+using NamedGraphs.GraphsExtensions: GraphsExtensions, arranged_edges, vertextype
+using NamedGraphs.PartitionedGraphs:
+    AbstractPartitionedGraph,
+    PartitionedGraphs,
+    departition,
+    partitioned_vertices,
+    partitionedgraph,
+    quotient_graph,
+    quotient_graph_type
+using .LazyNamedDimsArrays: lazy, Mul
+using DataGraphs: vertex_data_eltype, vertex_data, edge_data
+using DataGraphs.DataGraphsPartitionedGraphsExt
 
 function _TensorNetwork end
 
@@ -24,8 +35,14 @@ function _TensorNetwork(graph::AbstractGraph, tensors)
     return _TensorNetwork(graph, Dictionary(keys(tensors), values(tensors)))
 end
 
+function TensorNetwork{V, VD, UG, Tensors}(graph::UG) where {V, VD, UG <: AbstractGraph{V}, Tensors}
+    return _TensorNetwork(graph, Tensors())
+end
+
 DataGraphs.underlying_graph(tn::TensorNetwork) = getfield(tn, :underlying_graph)
 DataGraphs.vertex_data(tn::TensorNetwork) = getfield(tn, :tensors)
+DataGraphs.edge_data(tn::TensorNetwork) = Dictionary{edgetype(tn), Nothing}()
+DataGraphs.vertex_data_eltype(T::Type{<:TensorNetwork}) = eltype(fieldtype(T, :tensors))
 function DataGraphs.underlying_graph_type(type::Type{<:TensorNetwork})
     return fieldtype(type, :underlying_graph)
 end
@@ -70,7 +87,10 @@ function fix_links!(tn::AbstractTensorNetwork)
     for e in setdiff(arranged_edges(graph), tn_edges)
         insert_trivial_link!(tn, e)
     end
-    return tn
+    for edge in setdiff(arranged_edges(graph), arranged_edges(graph_structure))
+        insert_trivial_link!(network, edge)
+    end
+    return network
 end
 
 # Determine the graph structure from the tensors.
@@ -93,3 +113,56 @@ end
 
 NamedGraphs.convert_vertextype(::Type{V}, tn::TensorNetwork{V}) where {V} = tn
 NamedGraphs.convert_vertextype(V::Type, tn::TensorNetwork) = TensorNetwork{V}(tn)
+
+Graphs.connected_components(tn::TensorNetwork) = Graphs.connected_components(underlying_graph(tn))
+
+function Graphs.rem_edge!(tn::TensorNetwork, e)
+    if !has_edge(underlying_graph(tn), e)
+        return false
+    end
+    if !isempty(linkinds(tn, e))
+        throw(ArgumentError("cannot remove edge $e due to tensor indices existing on this edge."))
+    end
+    rem_edge!(underlying_graph(tn), e)
+    return true
+end
+
+function GraphsExtensions.graph_from_vertices(type::Type{<:TensorNetwork}, vertices)
+    DT = fieldtype(type, :tensors)
+    empty_dict = DT()
+    return TensorNetwork(similar_graph(underlying_graph_type(type), vertices), empty_dict)
+end
+
+## PartitionedGraphs
+function PartitionedGraphs.quotient_graph(tn::TensorNetwork)
+    ug = quotient_graph(underlying_graph(tn))
+    return TensorNetwork(ug, vertex_data(QuotientView(tn)))
+end
+function PartitionedGraphs.quotient_graph_type(type::Type{<:TensorNetwork})
+    UG = quotient_graph_type(underlying_graph_type(type))
+    VD = Vector{vertex_data_eltype(type)}
+    V = vertextype(UG)
+    return TensorNetwork{V, VD, UG, Dictionary{V, VD}}
+end
+
+function PartitionedGraphs.partitionedgraph(tn::TensorNetwork, parts)
+    pg = partitionedgraph(underlying_graph(tn), parts)
+    return TensorNetwork(pg, vertex_data(tn))
+end
+
+PartitionedGraphs.departition(tn::TensorNetwork) = tn
+function PartitionedGraphs.departition(
+        tn::TensorNetwork{<:Any, <:Any, <:AbstractPartitionedGraph}
+    )
+    return TensorNetwork(departition(underlying_graph(tn)), vertex_data(tn))
+end
+
+function DataGraphsPartitionedGraphsExt.to_quotient_vertex_data(::TensorNetwork, data)
+    return mapreduce(lazy, *, collect(last(data)))
+end
+
+function PartitionedGraphs.quotientview(tn::TensorNetwork)
+    qview = QuotientView(underlying_graph(tn))
+    tensors = vertex_data(QuotientView(tn))
+    return TensorNetwork(qview, tensors)
+end