diff --git a/Project.toml b/Project.toml index ce16e5d..bc6956e 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 69e2f30..d5f03e6 100644 --- a/src/TensorAlgebra.jl +++ b/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") diff --git a/src/dotperm.jl b/src/dotperm.jl new file mode 100644 index 0000000..20855c3 --- /dev/null +++ b/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 diff --git a/test/test_dotperm.jl b/test/test_dotperm.jl new file mode 100644 index 0000000..86bdd7d --- /dev/null +++ b/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 diff --git a/test/test_exports.jl b/test/test_exports.jl index ca60fdf..67e6b20 100644 --- a/test/test_exports.jl +++ b/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,