Add named directsum for tensors

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.12.0"
+version = "0.12.1"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -50,7 +50,7 @@ OMEinsumContractionOrders = "1.3"
 OrderedCollections = "1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.17.3"
+TensorAlgebra = "0.17.5"
 TensorKit = "0.17"
 TermInterface = "2"
 TupleTools = "1.6"

src/tensoralgebra.jl

@@ -149,6 +149,57 @@ function unmatricize_nameddims(na::AbstractNamedTensor, splitters::Vararg{Pair,
     return nameddims(a_split, names_split)
 end
 
+"""
+    directsum(A => inds_A, B => inds_B, ...)
+    directsum(out_inds, A => inds_A, B => inds_B, ...)
+
+Direct sum of the named tensors `A, B, …` over the indices paired with each. The remaining
+("shared") indices are common to every tensor; they are aligned and carried through unchanged,
+while the paired indices are concatenated block-diagonally. The result has the shared indices
+first and the summed indices trailing.
+
+The first form mints fresh summed indices and returns `S => out_inds`, mirroring the
+`tensor => indices` inputs. The second form takes the summed indices' names from `out_inds`
+(names or `NamedUnitRange`s) and returns just `S`; the summed axes themselves come from the
+direct sum, so only the names of `out_inds` are used.
+
+# Examples
+
+```jldoctest
+julia> using ITensorBase: Index
+
+julia> using TensorAlgebra: directsum
+
+julia> i, j, k = Index.((2, 2, 3));
+
+julia> a = randn(i, j);
+
+julia> b = randn(i, k);
+
+julia> s, (l,) = directsum(a => (j,), b => (k,));
+
+julia> length(l)
+5
+```
+"""
+function TA.directsum(
+        out_inds, pair1::Pair{<:AbstractNamedTensor}, pairs::Pair{<:AbstractNamedTensor}...
+    )
+    ps = (pair1, pairs...)
+    shared = namesetdiff(inds(first(pair1)), last(pair1))
+    summed_dims = length(shared) .+ eachindex(last(pair1))
+    aligned = map(p -> unname(first(p), [shared; collect(last(p))]), ps)
+    a = TA.directsum(summed_dims, aligned...)
+    return nameddims(a, [name.(shared); name.(collect(out_inds))])
+end
+function TA.directsum(
+        pair1::Pair{<:AbstractNamedTensor}, pairs::Pair{<:AbstractNamedTensor}...
+    )
+    out_names = [uniquename(dimnametype(first(pair1))) for _ in last(pair1)]
+    s = TA.directsum(out_names, pair1, pairs...)
+    return s => last(inds(s), length(out_names))
+end
+
 # Canonicalize a bond-name keyword to a `nametype -> name` minting function. A `NamedTuple` of
 # decoration is splatted into `uniquename` (kwargs such as `tags`/`plev`); a callable is used
 # as-is, for full control over how the name is minted. The default `(;)` reproduces the bare

test/Project.toml

@@ -41,7 +41,7 @@ Random = "1.10"
 SafeTestsets = "0.1"
 StableRNGs = "1"
 Suppressor = "0.2"
-TensorAlgebra = "0.17"
+TensorAlgebra = "0.17.5"
 TensorKit = "0.17"
 TermInterface = "2"
 Test = "1.10"

test/test_tensoralgebra.jl

@@ -5,7 +5,7 @@ using MatrixAlgebraKit: left_null, left_orth, left_polar, lq_compact, lq_full, q
     qr_full, right_null, right_orth, right_polar, svd_compact, svd_trunc, svd_vals
 using StableRNGs: StableRNG
 using TensorAlgebra.MatrixAlgebra: gram_eigh_full, gram_eigh_full_with_pinv
-using TensorAlgebra: TensorAlgebra, contract, matricize, project, trivialrange,
+using TensorAlgebra: TensorAlgebra, contract, directsum, matricize, project, trivialrange,
     unchecked_project, unmatricize
 using Test: @test, @test_broken, @testset
 
@@ -60,6 +60,34 @@ using Test: @test, @test_broken, @testset
             (unnamed(k), unnamed(i), unnamed(j), unnamed(l))
         )
     end
+    @testset "directsum" begin
+        i = namedoneto(2, "i")            # shared index, carried through unchanged
+        j1, j2 = namedoneto.((2, 3), ("j1", "j2"))
+        k1, k2 = namedoneto.((2, 3), ("k1", "k2"))
+        a = randn(elt, i, j1, k1)
+        b = randn(elt, i, j2, k2)
+        ref = cat(unname(a, (i, j1, k1)), unname(b, (i, j2, k2)); dims = (2, 3))
+        # Fresh output indices: the shared index is kept and the summed indices trail it,
+        # and the minted indices are returned as a second output.
+        s, summed = directsum(a => (j1, k1), b => (j2, k2))
+        @test eltype(s) ≡ elt
+        @test i in inds(s)
+        @test issetequal(inds(s), (i, summed...))
+        @test sort(length.(summed)) == [5, 5]
+        @test unname(s, (i, summed...)) == ref
+        # Explicit output indices name the summed dimensions.
+        o1, o2 = namedoneto.((5, 5), ("o1", "o2"))
+        s2 = directsum((o1, o2), a => (j1, k1), b => (j2, k2))
+        @test issetequal(inds(s2), (i, o1, o2))
+        @test unname(s2, (i, o1, o2)) == ref
+        # A single summed dimension.
+        u1, u2 = namedoneto.((2, 3), ("u1", "u2"))
+        c = randn(elt, i, u1)
+        d = randn(elt, i, u2)
+        sc, (su,) = directsum(c => (u1,), d => (u2,))
+        @test length(su) == 5
+        @test unname(sc, (i, su)) == cat(unname(c, (i, u1)), unname(d, (i, u2)); dims = 2)
+    end
     @testset "Matrix functions" begin
         for f in ITensorBase.MATRIX_FUNCTIONS
             f == :cbrt && elt <: Complex && continue