Introduce _contract_deltas

src/Graphs/abstractgraph.jl

@@ -18,11 +18,6 @@ function internal_edges(g::AbstractGraph)
   return filter(e -> !is_leaf_edge(g, e), edges(g))
 end
 
-"""Get all vertices which are leaves of a graph"""
-function leaf_vertices(g::AbstractGraph)
-  return vertices(g)[findall(==(1), [is_leaf(g, v) for v in vertices(g)])]
-end
-
 """Get distance of a vertex from a leaf"""
 function distance_to_leaf(g::AbstractGraph, v)
   leaves = leaf_vertices(g)

src/ITensorNetworks.jl

@@ -83,6 +83,7 @@ include("models.jl")
 include("tebd.jl")
 include("itensornetwork.jl")
 include("mincut.jl")
+include("contract_deltas.jl")
 include("binary_tree_partition.jl")
 include("utility.jl")
 include("specialitensornetworks.jl")

src/binary_tree_partition.jl

@@ -1,33 +1,3 @@
-"""
-Rewrite of the function
-  `DataStructures.root_union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}`.
-"""
-function _introot_union!(s::DataStructures.IntDisjointSets, x, y; left_root=true)
-  parents = s.parents
-  rks = s.ranks
-  @inbounds xrank = rks[x]
-  @inbounds yrank = rks[y]
-  if !left_root
-    x, y = y, x
-  end
-  @inbounds parents[y] = x
-  s.ngroups -= 1
-  return x
-end
-
-"""
-Rewrite of the function `DataStructures.root_union!(s::DisjointSet{T}, x::T, y::T)`.
-The difference is that in the output of `_root_union!`, x is guaranteed to be the root of y when
-setting `left_root=true`, and y will be the root of x when setting `left_root=false`.
-In `DataStructures.root_union!`, the root value cannot be specified.
-A specified root is useful in functions such as `_remove_deltas`, where when we union two
-indices into one disjointset, we want the index that is the outinds if the given tensor network
-to always be the root in the DisjointSets.
-"""
-function _root_union!(s::DisjointSets, x, y; left_root=true)
-  return s.revmap[_introot_union!(s.internal, s.intmap[x], s.intmap[y]; left_root=true)]
-end
-
 """
 Return the root vertex of a directed tree data graph
 """
@@ -56,98 +26,31 @@ Check if a data graph is a directed binary tree
 end
 
 """
-Partition the input network containing both `tn` and `deltas` (a vector of delta tensors)
-into two partitions, one adjacent to source_inds and the other adjacent to other external
-inds of the network.
+Partition the input network `tn` into two partitions, one adjacent to source_inds
+and the other adjacent to other external inds of the network.
 """
-function _binary_partition(
-  tn::ITensorNetwork, deltas::Vector{ITensor}, source_inds::Vector{<:Index}
-)
-  all_tensors = [Vector{ITensor}(tn)..., deltas...]
-  external_inds = noncommoninds(all_tensors...)
+function _binary_partition(tn::ITensorNetwork, source_inds::Vector{<:Index})
+  external_inds = noncommoninds(Vector{ITensor}(tn)...)
   # add delta tensor to each external ind
   external_sim_ind = [sim(ind) for ind in external_inds]
+  tn = map_data(t -> replaceinds(t, external_inds => external_sim_ind), tn; edges=[])
+  tn_wo_deltas = rename_vertices(v -> v[1], subgraph(v -> v[2] == 1, tn))
+  deltas = Vector{ITensor}(subgraph(v -> v[2] == 2, tn))
   new_deltas = [
     delta(external_inds[i], external_sim_ind[i]) for i in 1:length(external_inds)
   ]
-  deltas = map(t -> replaceinds(t, external_inds => external_sim_ind), deltas)
   deltas = [deltas..., new_deltas...]
-  tn = map_data(t -> replaceinds(t, external_inds => external_sim_ind), tn; edges=[])
+  tn = disjoint_union(tn_wo_deltas, ITensorNetwork(deltas))
   p1, p2 = _mincut_partition_maxweightoutinds(
-    disjoint_union(tn, ITensorNetwork(deltas)),
-    source_inds,
-    setdiff(external_inds, source_inds),
+    tn, source_inds, setdiff(external_inds, source_inds)
   )
