Add dot and dotperm for tensor inner products

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.23.1"
+version = "0.23.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/TensorAlgebra.jl

@@ -5,7 +5,7 @@ export contract, contract!, contractalign, dual, isdual, MatrixAlgebra
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
+            "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
         )
     )
 end
@@ -17,6 +17,7 @@ include("MatrixAlgebra.jl")
 include("bituple.jl")
 include("permutedimsadd.jl")
 include("matricize.jl")
+include("dotperm.jl")
 include("concatenate.jl")
 include("directsum.jl")
 include("dual.jl")

src/dotperm.jl

@@ -0,0 +1,81 @@
+using LinearAlgebra: LinearAlgebra
+
+"""
+    dotperm(a, b, perm)
+
+Compute the inner product of `a` and `b` after applying the full dimension permutation
+`perm` to `b`, using the codomain/domain split of `a`.
+The permuted axes must match: `axes(a, i) == axes(b, perm[i])`.
+
+# Examples
+
+```jldoctest
+julia> import TensorAlgebra
+
+julia> TensorAlgebra.dotperm([1 2; 3 4], [5 6; 7 8], (2, 1))
+69
+```
+
+See also [`dot`](@ref).
+"""
+function dotperm(a, b, perm)
+    check_input(dotperm, a, b, perm)
+    perm = NTuple{ndims(a)}(perm)
+    perm_codomain, perm_domain = bipartition(perm, Val(ndims_codomain(a)))
+    return LinearAlgebra.dot(
+        matricize(a, Val(ndims_codomain(a))), matricize(b, perm_codomain, perm_domain)
+    )
+end
+
+function check_input(::typeof(dotperm), a, b, perm)
+    ndims(a) == length(perm) == ndims(b) ||
+        throw(
+        ArgumentError("Permutation length must match both operands' dimension counts")
+    )
+    isperm(perm) || throw(ArgumentError("Invalid dimension permutation"))
+    for i in 1:ndims(a)
+        axes(a, i) == axes(b, perm[i]) || throw(
+            DimensionMismatch(
+                "Inner product axes do not match: axes(a, $i) = $(axes(a, i)), axes(b, $(perm[i])) = $(axes(b, perm[i]))"
+            )
+        )
+    end
+    return nothing
+end
+
+"""
+    TensorAlgebra.dot(a, labels_a, b, labels_b)
+
+Compute the inner product of `a` and `b` after permuting the dimensions of `b` to match
+the labels of `a`. Each operand must have one unique label per dimension, and both label
+sets must match.
+
+# Examples
+
+```jldoctest
+julia> import TensorAlgebra
+
+julia> a = [1 2; 3 4];
+       b = [5 6; 7 8];
+
+julia> TensorAlgebra.dot(a, (:i, :j), b, (:j, :i))
+69
+```
+
+See also [`dotperm`](@ref) and [`matricize`](@ref).
+"""
+function dot(a, labels_a, b, labels_b)
+    check_input(dot, a, labels_a, b, labels_b)
+    labels_a_tuple = NTuple{ndims(a)}(labels_a)
+    labels_b_tuple = NTuple{ndims(b)}(labels_b)
+    perm = tuple_indexin(labels_a_tuple, labels_b_tuple)
+    return dotperm(a, b, perm)
+end
+
+function check_input(::typeof(dot), a, labels_a, b, labels_b)
+    length(labels_a) == ndims(a) && length(labels_b) == ndims(b) ||
+        throw(ArgumentError("Each operand must have one label per dimension"))
+    allunique(labels_a) && allunique(labels_b) && issetequal(labels_a, labels_b) ||
+        throw(ArgumentError("Inner product labels must be unique and matching"))
+    return nothing
+end

test/test_dotperm.jl

