Refactor TensorNetwork type; add NormNetwork struct.

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
-version = "0.7.1"
+version = "0.8.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -27,7 +27,7 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 Adapt = "4.3"
 AlgorithmsInterface = "0.1"
 Combinatorics = "1"
-DataGraphs = "0.4"
+DataGraphs = "0.5"
 Dictionaries = "0.4.5"
 Graphs = "1.13.1"
 ITensorBase = "0.8"

docs/Project.toml

@@ -10,5 +10,5 @@ path = ".."
 [compat]
 Documenter = "1"
 ITensorFormatter = "0.2.27"
-ITensorNetworksNext = "0.7"
+ITensorNetworksNext = "0.8"
 Literate = "2"

examples/Project.toml

@@ -5,4 +5,4 @@ ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 path = ".."
 
 [compat]
-ITensorNetworksNext = "0.7"
+ITensorNetworksNext = "0.8"

src/ITensorNetworkGenerators/delta_network.jl

@@ -1,5 +1,5 @@
-using ..ITensorNetworksNext: ITensorNetwork
-using Graphs: AbstractGraph
+using ..ITensorNetworksNext: tensornetwork
+using Graphs: AbstractGraph, vertices
 using ITensorBase: NamedUnitRange, name, nameddims, unnamed
 using NamedGraphs.GraphsExtensions: incident_edges
 
