From 6af047ada7282b17cd602bb19ff16ca415a0326d Mon Sep 17 00:00:00 2001 From: Jack Dunham Date: Fri, 30 Oct 2026 13:29:43 -0400 Subject: [PATCH] Add `ITensorNetworkOperator` The network equivalent of `ITensorOperator`: a tensor network with a pairing of its dangling output and input names, each pair on one vertex. `operator(tn, outputnames, inputnames)` constructs one, and `operator_support` accepts it. Co-Authored-By: Claude Opus 5.5 --- src/ITensorNetworksNext.jl | 1 + src/abstracttensornetwork.jl | 8 +- src/itensornetworkoperator.jl | 119 ++++++++++++++++++++++++++++ test/test_itensornetworkoperator.jl | 115 +++++++++++++++++++++++++++ 4 files changed, 242 insertions(+), 1 deletion(-) create mode 100644 src/itensornetworkoperator.jl create mode 100644 test/test_itensornetworkoperator.jl diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 0298a85f..a8513d17 100644 --- a/src/ITensorNetworksNext.jl +++ b/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") diff --git a/src/abstracttensornetwork.jl b/src/abstracttensornetwork.jl index 7a6dcd82..98f5c03a 100644 --- a/src/abstracttensornetwork.jl +++ b/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 diff --git a/src/itensornetworkoperator.jl b/src/itensornetworkoperator.jl new file mode 100644 index 00000000..58e3e844 --- /dev/null +++ b/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 diff --git a/test/test_itensornetworkoperator.jl b/test/test_itensornetworkoperator.jl new file mode 100644 index 00000000..a2918297 --- /dev/null +++ b/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