Preserve operator structure under contraction

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.8.0"
+version = "0.8.1"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/itensoroperator.jl

@@ -23,14 +23,14 @@ get_codomain_name(a, i) = throw(MethodError(get_codomain_name, (a, i)))
 get_domain_name(a, i) = throw(MethodError(get_domain_name, (a, i)))
 
 function apply(x::AbstractITensor, y::AbstractITensor)
-    xy = x * y
+    xy = state(x) * state(y)
     return mapdimnames(xy) do i
         return get_domain_name(x, i)
     end
 end
 
 function apply_dag(x::AbstractITensor, y::AbstractITensor)
-    xy = x * y
+    xy = state(x) * state(y)
     return mapdimnames(xy) do i
         return get_codomain_name(y, i)
     end
@@ -131,20 +131,60 @@ function get_domain_name(a::ITensorOperator, i)
     return get(inverse(a.dimnames_bijection), i, i)
 end
 
+# `codomain` and `domain` may be given as dimension names or as named ranges
+# (such as `Index`es); `name` maps the latter to their names and leaves names as-is.
 # TODO: Unify these two functions.
 function operator(a::AbstractArray, codomain, domain)
+    codomain, domain = name.(codomain), name.(domain)
     na = nameddims(a, (codomain..., domain...))
     return operator(na, codomain, domain)
 end
 function operator(a::AbstractITensor, codomain, domain)
-    return ITensorOperator(a, codomain, domain)
+    return ITensorOperator(a, name.(codomain), name.(domain))
+end
+
+# Operator-preserving contraction. Contracting two named arrays sums over their
+# shared names, so the result keeps each operand's surviving codomain/domain
+# structure. A non-operator tensor contributes no pairs (all its names are
+# dangling from the operator point of view). The result is always an
+# `ITensorOperator`, even when its codomain and domain both come out empty, so
+# the product type does not depend on the runtime names being contracted.
+operator_pairs(a::ITensorOperator) = a.dimnames_bijection.domain_to_codomain
+operator_pairs(a::AbstractITensor) = ()
+
+# Compose the codomain/domain of `a * b`. The `domain => codomain` pairs of both
+# operands form a graph of maximum degree two (each name is paired at most once
+# per operand), so it is a disjoint union of simple paths. Each surviving domain
+# name reaches its surviving codomain partner by following the pairing through
+# any contracted (shared) names. A name whose chain dead-ends on a contracted
+# index is left dangling, so the result is well defined for any contraction.
+function product_codomain_domain(a::AbstractITensor, b::AbstractITensor)
+    shared = intersect(dimnames(a), dimnames(b))
+    pairs = collect(Iterators.flatten((operator_pairs(a), operator_pairs(b))))
+    forward = Dict(pairs)
+    domain = eltype(keys(forward))[]
+    codomain = eltype(values(forward))[]
+    for (d, c) in pairs
+        d in shared && continue
+        while c in shared && haskey(forward, c)
+            c = forward[c]
+        end
+        c in shared && continue
+        push!(domain, d)
+        push!(codomain, c)
+    end
+    return codomain, domain
+end
+
+function operator_product(a::AbstractITensor, b::AbstractITensor)
+    ab = state(a) * state(b)
+    codomain, domain = product_codomain_domain(a, b)
+    return operator(ab, codomain, domain)
 end
 
-# This helps preserve the ITensor type when multiplying,
-# for example when a ITensorOperator wraps an ITensor.
-Base.:*(a::ITensorOperator, b::ITensorOperator) = state(a) * state(b)
-Base.:*(a::ITensorOperator, b::AbstractITensor) = state(a) * state(b)
-Base.:*(a::AbstractITensor, b::ITensorOperator) = state(a) * state(b)
+Base.:*(a::ITensorOperator, b::ITensorOperator) = operator_product(a, b)
+Base.:*(a::ITensorOperator, b::AbstractITensor) = operator_product(a, b)
+Base.:*(a::AbstractITensor, b::ITensorOperator) = operator_product(a, b)
 
 # Operator-preserving broadcasting (the style struct and style-combination rules
 # live in `broadcast.jl`). An `ITensorOperator` broadcasts as itself, so `op .+ op`,

test/test_operator.jl

@@ -6,7 +6,7 @@ using Random: Random
 using StableRNGs: StableRNG
 using TensorAlgebra.MatrixAlgebra: gram_eigh_full, gram_eigh_full_with_pinv
 using TensorAlgebra: matricize