@@ -42,10 +42,11 @@ on each vertex. Link dimensions are defined using the function `f(e)` that shoul
 edge `e` as an input and should output the link index on that edge.
 """
 function delta_network(f, elt::Type, g::AbstractGraph)
-    return tn = ITensorNetwork(g) do v
+    tn = tensornetwork(vertices(g)) do v
         is = Tuple(f.(incident_edges(g, v)))
         return delta(elt, is)
     end
+    return tn
 end
 function delta_network(f, g::AbstractGraph)
     return delta_network(f, Float64, g)

src/ITensorNetworkGenerators/ising_network.jl

@@ -1,4 +1,4 @@
-using ..ITensorNetworksNext: @preserve_graph
+using ..ITensorNetworksNext
 using Graphs: degree, dst, edges, src
 using ITensorBase: apply, name, operator, uniquename
 using LinearAlgebra: Diagonal, eigen
@@ -29,20 +29,23 @@ function ising_network(
     )
     elt = typeof(β)
     l̃ = Dict(e => uniquename(f(e)) for e in edges(g))
-    f̃(e) = get(() -> l̃[reverse(e)], l̃, e)
-    tn = delta_network(f̃, elt, g)
+
+    fp(e) = get(() -> l̃[reverse(e)], l̃, e)
+    tn = delta_network(fp, elt, g)
     for v in sz_vertices
         tn[v] = diagonaltensor(elt[1, -1], axes(tn[v]))
     end
-    for e in edges(tn)
+
+    for e in edges(g)
         v1 = src(e)
         v2 = dst(e)
         deg1 = degree(tn, v1)
+
         deg2 = degree(tn, v2)
         m = sqrt_ising_bond(β; J, h, deg1, deg2)
-        t = operator(m, (name(f̃(e)),), (name(f(e)),))
-        @preserve_graph tn[v1] = apply(t, tn[v1])
-        @preserve_graph tn[v2] = apply(t, tn[v2])
+        t = operator(m, (name(fp(e)),), (name(f(e)),))
+        tn[v1] = apply(t, tn[v1])
+        tn[v2] = apply(t, tn[v2])
     end
     return tn
 end

src/ITensorNetworksNext.jl

@@ -3,7 +3,7 @@ module ITensorNetworksNext
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public apply_operator, apply_operators, beliefpropagation_normnetwork, identity_norm_message_env, normnetwork, norm_message_env, ones_norm_message_env, rand_norm_message_env, randn_norm_message_env, similar_norm_message_env"
+            "public apply_operator, apply_operators"
         )
     )
 end
@@ -12,6 +12,8 @@ include("select_algorithm.jl")
 include("AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl")
 include("abstracttensornetwork.jl")
 include("tensornetwork.jl")
+include("normnetwork.jl")
+include("normnetworkview.jl")
 include("ITensorNetworkGenerators/ITensorNetworkGenerators.jl")
 include("contract_network.jl")
 

src/abstracttensornetwork.jl

@@ -1,66 +1,54 @@
 using Adapt: Adapt, adapt
-using DataGraphs: DataGraphs, AbstractDataGraph, edge_data, set_vertex_data!,
-    underlying_graph, underlying_graph_type, vertex_data
+using DataGraphs: DataGraphs, AbstractDataGraph, AbstractVertexDataGraph, edge_data,
+    set_vertex_data!, underlying_graph, underlying_graph_type, vertex_data
 using Dictionaries: Dictionary
 using Graphs: Graphs, AbstractEdge, AbstractGraph, add_edge!, add_vertex!, dst, edges,
     edgetype, ne, neighbors, nv, rem_edge!, src, vertices
-using ITensorBase: dimnames, inds
+using ITensorBase: dimnames, inds, name, named, nametype, prime, uniquename, unnamedtype
 using LinearAlgebra: LinearAlgebra
 using MacroTools: @capture
 using NamedGraphs.GraphsExtensions: directed_graph, incident_edges, rem_edges!, vertextype
 using NamedGraphs.OrdinalIndexing: OrdinalSuffixedInteger
 using NamedGraphs: NamedGraphs, NamedGraph, not_implemented, similar_graph
+using TensorAlgebra: trivialrange
+
+abstract type AbstractITensorNetwork{T, V} <: AbstractVertexDataGraph{T, V} end
 
-abstract type AbstractITensorNetwork{V, VD} <: AbstractDataGraph{V, VD, Nothing} end
+# ====================================== Graphs.jl ======================================= #
 
 # Need to be careful about removing edges from tensor networks in case there is a bond
-Graphs.rem_edge!(::AbstractITensorNetwork, edge) = not_implemented()
+Graphs.rem_edge!(::AbstractITensorNetwork, _edge) = not_implemented()
 
-# Graphs.jl overloads
 function Graphs.weights(graph::AbstractITensorNetwork)
     V = vertextype(graph)
     es = Tuple.(edges(graph))
     ws = Dictionary{Tuple{V, V}, Float64}(es, undef)
     for e in edges(graph)
-        w = log2(dim(commoninds(graph, e)))
+        w = log2(dim(linkinds(graph, e)))
         ws[(src(e), dst(e))] = w
     end
     return ws
 end
 
-# Copy
-Base.copy(::AbstractITensorNetwork) = not_implemented()
-
-# Iteration
-Base.iterate(tn::AbstractITensorNetwork, args...) = iterate(vertex_data(tn), args...)
-Base.keys(tn::AbstractITensorNetwork) = vertices(tn)
-
-# TODO: This contrasts with the `DataGraphs.AbstractDataGraph` definition,
-# where it is defined as the `vertextype`. Does that cause problems or should it be changed?
-Base.eltype(tn::AbstractITensorNetwork) = eltype(vertex_data(tn))
-
 # Overload if needed
 Graphs.is_directed(::Type{<:AbstractITensorNetwork}) = false
 
-DataGraphs.underlying_graph(::AbstractITensorNetwork) = not_implemented()
-function NamedGraphs.vertex_positions(tn::AbstractITensorNetwork)
-    return NamedGraphs.vertex_positions(underlying_graph(tn))
-end
-function NamedGraphs.ordered_vertices(tn::AbstractITensorNetwork)
-    return NamedGraphs.ordered_vertices(underlying_graph(tn))
+# ==================================== NamedGraphs.jl ==================================== #
+
+function NamedGraphs.similar_graph(::AbstractITensorNetwork, VD::Type, vertices)
+    return ITensorNetwork{VD}(undef, collect(vertices))
 end
 
-function Adapt.adapt_structure(to, tn::AbstractITensorNetwork)
-    # TODO: Define and use:
-    #
-    # @preserve_graph map_vertex_data(adapt(to), tn)
-    #
-    # or just:
-    #
-    # @preserve_graph map(adapt(to), tn)
-    return map_vertex_data_preserve_graph(adapt(to), tn)
+# ==================================== DataGraphs.jl ===================================== #
+
+function DataGraphs.underlying_graph(tn::AbstractITensorNetwork)
+    ug = NamedGraph(vertices(tn))
+    add_edges!(ug, edges(tn))
+    return ug
 end
 
+# ====================================== interface ======================================= #
+
 linkinds(tn::AbstractGraph, edge::Pair) = linkinds(tn, edgetype(tn)(edge))
 # Pick the link indices from the `src` side, identified by name match with `dst`.
 # A range-strict intersection (`inds(src) ∩ inds(dst)`) would drop graded links
@@ -106,139 +94,40 @@ function sitenames(tn::AbstractGraph, v)
     return s
 end
 
-function setindex_preserve_graph!(tn::AbstractGraph, value, vertex)
-    set_vertex_data!(tn, value, vertex)
-    return tn
-end
+# Return the non-link vertices associated with an dim name
+function dimnamevertices(tn::AbstractGraph, name)
+    sites = vertextype(tn)[]
 
-function is_setindex!_expr(expr::Expr)
-    return is_assignment_expr(expr) && is_getindex_expr(first(expr.args))
-end
-is_setindex!_expr(x) = false
-is_getindex_expr(expr::Expr) = (expr.head === :ref)
-is_getindex_expr(x) = false
-is_assignment_expr(expr::Expr) = (expr.head === :(=))
-is_assignment_expr(expr) = false
-
-# TODO: Define this in terms of a function mapping
-# preserve_graph_function(::typeof(setindex!)) = setindex!_preserve_graph
-# preserve_graph_function(::typeof(map_vertex_data)) = map_vertex_data_preserve_graph
-# Also allow annotating codeblocks like `@views`.
-macro preserve_graph(expr)
-    if !is_setindex!_expr(expr)
-        error(
-            "preserve_graph must be used with setindex! syntax (as @preserve_graph a[i,j,...] = value)"
-        )
+    for v in vertices(tn)
+        if name ∈ dimnames(tn[v])
+            push!(sites, v)
+        end
     end
-    @capture(expr, array_[indices__] = value_)
-    return :(setindex_preserve_graph!($(esc(array)), $(esc(value)), $(esc.(indices)...)))
-end
 
-# Update the graph of the ITensorNetwork `tn` to include
-# edges that should exist based on the tensor connectivity.
-function add_missing_edges!(tn::AbstractGraph)
-    foreach(v -> add_missing_edges!(tn, v), vertices(tn))
-    return tn
+    return sites
 end
 
-# Update the graph of the ITensorNetwork `tn` to include
-# edges that should be incident to the vertex `v`
-# based on the tensor connectivity.
-function add_missing_edges!(tn::AbstractGraph, v)
-    for v′ in vertices(tn)
-        if v ≠ v′
-            e = v => v′
-            if !isempty(linknames(tn, e))
-                add_edge!(tn, e)
-            end
+function has_dimname(tn::AbstractGraph, name)
+    for v in vertices(tn)
+        if name ∈ dimnames(tn[v])
+            return true
         end
     end
-    return tn
+    return false
 end
 
-# Fix the edges of the ITensorNetwork `tn` to match
-# the tensor connectivity.
-function fix_edges!(tn::AbstractGraph)
-    foreach(v -> fix_edges!(tn, v), vertices(tn))
-    return tn
-end
-# Fix the edges of the ITensorNetwork `tn` to match
-# the tensor connectivity at vertex `v`.
-function fix_edges!(tn::AbstractGraph, v)
-    for e in incident_edges(tn, v)
-        # Remove an edge if there is no index on that edge.
-        if isempty(linknames(tn, e))
-            rem_edge!(tn, e)
-        end
-    end
-    add_missing_edges!(tn, v)
-    return tn
-end
+has_ind(tn::AbstractGraph, ind) = has_dimname(tn, name(ind))
 
-using ITensorBase: named, nametype, uniquename, unnamedtype
-using TensorAlgebra: trivialrange
-function insertlink!(tn, e)
-    add_edge!(tn, e)
+function insertlink!(tn::AbstractGraph, e)
     T = eltype(inds(tn[src(e)]))
-    l = named(trivialrange(unnamedtype(T)), uniquename(nametype(T)))
-    x = fill!(similar(tn[src(e)], (l,)), one(eltype(tn[src(e)])))
-    @preserve_graph tn[src(e)] = tn[src(e)] * x
-    @preserve_graph tn[dst(e)] = tn[dst(e)] * conj(x)
-    return tn
-end
 
-using ITensorBase: replacedimnames
-function randlinknames(tn)
-    new_tn = copy(tn)
-    for e in edges(new_tn)
-        u, v = src(e), dst(e)
-        for n in intersect(dimnames(new_tn[u]), dimnames(new_tn[v]))
-            n′ = uniquename(n)
-            new_tn[u] = replacedimnames(new_tn[u], n => n′)
-            new_tn[v] = replacedimnames(new_tn[v], n => n′)
-        end
-    end
-    return new_tn
-end
+    linkind = named(trivialrange(unnamedtype(T)), uniquename(nametype(T)))
 
-function Base.setindex!(tn::AbstractITensorNetwork, value, v)
-    @preserve_graph tn[v] = value
-    fix_edges!(tn, v)
-    return tn
-end
-# Fix ambiguity error.
-function Base.setindex!(
-        graph::AbstractITensorNetwork,
-        value,
-        vertex::OrdinalSuffixedInteger
-    )
-    graph[vertices(graph)[vertex]] = value
-    return graph
-end
-Base.setindex!(tn::AbstractITensorNetwork, value, edge::AbstractEdge) = not_implemented()
-Base.setindex!(tn::AbstractITensorNetwork, value, edge::Pair) = not_implemented()
-# Fix ambiguity error.
-function Base.setindex!(
-        tn::AbstractITensorNetwork,
-        value,
-        edge::Pair{<:OrdinalSuffixedInteger, <:OrdinalSuffixedInteger}
-    )
-    return not_implemented()
-end
+    x = similar(tn[src(e)], (linkind,))
+    fill!(x, true)
 
-function Base.show(io::IO, mime::MIME"text/plain", graph::AbstractITensorNetwork)
-    println(io, "$(typeof(graph)) with $(nv(graph)) vertices:")
-    show(io, mime, vertices(graph))
-    println(io, "\n")
-    println(io, "and $(ne(graph)) edge(s):")
-    for e in edges(graph)
-        show(io, mime, e)
-        println(io)
-    end
-    println(io)
-    println(io, "with vertex data:")
-    show(io, mime, axes.(vertex_data(graph)))
-    return nothing
-end
+    tn[src(e)] *= x
+    tn[dst(e)] *= conj(x)
 
-Base.show(io::IO, graph::AbstractITensorNetwork) = show(io, MIME"text/plain"(), graph)
+    return tn
+end

src/beliefpropagation/beliefpropagation.jl

@@ -5,7 +5,8 @@ using DataGraphs: edge_data
 using Graphs: AbstractEdge, edges, edgetype, has_edge, vertices
 using ITensorBase: AbstractITensor
 using LinearAlgebra: norm, normalize
-using NamedGraphs.GraphsExtensions: add_edges!, boundary_edges, subgraph
+using NamedGraphs.GraphsExtensions:
+    add_edges!, boundary_edges, forest_cover_edge_sequence, subgraph
 using NamedGraphs.PartitionedGraphs: quotientvertices
 
 # === Top-level user entry point ===
@@ -65,7 +66,10 @@ an explicit `AlgorithmsInterface.StoppingCriterion`.
 `message_update_algorithm` controls how a single message is recomputed
 from its incoming neighbours.
 """
-function beliefpropagation(
+function beliefpropagation(factors, messages; kwargs...)
+    return _beliefpropagation(factors, messages; kwargs...)
+end
+function _beliefpropagation(
         factors, messages;
         edges = default_beliefpropagation_edges(factors),
         stopping_criterion = nothing,