Allow linear combinations of operators and plain tensors

Project.toml

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

src/broadcast.jl

@@ -135,7 +135,8 @@ end
 # rules below enforce the input rules declaratively:
 #   - operator ⊗ operator → operator (preserved),
 #   - operator ⊗ scalar → operator (`2 .* op` stays an operator),
-#   - operator ⊗ non-operator tensor → error.
+#   - operator ⊗ non-operator tensor → operator (the tensor is a trivial, empty-pairing
+#     operator, so the result inherits the operator operand's pairing).
 # The `BroadcastStyle(::Type{<:NamedTensorOperator})` mapping and the operator-specific
 # `copy` (which unwraps, delegates to `NamedTensorStyle`, then rewraps) live in
 # `itensoroperator.jl`, where `NamedTensorOperator` is defined. `*` (contraction) is
@@ -158,15 +159,12 @@ function BC.BroadcastStyle(
     )
     return style
 end
-# operator ⊗ non-operator named tensor is type-nonsense and is rejected.
-function BC.BroadcastStyle(::NamedTensorOperatorStyle, ::NamedTensorStyle)
-    return throw(
-        ArgumentError(
-            "Cannot broadcast an `NamedTensorOperator` together with a non-operator " *
-                "tensor. Wrap the tensor as an operator first, or unwrap the " *
-                "operator with `state`."
-        )
-    )
+# operator ⊗ non-operator named tensor stays an operator: a plain tensor is a trivial
+# operator with no pairing, so `o - t` (etc.) combines the states elementwise and the
+# result inherits `o`'s output/input split (the split logic lives in
+# `broadcast_operator_output_input`).
+function BC.BroadcastStyle(style::NamedTensorOperatorStyle, ::NamedTensorStyle)
+    return style
 end
 
 # Reinterpret an operator-style `Broadcasted` under `NamedTensorStyle`, the broadcast

src/namedtensoroperator.jl

@@ -444,8 +444,8 @@ function BC.BroadcastStyle(arraytype::Type{<:NamedTensorOperator})
     return NamedTensorOperatorStyle{ndims(arraytype)}()
 end
 
-# Recover the output/input split shared by all operator operands of `bc`,
-# erroring if any two operators disagree.
+# Collect the operator operands of `bc`, skipping non-operator operands (plain tensors
+# and scalars), which contribute no pairing.
 operator_operands(bc::Broadcasted) = operator_operands(bc.args...)
 function operator_operands(arg::NamedTensorOperator, args...)
     return (arg, operator_operands(args...)...)
@@ -456,22 +456,30 @@ end
 operator_operands(arg, args...) = operator_operands(args...)
 operator_operands() = ()
 
+# The output/input split the broadcast result inherits from its operator operands. Each
+# operator contributes its (output, input) pairs; a non-operator operand contributes none
+# (a plain tensor is the trivial, empty-pairing operator), so combining an operator with a
+# plain tensor just inherits the operator's split. Combining operators requires them to
+# pair their shared names consistently: a name that appears in more than one distinct pair
+# (as an output or an input) is paired two different ways across the operands, which is an
+# error rather than a guess. An operand can be unwrapped with `state` to combine as a plain
+# tensor instead.
 function broadcast_operator_output_input(bc::Broadcasted)
     ops = operator_operands(bc)
-    op1 = first(ops)
-    out1 = outputnames(op1)
-    inp1 = inputnames(op1)
-    for op in Base.tail(ops)
-        (issetequal(outputnames(op), out1) && issetequal(inputnames(op), inp1)) ||
-            throw(
-            ArgumentError(
-                "Operator operands disagree on their output/input split: " *
-                    "$((out1, inp1)) vs $((outputnames(op), inputnames(op))). " *
-                    "Broadcasting operators requires a matching split."
-            )
-        )
+    DimName = eltype(outputnames(first(ops)))
+    pairs = Tuple{DimName, DimName}[]
+    for op in ops, pair in zip(outputnames(op), inputnames(op))
+        pair in pairs || push!(pairs, pair)
     end
-    return out1, inp1
+    outnames, innames = first.(pairs), last.(pairs)
+    allunique(outnames) && allunique(innames) && isdisjoint(outnames, innames) || throw(
+        ArgumentError(
+            "Operator operands pair a shared name two different ways; broadcasting " *
+                "operators requires each shared name to be paired the same way. Unwrap " *
+                "an operand with `state` to combine them as plain tensors instead."
+        )
+    )
+    return outnames, innames
 end
 
 function Base.copy(bc::Broadcasted{<:NamedTensorOperatorStyle})

test/test_operator.jl

@@ -343,16 +343,50 @@ end
     @test isempty(outputnames(oo))
     @test isempty(inputnames(oo))
 
-    # Operator combined with a non-operator tensor is rejected.
+    # Operator combined with a non-operator tensor: the tensor is a trivial
+    # (empty-pairing) operator, so the result stays an operator with `o`'s pairing.
     plain = NamedTensor(randn(2, 2), ("i'", "i"))
-    @test_throws ArgumentError o .+ plain
+    op = o .+ plain
+    @test op isa NamedTensorOperator
+    @test issetequal(outputnames(op), ("i'",))
+    @test issetequal(inputnames(op), ("i",))
+    @test unname(state(op), nms) ≈ unname(s, nms) .+ unname(plain, nms)
 
-    # Two operators whose name sets match but whose output/input split differs
-    # are rejected (the split would otherwise be ambiguous).
+    # Two operators that pair a shared name two different ways (here `i'` and `i` swap
+    # output/input roles) are an error, not a guess.
     o_swapped = operator(randn(2, 2), ("i",), ("i'",))
     @test_throws ArgumentError o .+ o_swapped
 end
 
+@testset "operator/state linear algebra" begin
+    # A plain tensor is a trivial (empty-pairing) operator, so combining it with an
+    # operator keeps the operator's pairing (the motivating `o - t` case).
+    o = operator(randn(2, 2), ("i",), ("j",))
+    t = NamedTensor(randn(2, 2), ("i", "j"))
+    for r in (o - t, t - o, o + t)
+        @test r isa NamedTensorOperator
+        @test issetequal(outputnames(r), ("i",))
+        @test issetequal(inputnames(r), ("j",))
+    end
+    @test unname(state(o - t), ("i", "j")) ≈
+        unname(state(o), ("i", "j")) - unname(t, ("i", "j"))
+    @test unname(state(t - o), ("i", "j")) ≈
+        unname(t, ("i", "j")) - unname(state(o), ("i", "j"))
+
+    # Two operators must pair their shared names consistently. `A` pairs i'->i and j'->j;
+    # `B` pairs i'->i (agreed) but j->j', so j' and j are paired two different ways and the
+    # whole combination errors, even though i'->i agrees.
+    A = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j"))
+    B = operator(randn(2, 2, 2, 2), ("i'", "j"), ("i", "j'"))
+    @test_throws ArgumentError A + B
+
+    # Operators that agree on every shared pairing do combine, keeping the pairing.
+    C = operator(randn(2, 2), ("i'",), ("i",))
+    D = operator(randn(2, 2), ("i'",), ("i",))
+    @test issetequal(outputnames(C + D), ("i'",))
+    @test issetequal(inputnames(C + D), ("i",))
+end
+
 @testset "operator-preserving contraction" begin
     # A shared *dangling* leg (in neither pairing) is summed away, and the
     # surviving output/input of each operand combine. This is the `c† * c`