@@ -0,0 +1,93 @@
+using LinearAlgebra: dot, norm
+using StableRNGs: StableRNG
+using TensorAlgebra: TensorAlgebra, dotperm
+using TensorKit: TensorKit, FermionParity, Rep, SU₂, U₁, Vect, randn, ←, ⊗
+using TensorOperations: TensorOperations
+using Test: @test, @test_throws, @testset
+
+@testset "dotperm dense (eltype=$elt)" for elt in (Float64, ComplexF64)
+    rng = StableRNG(123)
+    a = randn(rng, elt, 2, 3, 4)
+    b = randn(rng, elt, 3, 4, 2)
+    @test dotperm(a, a, (1, 2, 3)) ≈ norm(a)^2
+    @test dotperm(a, b, (3, 1, 2)) ≈ dot(a, permutedims(b, (3, 1, 2)))
+    @test dotperm(a, b, [3, 1, 2]) ≈ dot(a, permutedims(b, (3, 1, 2)))
+    @test dotperm(im * a, b, (3, 1, 2)) ≈ -im * dotperm(a, b, (3, 1, 2))
+    @test_throws ArgumentError dotperm(a, b, (3, 1))
+    @test_throws ArgumentError dotperm(a, b, [3, 1])
+    @test_throws ArgumentError dotperm(a, b, [3, 1, 2, 4])
+    @test_throws ArgumentError dotperm(a, b, (3, 1, 1))
+    @test_throws ArgumentError dotperm(a, b, (3, 1, 4))
+    @test_throws ArgumentError dotperm(a, vec(b), (3, 1, 2))
+    @test_throws DimensionMismatch dotperm(a, zeros(elt, 3, 4, 3), (3, 1, 2))
+    @test_throws DimensionMismatch dotperm(a, zeros(elt, 4, 3, 2), (3, 1, 2))
+    @test dotperm(fill(elt(2)), fill(elt(3)), ()) == 6
+end
+
+@testset "dotperm TensorKit (sector=$sector, eltype=$elt)" for sector in (U₁, SU₂),
+        elt in (Float64, ComplexF64)
+
+    rng = StableRNG(456)
+    V = sector === U₁ ? Rep[U₁](0 => 2, 1 => 1) : Rep[SU₂](0 => 2, 1 // 2 => 1)
+    a = randn(rng, elt, V ⊗ V ← V)
+    b = randn(rng, elt, V ⊗ V ← V)
+    @test dotperm(a, a, (1, 2, 3)) ≈ norm(a)^2
+    @test dotperm(a, b, (1, 2, 3)) ≈ dot(a, b)
+    @test dotperm(a, b, (2, 1, 3)) ≈ dot(a, TensorKit.permute(b, ((2, 1), (3,))))
+    @test_throws ArgumentError dotperm(a, b, (1, 2))
+end
+
+@testset "dotperm fermionic TensorKit rebasing" begin
+    rng = StableRNG(789)
+    V = Vect[FermionParity](0 => 2, 1 => 2)
+    a = randn(rng, ComplexF64, V ⊗ V ← V)
+    b = randn(rng, ComplexF64, V ⊗ V ← V)
+    @test dotperm(a, b, (2, 1, 3)) ≈ dot(a, TensorKit.permute(b, ((2, 1), (3,))))
+    b_rebased = TensorKit.permute(b, ((1,), (2, 3)))
+    @test dotperm(a, b_rebased, (1, 2, 3)) ≈ dot(a, b)
+    a_rebased = TensorKit.permute(a, ((1,), (2, 3)))
+    @test dotperm(a, a_rebased, (1, 2, 3)) ≈ norm(a)^2
+    for n in 0:3
+        pc, pd = Tuple(1:n), Tuple((n + 1):3)
+        a_split = TensorKit.permute(a, (pc, pd))
+        b_split = TensorKit.permute(b, (pc, pd))
+        @test dotperm(a_split, b_rebased, (1, 2, 3)) ≈ dot(a_split, b_split)
+    end
+    b_permuted = TensorKit.permute(b, ((2,), (1, 3)))
+    @test dotperm(a, b_permuted, (2, 1, 3)) ≈ dot(a, b)
+end
+
+@testset "labeled dot" begin
+    @test TensorAlgebra.dot !== dot
+    rng = StableRNG(321)
+    a = randn(rng, ComplexF64, 2, 3, 4)
+    b = randn(rng, ComplexF64, 3, 4, 2)
+    expected = dot(a, permutedims(b, (3, 1, 2)))
+    @test TensorAlgebra.dot(a, (:i, :j, :k), b, (:j, :k, :i)) ≈ expected
+    @test TensorAlgebra.dot(a, [10, 20, 30], b, [20, 30, 10]) ≈ expected
+    @test TensorAlgebra.dot(im * a, (:i, :j, :k), b, (:j, :k, :i)) ≈ -im * expected
+    @test_throws ArgumentError TensorAlgebra.dot(a, (:i, :j), b, (:j, :k, :i))
+    @test_throws ArgumentError TensorAlgebra.dot(a, [10, 20], b, [20, 30, 10])
+    @test_throws ArgumentError TensorAlgebra.dot(a, [10, 20, 30], b, [20, 30, 10, 40])
+    @test_throws ArgumentError TensorAlgebra.dot(a, (:i, :j, :k), b, (:j, :k))
+    @test_throws ArgumentError TensorAlgebra.dot(a, (:i, :i, :k), b, (:j, :k, :i))
+    @test_throws ArgumentError TensorAlgebra.dot(a, (:i, :j, :k), b, (:j, :j, :i))
+    @test_throws ArgumentError TensorAlgebra.dot(a, (:i, :j, :k), b, (:j, :k, :l))
+    @test_throws DimensionMismatch TensorAlgebra.dot(
+        a,
+        (:i, :j, :k),
+        zeros(ComplexF64, 4, 3, 2),
+        (:j, :k, :i)
+    )
+    @test TensorAlgebra.dot([1 2; 3 4], (:i, :j), [5 6; 7 8], (:j, :i)) == 69
+end
+
+@testset "labeled dot fermionic TensorKit rebasing" begin
+    rng = StableRNG(987)
+    V = Vect[FermionParity](0 => 2, 1 => 2)
+    a = randn(rng, ComplexF64, V ⊗ V ← V)
+    b = randn(rng, ComplexF64, V ⊗ V ← V)
+    b_rebased = TensorKit.permute(b, ((2,), (1, 3)))
+    @test TensorAlgebra.dot(a, (:i, :j, :k), b_rebased, (:j, :i, :k)) ≈ dot(a, b)
+    @test TensorAlgebra.dot(b_rebased, (:j, :i, :k), a, (:i, :j, :k)) ≈ conj(dot(a, b))
+end

test/test_exports.jl

@@ -26,7 +26,8 @@ using Test: @test, @testset
                 :contractperm!,
                 :contractpermadd!, :contractpermalign, :contractpermopadd!, :data,
                 :datatype,
-                :default_algorithm, :dims2cat, :directsum, :eig_full, :eig_trunc,
+                :default_algorithm, :dims2cat, :directsum, :dot, :dotperm, :eig_full,
+                :eig_trunc,
                 :eig_vals, :eigh_full, :eigh_trunc, :eigh_vals, :fill_map,
                 :flattenlinear, :has_bipartition, :infer_aux_space, :invsqrth_safe,
                 :is_output_view,