-using Test: @test, @testset
+using Test: @test, @test_throws, @testset
 
 @testset "operator" begin
     o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j"))
@@ -46,6 +46,15 @@ using Test: @test, @testset
     @test ov ≈ replacedimnames(o * v, "i'" => "i", "j'" => "j")
 end
 
+@testset "operator from named ranges" begin
+    # Codomain/domain may be given as named ranges, not just names.
+    i, ip = namedoneto(2, "i"), namedoneto(2, "i'")
+    o = operator(randn(2, 2), [ip], [i])
+    @test o isa ITensorOperator{String}
+    @test issetequal(codomainnames(o), ("i'",))
+    @test issetequal(domainnames(o), ("i",))
+end
+
 @testset "one(::ITensorOperator)" begin
     # Identity-operator construction: matricized form is the identity matrix.
     i, j, k, l = namedoneto.((2, 3, 2, 3), ("i", "j", "k", "l"))
@@ -113,8 +122,8 @@ end
 @testset "operator-preserving broadcasting" begin
     # `+`, `-`, and scalar multiplication lower to broadcasting. An operator
     # broadcasts as itself (it is not peeled to its `state`), so these operations
-    # preserve the `ITensorOperator` wrapper and its codomain/domain bijection. `*`
-    # (contraction) is unchanged and still decays to the underlying state.
+    # preserve the `ITensorOperator` wrapper and its codomain/domain bijection.
+    # (Contraction `*` is operator-preserving too, in its own testset below.)
     o = operator(randn(2, 2), ("i'",), ("i",))
     s = state(o)
     nms = ("i'", "i")
@@ -134,8 +143,14 @@ end
     @test unname(state(o .* 2), nms) ≈ 2 .* unname(s, nms)
     @test unname(state(o ./ 2), nms) ≈ unname(s, nms) ./ 2
 
-    # `*` (contraction) still decays to a plain tensor.
-    @test !((o * o) isa ITensorOperator)
+    # `o` shares both its names with itself, so `o * o` fully contracts to a
+    # scalar with no surviving codomain/domain. It is still an `ITensorOperator`
+    # (with empty codomain/domain), so the product type does not depend on which
+    # names happen to contract.
+    oo = o * o
+    @test oo isa ITensorOperator
+    @test isempty(codomainnames(oo))
+    @test isempty(domainnames(oo))
 
     # Operator combined with a non-operator tensor is rejected.
     plain = ITensor(randn(2, 2), ("i'", "i"))
@@ -147,6 +162,39 @@ end
     @test_throws ArgumentError o .+ o_swapped
 end
 
+@testset "operator-preserving contraction" begin
+    # A shared *dangling* leg (in neither bijection) is summed away, and the
+    # surviving codomain/domain of each operand combine. This is the `c† * c`
+    # hopping pattern: two operators paired over an auxiliary link.
+    a = operator(nameddims(randn(2, 2, 3), ("i'", "i", "aux")), ["i'"], ["i"])
+    b = operator(nameddims(randn(2, 2, 3), ("j'", "j", "aux")), ["j'"], ["j"])
+    ab = a * b
+    @test ab isa ITensorOperator
+    @test issetequal(codomainnames(ab), ("i'", "j'"))
+    @test issetequal(domainnames(ab), ("i", "j"))
+    @test !("aux" in dimnames(ab))
+    @test state(ab) ≈ state(a) * state(b)
+
+    # A shared *paired* index (a's domain equals b's codomain) chains through the
+    # contraction: a's codomain partner pairs with b's domain partner.
+    A = operator(randn(2, 2), ("a'",), ("m",))
+    B = operator(randn(2, 2), ("m",), ("b",))
+    AB = A * B
+    @test AB isa ITensorOperator
+    @test issetequal(codomainnames(AB), ("a'",))
+    @test issetequal(domainnames(AB), ("b",))
+
+    # Applying an operator to a plain state contracts the operator's domain and
+    # leaves its output dangling. The result stays an `ITensorOperator` with empty
+    # codomain/domain (the surviving `a'` leg is dangling, in neither).
+    v = nameddims(randn(2), ("m",))
+    Av = operator(randn(2, 2), ("a'",), ("m",)) * v
+    @test Av isa ITensorOperator
+    @test isempty(codomainnames(Av))
+    @test isempty(domainnames(Av))
+    @test issetequal(dimnames(Av), ("a'",))
+end
+
 @testset "gram_eigh_full on ITensorOperator" begin
     n = 5
     B = randn(n, n)