Add ITensorNetworkOperator

src/ITensorNetworksNext.jl

@@ -13,6 +13,7 @@ include("select_algorithm.jl")
 include("AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl")
 include("abstracttensornetwork.jl")
 include("tensornetwork.jl")
+include("itensornetworkoperator.jl")
 include("normnetwork.jl")
 include("normnetworkview.jl")
 include("ITensorNetworkGenerators/ITensorNetworkGenerators.jl")

src/abstracttensornetwork.jl

@@ -128,9 +128,15 @@ function insertlink!(tn::AbstractGraph, e)
 end
 
 function operator_support(tn::AbstractGraph, op::ITensorOperator)
+    return operator_support_names(tn, inputnames(op))
+end
+
+# Shared with the `ITensorNetworkOperator` method, which is defined alongside that type
+# because it is not yet known at this point in the include order.
+function operator_support_names(tn::AbstractGraph, opnames)
     support = Indices{vertextype(tn)}()
 
-    for name in inputnames(op)
+    for name in opnames
         vertices = dimnamevertices(tn, name)
 
         if length(vertices) > 1

src/itensornetworkoperator.jl

@@ -0,0 +1,119 @@
+using DataGraphs: DataGraphs, get_vertex_data, is_vertex_assigned
+using Dictionaries: Dictionaries, Dictionary
+using Graphs: Graphs, AbstractGraph, edges, vertices
+using ITensorBase: ITensorBase, NamedTensorOperator, inputnames, names, nametype, operator,
+    outputnames, state
+using NamedGraphs: NamedGraphs, decoded_vertex, encoded_graph, encoded_vertex
+
+"""
+    struct ITensorNetworkOperator{T, V, I, P} <: AbstractITensorNetwork{T, V}
+
+The network equivalent of `ITensorBase.ITensorOperator`: a tensor network of type
+`P <: AbstractITensorNetwork{T, V}` together with a pairing of its dangling index names,
+stored as a map from each output name to its input name. Applying the operator contracts over
+the input names and leaves the output names.
+
+The output and input of each pair must sit on the same vertex. Indexing returns the vertex
+tensor wrapped as an `ITensorOperator` carrying the pairs at that vertex.
+"""
+struct ITensorNetworkOperator{T, V, I, P <: AbstractITensorNetwork{T, V}} <:
+    AbstractITensorNetwork{T, V}
+    parent::P
+    pairing::Dictionary{I, I}
+    function ITensorNetworkOperator(
+            parent::AbstractITensorNetwork{T, V}, outputnames, inputnames
+        ) where {T, V}
+        I = nametype(T)
+        outputnames = collect(I, outputnames)
+        inputnames = collect(I, inputnames)
+        if length(outputnames) != length(inputnames)
+            throw(
+                ArgumentError(
+                    "Operator `outputnames` and `inputnames` must have equal length " *
+                        "(positional pairing), got $(length(outputnames)) and " *
+                        "$(length(inputnames))."
+                )
+            )
+        end
+        if !allunique(vcat(outputnames, inputnames))
+            throw(
+                ArgumentError("each operator output and input name must appear only once.")
+            )
+        end
+        for opname in Iterators.flatten((outputnames, inputnames))
+            nvertices = length(dimnamevertices(parent, opname))
+            if nvertices != 1
+                throw(
+                    ArgumentError(
+                        "operator dim name $opname is associated with $nvertices vertices " *
+                            "in the tensor network; an operator leg must be a dangling index."
+                    )
+                )
+            end
+        end
+        for (output, input) in zip(outputnames, inputnames)
+            if only(dimnamevertices(parent, output)) != only(dimnamevertices(parent, input))
+                throw(
+                    ArgumentError(
+                        "operator output $output and its paired input $input must sit on " *
+                            "the same vertex."
+                    )
+                )
+            end
+        end
+        pairing = Dictionary(outputnames, inputnames)
+        return new{T, V, I, typeof(parent)}(parent, pairing)
+    end
+end
+
+function ITensorBase.operator(tn::AbstractITensorNetwork, outputnames, inputnames)
+    return ITensorNetworkOperator(tn, outputnames, inputnames)
+end
+
+function Base.eltype(::Type{<:ITensorNetworkOperator{T, V, I}}) where {T, V, I}
+    return NamedTensorOperator{I, T}
+end
+
+# ====================================== Graphs.jl ======================================= #
+
+Graphs.edges(op::ITensorNetworkOperator) = edges(state(op))
+Graphs.vertices(op::ITensorNetworkOperator) = vertices(state(op))
+
+# ==================================== NamedGraphs.jl ==================================== #
+
+function NamedGraphs.encoded_vertex(op::ITensorNetworkOperator, vertex)
+    return encoded_vertex(state(op), vertex)
+end
+function NamedGraphs.decoded_vertex(op::ITensorNetworkOperator, code::Integer)
+    return decoded_vertex(state(op), code)
+end
+NamedGraphs.encoded_graph(op::ITensorNetworkOperator) = encoded_graph(state(op))
+
+# ==================================== DataGraphs.jl ===================================== #
+
+# Both names of a pair sit on one vertex, so the pairs whose output is on `vertex` are its pairing.
+function DataGraphs.get_vertex_data(op::ITensorNetworkOperator, vertex)
+    tensor = state(op)[vertex]
+    outputs = filter(name -> haskey(op.pairing, name), names(tensor))
+    return operator(tensor, outputs, [op.pairing[output] for output in outputs])
+end
+
+function DataGraphs.is_vertex_assigned(op::ITensorNetworkOperator, vertex)
+    return is_vertex_assigned(state(op), vertex)
+end
+
+# =================================== Dictionaries.jl ==================================== #
+
+Dictionaries.issettable(::ITensorNetworkOperator) = false
+Dictionaries.isinsertable(::ITensorNetworkOperator) = false
+
+# ====================================== interface ======================================= #
+
+ITensorBase.state(op::ITensorNetworkOperator) = op.parent
+Base.parent(op::ITensorNetworkOperator) = state(op)
+ITensorBase.outputnames(op::ITensorNetworkOperator) = collect(keys(op.pairing))
+ITensorBase.inputnames(op::ITensorNetworkOperator) = collect(op.pairing)
+
+function operator_support(tn::AbstractGraph, op::ITensorNetworkOperator)
+    return operator_support_names(tn, inputnames(op))
+end

