Validate NormNetwork bra names at construction, add layer accessors, interface improvements.

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
-version = "0.11.2"
+version = "0.12.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -23,6 +23,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [weakdeps]
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
@@ -48,4 +49,5 @@ Random = "1.10"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
 TensorAlgebra = "0.23.1"
+UUIDs = "1.10"
 julia = "1.10"

docs/Project.toml

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

examples/Project.toml

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

src/ITensorNetworkGenerators/ising_network.jl

@@ -1,6 +1,6 @@
 using ..ITensorNetworksNext
 using Graphs: degree, dst, edges, src
-using ITensorBase: NamedTensor, name, uniquename
+using ITensorBase: NamedTensor, inds, name, uniquename
 using LinearAlgebra: Diagonal, eigen
 using NamedGraphs: vertextype
 
@@ -33,7 +33,7 @@ function ising_network(
     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]))
+        tn[v] = diagonaltensor(elt[1, -1], axes(tn, v))
     end
 
     for e in edges(g)

src/abstracttensornetwork.jl

@@ -3,8 +3,8 @@ using DataGraphs: DataGraphs, AbstractDataGraph, AbstractVertexDataGraph, edge_d
     set_vertex_data!, underlying_graph, underlying_graph_type, vertex_data
 using Graphs: Graphs, AbstractEdge, AbstractGraph, add_edge!, add_vertex!, dst, edges,
     edgetype, ne, neighbors, nv, rem_edge!, src, vertices
-using ITensorBase: ITensorOperator, NamedUnitRange, inds, inputnames, name, names, nametype,
-    prime, uniquename, unnamedtype
+using ITensorBase: ITensorBase, ITensorOperator, NamedUnitRange, inds, inputnames, name,
+    names, nametype, prime, uniquename, unnamedtype
 using LinearAlgebra: LinearAlgebra
 using MacroTools: @capture
 using NamedGraphs:
@@ -37,67 +37,65 @@ end
 
 # ====================================== interface ======================================= #
 
+ITensorBase.inds(tn::AbstractITensorNetwork, v) = inds(tn[v])
+ITensorBase.axes(tn::AbstractITensorNetwork, v) = axes(tn[v])
+ITensorBase.names(tn::AbstractITensorNetwork, v) = names(tn[v])
+
 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
 # whose two endpoints carry dual-related ranges.
 function linkinds(tn::AbstractGraph, edge::AbstractEdge)
     ln = linknames(tn, edge)
-    return [i for i in inds(tn[src(edge)]) if name(i) in ln]
+    return [i for i in inds(tn, src(edge)) if name(i) in ln]
 end
 
 function linkaxes(tn::AbstractGraph, edge::Pair)
     return linkaxes(tn, edgetype(tn)(edge))
 end
 function linkaxes(tn::AbstractGraph, edge::AbstractEdge)
     ln = linknames(tn, edge)
-    return [ax for ax in axes(tn[src(edge)]) if name(ax) in ln]
+    return Tuple(ax for ax in axes(tn, src(edge)) if name(ax) in ln)
 end
 function linknames(tn::AbstractGraph, edge::Pair)
     return linknames(tn, edgetype(tn)(edge))
 end
 function linknames(tn::AbstractGraph, edge::AbstractEdge)
-    return names(tn[src(edge)]) ∩ names(tn[dst(edge)])
+    dstnames = names(tn, dst(edge))
+    return [n for n in names(tn, src(edge)) if n in dstnames]
 end
+eachlinkname(tn::AbstractGraph) = Iterators.flatten(linknames(tn, e) for e in edges(tn))
 
 function siteinds(tn::AbstractGraph, v)
-    s = inds(tn[v])
+    s = inds(tn, v)
     for v′ in neighbors(tn, v)
-        s = setdiff(s, inds(tn[v′]))
+        s = setdiff(s, inds(tn, v′))
     end
     return s
 end
 function siteaxes(tn::AbstractGraph, v)
-    s = axes(tn[v])
+    s = axes(tn, v)
     for v′ in neighbors(tn, v)
-        s = setdiff(s, axes(tn[v′]))
+        s = setdiff(s, axes(tn, v′))
     end
-    return s
+    return Tuple(s)
 end
 function sitenames(tn::AbstractGraph, v)
-    s = names(tn[v])
+    s = names(tn, v)
     for v′ in neighbors(tn, v)
-        s = setdiff(s, names(tn[v′]))
+        s = setdiff(s, names(tn, v′))
     end
     return s
 end
 
 # Return the non-link vertices associated with an dim name
 function dimnamevertices(tn::AbstractGraph, name)
-    sites = vertextype(tn)[]
-
-    for v in vertices(tn)
-        if name ∈ names(tn[v])
-            push!(sites, v)
-        end
-    end
-
-    return sites
+    return Indices(filter(v -> name ∈ names(tn, v), vertices(tn)))
 end
 
 function has_dimname(tn::AbstractGraph, name)
     for v in vertices(tn)
-        if name ∈ names(tn[v])
+        if name ∈ names(tn, v)
             return true
         end
     end
