diff --git a/Project.toml b/Project.toml index 81ae9fb..9747a3e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.17.0" +version = "0.17.1" authors = ["ITensor developers and contributors"] [workspace] diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 16d10e9..ebeda57 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -36,6 +36,10 @@ function TensorAlgebra.trivialrange(::Type{S}, n::Integer) where {S <: Elementar return TensorKit.oplus(ntuple(Returns(oneunit(S)), n)...) end +# The ungraded extent of a space-backed axis is the range over its dense dimension, dropping +# sectors and the arrow so `conj`-equal spaces share an ungraded value. +TensorAlgebra.ungrade(V::ElementarySpace) = Base.OneTo(dim(V)) + # Sum of the dense elements. Through the dense presentation rather than the block data: # for a non-abelian sector type the dense embedding expands each block by its fusion-tree # structure, so the block-data sum would differ. diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 9785f9f..725b378 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -9,12 +9,13 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, permutedims, permutedims!, to_range, zero!, scale!, permuteddims, PermutedDims" + "public biperm, bipartition, contractopadd!, data, datatype, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, ungrade, zero!, scale!, permuteddims, PermutedDims" ) ) end include("interface.jl") +include("datatype.jl") include("inplace.jl") include("MatrixAlgebra.jl") include("bituple.jl") diff --git a/src/datatype.jl b/src/datatype.jl new file mode 100644 index 0000000..2be8151 --- /dev/null +++ b/src/datatype.jl @@ -0,0 +1,46 @@ +""" + data(a) + +The underlying storage of `a`: the value reached by following `parent` to its fixed point +(an object that is its own `parent`). A wrapper returns the storage it ultimately wraps, and +a plain array returns itself. + +# Examples + +```jldoctest +julia> import TensorAlgebra + +julia> a = [1.0 2.0; 3.0 4.0]; + +julia> TensorAlgebra.data(transpose(a)) === a +true + +julia> TensorAlgebra.data(a) === a +true +``` +""" +function data(a) + parent_a = parent(a) + parent_a === a && return a + return data(parent_a) +end + +""" + datatype(a) -> Type + +The type of the underlying storage of `a`, i.e. `typeof(data(a))`, in contrast to +`scalartype`/`eltype`, which give its element type alone. + +# Examples + +```jldoctest +julia> import TensorAlgebra + +julia> TensorAlgebra.datatype([1.0 2.0; 3.0 4.0]) +Matrix{Float64} (alias for Array{Float64, 2}) + +julia> TensorAlgebra.datatype(transpose([1.0 2.0; 3.0 4.0])) +Matrix{Float64} (alias for Array{Float64, 2}) +``` +""" +datatype(a) = typeof(data(a)) diff --git a/src/factorizations.jl b/src/factorizations.jl index ef17ba8..85c2329 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -70,6 +70,48 @@ for f in ( end end +""" + TensorAlgebra.tr(A, labels_A, labels_codomain, labels_domain) + TensorAlgebra.tr(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) + TensorAlgebra.tr(A, ndims_codomain::Val) + +Trace of a generic N-dimensional array `A` interpreted as a linear map from its domain to its +codomain dimensions. The map is matricized into its square matrix, then the matrix trace is +taken, so the backend's own matrix `tr` (dense, graded, or `TensorMap`) does the work. The +partition is specified via labels, a bi-permutation, or directly as the codomain rank, matching +the factorization entry points. + +This is `TensorAlgebra`'s own function, distinct from `LinearAlgebra.tr`; the two-argument and +higher forms take a codomain/domain partition rather than a bare matrix. + +# Examples + +```jldoctest +julia> import TensorAlgebra + +julia> A = randn(2, 2, 2, 2); + +julia> TensorAlgebra.tr(A, (:i, :j, :k, :l), (:i, :k), (:j, :l)) ≈ + sum(A[i, i, k, k] for i in 1:2, k in 1:2) +true +``` +""" +function tr(style::FusionStyle, A, ndims_codomain::Val) + return LinearAlgebra.tr(matricize(style, A, ndims_codomain)) +end +function tr(A, ndims_codomain::Val) + return tr(FusionStyle(A), A, ndims_codomain) +end +function tr(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) + A_perm = bipermutedims(A, perm_codomain, perm_domain) + return tr(A_perm, Val(length(perm_codomain))) +end +function tr(A, labels_A, labels_codomain, labels_domain) + perm_codomain, perm_domain = + biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) + return tr(A, perm_codomain, perm_domain) +end + """ qr_compact(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R qr_compact(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R @@ -212,7 +254,7 @@ right_orth # Three-output SVD: `U` carries the codomain axes plus a trailing rank axis, `S` is the # rank × rank spectrum, and `Vᴴ` carries a leading rank axis plus the domain axes. -for f in (:svd_compact, :svd_full, :svd_trunc) +for f in (:svd_compact, :svd_full) @eval begin function $f(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) @@ -228,6 +270,22 @@ for f in (:svd_compact, :svd_full, :svd_trunc) end end +# `svd_trunc` matches the three-output SVD but additionally surfaces the truncation error +# `ϵ` (the 2-norm of the discarded singular values, computed by MatrixAlgebraKit without +# catastrophic cancellation), so it is spelled out here rather than sharing the loop above. +function svd_trunc(style::FusionStyle, A, ndims_codomain::Val; kwargs...) + A_mat = matricize(style, A, ndims_codomain) + U, S, Vᴴ, ϵ = MatrixAlgebraKit.svd_trunc(A_mat; kwargs...) + axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) + return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), + unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)), + unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain), + ϵ +end +function svd_trunc(A, ndims_codomain::Val; kwargs...) + return svd_trunc(FusionStyle(A), A, ndims_codomain; kwargs...) +end + # Eigendecomposition: `D` is the rank × rank spectrum, left as a matrix, while `V` carries # the codomain axes plus a trailing rank axis. for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) @@ -284,19 +342,38 @@ See also `MatrixAlgebraKit.svd_full!`. svd_full """ - svd_trunc(A, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ + svd_trunc(A, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ, ϵ + svd_trunc(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ, ϵ + svd_trunc(A, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ, ϵ Compute the truncated SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either via -labels or directly through a bi-permutation. +labels or directly through a bi-permutation. In addition to the factors, returns the +truncation error `ϵ`, the 2-norm of the discarded singular values. ## Keyword arguments - `trunc`: truncation strategy, passed on to `MatrixAlgebraKit.svd_trunc`. - Other keywords are passed on directly to MatrixAlgebraKit. +# Examples + +```jldoctest +julia> using TensorAlgebra: svd_trunc, contract + +julia> A = randn(4, 4); + +julia> U, S, Vᴴ, ϵ = svd_trunc(A, (:i, :j), (:i,), (:j,)); + +julia> SV = contract((:u, :j), S, (:u, :v), Vᴴ, (:v, :j)); + +julia> contract((:i, :j), U, (:i, :u), SV, (:u, :j)) ≈ A +true + +julia> isapprox(ϵ, 0; atol = 1e-10) +true +``` + See also `MatrixAlgebraKit.svd_trunc!`. """ svd_trunc @@ -741,6 +818,19 @@ function one!!(A, ndims_codomain::Val; kwargs...) return one!!(FusionStyle(A), A, ndims_codomain; kwargs...) end +# In-place identity fill: writes the identity into `A` and returns it. Matricizes `A`, fills the +# fused matrix with the identity, and — when the matricized form is a detached copy (a graded +# gather) rather than a view aliasing `A` (a dense reshape) — scatters it back with `unmatricize!`. +function one!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) + A_mat = matricize(style, A, ndims_codomain) + MatrixAlgebraKit.one!(A_mat) + Base.mightalias(A_mat, A) && return A + return unmatricize!(A, A_mat, ndims_codomain) +end +function one!(A, ndims_codomain::Val; kwargs...) + return one!(FusionStyle(A), A, ndims_codomain; kwargs...) +end + function one(style::FusionStyle, A, ndims_codomain::Val; kwargs...) return one!!(style, copy(A), ndims_codomain; kwargs...) end diff --git a/src/interface.jl b/src/interface.jl index 660e2be..9f7854d 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -14,9 +14,23 @@ function size end size(a) = Base.size(a) size(a, i::Int) = Base.size(a, i) -# The scalar held by a rank-0 tensor. The Base spelling is `a[]`, which a `TensorMap` -# with a nontrivial sector type does not support (TensorKit provides `scalar` instead). +""" + scalar(a) + +The single scalar held by a rank-0 (zero-dimensional) `a`, i.e. `a[]`. + +# Examples + +```jldoctest +julia> import TensorAlgebra + +julia> TensorAlgebra.scalar(fill(3.0)) +3.0 +``` +""" function scalar end +# The Base spelling is `a[]`, which a `TensorMap` with a nontrivial sector type does not +# support (TensorKit provides `scalar` instead). scalar(a) = a[] # The sum of the (dense) elements. A `TensorMap` is not iterable, so `Base.sum` does not diff --git a/src/matricize.jl b/src/matricize.jl index 2f1d11f..9113b7f 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -272,6 +272,21 @@ function unmatricizeperm!( return bipermutedims!(a_dest, a_perm, biperm_dest) end +# In-place split-axes counterpart of `unmatricize`, as `unmatricizeperm!` is of `unmatricizeperm`: +# scatter the fused matrix `m` back into `a_dest`'s existing storage across the codomain/domain +# split at `ndims_codomain`. The split applies no permutation, so this is `unmatricizeperm!` at the +# trivial bipermutation, reusing its in-place block scatter (no intermediate `unmatricize` copy). +function unmatricize!(a_dest, m, ndims_codomain::Val) + K = unval(ndims_codomain) + N = ndims(a_dest) + return unmatricizeperm!( + a_dest, + m, + ntuple(identity, Val(K)), + ntuple(i -> K + i, Val(N - K)) + ) +end + # Defaults to ReshapeFusion, a simple reshape struct ReshapeFusion <: FusionStyle end FusionStyle(::Type{<:AbstractArray}) = ReshapeFusion() diff --git a/src/to_range.jl b/src/to_range.jl index ee617ba..d31b100 100644 --- a/src/to_range.jl +++ b/src/to_range.jl @@ -13,3 +13,23 @@ that turns a vector of sector-to-multiplicity pairs into a graded range. """ to_range(space::AbstractUnitRange) = space to_range(space::Integer) = Base.OneTo(space) + +""" + TensorAlgebra.ungrade(r) + +Return the plain range underlying an axis, keeping only its extent and stripping any added +structure. On a plain `AbstractUnitRange` this is the identity (there is nothing to strip, and any +offset is preserved). Downstream packages extend it for richer axes: GradedArrays maps a graded +range to the `Base.OneTo` of its total dimension, and a native TensorKit space maps to the +`Base.OneTo` of its dimension. + +# Examples + +```jldoctest +julia> import TensorAlgebra + +julia> TensorAlgebra.ungrade(2:5) +2:5 +``` +""" +ungrade(r::AbstractUnitRange) = r diff --git a/test/test_datatype.jl b/test/test_datatype.jl new file mode 100644 index 0000000..a972c19 --- /dev/null +++ b/test/test_datatype.jl @@ -0,0 +1,31 @@ +using LinearAlgebra: Diagonal, transpose +using TensorAlgebra: data, datatype +using Test: @test, @testset + +@testset "data" begin + # A plain array is its own storage. + a = randn(2, 3) + @test data(a) === a + + # Wrappers recurse through `parent` to the underlying storage. + @test data(transpose(a)) === a + @test data(view(transpose(a), 1:2, 1:2)) === a +end + +@testset "datatype" begin + # A plain array is its own storage type. + a = randn(2, 3) + @test datatype(a) === Matrix{Float64} + + # Wrappers recurse through `parent` to the underlying storage. + @test datatype(transpose(a)) === Matrix{Float64} + @test datatype(view(a, 1:2, 1:2)) === Matrix{Float64} + @test datatype(Diagonal([1.0, 2.0])) === Vector{Float64} + + # A wrapper of a wrapper recurses all the way down; running to completion here is the + # check that the `parent(x) === x` base case terminates the recursion. + @test datatype(view(transpose(a), 1:2, 1:2)) === Matrix{Float64} + + # Element type is captured alongside the container. + @test datatype(randn(ComplexF32, 2)) === Vector{ComplexF32} +end diff --git a/test/test_exports.jl b/test/test_exports.jl index 67b2046..b7e516e 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -38,9 +38,9 @@ using Test: @test, @testset append!( exports, [ - :biperm, :bipartition, :contractopadd!, :label_type, :matricizeopperm, - :permutedims, :permutedims!, :to_range, :zero!, :scale!, :permuteddims, - :PermutedDims, + :biperm, :bipartition, :contractopadd!, :data, :datatype, :label_type, + :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map, + :to_range, :tr, :ungrade, :zero!, :scale!, :permuteddims, :PermutedDims, ] ) end diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index cdff41b..c6d7495 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -202,13 +202,15 @@ end _, S_untrunc, _ = svd_compact(A, labels_A, labels_U, labels_Vᴴ) trunc = truncrank(size(S_untrunc, 1) - 1) - U, S, Vᴴ = @constinferred svd_trunc(A, labels_A, labels_U, labels_Vᴴ; trunc) + U, S, Vᴴ, ϵ = @constinferred svd_trunc(A, labels_A, labels_U, labels_Vᴴ; trunc) @test A == Acopy # should not have altered initial array US, labels_US = contract(U, (labels_U..., :u), S, (:u, :v)) A′ = contract(labels_A, US, labels_US, Vᴴ, (:v, labels_Vᴴ...)) @test norm(A - A′) ≈ S_untrunc[end] @test size(S, 1) == size(S_untrunc, 1) - 1 + # `ϵ` is the 2-norm of the discarded singular values (here the single dropped value). + @test ϵ ≈ S_untrunc[end] end @testset "Nullspace ($T)" for T in elts @@ -383,4 +385,32 @@ end @test TensorAlgebra.matricize(Id_perm, Val(2)) ≈ I # Perm- and biperm-tuple forms agree with the label form. @test TensorAlgebra.one(B, (1, 3), (2, 4)) ≈ Id_perm + + # In-place `one!` fills the identity into its argument and returns it. + C = randn(T, 2, 3, 2, 3) + Cret = @constinferred TensorAlgebra.one!(C, Val(2)) + @test Cret === C + @test TensorAlgebra.matricize(C, Val(2)) ≈ I + @test C ≈ TensorAlgebra.one(A, Val(2)) + + # `unmatricize!` scatters a fused matrix back into an existing array. + D = randn(T, 2, 3, 2, 3) + Dmat = TensorAlgebra.matricize(D, Val(2)) + E = similar(D) + Eret = TensorAlgebra.unmatricize!(E, Dmat, Val(2)) + @test Eret === E + @test E ≈ D +end + +# Trace +# ----- +@testset "tr ($T)" for T in elts + A = randn(T, 2, 3, 2, 3) + m = reshape(A, 6, 6) + # The labels, bi-permutation, and codomain-rank forms all agree with the matrix trace of + # the matricized map. + @test TensorAlgebra.tr(A, (:i, :j, :ip, :jp), (:i, :j), (:ip, :jp)) ≈ + LinearAlgebra.tr(m) + @test TensorAlgebra.tr(A, (1, 2), (3, 4)) ≈ LinearAlgebra.tr(m) + @test TensorAlgebra.tr(A, Val(2)) ≈ LinearAlgebra.tr(m) end diff --git a/test/test_scalar.jl b/test/test_scalar.jl new file mode 100644 index 0000000..89098c5 --- /dev/null +++ b/test/test_scalar.jl @@ -0,0 +1,10 @@ +using TensorAlgebra: scalar +using Test: @test, @testset + +@testset "scalar" begin + @test scalar(fill(3.0)) === 3.0 + a = Array{Float64, 0}(undef) + a[] = 5.0 + @test scalar(a) === 5.0 + @test scalar(fill(2 + 3im)) === 2 + 3im +end diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index 6e69b61..2545082 100644 --- a/test/test_tensorkitext.jl +++ b/test/test_tensorkitext.jl @@ -1,10 +1,10 @@ using Base.Broadcast: broadcasted -using LinearAlgebra: norm +using LinearAlgebra: LinearAlgebra, norm using StableRNGs: StableRNG using TensorAlgebra: TensorAlgebra, contract, matricize, project, projectto!, rand_map, randn_map, similar_map, tryflattenlinear, tryproject, unchecked_project, unmatricize, zeros_map -using TensorKit: @tensor, AbstractTensorMap, Rep, SU₂, TensorMap, U₁, dual, fuse, +using TensorKit: @tensor, AbstractTensorMap, Rep, SU₂, TensorMap, U₁, dim, dual, fuse, isomorphism, randn, space, storagetype, ←, ⊗ using Test: @test, @test_throws, @testset @@ -299,4 +299,17 @@ using Test: @test, @test_throws, @testset @test isnothing(tryflattenlinear(broadcasted(*, a, b))) @test_throws ErrorException copy(broadcasted(*, a, b)) end + + @testset "ungrade / tr" begin + W = Rep[U₁](0 => 2, 1 => 1) + X = Rep[U₁](0 => 1, 1 => 2) + # `ungrade` drops sectors and the arrow, so a space and its dual share the ungraded extent. + @test TensorAlgebra.ungrade(W) == Base.OneTo(dim(W)) + @test TensorAlgebra.ungrade(dual(W)) == TensorAlgebra.ungrade(W) + + # `tr` over a codomain/domain bipartition matches TensorKit's native trace of the endomorphism. + t = randn(rng, elt, W ⊗ X, W ⊗ X) + @test TensorAlgebra.tr(t, (:i, :j, :ip, :jp), (:i, :j), (:ip, :jp)) ≈ + LinearAlgebra.tr(t) + end end diff --git a/test/test_to_range.jl b/test/test_to_range.jl index 20429ef..73d8ed9 100644 --- a/test/test_to_range.jl +++ b/test/test_to_range.jl @@ -13,3 +13,10 @@ using Test: @test, @testset end end end + +@testset "ungrade" begin + # On a plain range the ungraded range is the identity, offset preserved. + for r in (Base.OneTo(4), 2:5) + @test TensorAlgebra.ungrade(r) === r + end +end