diff --git a/Project.toml b/Project.toml index 718d95a..9628e9a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.16.0" +version = "0.16.1" authors = ["ITensor developers and contributors"] [workspace] @@ -16,10 +16,12 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" [weakdeps] Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" [extensions] TensorAlgebraMooncakeExt = "Mooncake" +TensorAlgebraTensorKitExt = ["TensorKit", "TensorOperations"] TensorAlgebraTensorOperationsExt = "TensorOperations" [compat] @@ -29,6 +31,7 @@ MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6" Mooncake = "0.4.202, 0.5" Strided = "2.6" StridedViews = "0.5" +TensorKit = "0.17" TensorOperations = "5" TupleTools = "1.6" julia = "1.10" diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl new file mode 100644 index 0000000..8e5c8a8 --- /dev/null +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -0,0 +1,87 @@ +module TensorAlgebraTensorKitExt + +using TensorAlgebra: TensorAlgebra +using TensorKit: TensorKit, AbstractTensorMap, ProductSpace, numind, permute, space, + spacetype, zerovector!, ← +using TensorOperations: TensorOperations as TO + +# ============================ AbstractArray-vocabulary bridge ============================ +# TensorAlgebra's generic orchestration describes operands in the `AbstractArray` vocabulary +# (`TensorAlgebra.ndims`, `TensorAlgebra.axes`), while a `TensorMap` speaks `numind`/`space`. +# Overload the TensorAlgebra-owned accessors so a `TensorMap` flows through the generic code +# unchanged. The `i`-th "axis" of a `TensorMap` is its `i`-th index space `space(t, i)`, which +# for domain indices is already dualized. +TensorAlgebra.ndims(t::AbstractTensorMap) = numind(t) +TensorAlgebra.axes(t::AbstractTensorMap, i::Int) = space(t, i) +TensorAlgebra.axes(t::AbstractTensorMap) = ntuple(i -> space(t, i), numind(t)) + +# ===================================== similar_map ======================================= +# `similar_map` takes the codomain/domain axes in codomain-facing (un-dualized) form, which is +# exactly what TensorKit's `similar(t, T, codomain, domain)` wants, so build the two +# `ProductSpace`s directly. +function TensorAlgebra.similar_map( + a::AbstractTensorMap, ::Type{T}, codomain_axes, domain_axes + ) where {T} + S = spacetype(a) + return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...)) +end + +# ================================ bipermutedimsopadd! ===================================== +# `dest = β * dest + α * permutedims(op.(src), (perm_codomain, perm_domain))`. Delegate to +# TensorKit's TensorOperations interface: `tensoradd!` realizes the permutation, the `op === conj` +# data conjugation (via `adjoint` internally), and the `α`/`β` scaling in one call. +function TensorAlgebra.bipermutedimsopadd!( + dest::AbstractTensorMap, op, src::AbstractTensorMap, + perm_codomain, perm_domain, α::Number, β::Number + ) + conjA = op === conj + (op === identity || conjA) || + throw(ArgumentError("`op` must be `identity` or `conj`, got `$op`")) + TO.tensoradd!(dest, src, (perm_codomain, perm_domain), conjA, α, β) + return dest +end + +# ================================== matricize / unmatricize ============================== +# A `TensorMap` is already a linear map codomain ← domain, so "matricizing" is just regrouping +# its indices into the requested codomain/domain bipartition (`permute`). No fusion or copy of +# the array vocabulary is needed: MatrixAlgebraKit factorizes the regrouped `TensorMap` directly. +struct TensorKitFusion <: TensorAlgebra.FusionStyle end +TensorAlgebra.FusionStyle(::Type{<:AbstractTensorMap}) = TensorKitFusion() + +function TensorAlgebra.matricize( + ::TensorKitFusion, t::AbstractTensorMap, ndims_codomain::Val{K} + ) where {K} + N = numind(t) + return permute(t, (ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K)))) +end + +# `unmatricize` reconstructs the codomain/domain axes from the matrix `m`. A `TensorMap` already +# is the linear map its space describes, so the only valid request is the one whose codomain/domain +# split matches `m`'s own space, and `unmatricize` returns `m` unchanged. The domain axes arrive +# codomain-facing (un-dualized), which is exactly TensorKit's domain convention, so they build the +# domain `ProductSpace` directly. +function TensorAlgebra.unmatricize( + ::TensorKitFusion, m::AbstractTensorMap, codomain_axes, domain_axes + ) + S = spacetype(m) + dest = ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...) + space(m) == dest || + throw(ArgumentError("`unmatricize` space `$dest` does not match `$(space(m))`")) + return m +end + +# ====================================== contract ========================================= +# Contraction of `TensorMap`s is index regrouping plus a matrix product, which TensorKit +# already implements through its TensorOperations interface. Route the generic `contract` +# there: `zero!` clears the `similar_map`-allocated destination, and the default algorithm +# hands the in-place contraction to the TensorOperations backend (see the TensorOperations +# extension's `contractopadd!`). +TensorAlgebra.zero!(t::AbstractTensorMap) = zerovector!(t) + +function TensorAlgebra.default_contract_algorithm( + ::Type{<:AbstractTensorMap}, ::Type{<:AbstractTensorMap} + ) + return TensorAlgebra.ContractAlgorithm(TO.DefaultBackend()) +end + +end diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 92bbcd2..cb4265a 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -13,6 +13,7 @@ if VERSION >= v"1.11.0-DEV.469" ) end +include("interface.jl") include("inplace.jl") include("MatrixAlgebra.jl") include("bituple.jl") diff --git a/src/interface.jl b/src/interface.jl new file mode 100644 index 0000000..50c8912 --- /dev/null +++ b/src/interface.jl @@ -0,0 +1,11 @@ +# TensorAlgebra's generic operations describe their operands in the `AbstractArray` vocabulary of +# `ndims` and `axes`. These are TensorAlgebra-owned functions, distinct from `Base.ndims`/ +# `Base.axes`, that forward to Base by default. A backend for a non-`AbstractArray` tensor type +# (such as a `TensorMap`, whose "axes" are its index spaces) overloads these instead of committing +# type piracy on the `Base` functions. +function ndims end +ndims(a) = Base.ndims(a) + +function axes end +axes(a) = Base.axes(a) +axes(a, i::Int) = Base.axes(a, i) diff --git a/test/Project.toml b/test/Project.toml index 963ee89..5b304de 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -14,6 +14,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" +TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" @@ -37,6 +38,7 @@ SafeTestsets = "0.1" StableRNGs = "1.0.2" Suppressor = "0.2" TensorAlgebra = "0.16" +TensorKit = "0.17" TensorOperations = "5.1.4" Test = "1.10" TestExtras = "0.3.1" diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl new file mode 100644 index 0000000..a9799d0 --- /dev/null +++ b/test/test_tensorkitext.jl @@ -0,0 +1,102 @@ +using StableRNGs: StableRNG +using TensorAlgebra: contract, matricize, similar_map, unmatricize +using TensorKit: @tensor, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗ +using Test: @test, @test_throws, @testset + +# A shared bond contracts when it sits in one operand's domain and the other's codomain, i.e. +# `space(a, ka) == dual(space(b, kb))`, exactly as it would in a TensorKit tensor network. +@testset "TensorKitExt (eltype = $elt)" for elt in (Float64, ComplexF64) + rng = StableRNG(1234) + + @testset "contract abelian, rank-2 bond" begin + W = Rep[U₁](0 => 2, 1 => 1) + X = Rep[U₁](0 => 1, 1 => 2) + Y = Rep[U₁](-1 => 1, 0 => 2) + a = randn(rng, elt, W, X) + b = randn(rng, elt, X, Y) + c, labels = contract(a, (:i, :j), b, (:j, :k)) + @test labels == [:i, :k] + @test c ≈ a * b + end + + @testset "contract abelian, rank-3 with permuted output" begin + A1 = Rep[U₁](0 => 2, 1 => 1) + A2 = Rep[U₁](0 => 1, 1 => 1) + B = Rep[U₁](0 => 1, -1 => 2) + C1 = Rep[U₁](0 => 2) + C2 = Rep[U₁](1 => 1, 0 => 1) + a = randn(rng, elt, A1 ⊗ A2, B) + b = randn(rng, elt, B, C1 ⊗ C2) + c, labels = contract(a, (:i, :j, :m), b, (:m, :k, :l)) + @test labels == [:i, :j, :k, :l] + @tensor ref[i, j; k, l] := a[i, j, m] * b[m, k, l] + @test c ≈ ref + end + + @testset "contract non-abelian (SU2)" begin + P = Rep[SU₂](1 // 2 => 1) + Q = Rep[SU₂](0 => 1, 1 => 1) + R = Rep[SU₂](1 // 2 => 2) + s = randn(rng, elt, P ⊗ Q, R) + w = randn(rng, elt, R, P) + c, labels = contract(s, (:i, :j, :m), w, (:m, :k)) + @test labels == [:i, :j, :k] + @tensor ref[i, j; k] := s[i, j, m] * w[m, k] + @test c ≈ ref + end + + @testset "matricize / unmatricize round-trip" begin + A1 = Rep[U₁](0 => 2, 1 => 1) + A2 = Rep[U₁](0 => 1, 1 => 1) + B = Rep[U₁](0 => 1, -1 => 2) + C1 = Rep[U₁](0 => 2) + t = randn(rng, elt, A1 ⊗ A2, B ⊗ C1) + codomain_axes = (space(t, 1), space(t, 2)) + # `unmatricize` takes the domain axes codomain-facing (un-dualized), so pass `B`, `C1` + # directly rather than the dualized `space(t, 3)`, `space(t, 4)`. + domain_axes = (B, C1) + m = matricize(t, Val(2)) + @test space(m) == space(t) + back = unmatricize(m, codomain_axes, domain_axes) + @test back ≈ t + end + + @testset "unmatricize rejects a mismatched split" begin + for (V1, V2, U) in ( + (Rep[U₁](0 => 1, 1 => 2), Rep[U₁](0 => 2, -1 => 1), Rep[U₁](0 => 1, 1 => 1)), + ( + Rep[SU₂](1 // 2 => 1, 0 => 1), + Rep[SU₂](1 // 2 => 2), + Rep[SU₂](0 => 1, 1 => 1), + ), + ) + a = randn(rng, elt, V1 ⊗ V2, U) + # Fuse the codomain into a single space so the matrix codomain no longer splits into + # `(V1, V2)`; `unmatricize` is a strict no-op, so it rejects the fused split rather + # than rewrapping the data. + m = isomorphism(elt, fuse(V1, V2), V1 ⊗ V2) * a + @test space(m) != space(a) + @test_throws ArgumentError unmatricize(m, (V1, V2), (U,)) + end + end + + @testset "similar_map space convention" begin + A1 = Rep[U₁](0 => 2, 1 => 1) + A2 = Rep[U₁](0 => 1, 1 => 1) + B = Rep[U₁](0 => 1, -1 => 2) + C1 = Rep[U₁](0 => 2) + t = randn(rng, elt, A1 ⊗ A2, B ⊗ C1) + sm = similar_map(t, elt, (A1, A2), (B, C1)) + @test space(sm) == space(t) + + # An all-codomain `TensorMap` (empty domain) is how ITensorBase direct-wraps a + # `TensorMap`, so `similar_map` must handle empty axis tuples on either side. + t_codomain = randn(rng, elt, (A1 ⊗ A2) ← one(A1)) + sm_codomain = similar_map(t_codomain, elt, (A1, A2), ()) + @test space(sm_codomain) == space(t_codomain) + + t_domain = randn(rng, elt, one(A1) ← (A1 ⊗ A2)) + sm_domain = similar_map(t_domain, elt, (), (A1, A2)) + @test space(sm_domain) == space(t_domain) + end +end