Drop _TreeTensorNetwork inner ctor; share unregister-inds helper

mtfishman · claude · mtfishman · commit dc22cf8b301a · 2026-05-12T13:53:36.000-04:00
- TreeTensorNetwork now uses the auto-generated all-fields constructor
  plus a single outer ctor that performs the is_tree check, matching
  the ITensorNetwork constructor design.
- Add ITensorNetwork{V}() empty ctor used as the seed for the
  tensor-collection constructor, replacing the explicit three-field
  call.
- Extract _unregister_inds! so _set_vertex_data! and _rem_vertex!
  share their reverse-map cleanup, without altering the in-place
  vertex_data update path.
- Build dictionaries via map(Indices(...)) in test/utils.jl.
- Fix the TreeTensorNetwork jldoctest to use a properly connected
  two-vertex example.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/itensornetwork.jl b/src/itensornetwork.jl
@@ -1,9 +1,9 @@
 using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data
 using Dictionaries: Dictionaries, Dictionary
-using Graphs: Graphs, add_edge!, add_vertex!, edges, has_edge, has_vertex, neighbors,
-    rem_edge!, rem_vertex!, vertices
-using ITensors: ITensors, ITensor, Index, inds
-using NamedGraphs: NamedGraphs, NamedEdge, NamedGraph, vertextype
+using Graphs:
+    Graphs, add_edge!, add_vertex!, has_edge, has_vertex, neighbors, rem_edge!, rem_vertex!
+using ITensors: ITensor, Index, inds
+using NamedGraphs: NamedGraphs, NamedGraph, vertextype
 
 """
     ITensorNetwork{V}
@@ -69,14 +69,22 @@ end
 # Constructors
 #
 
+# Empty network with no vertices, used as the starting point for the
+# tensor-collection constructor below.
+function ITensorNetwork{V}() where {V}
+    return ITensorNetwork{V}(
+        NamedGraph{V}(), Dictionary{V, ITensor}(), Dict{Index, Set{V}}()
+    )
+end
+
 # Construct by feeding `tensors` through `set_vertex_data!` one vertex
 # at a time — this centralizes the reverse-map registration, edge
 # inference, and hypergraph check in a single place (the `setindex!`
 # code path). Walking `keys(tensors)` in order makes the resulting
 # `neighbors(g, v)` / `edges(g)` iteration order deterministic in the
 # input order.
 function ITensorNetwork{V}(tensors) where {V}
-    tn = ITensorNetwork{V}(NamedGraph{V}(), Dictionary{V, ITensor}(), Dict{Index, Set{V}}())
+    tn = ITensorNetwork{V}()
     for v in keys(tensors)
         set_vertex_data!(tn, tensors[v], v)
     end
@@ -100,33 +108,43 @@ Base.copy(tn::ITensorNetwork) = ITensorNetwork(map(copy, vertex_data(tn)))
 # Mutation: keep `graph`, `vertex_data`, and `ind_to_vertices` in sync.
 #
 
+# Drop the inds of `vertex_data[v]` from the reverse map, leaving
+# `vertex_data` and `graph` themselves untouched. Used both as a
+# prelude to overwriting `v` and as a step in `_rem_vertex!`.
+function _unregister_inds!(
+        vertex_data::Dictionary{V, ITensor},
+        ind_to_vertices::Dict{Index, Set{V}},
+        v
+    ) where {V}
+    haskey(vertex_data, v) || return nothing
+    for i in inds(vertex_data[v])
+        owners = ind_to_vertices[i]
+        delete!(owners, v)
+        isempty(owners) && delete!(ind_to_vertices, i)
+    end
+    return nothing
+end
+
 # Write `value` to vertex `v`, updating the reverse map and reconciling
-# edges so the graph-edge ↔ shared-`Index` invariant holds. Cost is
-# O(deg(v) + |inds(value)|). If `v` isn't already in the network, it's
-# added — so this is also the natural way to grow the network one tensor
-# at a time without a separate `add_vertex!` step. Operates on raw
-# storage so `ITensorNetwork` and `TreeTensorNetwork` can share it.
+# edges so the graph-edge ↔ shared-`Index` invariant holds. If `v` is
+# new it's added to the graph — so this is also the natural way to grow
+# the network one tensor at a time without a separate `add_vertex!`
+# step. Operates on raw storage so `ITensorNetwork` and
+# `TreeTensorNetwork` can share it.
 function _set_vertex_data!(
         graph::NamedGraph{V},
         vertex_data::Dictionary{V, ITensor},
         ind_to_vertices::Dict{Index, Set{V}},
         value,
         v
     ) where {V}
-    # Add the vertex to the graph if it's new.
     has_vertex(graph, v) || add_vertex!(graph, v)
-    # Unregister old inds of `vertex_data[v]` from the reverse map.
-    if haskey(vertex_data, v)
-        for i in inds(vertex_data[v])
-            owners = ind_to_vertices[i]
-            delete!(owners, v)
-            isempty(owners) && delete!(ind_to_vertices, i)
-        end
-    end
-    # Write the new tensor. `Dictionaries.set!` inserts or updates;
-    # plain `setindex!` would error on a vertex not already in the dict.
+    _unregister_inds!(vertex_data, ind_to_vertices, v)
+    # `set!` updates in place when `v` is already present, preserving
+    # the insertion order of `vertex_data`. Plain `setindex!` would
+    # error on a missing key, and `insert!` would error on an existing
+    # one — `set!` handles both branches.
     Dictionaries.set!(vertex_data, v, value)
-    # Register new inds.
     for i in inds(value)
         push!(get!(ind_to_vertices, i, Set{V}()), v)
         length(ind_to_vertices[i]) <= 2 || error(
@@ -162,14 +180,8 @@ function _rem_vertex!(
         ind_to_vertices::Dict{Index, Set{V}},
         v
     ) where {V}
-    if haskey(vertex_data, v)
-        for i in inds(vertex_data[v])
-            owners = ind_to_vertices[i]
-            delete!(owners, v)
-            isempty(owners) && delete!(ind_to_vertices, i)
-        end
-        delete!(vertex_data, v)
-    end
+    _unregister_inds!(vertex_data, ind_to_vertices, v)
+    haskey(vertex_data, v) && delete!(vertex_data, v)
     rem_vertex!(graph, v)
     return nothing
 end
diff --git a/src/treetensornetworks/treetensornetwork.jl b/src/treetensornetworks/treetensornetwork.jl
@@ -1,6 +1,6 @@
 using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data
-using Dictionaries: Dictionaries, Dictionary, Indices
-using Graphs: Graphs, add_vertex!, has_vertex, is_tree, rem_vertex!, vertices
+using Dictionaries: Dictionary, Indices
+using Graphs: Graphs, is_tree, rem_vertex!, vertices
 using ITensors: ITensor, Index
 using NamedGraphs.GraphsExtensions: vertextype
 using NamedGraphs: NamedGraph
@@ -24,27 +24,6 @@ struct TreeTensorNetwork{V} <: AbstractTreeTensorNetwork{V}
     vertex_data::Dictionary{V, ITensor}
     ind_to_vertices::Dict{Index, Set{V}}
     ortho_region::Indices{V}
-    global function _TreeTensorNetwork(
-            graph::NamedGraph{V},
-            vertex_data::Dictionary{V, ITensor},
-            ind_to_vertices::Dict{Index, Set{V}},
-            ortho_region::Indices{V}
-        ) where {V}
-        @assert is_tree(graph)
-        return new{V}(graph, vertex_data, ind_to_vertices, ortho_region)
-    end
-end
-
-function _TreeTensorNetwork(tn::ITensorNetwork{V}, ortho_region::Indices{V}) where {V}
-    return _TreeTensorNetwork(tn.graph, tn.vertex_data, tn.ind_to_vertices, ortho_region)
-end
-
-function _TreeTensorNetwork(tn::ITensorNetwork{V}, ortho_region) where {V}
-    return _TreeTensorNetwork(tn, Indices{V}(ortho_region))
-end
-
-function _TreeTensorNetwork(tn::ITensorNetwork)
-    return _TreeTensorNetwork(tn, vertices(tn))
 end
 
 """