test/test_itensornetworkoperator.jl

@@ -0,0 +1,115 @@
+using DataGraphs: is_vertex_assigned
+using Dictionaries: isinsertable, issettable
+using Graphs: edges, vertices
+using ITensorBase: ITensor, Index, inputnames, name, operator, outputnames, state
+using ITensorNetworksNext:
+    ITensorNetwork, ITensorNetworkOperator, operator_support, tensornetwork
+using NamedGraphs: incident_edges, named_path_graph
+using Test: @test, @test_throws, @testset
+
+# A bondless operator network on `g`, with one input/output pair per vertex.
+function local_operator(g; d = 2)
+    inp = Dict(v => Index(d) for v in vertices(g))
+    out = Dict(v => Index(d) for v in vertices(g))
+    tn = tensornetwork(vertices(g)) do v
+        return randn((out[v], inp[v]))
+    end
+    vs = collect(vertices(g))
+    op = operator(tn, [name(out[v]) for v in vs], [name(inp[v]) for v in vs])
+    return op, inp, out
+end
+
+@testset "`ITensorNetworkOperator`" begin
+    @testset "Basics" begin
+        g = named_path_graph(3)
+        op, inp, out = local_operator(g)
+
+        @test op isa ITensorNetworkOperator
+
+        # The operator shares the graph structure of the network it wraps.
+        @test issetequal(vertices(op), vertices(state(op)))
+        @test issetequal(edges(op), edges(state(op)))
+        @test parent(op) === state(op)
+
+        # The pairing is positional: `outputnames[i]` goes with `inputnames[i]`.
+        vs = collect(vertices(g))
+        @test outputnames(op) == [name(out[v]) for v in vs]
+        @test inputnames(op) == [name(inp[v]) for v in vs]
+
+        @test is_vertex_assigned(op, first(vs))
+
+        # It is a lazy wrapper, so it is neither settable nor insertable.
+        @test !issettable(op)
+        @test !isinsertable(op)
+    end
+
+    @testset "`getindex` wraps the vertex tensor" begin
+        g = named_path_graph(3)
+        op, inp, out = local_operator(g)
+
+        # Both halves of each pair sit on the same vertex here, so the wrapper carries them.
+        for v in vertices(g)
+            @test outputnames(op[v]) == [name(out[v])]
+            @test inputnames(op[v]) == [name(inp[v])]
+            @test state(op[v]) === state(op)[v]
+        end
+
+        # `eltype` is the type of the wrapped vertex data.
+        @test eltype(op) === typeof(op[first(vertices(g))])
+    end
+
+    @testset "a pair must sit on one vertex" begin
+        # A swap pairs the output at vertex 1 with the input at vertex 2, and vice versa.
+        a, b = Index(2), Index(2)
+        a′, b′ = Index(2), Index(2)
+        tn = ITensorNetwork(Dict(1 => randn((a′, a)), 2 => randn((b′, b))))
+        @test_throws ArgumentError operator(tn, [name(a′), name(b′)], [name(b), name(a)])
+    end
+
+    @testset "constructor validation" begin
+        g = named_path_graph(2)
+        op, inp, out = local_operator(g)
+        tn = state(op)
+        vs = collect(vertices(g))
+
+        # The pairing is a bijection, so the two name lists must have equal length.
+        @test_throws ArgumentError operator(tn, [name(out[vs[1]])], [])
+        @test_throws ArgumentError operator(
+            tn, [name(out[v]) for v in vs], [name(inp[vs[1]])]
+        )
+        # An output name can be paired with only one input.
+        @test_throws ArgumentError operator(
+            tn, [name(out[vs[1]]), name(out[vs[1]])], [name(inp[vs[1]]), name(inp[vs[1]])]
+        )
+        # An input name can be paired with only one output, and no name is both.
+        @test_throws ArgumentError operator(
+            tn, [name(out[vs[1]]), name(out[vs[2]])], [name(inp[vs[1]]), name(inp[vs[1]])]
+        )
+        @test_throws ArgumentError operator(
+            tn, [name(out[vs[1]]), name(inp[vs[1]])], [name(inp[vs[1]]), name(out[vs[2]])]
+        )
+
+        # An operator leg must be dangling: a link name touches two vertices.
+        l = Index(2)
+        linked = tensornetwork(vertices(g)) do v
+            return randn((inp[v], l))
+        end
+        @test_throws ArgumentError operator(linked, [name(l)], [name(inp[vs[1]])])
+    end
+
+    @testset "`operator_support`" begin
+        g = named_path_graph(3)
+        s = Dict(v => Index(2) for v in vertices(g))
+        state_tn = tensornetwork(vertices(g)) do v
+            is = map(e -> Index(2), incident_edges(g, v))
+            return randn((s[v], is...))
+        end
+
+        # An operator acting on sites 1 and 3 is supported on exactly those vertices.
+        out1, out3 = Index(2), Index(2)
+        optn = ITensorNetwork(Dict(1 => randn((out1, s[1])), 3 => randn((out3, s[3]))))
+        op = operator(optn, [name(out1), name(out3)], [name(s[1]), name(s[3])])
+
+        @test issetequal(operator_support(state_tn, op), Set([1, 3]))
+    end
+end