@@ -107,7 +105,7 @@ end
 has_ind(tn::AbstractGraph, ind) = has_dimname(tn, name(ind))
 
 function insertlink!(tn::AbstractGraph, e)
-    T = eltype(inds(tn[src(e)]))
+    T = eltype(inds(tn, src(e)))
 
     linkind = NamedUnitRange(trivialrange(unnamedtype(T)), uniquename(nametype(T)))
 

src/beliefpropagation/beliefpropagation.jl

@@ -272,9 +272,9 @@ end
 # the message is positive semidefinite and its trace is a positive normalization.
 function message_update!(algorithm::SimpleMessageUpdate, cache, factors::NormNetwork, edge)
     new_tensor = updated_message(algorithm, cache, factors, edge)
-    new_message = operator(
-        new_tensor, linknames(BraView(factors), edge), linknames(KetView(factors), edge)
-    )
+    branames = linknames(branetwork(factors), edge)
+    ketnames = linknames(ketnetwork(factors), edge)
+    new_message = operator(new_tensor, branames, ketnames)
     if algorithm.normalize
         message_norm = tr(new_message)
         iszero(message_norm) || (new_message /= message_norm)

src/beliefpropagation/messagecache.jl

@@ -191,16 +191,16 @@ bethe_free_energy(factors, messages) = -bethe_free_entropy(factors, messages)
 function similar_message_environment(nn::NormNetwork)
     messages = mapmany(vertices(nn)) do vertex
         return map(in_incident_edges(nn, vertex)) do edge
-            braview = BraView(nn)
-            ketview = KetView(nn)
+            bra = branetwork(nn)
+            ket = ketnetwork(nn)
 
-            ketnames = linknames(ketview, edge)
-            branames = linknames(braview, edge)
-            braaxis = unnamed.(linkaxes(braview, edge))
+            ketnames = linknames(ket, edge)
+            branames = linknames(bra, edge)
+            braaxis = unnamed.(linkaxes(bra, edge))
 
             # Bra leg = operator output, ket leg = input, the bipartition in which the message
             # is positive semidefinite.
-            message = similar_operator(ketview[vertex], braaxis, branames, ketnames)
+            message = similar_operator(ket[vertex], braaxis, branames, ketnames)
 
             return edge => message
         end

src/normnetwork.jl