@@ -61,31 +40,27 @@ Throws an error if the underlying graph of `tn` is not a tree.
 # Example
 
 ```jldoctest
-julia> using NamedGraphs.NamedGraphGenerators: named_comb_tree
-
-julia> using NamedGraphs: NamedGraph
+julia> using ITensors: Index, ITensor
 
 julia> using Graphs: vertices
 
-julia> using ITensors: ITensor
+julia> i, j, k = Index(2, "i"), Index(2, "j"), Index(2, "k");
 
-julia> g = named_comb_tree((2, 2));
+julia> itn = ITensorNetwork([ITensor(i, j), ITensor(j, k)]);
 
-julia> s = siteinds("S=1/2", g);
-
-julia> tensors = Dict(v => ITensor(s[v]...) for v in vertices(g));
-
-julia> itn = ITensorNetwork(tensors);
-
-julia> ttn_state = TreeTensorNetwork(itn; ortho_region = [first(vertices(itn))]);
+julia> ttn = TreeTensorNetwork(itn; ortho_region = [first(vertices(itn))]);
 
 ```
 
 See also: [`ITensorNetwork`](@ref), [`orthogonalize`](@ref).
 """
-function TreeTensorNetwork(tn::ITensorNetwork; ortho_region = vertices(tn))
-    return _TreeTensorNetwork(tn, ortho_region)
+function TreeTensorNetwork(tn::ITensorNetwork{V}; ortho_region = vertices(tn)) where {V}
+    @assert is_tree(tn)
+    return TreeTensorNetwork{V}(
+        tn.graph, tn.vertex_data, tn.ind_to_vertices, Indices{V}(ortho_region)
+    )
 end
+
 function TreeTensorNetwork{V}(tn::ITensorNetwork) where {V}
     return TreeTensorNetwork(ITensorNetwork{V}(tn))
 end
@@ -132,24 +107,19 @@ function Graphs.rem_vertex!(tn::TTN, v)
     return tn
 end
 
-function Base.copy(tn::TTN)
-    V = vertextype(tn)
-    return _TreeTensorNetwork(
+function Base.copy(tn::TTN{V}) where {V}
+    return TreeTensorNetwork{V}(
         copy(tn.graph),
         map(copy, tn.vertex_data),
         Dict{Index, Set{V}}(i => copy(vs) for (i, vs) in tn.ind_to_vertices),
         copy(tn.ortho_region)
     )
 end
 
-#
-# Constructor
-#
-
 # set_ortho_region: low-level update of the ortho_region metadata only,
 # without any gauge transformations. To move the orthogonality center use orthogonalize.
-function set_ortho_region(tn::TTN, ortho_region)
-    return _TreeTensorNetwork(
-        tn.graph, tn.vertex_data, tn.ind_to_vertices, Indices{vertextype(tn)}(ortho_region)
+function set_ortho_region(tn::TTN{V}, ortho_region) where {V}
+    return TreeTensorNetwork{V}(
+        tn.graph, tn.vertex_data, tn.ind_to_vertices, Indices{V}(ortho_region)
     )
 end
diff --git a/test/utils.jl b/test/utils.jl
@@ -4,7 +4,7 @@
 # inside its gensym module.
 
 using DataGraphs: underlying_graph, vertex_data
-using Dictionaries: Dictionary
+using Dictionaries: Indices
 using Graphs: AbstractGraph, dst, edges, src, vertices
 using ITensorNetworks: ITensorNetwork, IndsNetwork
 using ITensors.NDTensors: dim
@@ -36,15 +36,15 @@ function random_tensornetwork(
     g = NamedGraph(graph)
     links = Dict(e => Index(link_space, "Link") for e in edges(g))
     links = merge(links, Dict(reverse(e) => links[e] for e in edges(g)))
-    # Use a `Dictionary` (insertion-ordered) so the constructed
-    # `ITensorNetwork`'s vertex / edge order tracks `vertices(g)`.
-    vs = collect(vertices(g))
-    ts = map(vs) do v
+    # `Indices`-keyed `map` returns a `Dictionary` (insertion-ordered),
+    # so the constructed `ITensorNetwork`'s vertex / edge order tracks
+    # `vertices(g)`.
+    ts = map(Indices(vertices(g))) do v
         link_v = [links[e] for e in incident_edges(g, v)]
         inds_v = [siteinds[v]; link_v]
         return itensor(randn(rng, eltype, dim.(inds_v)...), inds_v)
     end
-    return ITensorNetwork(Dictionary(vs, ts))
+    return ITensorNetwork(ts)
 end
 
 # `IndsNetwork`: extract site inds (`Index[]` where unassigned).
@@ -108,13 +108,12 @@ function productstate(elt::Type, state::Function, s::IndsNetwork)
     return productstate(elt, Dict(v => state(v) for v in vertices(s)), s)
 end
 function productstate(elt::Type, state, s::IndsNetwork)
-    vs = collect(vertices(s))
-    ts = map(vs) do v
+    ts = map(Indices(vertices(s))) do v
         site_v = isassigned(vertex_data(s), v) ? s[v] : Index[]
         t = ITensors.state(state[v], only(site_v))
         return ITensors.convert_eltype(elt, t)
     end
-    tn = ITensorNetwork(Dictionary(vs, ts))
+    tn = ITensorNetwork(ts)
     for e in edges(s)
         _add_edge!(elt, tn, e)
     end