-  tn_vs = [v[1] for v in p1 if v[2] == 1]
-  source_tn = subgraph(tn, tn_vs)
-  delta_indices = [v[1] for v in p1 if v[2] == 2]
-  source_deltas = Vector{ITensor}([deltas[i] for i in delta_indices])
-  source_tn, source_deltas = _remove_deltas(source_tn, source_deltas)
-  tn_vs = [v[1] for v in p2 if v[2] == 1]
-  remain_tn = subgraph(tn, tn_vs)
-  delta_indices = [v[1] for v in p2 if v[2] == 2]
-  remain_deltas = Vector{ITensor}([deltas[i] for i in delta_indices])
-  remain_tn, remain_deltas = _remove_deltas(remain_tn, remain_deltas)
+  source_tn = _contract_deltas(subgraph(tn, p1))
+  remain_tn = _contract_deltas(subgraph(tn, p2))
   @assert (
-    length(noncommoninds(all_tensors...)) == length(
-      noncommoninds(
-        Vector{ITensor}(source_tn)...,
-        source_deltas...,
-        Vector{ITensor}(remain_tn)...,
-        remain_deltas...,
-      ),
-    )
-  )
-  return source_tn, source_deltas, remain_tn, remain_deltas
-end
-
-"""
-Given an input tensor network containing tensors in the input `tn`` and
-tensors in `deltas``, remove redundent delta tensors in `deltas` and change
-inds accordingly to make the output `tn` and `out_deltas` represent the same
-tensor network but with less delta tensors.
-Note: inds of tensors in `tn` and `deltas` may be changed, and `out_deltas`
-  may still contain necessary delta tensors.
-
-========
-Example:
-  julia> is = [Index(2, "i") for i in 1:6]
-  julia> a = ITensor(is[1], is[2])
-  julia> b = ITensor(is[2], is[3])
-  julia> delta1 = delta(is[3], is[4])
-  julia> delta2 = delta(is[5], is[6])
-  julia> tn = ITensorNetwork([a,b])
-  julia> tn, out_deltas = ITensorNetworks._remove_deltas(tn, [delta1, delta2])
-  julia> noncommoninds(Vector{ITensor}(tn)...)
-  2-element Vector{Index{Int64}}:
-   (dim=2|id=339|"1")
-   (dim=2|id=489|"4")
-  julia> length(out_deltas)
-  1
-"""
-function _remove_deltas(tn::ITensorNetwork, deltas::Vector{ITensor})
-  out_delta_inds = Vector{Pair}()
-  network = [Vector{ITensor}(tn)..., deltas...]
-  outinds = noncommoninds(network...)
-  inds_list = map(t -> collect(inds(t)), deltas)
-  deltainds = collect(Set(vcat(inds_list...)))
-  ds = DisjointSets(deltainds)
-  for t in deltas
-    i1, i2 = inds(t)
-    if find_root!(ds, i1) in outinds && find_root!(ds, i2) in outinds
-      push!(out_delta_inds, find_root!(ds, i1) => find_root!(ds, i2))
-    end
-    if find_root!(ds, i1) in outinds
-      _root_union!(ds, find_root!(ds, i1), find_root!(ds, i2))
-    else
-      _root_union!(ds, find_root!(ds, i2), find_root!(ds, i1))
-    end
-  end
-  tn = map_data(
-    t -> replaceinds(t, deltainds => [find_root!(ds, i) for i in deltainds]), tn; edges=[]
+    length(external_inds) ==
+    length(noncommoninds(Vector{ITensor}(source_tn)..., Vector{ITensor}(remain_tn)...))
   )