@@ -1,5 +1,5 @@
-using Dictionaries: Dictionary, dictionary
-using ITensorBase: rename, setname, similar_operator, uniquename
+using Dictionaries: Dictionary, Indices, dictionary
+using ITensorBase: ITensorBase, conj, nametype, rename, similar_operator
 using ITensorNetworksNext
 
 """
@@ -16,14 +16,24 @@ struct NormNetwork{T, V, I} <: AbstractITensorNetwork{T, V}
             ket::ITensorNetwork{T, V, I},
             map::Dictionary{I, I}
         ) where {T, V, I}
-        braname = Dictionary{I, I}()
-        for (name, vertices) in pairs(ket.dimname_vertices)
-            if length(vertices) == 2
-                insert!(braname, name, map[name])
-            end
+        check_input(NormNetwork, ket, map)
+        return new{T, V, I}(ket, map)
+    end
+end
+
+function check_input(::Type{<:NormNetwork}, ket, braname)
+    if !issubset(eachlinkname(ket), keys(braname))
+        error("every link name of the ket layer must have a bra name.")
+    end
+    if !allunique(braname)
+        error("the bra names of a norm network must be distinct.")
+    end
+    for name in braname
+        if has_dimname(ket, name)
+            error("bra name $name is already an index name of the ket layer.")
         end
-        return new{T, V, I}(ket, braname)
     end
+    return nothing
 end
 
 function Base.eltype(::Type{<:NormNetwork})
@@ -32,20 +42,20 @@ function Base.eltype(::Type{<:NormNetwork})
     )
 end
 
-function NormNetwork(tn::ITensorNetwork)
-    return NormNetwork(tn, map(uniquename, keys(tn.dimname_vertices)))
-end
-
 # ====================================== Graphs.jl ======================================= #
 
-Graphs.edges(nn::NormNetwork) = edges(nn.ket)
-Graphs.vertices(nn::NormNetwork) = vertices(nn.ket)
+Graphs.edges(nn::NormNetwork) = edges(ketnetwork(nn))
+Graphs.vertices(nn::NormNetwork) = vertices(ketnetwork(nn))
 
 # ==================================== NamedGraphs.jl ==================================== #
 
-NamedGraphs.encoded_vertex(nn::NormNetwork, vertex) = encoded_vertex(nn.ket, vertex)
-NamedGraphs.decoded_vertex(nn::NormNetwork, code::Integer) = decoded_vertex(nn.ket, code)
-NamedGraphs.encoded_graph(nn::NormNetwork) = encoded_graph(nn.ket)
+function NamedGraphs.encoded_vertex(nn::NormNetwork, vertex)
+    return encoded_vertex(ketnetwork(nn), vertex)
+end
+function NamedGraphs.decoded_vertex(nn::NormNetwork, code::Integer)
+    return decoded_vertex(ketnetwork(nn), code)
+end
+NamedGraphs.encoded_graph(nn::NormNetwork) = encoded_graph(ketnetwork(nn))
 
 # ==================================== DataGraphs.jl ===================================== #
 
@@ -56,31 +66,74 @@ function DataGraphs.get_vertex_data(nn::NormNetwork, vertex)
 end
 
 function DataGraphs.is_vertex_assigned(nn::NormNetwork, vertex)
-    return isassigned(nn.ket, vertex)
+    return isassigned(ketnetwork(nn), vertex)
 end
+
 # =================================== Dictionaries.jl ==================================== #
 
 Dictionaries.issettable(::NormNetwork) = false
 Dictionaries.isinsertable(::NormNetwork) = false
 
 # ====================================== interface ======================================= #
 
+# The name type of the indices, the same as the ket network's.
+ITensorBase.nametype(::Type{<:NormNetwork{T, V, I}}) where {T, V, I} = I
+ITensorBase.nametype(nn::NormNetwork) = nametype(typeof(nn))
+
+"""
+    ketnetwork(nn::NormNetwork)
+
+The ket-layer network of `nn`.
+"""
+ketnetwork(nn::NormNetwork) = nn.ket
+
+"""
+    branetwork(nn::NormNetwork)
+
+The bra-layer network of `nn`, a `BraView` whose tensors are built by [`bratensor`](@ref) when
+accessed.
+"""
+branetwork(nn::NormNetwork) = BraView(nn)
+
+"""
+    branamemap(nn::NormNetwork)
+
+The ket→bra name map, holding a bra name for each ket index name that has a separate bra copy.
+"""
+branamemap(nn::NormNetwork) = nn.braname
+
+"""
+    braname(nn::NormNetwork, name)
+
+The bra-layer index name corresponding to the ket-layer index name `name`.
+"""
 function braname(nn::NormNetwork, name)
-    if !has_dimname(nn.ket, name)
+    if !has_dimname(ketnetwork(nn), name)
         error("index name $name not found underlying tensor network.")
     end
-    # The indices not stored in `nn.braname` are precisely the site indices, which
-    # get mapped to themselves.
-    return get(nn.braname, name, name)
+    # The names absent from the map are the site indices, which map to themselves.
+    return get(branamemap(nn), name, name)
 end
 
-indmap(nn::NormNetwork, ind) = setname(conj(ind), braname(nn, name(ind)))
+# The default ket→bra name map
+derive_braname(name) = prime(name)
+
+"""
+    kettensor(nn::NormNetwork, vertex)
+
+The ket-layer tensor at `vertex`.
+"""
+kettensor(nn::NormNetwork, vertex) = ketnetwork(nn)[vertex]
 
-kettensor(nn::NormNetwork, vertex) = nn.ket[vertex]
 function conj_bratensor(nn::NormNetwork, vertex)
     return rename(n -> braname(nn, n), kettensor(nn, vertex))
 end
 
+"""
+    bratensor(nn::NormNetwork, vertex)
+
+The bra-layer tensor at `vertex`.
+"""
 bratensor(nn::NormNetwork, vertex) = conj(conj_bratensor(nn, vertex))
 
 """
@@ -102,12 +155,22 @@ function flatten_network(nn::NormNetwork)
 end
 
 """
-    normnetwork(tn::ITensorNetwork, [braname]) -> NormNetwork
-
-Build the double-layer norm network `⟨tn|tn⟩`, represented lazily as a `NomnNetwork` object.
-The optional second argument `braname` should implement `braname[ketdimname] = bradimname` for
-every link dimension name `ketdimname` in `tn`. If this is not specified, then a name is
-generated via the `ITensorBase.uniquename` function.
+    normnetwork([f::Function], tn::ITensorNetwork, [dangling]) -> NormNetwork
+    normnetwork(tn::ITensorNetwork, braname::Dictionary) -> NormNetwork
+
+Build the double-layer norm network `⟨tn|tn⟩`, represented lazily as a `NormNetwork` object.
+Each ket index name in `union(dangling, eachlinkname(tn))` gets the bra name `f(name)`,
+where `f` defaults to `derive_braname`, which primes the name, and `dangling` is empty by
+default.
+Alternatively pass the ket→bra name map `braname` directly, which must include every link
+name of `tn` and may include others. Either way an error is thrown if a link name is missing,
+two bra names coincide or a bra name is already an index name of `tn`.
 """
-normnetwork(tn::ITensorNetwork) = NormNetwork(tn)
-normnetwork(tn::ITensorNetwork, braname) = NormNetwork(tn, braname)
+normnetwork(tn::ITensorNetwork, braname::Dictionary) = NormNetwork(tn, braname)
+function normnetwork(f::Function, tn::ITensorNetwork, dangling = nametype(tn)[])
+    names = Indices{nametype(tn)}(union(eachlinkname(tn), dangling))
+    return normnetwork(tn, narrow_map(f, names))
+end
+function normnetwork(tn::ITensorNetwork, dangling = nametype(tn)[])
+    return normnetwork(derive_braname, tn, dangling)
+end