diff --git a/Project.toml b/Project.toml index 59ad31f..48d215c 100644 --- a/Project.toml +++ b/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 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" diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index ef08a0a..b89f767 100644 --- a/src/tensoralgebra.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index fdba84e..2093879 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_tensoralgebra.jl b/test/test_tensoralgebra.jl index 601c076..b1672b4 100644 --- a/test/test_tensoralgebra.jl +++ b/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