Add datatype/data, one!/unmatricize!, ungrade, matricize-based tr, and svd_trunc error

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 <support@itensor.org> and contributors"]
 
 [workspace]

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.

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")

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))

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

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

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()

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

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

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!, :datatype, :label_type,
+                :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map,
+                :to_range, :tr, :ungrade, :zero!, :scale!, :permuteddims, :PermutedDims,
             ]
         )
     end

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