-  out_deltas = Vector{ITensor}([delta(i.first, i.second) for i in out_delta_inds])
-  return tn, out_deltas
+  return source_tn, remain_tn
 end
 
 """
@@ -170,36 +73,30 @@ function partition(
   @assert _is_directed_binary_tree(inds_btree)
   output_tns = Vector{ITensorNetwork}()
   output_deltas_vector = Vector{Vector{ITensor}}()
-  # Mapping each vertex of the binary tree to a tn and a vector of deltas
-  # representing the partition of the subtree containing this vertex and
-  # its descendant vertices.
+  # Mapping each vertex of the binary tree to a tn representing the partition
+  # of the subtree containing this vertex and its descendant vertices.
   leaves = leaf_vertices(inds_btree)
   root = _root(inds_btree)
-  v_to_subtree_tn_deltas = Dict{vertextype(inds_btree),Tuple}()
-  v_to_subtree_tn_deltas[root] = (tn, Vector{ITensor}())
+  v_to_subtree_tn = Dict{vertextype(inds_btree),ITensorNetwork}()
+  v_to_subtree_tn[root] = disjoint_union(tn, ITensorNetwork())
   for v in pre_order_dfs_vertices(inds_btree, root)
-    @assert haskey(v_to_subtree_tn_deltas, v)
-    input_tn, input_deltas = v_to_subtree_tn_deltas[v]
-    if is_leaf(inds_btree, v)
-      push!(output_tns, input_tn)
-      push!(output_deltas_vector, input_deltas)
-      continue
+    @assert haskey(v_to_subtree_tn, v)
+    input_tn = v_to_subtree_tn[v]
+    if !is_leaf(inds_btree, v)
+      c1, c2 = child_vertices(inds_btree, v)
+      descendant_c1 = pre_order_dfs_vertices(inds_btree, c1)
+      indices = [inds_btree[l] for l in intersect(descendant_c1, leaves)]
+      tn1, input_tn = _binary_partition(input_tn, indices)
+      v_to_subtree_tn[c1] = tn1
+      descendant_c2 = pre_order_dfs_vertices(inds_btree, c2)
+      indices = [inds_btree[l] for l in intersect(descendant_c2, leaves)]
+      tn1, input_tn = _binary_partition(input_tn, indices)
+      v_to_subtree_tn[c2] = tn1
     end
-    c1, c2 = child_vertices(inds_btree, v)
-    descendant_c1 = pre_order_dfs_vertices(inds_btree, c1)
-    indices = [inds_btree[l] for l in intersect(descendant_c1, leaves)]
-    tn1, deltas1, input_tn, input_deltas = _binary_partition(
-      input_tn, input_deltas, indices
-    )
-    v_to_subtree_tn_deltas[c1] = (tn1, deltas1)
-    descendant_c2 = pre_order_dfs_vertices(inds_btree, c2)
-    indices = [inds_btree[l] for l in intersect(descendant_c2, leaves)]
-    tn1, deltas1, input_tn, input_deltas = _binary_partition(
-      input_tn, input_deltas, indices
-    )
-    v_to_subtree_tn_deltas[c2] = (tn1, deltas1)
-    push!(output_tns, input_tn)
-    push!(output_deltas_vector, input_deltas)
+    tn = rename_vertices(u -> u[1], subgraph(u -> u[2] == 1, input_tn))
+    deltas = Vector{ITensor}(subgraph(u -> u[2] == 2, input_tn))
+    push!(output_tns, tn)
+    push!(output_deltas_vector, deltas)
   end
   # In subgraph_vertices, each element is a vector of vertices to be
   # grouped in one partition.

src/contract_deltas.jl

@@ -0,0 +1,160 @@
+"""
+Rewrite of the function
+  `DataStructures.root_union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}`.
+"""
+function _introot_union!(s::DataStructures.IntDisjointSets, x, y; left_root=true)
+  parents = s.parents
+  rks = s.ranks
+  @inbounds xrank = rks[x]
+  @inbounds yrank = rks[y]
+  if !left_root
+    x, y = y, x
+  end
+  @inbounds parents[y] = x
+  s.ngroups -= 1
+  return x
+end
+
+"""
+Rewrite of the function `DataStructures.root_union!(s::DisjointSet{T}, x::T, y::T)`.
+The difference is that in the output of `_root_union!`, x is guaranteed to be the root of y when
+setting `left_root=true`, and y will be the root of x when setting `left_root=false`.
+In `DataStructures.root_union!`, the root value cannot be specified.
+A specified root is useful in functions such as `_remove_deltas`, where when we union two
+indices into one disjointset, we want the index that is the outinds if the given tensor network
+to always be the root in the DisjointSets.
+"""
+function _root_union!(s::DisjointSets, x, y; left_root=true)
+  return s.revmap[_introot_union!(s.internal, s.intmap[x], s.intmap[y]; left_root=true)]
+end
+
+"""
+Given a list of delta tensors `deltas`, return a `DisjointSets` of all its indices
+such that each pair of indices adjacent to any delta tensor must be in the same disjoint set.
+If a disjoint set contains indices in `rootinds`, then one of such indices in `rootinds`
+must be the root of this set.
+"""
+function _delta_inds_disjointsets(deltas::Vector{<:ITensor}, rootinds::Vector{<:Index})
+  inds_list = map(t -> collect(inds(t)), deltas)
+  deltainds = collect(Set(vcat(inds_list...)))
+  ds = DisjointSets(deltainds)
+  for t in deltas
+    i1, i2 = inds(t)
+    if find_root!(ds, i1) in rootinds
+      _root_union!(ds, find_root!(ds, i1), find_root!(ds, i2))
+    else
+      _root_union!(ds, find_root!(ds, i2), find_root!(ds, i1))
+    end
+  end
+  return ds
+end
+
+"""
+Given an input tensor network `tn`, remove redundent delta tensors
+in `tn` and change inds accordingly to make the output `tn` represent
+the same tensor network but with less delta tensors.
+
+========
+Example:
+  julia> is = [Index(2, string(i)) for i in 1:6]
+  julia> a = ITensor(is[1], is[2])
+  julia> b = ITensor(is[2], is[3])
+  julia> delta1 = delta(is[3], is[4])
+  julia> delta2 = delta(is[5], is[6])
+  julia> tn = ITensorNetwork([a, b, delta1, delta2])
+  julia> ITensorNetworks._contract_deltas(tn)
+  ITensorNetwork{Int64} with 3 vertices:
+  3-element Vector{Int64}:
+   1
+   2
+   4
+
+  and 1 edge(s):
+  1 => 2
+
+  with vertex data:
+  3-element Dictionaries.Dictionary{Int64, Any}
+   1 │ ((dim=2|id=457|"1"), (dim=2|id=296|"2"))
+   2 │ ((dim=2|id=296|"2"), (dim=2|id=613|"4"))
+   4 │ ((dim=2|id=626|"6"), (dim=2|id=237|"5"))
+"""
+function _contract_deltas(tn::ITensorNetwork)
+  tn = copy(tn)
+  network = Vector{ITensor}(tn)
+  deltas = filter(t -> is_delta(t), network)
+  outinds = noncommoninds(network...)
+  ds = _delta_inds_disjointsets(deltas, outinds)
+  deltainds = [ds...]
+  sim_deltainds = [find_root!(ds, i) for i in deltainds]
+  # `rem_vertex!(tn, v)` changes `vertices(tn)` in place.
+  # We copy it here so that the enumeration won't be affected.
+  for v in copy(vertices(tn))
+    if !is_delta(tn[v])
+      tn[v] = replaceinds(tn[v], deltainds, sim_deltainds)
+      continue
+    end
+    i1, i2 = inds(tn[v])
+    root = find_root!(ds, i1)
+    @assert root === find_root!(ds, i2)
+    if i1 != root && i1 in outinds
+      tn[v] = delta(i1, root)
+    elseif i2 != root && i2 in outinds
+      tn[v] = delta(i2, root)
+    else
+      rem_vertex!(tn, v)
+    end
+  end
+  return tn
+end
+
+"""
+  TODO: do we want to make it a public function?
+"""
+function _noncommoninds(partition::DataGraph)
+  networks = [Vector{ITensor}(partition[v]) for v in vertices(partition)]
+  network = vcat(networks...)
+  return noncommoninds(network...)
+end
+
+"""
+Given an input `partition`, contract redundent delta tensors in `partition`
+without changing the tensor network value. `root` is the root of the
+dfs_tree that defines the leaves.
+Note: for each vertex `v` of `partition`, the number of non-delta tensors
+  in `partition[v]` will not be changed.
+Note: only delta tensors of non-leaf vertices will be contracted.
+"""
+function _contract_deltas(partition::DataGraph; root=1)
+  partition = copy(partition)
+  leaves = leaf_vertices(dfs_tree(partition, root))
+  # We only remove deltas in non-leaf vertices
+  nonleaf_vertices = setdiff(vertices(partition), leaves)
+  rootinds = _noncommoninds(subgraph(partition, nonleaf_vertices))
+  all_deltas = mapreduce(
+    tn_v -> [
+      partition[tn_v][v] for v in vertices(partition[tn_v]) if is_delta(partition[tn_v][v])
+    ],
+    vcat,
+    nonleaf_vertices,
+  )
+  if length(all_deltas) == 0
+    return partition
+  end
+  ds = _delta_inds_disjointsets(Vector{ITensor}(all_deltas), rootinds)
+  deltainds = [ds...]
+  sim_deltainds = [find_root!(ds, ind) for ind in deltainds]
+  for tn_v in nonleaf_vertices
+    tn = partition[tn_v]
+    nondelta_vertices = [v for v in vertices(tn) if !is_delta(tn[v])]
+    tn = subgraph(tn, nondelta_vertices)
+    partition[tn_v] = map_data(t -> replaceinds(t, deltainds, sim_deltainds), tn; edges=[])
+  end
+  # Note: we also need to change inds in the leaves since they can be connected by deltas
+  # in nonleaf vertices
+  for tn_v in leaves
+    partition[tn_v] = map_data(
+      t -> replaceinds(t, deltainds, sim_deltainds), partition[tn_v]; edges=[]
+    )
+  end
+  return partition
+end

src/itensors.jl

@@ -41,3 +41,9 @@ trivial_space(x::Index) = _trivial_space(x)
 trivial_space(x::Vector{<:Index}) = _trivial_space(x)
 trivial_space(x::ITensor) = trivial_space(inds(x))
 trivial_space(x::Tuple{Vararg{Index}}) = trivial_space(first(x))
+
+is_delta(it::ITensor) = is_delta(ITensors.tensor(it))
+is_delta(t::ITensors.Tensor) = false
+function is_delta(t::ITensors.NDTensors.UniformDiagTensor)
+  return isone(ITensors.NDTensors.getdiagindex(t, 1))
+end