From 83cd8c12ec29ad22a87dd94baaa57e8fb6338e44 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sun, 5 Jul 2026 15:46:52 -0400 Subject: [PATCH 01/15] Add Hermitian-safe matrix functions and symmetric-backend factorization glue Adds a Hermitian-safe matrix function family to `MatrixAlgebra` (`sqrth_safe`, `invsqrth_safe`, and `sqrth_invsqrth_safe`) that projects onto the Hermitian part and clamps eigenvalues below tolerance in the Moore-Penrose convention, with labeled, permutation, and `Val` tensor forms that work on dense, graded, and `TensorMap` matricizations. Adds a tensor-level `project_hermitian` matching the `MatrixAlgebraKit` function. Factorization outputs now unfuse through `unmatricize_codomain` and `unmatricize_domain`, which backends whose `matricize` does not fuse (TensorKit) override to validate and return the factor unchanged. `project` tolerates trailing length-1 axes on size-strict backends. Structured matrices such as `Diagonal` flow through `bipermutedimsopadd!` via `permuteddims` and the `isstrided` trait. Raises the `MatrixAlgebraKit` compat floor to `0.6`, the first version with `project_hermitian`. Co-authored-by: Claude Fable 5 --- Project.toml | 4 +- docs/Project.toml | 2 +- examples/Project.toml | 2 +- ext/TensorAlgebraTensorKitExt.jl | 157 +++++++++++++++++++++++++++---- src/MatrixAlgebra.jl | 81 ++++++++++++++-- src/TensorAlgebra.jl | 7 +- src/diagonal.jl | 8 +- src/factorizations.jl | 146 ++++++++++++++++++++++++++-- src/interface.jl | 32 ++++++- src/permutedimsadd.jl | 11 ++- src/projectto.jl | 152 ++++++++++++++++++++---------- test/Project.toml | 2 +- test/test_exports.jl | 5 + test/test_projectto.jl | 117 ++++++++++++++--------- test/test_tensorkitext.jl | 107 +++++++++++++++++---- 15 files changed, 670 insertions(+), 163 deletions(-) diff --git a/Project.toml b/Project.toml index 4f1904a..81ae9fb 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.16.5" +version = "0.17.0" authors = ["ITensor developers and contributors"] [workspace] @@ -28,7 +28,7 @@ TensorAlgebraTensorOperationsExt = "TensorOperations" [compat] EllipsisNotation = "1.8" LinearAlgebra = "1.10" -MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6" +MatrixAlgebraKit = "0.6" Mooncake = "0.4.202, 0.5" Random = "1.10" Strided = "2.6" diff --git a/docs/Project.toml b/docs/Project.toml index 065b705..de734b6 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -11,4 +11,4 @@ path = ".." Documenter = "1.8.1" ITensorFormatter = "0.2.27" Literate = "2.20.1" -TensorAlgebra = "0.16" +TensorAlgebra = "0.17" diff --git a/examples/Project.toml b/examples/Project.toml index 480a18f..70b75c7 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" path = ".." [compat] -TensorAlgebra = "0.16" +TensorAlgebra = "0.17" diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index b49ecb7..e41f7d0 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -2,9 +2,9 @@ module TensorAlgebraTensorKitExt using Random: AbstractRNG using TensorAlgebra: TensorAlgebra -using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, - TensorMapWithStorage, numind, permute, project_symmetric!, space, spacetype, - zerovector!, ← +using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, TensorMap, + TensorMapWithStorage, codomain, dim, domain, dual, fuse, numind, permute, + project_symmetric!, sectors, space, spacetype, zerovector!, ← using TensorOperations: TensorOperations as TO # ============================ AbstractArray-vocabulary bridge ============================ @@ -16,6 +16,35 @@ using TensorOperations: TensorOperations as TO 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)) +# A `TensorMap` has no `Base.size`; its dense size is the per-index space dimension. +TensorAlgebra.size(t::AbstractTensorMap, i::Int) = dim(space(t, i)) +TensorAlgebra.size(t::AbstractTensorMap) = ntuple(i -> dim(space(t, i)), numind(t)) + +# TensorKit deliberately defines no `Base.conj` for tensors; its native conjugation is +# `adjoint`, which also swaps the codomain and domain. `TensorAlgebra.conjugate` +# (conjugated elements on dualized spaces, leg order unchanged) is the adjoint followed by +# the block swap returning each leg to its original flat position — the inverse of +# TensorKit's `adjointtensorindex` convention (`i <= numout ? numin + i : i - numout`), +# the same lowering its TensorOperations interface uses for `conj` arguments. +function TensorAlgebra.conjugate(t::AbstractTensorMap) + m, n = TensorKit.numout(t), TensorKit.numin(t) + p = ntuple(i -> i <= m ? n + i : i - m, m + n) + return permute(adjoint(t), (p[1:m], p[(m + 1):(m + n)])) +end + +# `t[]` on a rank-0 `TensorMap` requires a trivial sector type; `TensorKit.scalar` is the +# general spelling. +TensorAlgebra.scalar(t::AbstractTensorMap) = TensorKit.scalar(t) + +# The trivial length-1 axis of a space is its unit space (`oneunit`), the trivial-sector +# one-dimensional space. +TensorAlgebra.trivialrange(V::ElementarySpace) = oneunit(V) +TensorAlgebra.trivialrange(::Type{S}) where {S <: ElementarySpace} = oneunit(S) + +# 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. +TensorAlgebra.sum(t::AbstractTensorMap; kwargs...) = Base.sum(convert(Array, t); kwargs...) # ===================================== similar_map ======================================= # `similar_map` takes the codomain/domain axes in codomain-facing (un-dualized) form, which is @@ -99,24 +128,73 @@ end # to TensorKit's `project_symmetric!`, which fills the symmetry-allowed blocks from the dense data # and discards any component outside the block structure. Composed with the map constructors above, # this makes `project(dense, codomain_axes, domain_axes)` build a `TensorMap` from a dense matrix. +# `project_symmetric!` requires a matching dense size, so reshape `src` to `size(dest)` first (a +# no-op when the ranks already match); this lets a lower-rank `src` omit trailing length-1 axes, +# matching the generic `projectto!`. function TensorAlgebra.projectto!(dest::AbstractTensorMap, src::AbstractArray) - return project_symmetric!(dest, src) -end - -# The generic `checked_projectto!` verifies the projection with `isapprox(src, dest)`, but a -# `TensorMap` `dest` is not elementwise-comparable to the dense `src`. Densify `dest` with -# `convert(Array, ...)` so the check is the same elementwise `isapprox(src, dest)` as the dense path, -# keeping one `InexactError`/`kwargs` contract across backends rather than TensorKit's own -# residual-norm `tol`/`ArgumentError` check. -function TensorAlgebra.checked_projectto!( - dest::AbstractTensorMap, - src::AbstractArray; - kwargs... - ) - TensorAlgebra.projectto!(dest, src) - isapprox(src, convert(Array, dest); kwargs...) || - throw(InexactError(:checked_projectto!, typeof(dest), src)) - return dest + return project_symmetric!(dest, reshape(src, TensorAlgebra.size(dest))) +end + +# The `is_projected` check compares through `convert(Array, dest)`, which TensorKit already +# defines for an `AbstractTensorMap`, so no dedicated method is needed here. + +# ============================= allocate_project (aux-leg derivation) ===================== +# `allocate_project` with `TensorMap` spaces: the destination allocation, which is where a +# trailing surplus axis in `raw` — an auxiliary leg appended as the last domain axis, matching +# the shape of `raw` — gets its space derived so the result is symmetry-allowed. The candidate +# irreps are the operator content `codomain ⊗ conj(domain)`; scan the aux axis in the content's +# canonical (sorted) sector order, consuming contiguous width-`dim(s)` slices that are covariant +# with irrep `s` (`TensorMap(data, space)` does the projection and the covariance check per +# slice). The result is a possibly multi-sector (direct-sum) aux, e.g. an MPO-style virtual leg; +# a single irrep and the abelian single-charge case fall out. The slice order must match the +# canonical sector order, since a `GradedSpace` sorts its sectors and the dense layout follows +# it. The fill onto the derived axes (`projectto!`) is magnitude-blind; the `project` wrapper +# verifies that nothing was discarded. With no surplus axis (a lower-rank `raw` omits trailing +# length-1 axes, which the `projectto!` reshape pads), this is plain `similar_map`. +function TensorAlgebra.allocate_project( + raw::AbstractArray, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} + ) where {S <: ElementarySpace} + nphys = length(codomain_axes) + length(domain_axes) + if ndims(raw) > nphys + ndims(raw) == nphys + 1 || throw( + ArgumentError( + "`project`: expected at most one trailing auxiliary axis beyond the \ + $nphys given axes, got a rank-$(ndims(raw)) input" + ) + ) + auxdim = size(raw, nphys + 1) + content = fuse(codomain_axes..., dual.(domain_axes)...) + function slice_is_covariant(r, s) + target = _map_homspace(S, codomain_axes, (domain_axes..., S(s => 1))) + return try + TensorMap(selectdim(raw, nphys + 1, r), target) + true + catch e + e isa Union{ArgumentError, DimensionMismatch} || rethrow() + false + end + end + seccounts = Pair{TensorKit.sectortype(S), Int}[] + pos = 1 + for s in sectors(content) + d = dim(S(s => 1)) + m = 0 + while pos + d - 1 <= auxdim && slice_is_covariant(pos:(pos + d - 1), s) + m += 1 + pos += d + end + m > 0 && push!(seccounts, s => m) + end + pos == auxdim + 1 || throw( + ArgumentError( + "`project`: could not derive a covariant auxiliary space for the surplus axis \ + of dimension $auxdim; the aux slices must be ordered by the canonical (sorted) \ + sector order of the derived space" + ) + ) + domain_axes = (domain_axes..., S(seccounts...)) + end + return TensorAlgebra.similar_map(raw, codomain_axes, domain_axes) end # ================================ bipermutedimsopadd! ===================================== @@ -148,6 +226,14 @@ function TensorAlgebra.matricize( return permute(t, (ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K)))) end +# The identity fill on the regrouped map is TensorKit's own `one!` (MatrixAlgebraKit's +# `one!` speaks `AbstractMatrix` only). +function TensorAlgebra.one!!( + style::TensorKitFusion, A::AbstractTensorMap, ndims_codomain::Val; kwargs... + ) + return TensorKit.one!(TensorAlgebra.matricize(style, A, ndims_codomain)) +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 @@ -163,6 +249,37 @@ function TensorAlgebra.unmatricize( return m end +# A factorization factor needs no unfusing here (`matricize` above does not fuse, so the +# factors come back from MatrixAlgebraKit with their original legs plus the bond). The +# generic defaults would reconstruct the bond axis with the matrix-flavored `axes(X, 2)`, +# which on an unfused factor reads a codomain leg instead of the bond (the two only +# coincide when the bond carries the trivial sector). Validate the known side and return +# the factor unchanged. +function TensorAlgebra.unmatricize_codomain( + ::TensorKitFusion, X::AbstractTensorMap, codomain_axes + ) + S = spacetype(X) + dest = ProductSpace{S}(codomain_axes...) + codomain(X) == dest || throw( + ArgumentError( + "`unmatricize_codomain` space `$dest` does not match `$(codomain(X))`" + ) + ) + return X +end +function TensorAlgebra.unmatricize_domain( + ::TensorKitFusion, Y::AbstractTensorMap, domain_axes + ) + S = spacetype(Y) + # The requested axes arrive codomain-facing (un-dualized), which is TensorKit's + # domain convention, so they build the domain `ProductSpace` directly. + dest = ProductSpace{S}(domain_axes...) + domain(Y) == dest || throw( + ArgumentError("`unmatricize_domain` space `$dest` does not match `$(domain(Y))`") + ) + return Y +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` diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index d3ed5e6..d7483c9 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -7,7 +7,8 @@ export gram_eigh_full, pow_diag_safe, powh_safe, sqrt_diag_safe, - sqrth_safe + sqrth_safe, + sqrth_invsqrth_safe using LinearAlgebra: LinearAlgebra, Diagonal, isdiag, norm using MatrixAlgebraKit: MatrixAlgebraKit as MAK @@ -54,6 +55,37 @@ function pow_diag_safe( return MAK.diagonal(map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ)) end +# Non-`AbstractMatrix` matrix-like backends (e.g. a `TensorMap`): `map` over the +# `MAK.diagview` may not preserve the backend's block structure, so write the mapped +# diagonal back into a `copy` through the (aliasing) `diagview` instead of rebuilding +# with `MAK.diagonal`. The tolerance uses `norm(D, Inf)` (the largest-magnitude entry), +# which equals the largest-magnitude diagonal entry under the `isdiag` guard. +function pow_diag_safe( + D, p; + atol = zero(real(eltype(D))), + rtol = iszero(atol) ? eps(real(eltype(D)))^(3 // 4) : + zero(real(eltype(D))) + ) + isdiag(D) || throw( + ArgumentError("pow_diag_safe requires a diagonal-structured matrix") + ) + tol = max(atol, rtol * norm(D, Inf)) + Dp = copy(D) + _pow_diag!(MAK.diagview(Dp), p, tol) + return Dp +end + +function _pow_diag!(σ, p, tol) + σ .= map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ) + return σ +end +# A backend's `diagview` may be a dict of per-block diagonal views (e.g. a `TensorMap`, +# keyed by sector) rather than a single vector view. +function _pow_diag!(σ::AbstractDict, p, tol) + foreach(v -> _pow_diag!(v, p, tol), values(σ)) + return σ +end + """ sqrt_diag_safe(D::AbstractMatrix; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^(1//2) @@ -80,12 +112,14 @@ $(_clamp_kwargs_doc("D")) invsqrt_diag_safe(D::AbstractMatrix; kwargs...) = pow_diag_safe(D, -1 // 2; kwargs...) """ - powh_safe(M::AbstractMatrix, p; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^p + powh_safe(M, p; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^p Raise an approximately Hermitian positive semi-definite matrix to the power `p`. For diagonal-structured `M` (`isdiag(M) == true`), dispatches to [`pow_diag_safe`](@ref) and skips the eigendecomposition. Otherwise, -computes via `M = V * D * V'` as `V * pow_diag_safe(D, p; atol, rtol) * V'`. +projects `M` onto its Hermitian part `(M + M') / 2` (so input that is +Hermitian only up to numerical noise is accepted) and computes via +`M = V * D * V'` as `V * pow_diag_safe(D, p; atol, rtol) * V'`. ## Keyword arguments @@ -93,14 +127,19 @@ computes via `M = V * D * V'` as `V * pow_diag_safe(D, p; atol, rtol) * V'`. $(_clamp_kwargs_doc("M")) """ -function powh_safe(M::AbstractMatrix, p; alg = nothing, kwargs...) +function powh_safe(M, p; alg = nothing, kwargs...) isdiag(M) && return pow_diag_safe(M, p; kwargs...) - D, V = MAK.eigh_full(M, MAK.select_algorithm(MAK.eigh_full, M, alg)) + D, V = _eigh_full_hermitianpart(M, alg) return V * pow_diag_safe(D, p; kwargs...) * V' end +function _eigh_full_hermitianpart(M, alg) + M = MAK.project_hermitian(M) + return MAK.eigh_full(M, MAK.select_algorithm(MAK.eigh_full, M, alg)) +end + """ - sqrth_safe(M::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2) + sqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2) Square root of an approximately Hermitian positive semi-definite matrix. Equivalent to `powh_safe(M, 1//2; alg, atol, rtol)`. @@ -111,10 +150,10 @@ Equivalent to `powh_safe(M, 1//2; alg, atol, rtol)`. $(_clamp_kwargs_doc("M")) """ -sqrth_safe(M::AbstractMatrix; kwargs...) = powh_safe(M, 1 // 2; kwargs...) +sqrth_safe(M; kwargs...) = powh_safe(M, 1 // 2; kwargs...) """ - invsqrth_safe(M::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(-1//2) + invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(-1//2) Inverse square root of an approximately Hermitian positive semi-definite matrix. Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. @@ -125,7 +164,31 @@ matrix. Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. $(_clamp_kwargs_doc("M")) """ -invsqrth_safe(M::AbstractMatrix; kwargs...) = powh_safe(M, -1 // 2; kwargs...) +invsqrth_safe(M; kwargs...) = powh_safe(M, -1 // 2; kwargs...) + +""" + sqrth_invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2), M^(-1//2) + +Square root and pseudo-inverse square root of an approximately Hermitian +positive semi-definite matrix, from a single eigendecomposition. Equivalent +to `(sqrth_safe(M; ...), invsqrth_safe(M; ...))` but with the +eigendecomposition computed once. Eigenvalues below tolerance are clamped +to zero in both factors (Moore-Penrose convention for the inverse). + +## Keyword arguments + + - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. + +$(_clamp_kwargs_doc("M")) +""" +function sqrth_invsqrth_safe(M; alg = nothing, kwargs...) + if isdiag(M) + return pow_diag_safe(M, 1 // 2; kwargs...), pow_diag_safe(M, -1 // 2; kwargs...) + end + D, V = _eigh_full_hermitianpart(M, alg) + return V * pow_diag_safe(D, 1 // 2; kwargs...) * V', + V * pow_diag_safe(D, -1 // 2; kwargs...) * V' +end for (gram, gram_with_pinv, eigh_full) in ( (:gram_eigh_full, :gram_eigh_full_with_pinv, :eigh_full), diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 35317cd..ce3d708 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -1,9 +1,10 @@ module TensorAlgebra export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, - eigh_vals, gram_eigh_full, gram_eigh_full_with_pinv, left_null, left_orth, - left_polar, lq_compact, lq_full, qr_compact, qr_full, right_null, right_orth, - right_polar, svd_compact, svd_full, svd_trunc, svd_vals + eigh_vals, gram_eigh_full, gram_eigh_full_with_pinv, invsqrth_safe, left_null, + left_orth, left_polar, lq_compact, lq_full, project_hermitian, qr_compact, + qr_full, right_null, right_orth, right_polar, sqrth_invsqrth_safe, sqrth_safe, + svd_compact, svd_full, svd_trunc, svd_vals if VERSION >= v"1.11.0-DEV.469" eval( diff --git a/src/diagonal.jl b/src/diagonal.jl index f836d93..48d03bd 100644 --- a/src/diagonal.jl +++ b/src/diagonal.jl @@ -10,8 +10,12 @@ using LinearAlgebra: Diagonal # (the latter falls back to Base's dense `similar`, since `contract` allocates from a flat # axis tuple, not a `BiTuple`). -# Permuting the two axes of a square `Diagonal` (identity or transpose) leaves the stored -# diagonal unchanged, so accumulate straight onto it. +# Permuting the two axes of a square `Diagonal` (identity or transpose) leaves it +# unchanged, so the lazy permutation is the matrix itself. +permuteddims(a::Diagonal, perm) = a + +# Same reasoning for the in-place accumulate: skip the permutation and accumulate +# straight onto the destination. function bipermutedimsopadd!( dest::Diagonal, op, src::Diagonal, perm_codomain, perm_domain, diff --git a/src/factorizations.jl b/src/factorizations.jl index 0cf471c..14ef186 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -1,6 +1,20 @@ using LinearAlgebra: LinearAlgebra using MatrixAlgebraKit: MatrixAlgebraKit +# Unfuse one side of a factorization factor, keeping its other (bond) axis as is: a +# factor `X` carries the fused codomain (or domain) group of the input on one axis and +# the new bond on the other, and only the fused group needs reconstructing. For fusing +# backends the bond axis is read off the factor itself (`axes(X, 2)`, dualized to stay +# codomain-facing). A backend whose `matricize` does not fuse (e.g. a `TensorMap`) +# overrides these to return the factor unchanged: its factors keep their original legs, +# so `axes(X, 2)` does not address the bond. +function unmatricize_codomain(style::FusionStyle, X, axes_codomain) + return unmatricize(style, X, axes_codomain, (conj(axes(X, 2)),)) +end +function unmatricize_domain(style::FusionStyle, Y, axes_domain) + return unmatricize(style, Y, (axes(Y, 1),), axes_domain) +end + # Two-output factorizations: the first factor `X` has the codomain axes plus a trailing # rank axis, the second factor `Y` has a leading rank axis plus the domain axes. for f in ( @@ -12,8 +26,8 @@ for f in ( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize(style, X, axes_codomain, (conj(axes(X, 2)),)), - unmatricize(style, Y, (axes(Y, 1),), axes_domain) + return unmatricize_codomain(style, X, axes_codomain), + unmatricize_domain(style, Y, axes_domain) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -27,6 +41,7 @@ for f in ( :svd_compact, :svd_full, :svd_trunc, :svd_vals, :eigh_full, :eig_full, :eigh_trunc, :eig_trunc, :eigh_vals, :eig_vals, :left_null, :right_null, :gram_eigh_full, :gram_eigh_full_with_pinv, :one, + :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, :project_hermitian, ) @eval begin function $f( @@ -210,9 +225,9 @@ for f in (:svd_compact, :svd_full, :svd_trunc) A_mat = matricize(style, A, ndims_codomain) U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize(style, U, axes_codomain, (conj(axes(U, 2)),)), + return unmatricize_codomain(style, U, axes_codomain), unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)), - unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) + unmatricize_domain(style, Vᴴ, axes_domain) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -228,7 +243,7 @@ for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) A_mat = matricize(style, A, ndims_codomain) D, V = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return D, unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),)) + return D, unmatricize_codomain(style, V, axes_codomain) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -406,7 +421,7 @@ function left_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize(style, N, axes_codomain, (conj(axes(N, 2)),)) + return unmatricize_codomain(style, N, axes_codomain) end function left_null!!(A, ndims_codomain::Val; kwargs...) return left_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -443,7 +458,7 @@ function right_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) _, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) + return unmatricize_domain(style, Nᴴ, axes_domain) end function right_null!!(A, ndims_codomain::Val; kwargs...) return right_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -499,7 +514,7 @@ function gram_eigh_full!!( A_mat = matricize(style, A, ndims_codomain) X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize(style, X, axes_codomain, (conj(axes(X, 2)),)) + return unmatricize_codomain(style, X, axes_codomain) end function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -560,8 +575,8 @@ function gram_eigh_full_with_pinv!!( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize(style, X, axes_codomain, (conj(axes(X, 2)),)), - unmatricize(style, Y, (axes(Y, 1),), axes_codomain) + return unmatricize_codomain(style, X, axes_codomain), + unmatricize_domain(style, Y, axes_codomain) end function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -576,6 +591,117 @@ function gram_eigh_full_with_pinv(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(copy(A), ndims_codomain; kwargs...) end +""" + sqrth_safe(A, labels_A, labels_codomain, labels_domain; kwargs...) -> P + sqrth_safe(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P + sqrth_safe(A, ndims_codomain::Val; kwargs...) -> P + +Square root of a generic N-dimensional array, interpreting it as an +approximately Hermitian positive semi-definite linear map from the domain +to the codomain dimensions. The result carries the same codomain and +domain axes as `A`. Eigenvalues below tolerance are clamped to zero. + +## Keyword arguments + + - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. + +$(MatrixAlgebra._clamp_kwargs_doc("A")) + +See also [`invsqrth_safe`](@ref), [`sqrth_invsqrth_safe`](@ref), and +[`MatrixAlgebra.sqrth_safe`](@ref). +""" +sqrth_safe + +""" + invsqrth_safe(A, labels_A, labels_codomain, labels_domain; kwargs...) -> P + invsqrth_safe(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P + invsqrth_safe(A, ndims_codomain::Val; kwargs...) -> P + +Pseudo-inverse square root of a generic N-dimensional array, interpreting +it as an approximately Hermitian positive semi-definite linear map from +the domain to the codomain dimensions. The result carries the same +codomain and domain axes as `A`. Eigenvalues below tolerance are clamped +to zero (Moore-Penrose convention). + +## Keyword arguments + + - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. + +$(MatrixAlgebra._clamp_kwargs_doc("A")) + +See also [`sqrth_safe`](@ref), [`sqrth_invsqrth_safe`](@ref), and +[`MatrixAlgebra.invsqrth_safe`](@ref). +""" +invsqrth_safe + +for f in (:sqrth_safe, :invsqrth_safe) + @eval begin + function $f(style::FusionStyle, A, ndims_codomain::Val; kwargs...) + A_mat = matricize(style, A, ndims_codomain) + P_mat = MatrixAlgebra.$f(A_mat; kwargs...) + axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) + return unmatricize(style, P_mat, axes_codomain, axes_domain) + end + function $f(A, ndims_codomain::Val; kwargs...) + return $f(FusionStyle(A), A, ndims_codomain; kwargs...) + end + end +end + +""" + project_hermitian(A, labels_A, labels_codomain, labels_domain; kwargs...) -> H + project_hermitian(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> H + project_hermitian(A, ndims_codomain::Val; kwargs...) -> H + +Hermitian part `(M + M') / 2` of a generic N-dimensional array, interpreting +it as a linear map `M` from the domain to the codomain dimensions. The result +carries the same codomain and domain axes as `A`. + +See also `MatrixAlgebraKit.project_hermitian`. +""" +project_hermitian + +function project_hermitian(style::FusionStyle, A, ndims_codomain::Val; kwargs...) + A_mat = matricize(style, A, ndims_codomain) + H_mat = MatrixAlgebraKit.project_hermitian(A_mat; kwargs...) + axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) + return unmatricize(style, H_mat, axes_codomain, axes_domain) +end +function project_hermitian(A, ndims_codomain::Val; kwargs...) + return project_hermitian(FusionStyle(A), A, ndims_codomain; kwargs...) +end + +""" + sqrth_invsqrth_safe(A, labels_A, labels_codomain, labels_domain; kwargs...) -> P, Pinv + sqrth_invsqrth_safe(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P, Pinv + sqrth_invsqrth_safe(A, ndims_codomain::Val; kwargs...) -> P, Pinv + +Square root and pseudo-inverse square root of a generic N-dimensional +array (see [`sqrth_safe`](@ref) and [`invsqrth_safe`](@ref)), from a +single eigendecomposition. Both results carry the same codomain and +domain axes as `A`. + +## Keyword arguments + + - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. + +$(MatrixAlgebra._clamp_kwargs_doc("A")) + +See also [`MatrixAlgebra.sqrth_invsqrth_safe`](@ref). +""" +sqrth_invsqrth_safe + +function sqrth_invsqrth_safe(style::FusionStyle, A, ndims_codomain::Val; kwargs...) + A_mat = matricize(style, A, ndims_codomain) + P_mat, Pinv_mat = MatrixAlgebra.sqrth_invsqrth_safe(A_mat; kwargs...) + axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) + return unmatricize(style, P_mat, axes_codomain, axes_domain), + unmatricize(style, Pinv_mat, axes_codomain, axes_domain) +end +function sqrth_invsqrth_safe(A, ndims_codomain::Val; kwargs...) + return sqrth_invsqrth_safe(FusionStyle(A), A, ndims_codomain; kwargs...) +end + """ TensorAlgebra.one(A, labels_A, labels_codomain, labels_domain) -> Id TensorAlgebra.one(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id diff --git a/src/interface.jl b/src/interface.jl index 50c8912..932eff5 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -1,11 +1,35 @@ # 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. +# `ndims`, `axes`, and `size`. These are TensorAlgebra-owned functions, distinct from `Base.ndims`/ +# `Base.axes`/`Base.size`, 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) + +function size end +size(a) = Base.size(a) +size(a, i::Int) = Base.size(a, i) + +# Eager conjugation in the same vocabulary: conjugated elements on conjugated (dualized) +# axes, leg order unchanged. Backends whose native conjugation is `adjoint`-shaped (such as +# a `TensorMap`, where the adjoint also swaps codomain and domain) overload this to restore +# the original leg order. Named `conjugate` (the eager companion of the lazy `conjed`) +# rather than `conj`: `Base.conj` is passed around as an `op` value inside this package +# (`op === conj` guards, `::typeof(conj)` dispatch), so a module-level `conj` binding +# would silently change what those mean. +function conjugate end +conjugate(a) = Base.conj(a) + +# 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). +function scalar end +scalar(a) = a[] + +# The sum of the (dense) elements. A `TensorMap` is not iterable, so `Base.sum` does not +# apply to it directly. +function sum end +sum(a; kwargs...) = Base.sum(a; kwargs...) diff --git a/src/permutedimsadd.jl b/src/permutedimsadd.jl index 8995baf..dd85d66 100644 --- a/src/permutedimsadd.jl +++ b/src/permutedimsadd.jl @@ -89,9 +89,14 @@ function bipermutedimsopadd!( perm = (perm_codomain..., perm_domain...) check_input(bipermutedimsopadd!, dest, op, src, perm_codomain, perm_domain) - dest′ = SV.StridedView(dest) - src′ = Base.permutedims(SV.StridedView(src), perm) - _opadd!(dest′, op, src′, α, β) + src′ = permuteddims(src, perm) + if SV.isstrided(dest) && SV.isstrided(src′) + _opadd!(SV.StridedView(dest), op, SV.StridedView(src′), α, β) + else + # Non-strided operands (e.g. structured matrices like `Diagonal`) accumulate + # through generic broadcasting. + _opadd!(dest, op, src′, α, β) + end return dest end diff --git a/src/projectto.jl b/src/projectto.jl index 15649c6..45769c4 100644 --- a/src/projectto.jl +++ b/src/projectto.jl @@ -2,78 +2,134 @@ projectto!(dest, src) -> dest Project `src` into the restricted space of `dest` without checking which -components may have been projected out. Defaults to `copyto!`. See -[`checked_projectto!`](@ref) for a checked version, and [`project`](@ref) -for the allocating form. +components may have been projected out. Defaults to `copyto!`, which copies +by linear index, so a lower-rank `src` may omit trailing length-1 axes (e.g. +an auxiliary flux-canceling leg a codomain/domain split introduces on a +symmetric state). A size-strict backend overloads this to reshape `src` to +`size(dest)` for the same effect. This is the in-place fill primitive that +[`unchecked_project`](@ref) allocates a destination for. """ projectto!(dest, src) = copyto!(dest, src) """ - checked_projectto!(dest, src; kwargs...) -> dest + allocate_project(raw, codomain_axes, domain_axes) -> dest -Project `src` into the restricted space of `dest` via [`projectto!`](@ref) -and verify via `isapprox(src, dest; kwargs...)` that the discarded -component is within tolerance. Keyword arguments are forwarded to -`isapprox`. The default tolerances are subject to change in future -versions. +Allocate the destination that projecting `raw` onto +`codomain_axes`/`domain_axes` fills. The generic method is +`similar_map(raw, codomain_axes, domain_axes)`. This is a backend +customization point (with [`projectto!`](@ref) and [`is_projected`](@ref)): +the allocation may depend on the data, since a symmetric backend derives the +space of one trailing surplus axis in `raw` (an auxiliary leg appended as +the last domain axis, e.g. a flux-canceling leg for a charge-shifting +operator) before allocating. """ -function checked_projectto!(dest, src; kwargs...) - projectto!(dest, src) - isapprox(src, dest; kwargs...) || - throw(InexactError(:checked_projectto!, typeof(dest), src)) - return dest +function allocate_project(raw, codomain_axes, domain_axes) + return similar_map(raw, codomain_axes, domain_axes) end """ - project_map(raw, codomain_axes, domain_axes) -> dest + unchecked_project(raw, codomain_axes, domain_axes) -> dest + unchecked_project(raw, axes) -> dest -Allocate a map-shaped array via [`similar_map`](@ref) and project `raw` -into it with [`projectto!`](@ref). This is the strict form that takes an -explicit codomain/domain split; [`project`](@ref) is the convenience entry -point that also accepts a flat list of axes. See -[`checked_project_map`](@ref) for a checked version. +Project `raw` into a symmetry-restricted array shaped as a map from +`domain_axes` to `codomain_axes`, without checking which components are +discarded: entries of `raw` outside the symmetry-allowed structure are +dropped without inspection. Most callers want [`project`](@ref), which +verifies that nothing was discarded, or [`tryproject`](@ref), its nullable +sibling. All three derive from the backend customization points: this one is +`projectto!(allocate_project(raw, codomain_axes, domain_axes), raw)`. The +two-argument form takes a flat list of `axes` and is equivalent to an empty +domain. """ -function project_map(raw, codomain_axes, domain_axes) - return projectto!(similar_map(raw, codomain_axes, domain_axes), raw) +function unchecked_project(raw, codomain_axes, domain_axes) + return projectto!(allocate_project(raw, codomain_axes, domain_axes), raw) end +# Forward to the three-argument form so a backend's surplus-axis derivation also +# applies to the flat all-codomain (state) form. +unchecked_project(raw, axes) = unchecked_project(raw, axes, ()) """ - checked_project_map(raw, codomain_axes, domain_axes; kwargs...) -> dest + is_projected(dest, src; kwargs...) -> Bool -Allocate a map-shaped array via [`similar_map`](@ref) and project `raw` -into it with [`checked_projectto!`](@ref). Keyword arguments are forwarded -to [`checked_projectto!`](@ref). +Whether the projected `dest` still represents `src` within the `isapprox` +tolerance, i.e. whether the projection that produced `dest` discarded only a +negligible component of `src`. Keyword arguments are forwarded to `isapprox`. + +Together with [`unchecked_project`](@ref) this is the backend customization +point ([`project`](@ref) and [`tryproject`](@ref) derive from the two). The +generic method reshapes `src` to `size(dest)` (so a lower-rank `src` that +omits trailing length-1 axes lines up) and compares against +`convert(Array, dest)`, so a backend whose arrays are not elementwise +comparable to a dense array (opaque block storage, a `TensorMap`) only needs +that conversion. """ -function checked_project_map(raw, codomain_axes, domain_axes; kwargs...) - return checked_projectto!( - similar_map(raw, codomain_axes, domain_axes), raw; kwargs... - ) +function is_projected(dest, src; kwargs...) + return isapprox(reshape(src, size(dest)), convert(Array, dest); kwargs...) end """ - project(raw, codomain_axes, domain_axes) -> dest - project(raw, axes) -> dest + project!(dest, src; kwargs...) -> dest -Project `raw` into a symmetry-restricted array. The three-argument form -takes an explicit codomain/domain split; the two-argument form takes a -flat list of `axes` and is equivalent to an empty domain. Both forward to -[`project_map`](@ref). See [`checked_project`](@ref) for a checked version. +In-place checked projection: project `src` into the restricted space of +`dest` via [`projectto!`](@ref) and verify with [`is_projected`](@ref) that +only a negligible component was discarded, throwing an `InexactError` +otherwise (keyword arguments are forwarded to the `isapprox` tolerance +check). This is the checked sibling of the [`projectto!`](@ref) primitive, +in the way `copy!` relates to `copyto!`; see [`project`](@ref) for the +allocating form. """ -project(raw, codomain_axes, domain_axes) = project_map(raw, codomain_axes, domain_axes) -project(raw, axes) = project_map(raw, axes, ()) +function project!(dest, src; kwargs...) + projectto!(dest, src) + is_projected(dest, src; kwargs...) || + throw(InexactError(:project!, typeof(dest), src)) + return dest +end """ - checked_project(raw, codomain_axes, domain_axes; kwargs...) -> dest - checked_project(raw, axes; kwargs...) -> dest + project(raw, codomain_axes, domain_axes; kwargs...) -> dest + project(raw, axes; kwargs...) -> dest + +Project `raw` into a symmetry-restricted array shaped as a map from +`domain_axes` to `codomain_axes`, verifying that only a negligible component +of `raw` is discarded and throwing an `InexactError` otherwise (keyword +arguments are forwarded to the `isapprox` tolerance check; the default +tolerances are subject to change in future versions). See +[`tryproject`](@ref) for a nullable version and [`unchecked_project`](@ref) +for the unchecked projection this derives from. -Checked form of [`project`](@ref): projects `raw` via -[`checked_project_map`](@ref), verifying that the discarded component is -within tolerance. Keyword arguments are forwarded to -[`checked_project_map`](@ref). +When `raw` has one axis more than the given axes account for, that trailing +surplus axis is an auxiliary leg whose space a symmetric backend derives so +the result is symmetry-allowed (e.g. a flux-canceling leg for a +charge-shifting operator); the result's shape matches `raw`'s shape. The +derivation is backend-internal: a graded backend reads the sector, the +`TensorMap` backend projects over the `codomain ⊗ conj(domain)` content. The +two-argument form takes a flat list of `axes` and is equivalent to an empty +domain. """ -function checked_project(raw, codomain_axes, domain_axes; kwargs...) - return checked_project_map(raw, codomain_axes, domain_axes; kwargs...) +function project(raw, codomain_axes, domain_axes; kwargs...) + dest = unchecked_project(raw, codomain_axes, domain_axes) + is_projected(dest, raw; kwargs...) || + throw(InexactError(:project, typeof(dest), raw)) + return dest end -function checked_project(raw, axes; kwargs...) - return checked_project_map(raw, axes, (); kwargs...) +project(raw, axes; kwargs...) = project(raw, axes, (); kwargs...) + +""" + tryproject(raw, codomain_axes, domain_axes; kwargs...) -> Union{dest, Nothing} + tryproject(raw, axes; kwargs...) -> Union{dest, Nothing} + +Like [`project`](@ref), but return `nothing` instead of throwing when more +than a negligible component of `raw` would be discarded. Useful for +branching on whether `raw` is symmetry-allowed in the given axes, e.g. +projecting a state as invariant and falling back to deriving an auxiliary +flux-carrying leg: + + @something tryproject(v, (cod,)) project(reshape(v, (length(v), 1)), (cod,)) + +Keyword arguments are forwarded to the `isapprox` tolerance check. +""" +function tryproject(raw, codomain_axes, domain_axes; kwargs...) + dest = unchecked_project(raw, codomain_axes, domain_axes) + return is_projected(dest, raw; kwargs...) ? dest : nothing end +tryproject(raw, axes; kwargs...) = tryproject(raw, axes, (); kwargs...) diff --git a/test/Project.toml b/test/Project.toml index 5b304de..469b463 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -37,7 +37,7 @@ Random = "1.10" SafeTestsets = "0.1" StableRNGs = "1.0.2" Suppressor = "0.2" -TensorAlgebra = "0.16" +TensorAlgebra = "0.17" TensorKit = "0.17" TensorOperations = "5.1.4" Test = "1.10" diff --git a/test/test_exports.jl b/test/test_exports.jl index 8828349..cafc339 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -14,16 +14,20 @@ using Test: @test, @testset :eigh_vals, :gram_eigh_full, :gram_eigh_full_with_pinv, + :invsqrth_safe, :left_null, :left_orth, :left_polar, :lq_compact, :lq_full, + :project_hermitian, :qr_compact, :qr_full, :right_null, :right_orth, :right_polar, + :sqrth_invsqrth_safe, + :sqrth_safe, :svd_compact, :svd_full, :svd_trunc, @@ -51,6 +55,7 @@ using Test: @test, @testset :pow_diag_safe, :powh_safe, :sqrt_diag_safe, + :sqrth_invsqrth_safe, :sqrth_safe, ] @test issetequal(names(TensorAlgebra.MatrixAlgebra), exports) diff --git a/test/test_projectto.jl b/test/test_projectto.jl index 4bbbd4e..159dafb 100644 --- a/test/test_projectto.jl +++ b/test/test_projectto.jl @@ -1,70 +1,101 @@ -using TensorAlgebra: TensorAlgebra, checked_project, checked_project_map, - checked_projectto!, project, project_map, projectto! +using TensorAlgebra: TensorAlgebra, is_projected, project, project!, projectto!, tryproject, + unchecked_project using Test: @test, @test_throws, @testset const elts = (Float32, Float64, ComplexF32, ComplexF64) -# `projectto!` defaults to `copyto!` and so cannot itself produce a -# discrepancy that `checked_projectto!` would reject. To exercise the -# tolerance check we wrap the destination in a custom array type whose -# `projectto!` method rounds to one decimal place. +# A dense projection is exact, so it cannot itself trip the `is_projected` check. To +# exercise the checked verbs generically we wrap the source in a custom array type whose +# `projectto!` method rounds to one decimal place, discarding information. struct Rounded{T, A <: AbstractArray{T}} <: AbstractArray{T, 2} data::A end Base.size(r::Rounded) = size(r.data) Base.getindex(r::Rounded, I...) = r.data[I...] -Base.setindex!(r::Rounded, v, I...) = (r.data[I...] = v; r) -function TensorAlgebra.projectto!(dest::Rounded, src::AbstractArray) - dest.data .= round.(src; digits = 1) +function TensorAlgebra.projectto!(dest::AbstractArray, src::Rounded) + dest .= round.(src.data; digits = 1) return dest end @testset "projectto!/project ($T)" for T in elts src = randn(T, 2, 3) - # `projectto!` defaults to `copyto!`. + # `projectto!` is the in-place fill primitive and defaults to `copyto!`. dest = similar(src) @test projectto!(dest, src) === dest @test dest == src - # `checked_projectto!` accepts when the projection is exact. - dest2 = similar(src) - @test checked_projectto!(dest2, src) === dest2 - @test dest2 == src + # `project!` is its checked sibling (as `copy!` is to `copyto!`). + dest1 = similar(src) + @test project!(dest1, src) === dest1 + @test dest1 == src - # `checked_projectto!` rejects when the custom `projectto!` discards - # information beyond `isapprox`'s tolerance. - if T <: Real - rough_src = randn(T, 2, 3) - dest3 = Rounded(similar(rough_src)) - @test_throws InexactError checked_projectto!( - dest3, - rough_src; - atol = 0.0, - rtol = 0.0 - ) - end - - # `project_map` allocates an `(cod..., conj.(dom)...)`-shaped buffer and - # projects `raw` into it. + # `project` projects into a `(cod..., conj.(dom)...)`-shaped destination (allocated by + # `allocate_project`, filled by `projectto!`) and verifies via `is_projected` that + # nothing was discarded. A dense projection is exact, so it always accepts here; the + # symmetry-rejection cases are exercised on the GradedArrays and TensorKit backends. raw = randn(T, 2, 3, 2, 3) cod = (Base.OneTo(2), Base.OneTo(3)) dom = (Base.OneTo(2), Base.OneTo(3)) - Mmap = project_map(raw, cod, dom) - @test eltype(Mmap) === T - @test size(Mmap) == (2, 3, 2, 3) - @test Mmap == raw - @test checked_project_map(raw, cod, dom) == raw + M = project(raw, cod, dom) + @test eltype(M) === T + @test size(M) == (2, 3, 2, 3) + @test M == raw - # `project` forwards to `project_map`: the three-argument form takes the - # explicit split, the two-argument form takes a flat list (empty domain). - @test project(raw, cod, dom) == raw - @test checked_project(raw, cod, dom) == raw + # `unchecked_project` is the same projection without the check. + @test unchecked_project(raw, cod, dom) == raw + # the two-argument forms take a flat list of axes (empty domain) flat = randn(T, 2, 3) - M = project(flat, (Base.OneTo(2), Base.OneTo(3))) - @test eltype(M) === T - @test size(M) == (2, 3) - @test M == flat - @test checked_project(flat, (Base.OneTo(2), Base.OneTo(3))) == flat + Mflat = project(flat, (Base.OneTo(2), Base.OneTo(3))) + @test eltype(Mflat) === T + @test size(Mflat) == (2, 3) + @test Mflat == flat + @test unchecked_project(flat, (Base.OneTo(2), Base.OneTo(3))) == flat + + # A lossy `projectto!` (the `Rounded` fixture) flips the outcome of each verb: + # `unchecked_project` silently returns the truncated result, `is_projected` reports + # the discard, `project` throws, and `tryproject` gives `nothing`. + if T <: Real + rough = Rounded(randn(T, 2, 3)) + lossy = unchecked_project(rough, (Base.OneTo(2), Base.OneTo(3))) + @test lossy == round.(rough.data; digits = 1) + @test !is_projected(lossy, rough; atol = 0, rtol = 0) + @test_throws InexactError project( + rough, (Base.OneTo(2), Base.OneTo(3)); atol = 0, rtol = 0 + ) + @test isnothing( + tryproject(rough, (Base.OneTo(2), Base.OneTo(3)); atol = 0, rtol = 0) + ) + @test_throws InexactError project!(similar(rough.data), rough; atol = 0, rtol = 0) + end +end + +@testset "tryproject ($T)" for T in elts + # `tryproject` is the nullable sibling of `project`: same projection and check, but + # `nothing` instead of an `InexactError` when weight is discarded. A dense projection + # is exact, so it always succeeds here; the failure cases are exercised on symmetric + # backends (see the GradedArrays and TensorKit tests) and via the `Rounded` fixture. + raw = randn(T, 2, 3) + t = tryproject(raw, (Base.OneTo(2), Base.OneTo(3))) + @test t == raw + @test tryproject(raw, (Base.OneTo(2),), (Base.OneTo(3),)) == raw +end + +@testset "project pads trailing length-1 axes ($T)" for T in elts + # A caller may pass the dense data over the non-trivial axes and let `project` + # supply the length-1 axes a split introduces (e.g. an auxiliary flux-canceling + # leg on a symmetric state); `size(raw, d)` is 1 past `ndims(raw)`, matching + # `raw[.., 1:1]` slicing. + flat = randn(T, 2, 3) + + # a trailing length-1 axis absent from `raw`'s rank + M = project(flat, (Base.OneTo(2), Base.OneTo(3), Base.OneTo(1))) + @test size(M) == (2, 3, 1) + @test vec(M) == vec(flat) + + # the length-1 axis may sit in the domain half of an explicit split + Msplit = project(flat, (Base.OneTo(2), Base.OneTo(3)), (Base.OneTo(1),)) + @test size(Msplit) == (2, 3, 1) + @test vec(Msplit) == vec(flat) end diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index 3385532..6e69b61 100644 --- a/test/test_tensorkitext.jl +++ b/test/test_tensorkitext.jl @@ -1,9 +1,9 @@ using Base.Broadcast: broadcasted using LinearAlgebra: norm using StableRNGs: StableRNG -using TensorAlgebra: TensorAlgebra, checked_project, contract, matricize, project, - project_map, projectto!, rand_map, randn_map, similar_map, tryflattenlinear, - unmatricize, zeros_map +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, isomorphism, randn, space, storagetype, ←, ⊗ using Test: @test, @test_throws, @testset @@ -133,11 +133,11 @@ using Test: @test, @test_throws, @testset @test space(zd) == (one(A1) ← (A1 ⊗ B)) end - # `project` builds a `TensorMap` from a dense matrix: `similar_map` allocates a same-device - # buffer (its block storage type follows the dense prototype) and `projectto!` fills the - # symmetry-allowed blocks via `project_symmetric!`, discarding the rest. A charge-preserving - # matrix survives; a charge-breaking one is projected away, and `checked_project` rejects that - # loss. + # `project` builds a `TensorMap` from a dense matrix: `allocate_project` allocates a + # same-device buffer (its block storage type follows the dense prototype) and `projectto!` + # fills the symmetry-allowed blocks via `project_symmetric!`, discarding the rest. A + # charge-preserving matrix survives; a charge-breaking one is projected away, which + # `unchecked_project` allows silently and `project` rejects. @testset "project a dense matrix into a TensorMap" begin W = Rep[U₁](0 => 1, 1 => 1) Sz = elt[0.5 0; 0 -0.5] @@ -147,19 +147,16 @@ using Test: @test, @test_throws, @testset @test pz isa AbstractTensorMap @test space(pz) == (W ← W) @test pz ≈ TensorMap(Sz, W ← W) - # `project` forwards to the `project_map` hook - @test project_map(Sz, (W,), (W,)) ≈ pz # the block storage type follows the dense prototype's array type (device-preserving) @test storagetype(pz) == Vector{elt} # `projectto!` into a same-space buffer agrees with `project` @test projectto!(similar_map(Sz, elt, (W,), (W,)), Sz) ≈ pz - # `checked_project` accepts the charge-preserving matrix (nothing discarded) - @test checked_project(Sz, (W,), (W,)) ≈ pz - # a charge-breaking matrix is projected to zero; `checked_project` rejects the discard - @test norm(project(Sx, (W,), (W,))) == 0 - @test_throws InexactError checked_project(Sx, (W,), (W,); atol = 0, rtol = 0) + # a charge-breaking matrix is projected to zero by `unchecked_project`; the checked + # `project` rejects the discard + @test norm(unchecked_project(Sx, (W,), (W,))) == 0 + @test_throws InexactError project(Sx, (W,), (W,); atol = 0, rtol = 0) # the flat two-argument form builds an all-codomain `TensorMap` (empty domain): only # the trivial-charge component of a dense vector survives the projection @@ -167,7 +164,85 @@ using Test: @test, @test_throws, @testset @test pv isa AbstractTensorMap @test space(pv) == (W ← one(W)) @test norm(pv) ≈ 1 - @test norm(project(elt[0, 1], (W,))) == 0 + @test norm(unchecked_project(elt[0, 1], (W,))) == 0 + + # a flat vector may omit a trailing length-1 auxiliary axis from its rank: + # `project` appends it. An aux index carrying the canceling charge lets a + # charged (equivariant) component survive the projection. + aux = Rep[U₁](-1 => 1) + pa = project(elt[0, 1], (W, aux)) + @test pa isa AbstractTensorMap + @test space(pa) == ((W ⊗ aux) ← one(W)) + @test norm(pa) ≈ 1 + + # only trailing *length-1* axes may be omitted: a longer trailing axis is a genuine + # size mismatch that the block projection still rejects + @test_throws DimensionMismatch project(elt[0, 1], (W, Rep[U₁](-1 => 3))) + end + + # When `raw` has one trailing axis more than the given codomain/domain account for, that + # surplus axis is an auxiliary leg (appended as the last domain axis, matching the shape of + # `raw`) whose space `project` derives, so the result is symmetry-allowed. The derivation + # scans the aux axis against the operator content and works for non-abelian symmetries and + # multi-sector (direct-sum) auxes; the abelian single-charge case falls out. + @testset "project derives the auxiliary space" begin + # SU(2): a spin operator (aux = spin-1, dim 3) is recovered from its dense components. + Ws = Rep[SU₂](1 // 2 => 1) + ts = randn(rng, elt, Ws, Ws ⊗ Rep[SU₂](1 => 1)) + rs = project(convert(Array, ts), (Ws,), (Ws,)) + @test space(rs) == (Ws ← (Ws ⊗ Rep[SU₂](1 => 1))) + @test rs ≈ ts + + # U(1): a charge-shifting operator (non-self-dual aux = charge +1) is recovered. + Wu = Rep[U₁](0 => 1, 1 => 1) + tu = randn(rng, elt, Wu, Wu ⊗ Rep[U₁](1 => 1)) + ru = project(convert(Array, tu), (Wu,), (Wu,)) + @test space(ru) == (Wu ← (Wu ⊗ Rep[U₁](1 => 1))) + @test ru ≈ tu + + # U(1) direct sum (an MPO-style virtual leg): each slice carries its own charge. + tds = randn(rng, elt, Wu, Wu ⊗ Rep[U₁](1 => 1, -1 => 1)) + rds = project(convert(Array, tds), (Wu,), (Wu,)) + @test space(rds) == (Wu ← (Wu ⊗ Rep[U₁](1 => 1, -1 => 1))) + @test rds ≈ tds + + # SU(2) direct sum of different irreps (scalar ⊕ vector part, dim 4). + tmx = randn(rng, elt, Ws, Ws ⊗ Rep[SU₂](0 => 1, 1 => 1)) + rmx = project(convert(Array, tmx), (Ws,), (Ws,)) + @test space(rmx) == (Ws ← (Ws ⊗ Rep[SU₂](0 => 1, 1 => 1))) + @test rmx ≈ tmx + + # SU(2) multiplicity > 1: two spin-1 copies (dim 6). + tm2 = randn(rng, elt, Ws, Ws ⊗ Rep[SU₂](1 => 2)) + rm2 = project(convert(Array, tm2), (Ws,), (Ws,)) + @test space(rm2) == (Ws ← (Ws ⊗ Rep[SU₂](1 => 2))) + @test rm2 ≈ tm2 + + # data not covariant with any aux decomposition of the surplus axis is rejected + @test_throws ArgumentError project(randn(rng, elt, 2, 2, 3), (Ws,), (Ws,)) + + # only one trailing surplus axis is supported: more is an error, not a flattening + @test_throws ArgumentError project( + reshape(convert(Array, ts), 2, 2, 3, 1), (Ws,), (Ws,) + ) + + # a lower-rank `raw` that omits explicitly-given trailing length-1 axes is the + # trailing-axes tolerance, not a surplus axis: it pads, it does not derive. A domain + # aux of charge +1 admits the charge-1 component. + aux1 = Rep[U₁](1 => 1) + po = project(elt[0, 1], (Wu,), (aux1,)) + @test space(po) == (Wu ← aux1) + @test norm(po) ≈ 1 + + # the flat all-codomain (state) form also derives: a stack of basis states with + # different charges gets a multi-sector aux + ps = project(elt[1 0; 0 1], (Wu,)) + @test space(ps) == (Wu ← Rep[U₁](0 => 1, 1 => 1)) + + # `tryproject` gives `nothing` instead of throwing when the data is not invariant + # in the given (all-given) axes — the branch-and-fall-back-to-derivation idiom + @test isnothing(tryproject(elt[0, 1], (Wu,))) + @test tryproject(elt[1, 0], (Wu,)) isa AbstractTensorMap end # `permutedims` reorders a `TensorMap`'s indices; the flat form gives an all-codomain From 81d2c3fbc77cf0a6f5c08e3ce927c90849152d08 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sun, 5 Jul 2026 17:13:05 -0400 Subject: [PATCH 02/15] Add a length-n form of trivialrange `trivialrange(r, n)` returns the n-dimensional identity range of a range family, defaulting to `Base.OneTo(n)`, with the TensorKit extension returning the direct sum of n unit spaces. Downstream packages overload the type-level method the way they already do for the length-1 form. --- ext/TensorAlgebraTensorKitExt.jl | 8 +++++++- src/matricize.jl | 11 ++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index e41f7d0..8e4aee5 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -37,9 +37,15 @@ end TensorAlgebra.scalar(t::AbstractTensorMap) = TensorKit.scalar(t) # The trivial length-1 axis of a space is its unit space (`oneunit`), the trivial-sector -# one-dimensional space. +# one-dimensional space; the length-`n` form is the direct sum of `n` unit spaces. TensorAlgebra.trivialrange(V::ElementarySpace) = oneunit(V) TensorAlgebra.trivialrange(::Type{S}) where {S <: ElementarySpace} = oneunit(S) +function TensorAlgebra.trivialrange(V::ElementarySpace, n::Integer) + return TensorAlgebra.trivialrange(typeof(V), n) +end +function TensorAlgebra.trivialrange(::Type{S}, n::Integer) where {S <: ElementarySpace} + return TensorKit.oplus(ntuple(Returns(oneunit(S)), n)...) +end # 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 diff --git a/src/matricize.jl b/src/matricize.jl index d6c6a82..2f1d11f 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -12,17 +12,22 @@ FusionStyle(style1::FusionStyle, style2::FusionStyle) = ReshapeFusion() # ======================================= misc ======================================== """ - TensorAlgebra.trivialrange(R::Type{<:AbstractUnitRange}) - TensorAlgebra.trivialrange(r::AbstractUnitRange) + TensorAlgebra.trivialrange(R::Type{<:AbstractUnitRange}[, n::Integer]) + TensorAlgebra.trivialrange(r::AbstractUnitRange[, n::Integer]) Return the identity range for fusing ranges of type `R`: a one-dimensional range `t` for which fusing `t` with any other range of the same family leaves that range unchanged. Defaults to `Base.OneTo(1)`. Downstream packages overload the type-level -method to return their own identity (for example, a charge-0 one-dimensional sector +methods to return their own identity (for example, a charge-0 one-dimensional sector for a graded range). + +With a length `n`, return the `n`-dimensional analogue: `n` copies of the identity +range stacked into one range (for a graded range, a charge-0 sector of dimension `n`). """ trivialrange(r::AbstractUnitRange) = trivialrange(typeof(r)) trivialrange(::Type{<:AbstractUnitRange}) = Base.OneTo(1) +trivialrange(r::AbstractUnitRange, n::Integer) = trivialrange(typeof(r), n) +trivialrange(::Type{<:AbstractUnitRange}, n::Integer) = Base.OneTo(n) """ permutedimsop(op, src, perm_codomain, perm_domain) From 3126e1116645aceb8d286c62d368321a109ba5dd Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sun, 5 Jul 2026 19:32:44 -0400 Subject: [PATCH 03/15] Preserve the diagonal eltype in pow_diag_safe The AbstractMatrix method rebuilt the diagonal from a map whose clamp branch returned `zero(d)` and whose power branch returned `real(d)^p`, so the first below-tolerance entry of a complex diagonal widened the result to an abstract eltype. Convert both branches with `oftype`, matching what the in-place variant already does through its broadcast assignment. --- src/MatrixAlgebra.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index d7483c9..51180b2 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -52,7 +52,10 @@ function pow_diag_safe( ) σ = MAK.diagview(D) tol = max(atol, rtol * maximum(abs, σ; init = zero(real(eltype(D))))) - return MAK.diagonal(map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ)) + # `oftype` keeps the entry type: the clamp branch (`zero(d)`) and the power branch + # (`real(d)^p`) otherwise differ (e.g. `ComplexF32` vs `Float64`) and `map` would + # widen the diagonal to an abstract eltype. + return MAK.diagonal(map(d -> oftype(d, abs(d) < tol ? zero(d) : real(d)^p), σ)) end # Non-`AbstractMatrix` matrix-like backends (e.g. a `TensorMap`): `map` over the From 44c10ba5deccd8a2e81dfd1fa0fc52d2bcc52303 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sun, 5 Jul 2026 21:31:02 -0400 Subject: [PATCH 04/15] Require Hermitian input in the safe spectral functions powh_safe, sqrth_invsqrth_safe, and the gram functions call MatrixAlgebraKit.eigh_full directly instead of projecting onto the Hermitian part internally. Inputs that are Hermitian only up to numerical noise are the caller's responsibility: project with MatrixAlgebraKit.project_hermitian first, as the docstrings now state. Co-authored-by: Claude --- src/MatrixAlgebra.jl | 41 +++++++++++++++++++++-------------------- src/factorizations.jl | 19 +++++++++++-------- 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index 51180b2..570f2c1 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -117,12 +117,14 @@ invsqrt_diag_safe(D::AbstractMatrix; kwargs...) = pow_diag_safe(D, -1 // 2; kwar """ powh_safe(M, p; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^p -Raise an approximately Hermitian positive semi-definite matrix to the -power `p`. For diagonal-structured `M` (`isdiag(M) == true`), dispatches -to [`pow_diag_safe`](@ref) and skips the eigendecomposition. Otherwise, -projects `M` onto its Hermitian part `(M + M') / 2` (so input that is -Hermitian only up to numerical noise is accepted) and computes via -`M = V * D * V'` as `V * pow_diag_safe(D, p; atol, rtol) * V'`. +Raise a Hermitian positive semi-definite matrix to the power `p`. For +diagonal-structured `M` (`isdiag(M) == true`), dispatches to +[`pow_diag_safe`](@ref) and skips the eigendecomposition. Otherwise +computes via `M = V * D * V'` as `V * pow_diag_safe(D, p; atol, rtol) * V'`. + +The input must be Hermitian (as for `MatrixAlgebraKit.eigh_full`): project +with `MatrixAlgebraKit.project_hermitian` first if it is Hermitian only up +to numerical noise. ## Keyword arguments @@ -132,19 +134,14 @@ $(_clamp_kwargs_doc("M")) """ function powh_safe(M, p; alg = nothing, kwargs...) isdiag(M) && return pow_diag_safe(M, p; kwargs...) - D, V = _eigh_full_hermitianpart(M, alg) + D, V = MAK.eigh_full(M; alg) return V * pow_diag_safe(D, p; kwargs...) * V' end -function _eigh_full_hermitianpart(M, alg) - M = MAK.project_hermitian(M) - return MAK.eigh_full(M, MAK.select_algorithm(MAK.eigh_full, M, alg)) -end - """ sqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2) -Square root of an approximately Hermitian positive semi-definite matrix. +Square root of a Hermitian positive semi-definite matrix. Equivalent to `powh_safe(M, 1//2; alg, atol, rtol)`. ## Keyword arguments @@ -158,8 +155,8 @@ sqrth_safe(M; kwargs...) = powh_safe(M, 1 // 2; kwargs...) """ invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(-1//2) -Inverse square root of an approximately Hermitian positive semi-definite -matrix. Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. +Inverse square root of a Hermitian positive semi-definite matrix. +Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. ## Keyword arguments @@ -172,12 +169,16 @@ invsqrth_safe(M; kwargs...) = powh_safe(M, -1 // 2; kwargs...) """ sqrth_invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2), M^(-1//2) -Square root and pseudo-inverse square root of an approximately Hermitian -positive semi-definite matrix, from a single eigendecomposition. Equivalent +Square root and pseudo-inverse square root of a Hermitian positive +semi-definite matrix, from a single eigendecomposition. Equivalent to `(sqrth_safe(M; ...), invsqrth_safe(M; ...))` but with the eigendecomposition computed once. Eigenvalues below tolerance are clamped to zero in both factors (Moore-Penrose convention for the inverse). +The input must be Hermitian (as for `MatrixAlgebraKit.eigh_full`): project +with `MatrixAlgebraKit.project_hermitian` first if it is Hermitian only up +to numerical noise. + ## Keyword arguments - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. @@ -188,7 +189,7 @@ function sqrth_invsqrth_safe(M; alg = nothing, kwargs...) if isdiag(M) return pow_diag_safe(M, 1 // 2; kwargs...), pow_diag_safe(M, -1 // 2; kwargs...) end - D, V = _eigh_full_hermitianpart(M, alg) + D, V = MAK.eigh_full(M; alg) return V * pow_diag_safe(D, 1 // 2; kwargs...) * V', V * pow_diag_safe(D, -1 // 2; kwargs...) * V' end @@ -199,11 +200,11 @@ for (gram, gram_with_pinv, eigh_full) in ( ) @eval begin function $gram(A::AbstractMatrix; alg = nothing, kwargs...) - D, V = MAK.$eigh_full(A, MAK.select_algorithm(MAK.$eigh_full, A, alg)) + D, V = MAK.$eigh_full(A; alg) return V * sqrth_safe(D; kwargs...) end function $gram_with_pinv(A::AbstractMatrix; alg = nothing, kwargs...) - D, V = MAK.$eigh_full(A, MAK.select_algorithm(MAK.$eigh_full, A, alg)) + D, V = MAK.$eigh_full(A; alg) return V * sqrth_safe(D; kwargs...), invsqrth_safe(D; kwargs...) * V' end end diff --git a/src/factorizations.jl b/src/factorizations.jl index 14ef186..b5a11cd 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -596,10 +596,12 @@ end sqrth_safe(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P sqrth_safe(A, ndims_codomain::Val; kwargs...) -> P -Square root of a generic N-dimensional array, interpreting it as an -approximately Hermitian positive semi-definite linear map from the domain -to the codomain dimensions. The result carries the same codomain and -domain axes as `A`. Eigenvalues below tolerance are clamped to zero. +Square root of a generic N-dimensional array, interpreting it as a +Hermitian positive semi-definite linear map from the domain to the +codomain dimensions. The result carries the same codomain and domain axes +as `A`. Eigenvalues below tolerance are clamped to zero. The input must be +Hermitian: project with `project_hermitian` first if it is Hermitian only +up to numerical noise. ## Keyword arguments @@ -618,10 +620,11 @@ sqrth_safe invsqrth_safe(A, ndims_codomain::Val; kwargs...) -> P Pseudo-inverse square root of a generic N-dimensional array, interpreting -it as an approximately Hermitian positive semi-definite linear map from -the domain to the codomain dimensions. The result carries the same -codomain and domain axes as `A`. Eigenvalues below tolerance are clamped -to zero (Moore-Penrose convention). +it as a Hermitian positive semi-definite linear map from the domain to the +codomain dimensions. The result carries the same codomain and domain axes +as `A`. Eigenvalues below tolerance are clamped to zero (Moore-Penrose +convention). The input must be Hermitian: project with `project_hermitian` +first if it is Hermitian only up to numerical noise. ## Keyword arguments From d14fb7920f00e0b81f746d733967e09a46d893a9 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 09:45:03 -0400 Subject: [PATCH 05/15] Simplify pow_diag_safe and the TensorKit aux-space scan pow_diag_safe now has one implementation for every backend: write the clamped powers back through MAK.diagview onto a copy of the input, dropping the separate AbstractMatrix method that rebuilt through MAK.diagonal. The result keeps the input's type and structure, and block-structured diagonal views (graded, TensorMap) come along through their blockwise copyto!. The TensorMap allocate_project aux-space scan tests each candidate slice with tryproject instead of a try/catch around the TensorMap constructor, so covariance is checked through the same verify path the rest of project uses rather than through TensorKit's constructor error types. Co-authored-by: Claude --- ext/TensorAlgebraTensorKitExt.jl | 21 +++++++-------- src/MatrixAlgebra.jl | 46 ++++++++++---------------------- test/test_matrixalgebra.jl | 6 +++-- 3 files changed, 28 insertions(+), 45 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 8e4aee5..f257bd9 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -2,7 +2,7 @@ module TensorAlgebraTensorKitExt using Random: AbstractRNG using TensorAlgebra: TensorAlgebra -using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, TensorMap, +using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, TensorMapWithStorage, codomain, dim, domain, dual, fuse, numind, permute, project_symmetric!, sectors, space, spacetype, zerovector!, ← using TensorOperations: TensorOperations as TO @@ -150,8 +150,8 @@ end # the shape of `raw` — gets its space derived so the result is symmetry-allowed. The candidate # irreps are the operator content `codomain ⊗ conj(domain)`; scan the aux axis in the content's # canonical (sorted) sector order, consuming contiguous width-`dim(s)` slices that are covariant -# with irrep `s` (`TensorMap(data, space)` does the projection and the covariance check per -# slice). The result is a possibly multi-sector (direct-sum) aux, e.g. an MPO-style virtual leg; +# with irrep `s` (`tryproject` at default tolerances does the projection and the covariance check +# per slice). The result is a possibly multi-sector (direct-sum) aux, e.g. an MPO-style virtual leg; # a single irrep and the abelian single-charge case fall out. The slice order must match the # canonical sector order, since a `GradedSpace` sorts its sectors and the dense layout follows # it. The fill onto the derived axes (`projectto!`) is magnitude-blind; the `project` wrapper @@ -170,15 +170,14 @@ function TensorAlgebra.allocate_project( ) auxdim = size(raw, nphys + 1) content = fuse(codomain_axes..., dual.(domain_axes)...) + # A slice keeps the aux axis (width `dim(s)`), so its rank matches the candidate + # axes exactly and `tryproject` allocates, fills, and round-trip-verifies without + # re-entering the derivation branch. function slice_is_covariant(r, s) - target = _map_homspace(S, codomain_axes, (domain_axes..., S(s => 1))) - return try - TensorMap(selectdim(raw, nphys + 1, r), target) - true - catch e - e isa Union{ArgumentError, DimensionMismatch} || rethrow() - false - end + slice = selectdim(raw, nphys + 1, r) + return !isnothing( + TensorAlgebra.tryproject(slice, codomain_axes, (domain_axes..., S(s => 1))) + ) end seccounts = Pair{TensorKit.sectortype(S), Int}[] pos = 1 diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index 570f2c1..17a58f8 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -23,17 +23,19 @@ function _clamp_kwargs_doc(arg::AbstractString) end """ - pow_diag_safe(D::AbstractMatrix, p; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^p + pow_diag_safe(D, p; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^p Raise a diagonal-structured matrix `D` to the power `p`. Diagonal entries -`d` of `MAK.diagview(D)` with `abs(d) < tol` are clamped to zero before -exponentiation, where `tol = max(atol, rtol * maximum(abs, diagview(D)))`. +`d` with `abs(d) < tol` are clamped to zero before exponentiation, where +`tol = max(atol, rtol * norm(D, Inf))` (the largest-magnitude entry, which +is the largest-magnitude diagonal entry for a diagonal-structured matrix). Negative `d` above `tol` cause `d^p` to error for fractional `p` (e.g. `p = 1//2`) and pass through for integer `p`, so the operation itself enforces the PSD precondition per-power. Errors if `isdiag(D)` is `false`. -The implementation extracts entries via `MAK.diagview` and rebuilds via -`MAK.diagonal`, so types extending those (e.g. graded or block diagonal) +The implementation writes the clamped powers back through `MAK.diagview` +onto a `copy` of `D`, so the result has the input's type and structure, and +types extending `diagview` (e.g. graded or block diagonal, a `TensorMap`) automatically extend [`sqrt_diag_safe`](@ref), [`invsqrt_diag_safe`](@ref), and the [`powh_safe`](@ref) family. @@ -41,28 +43,6 @@ and the [`powh_safe`](@ref) family. $(_clamp_kwargs_doc("D")) """ -function pow_diag_safe( - D::AbstractMatrix, p; - atol = zero(real(eltype(D))), - rtol = iszero(atol) ? eps(real(eltype(D)))^(3 // 4) : - zero(real(eltype(D))) - ) - isdiag(D) || throw( - ArgumentError("pow_diag_safe requires a diagonal-structured matrix") - ) - σ = MAK.diagview(D) - tol = max(atol, rtol * maximum(abs, σ; init = zero(real(eltype(D))))) - # `oftype` keeps the entry type: the clamp branch (`zero(d)`) and the power branch - # (`real(d)^p`) otherwise differ (e.g. `ComplexF32` vs `Float64`) and `map` would - # widen the diagonal to an abstract eltype. - return MAK.diagonal(map(d -> oftype(d, abs(d) < tol ? zero(d) : real(d)^p), σ)) -end - -# Non-`AbstractMatrix` matrix-like backends (e.g. a `TensorMap`): `map` over the -# `MAK.diagview` may not preserve the backend's block structure, so write the mapped -# diagonal back into a `copy` through the (aliasing) `diagview` instead of rebuilding -# with `MAK.diagonal`. The tolerance uses `norm(D, Inf)` (the largest-magnitude entry), -# which equals the largest-magnitude diagonal entry under the `isdiag` guard. function pow_diag_safe( D, p; atol = zero(real(eltype(D))), @@ -78,8 +58,10 @@ function pow_diag_safe( return Dp end +# `copyto!` rather than `.=`: block-structured diagonal views (e.g. a graded fused +# vector) have a blockwise `map` and `copyto!` but no broadcast support. function _pow_diag!(σ, p, tol) - σ .= map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ) + copyto!(σ, map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ)) return σ end # A backend's `diagview` may be a dict of per-block diagonal views (e.g. a `TensorMap`, @@ -90,7 +72,7 @@ function _pow_diag!(σ::AbstractDict, p, tol) end """ - sqrt_diag_safe(D::AbstractMatrix; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^(1//2) + sqrt_diag_safe(D; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^(1//2) Square root of a diagonal-structured matrix `D`, equivalent to `pow_diag_safe(D, 1//2; atol, rtol)`. @@ -99,10 +81,10 @@ Square root of a diagonal-structured matrix `D`, equivalent to $(_clamp_kwargs_doc("D")) """ -sqrt_diag_safe(D::AbstractMatrix; kwargs...) = pow_diag_safe(D, 1 // 2; kwargs...) +sqrt_diag_safe(D; kwargs...) = pow_diag_safe(D, 1 // 2; kwargs...) """ - invsqrt_diag_safe(D::AbstractMatrix; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^(-1//2) + invsqrt_diag_safe(D; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^(-1//2) Inverse square root of a diagonal-structured matrix `D`, treating diagonal entries below tolerance as zero (Moore-Penrose convention). Equivalent to @@ -112,7 +94,7 @@ entries below tolerance as zero (Moore-Penrose convention). Equivalent to $(_clamp_kwargs_doc("D")) """ -invsqrt_diag_safe(D::AbstractMatrix; kwargs...) = pow_diag_safe(D, -1 // 2; kwargs...) +invsqrt_diag_safe(D; kwargs...) = pow_diag_safe(D, -1 // 2; kwargs...) """ powh_safe(M, p; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^p diff --git a/test/test_matrixalgebra.jl b/test/test_matrixalgebra.jl index 2ecf63e..68a61e3 100644 --- a/test/test_matrixalgebra.jl +++ b/test/test_matrixalgebra.jl @@ -202,11 +202,13 @@ elts = (Float32, Float64, ComplexF32, ComplexF64) # `isdiag` fast-path: a dense matrix that happens to be diagonal- # structured goes through `pow_diag_safe`, not `eigh_full`, and - # the result is still a `Diagonal` (from `MAK.diagonal`). + # the result keeps the input's type (the clamped powers are written + # back onto a `copy` through `MAK.diagview`). Mdiag = Matrix(D) sqrtMdiag = MatrixAlgebra.powh_safe(Mdiag, 1 // 2) - @test sqrtMdiag isa Diagonal + @test sqrtMdiag isa Matrix @test sqrtMdiag ≈ MatrixAlgebra.pow_diag_safe(D, 1 // 2) + @test MatrixAlgebra.pow_diag_safe(D, 1 // 2) isa Diagonal # `pow_diag_safe` directly on a diagonal-structured `AbstractMatrix`. @test MatrixAlgebra.pow_diag_safe(Mdiag, 1 // 2) ≈ From c7bb4dd8448937ec0c519f50ec0e2e5016d6c54b Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 13:11:56 -0400 Subject: [PATCH 06/15] Reconstruct factorization factors on any backend without per-backend overrides The generic factor-reconstruction helpers located a factor's new bond axis with `axes(X, 2)`, which addresses the bond only when `matricize` fuses the codomain into one axis. A non-fusing backend such as a `TensorMap` keeps the group's original legs plus the bond, so `axes(X, 2)` read a physical leg instead, and the TensorKit extension had to override both helpers. The bond is the factor's last axis (codomain factor) or first axis (domain factor) on every backend, so locating it with `axes(X, ndims(X))` makes the helpers backend-agnostic and lets the overrides go. Also trims the `allocate_project` aux-derivation comment. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 49 ++++---------------------------- src/factorizations.jl | 16 +++++------ 2 files changed, 14 insertions(+), 51 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index f257bd9..97ceaa8 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -145,18 +145,12 @@ end # defines for an `AbstractTensorMap`, so no dedicated method is needed here. # ============================= allocate_project (aux-leg derivation) ===================== -# `allocate_project` with `TensorMap` spaces: the destination allocation, which is where a -# trailing surplus axis in `raw` — an auxiliary leg appended as the last domain axis, matching -# the shape of `raw` — gets its space derived so the result is symmetry-allowed. The candidate -# irreps are the operator content `codomain ⊗ conj(domain)`; scan the aux axis in the content's -# canonical (sorted) sector order, consuming contiguous width-`dim(s)` slices that are covariant -# with irrep `s` (`tryproject` at default tolerances does the projection and the covariance check -# per slice). The result is a possibly multi-sector (direct-sum) aux, e.g. an MPO-style virtual leg; -# a single irrep and the abelian single-charge case fall out. The slice order must match the -# canonical sector order, since a `GradedSpace` sorts its sectors and the dense layout follows -# it. The fill onto the derived axes (`projectto!`) is magnitude-blind; the `project` wrapper -# verifies that nothing was discarded. With no surplus axis (a lower-rank `raw` omits trailing -# length-1 axes, which the `projectto!` reshape pads), this is plain `similar_map`. +# `allocate_project` for `TensorMap` spaces. An optional trailing surplus axis in `raw` (an +# auxiliary leg appended as the last domain axis) has its space derived here so the result is +# symmetry-allowed; without one this is plain `similar_map`. Candidates are the operator content +# `codomain ⊗ conj(domain)`, scanned in canonical (sorted) sector order — a `GradedSpace` sorts +# its sectors and the dense layout follows, so the aux slices must appear in that order. The +# derived aux may span several sectors (a direct-sum, MPO-style virtual leg). function TensorAlgebra.allocate_project( raw::AbstractArray, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} ) where {S <: ElementarySpace} @@ -254,37 +248,6 @@ function TensorAlgebra.unmatricize( return m end -# A factorization factor needs no unfusing here (`matricize` above does not fuse, so the -# factors come back from MatrixAlgebraKit with their original legs plus the bond). The -# generic defaults would reconstruct the bond axis with the matrix-flavored `axes(X, 2)`, -# which on an unfused factor reads a codomain leg instead of the bond (the two only -# coincide when the bond carries the trivial sector). Validate the known side and return -# the factor unchanged. -function TensorAlgebra.unmatricize_codomain( - ::TensorKitFusion, X::AbstractTensorMap, codomain_axes - ) - S = spacetype(X) - dest = ProductSpace{S}(codomain_axes...) - codomain(X) == dest || throw( - ArgumentError( - "`unmatricize_codomain` space `$dest` does not match `$(codomain(X))`" - ) - ) - return X -end -function TensorAlgebra.unmatricize_domain( - ::TensorKitFusion, Y::AbstractTensorMap, domain_axes - ) - S = spacetype(Y) - # The requested axes arrive codomain-facing (un-dualized), which is TensorKit's - # domain convention, so they build the domain `ProductSpace` directly. - dest = ProductSpace{S}(domain_axes...) - domain(Y) == dest || throw( - ArgumentError("`unmatricize_domain` space `$dest` does not match `$(domain(Y))`") - ) - return Y -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` diff --git a/src/factorizations.jl b/src/factorizations.jl index b5a11cd..ea19715 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -1,15 +1,15 @@ using LinearAlgebra: LinearAlgebra using MatrixAlgebraKit: MatrixAlgebraKit -# Unfuse one side of a factorization factor, keeping its other (bond) axis as is: a -# factor `X` carries the fused codomain (or domain) group of the input on one axis and -# the new bond on the other, and only the fused group needs reconstructing. For fusing -# backends the bond axis is read off the factor itself (`axes(X, 2)`, dualized to stay -# codomain-facing). A backend whose `matricize` does not fuse (e.g. a `TensorMap`) -# overrides these to return the factor unchanged: its factors keep their original legs, -# so `axes(X, 2)` does not address the bond. +# Reconstruct one side of a factorization factor, keeping its other (bond) axis as is: a +# factor carries the codomain (or domain) group of the input on one side and the new bond +# on the other, and only that group needs reconstructing. The bond is the factor's last +# axis (codomain factor) or first axis (domain factor) on every backend: a fusing backend +# returns a rank-2 factor `[fused group, bond]`, a non-fusing backend (e.g. a `TensorMap`) +# returns the group's original legs plus the bond. So `unmatricize` reconstructs both with +# no per-backend override; the bond axis is dualized to stay codomain-facing. function unmatricize_codomain(style::FusionStyle, X, axes_codomain) - return unmatricize(style, X, axes_codomain, (conj(axes(X, 2)),)) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)) end function unmatricize_domain(style::FusionStyle, Y, axes_domain) return unmatricize(style, Y, (axes(Y, 1),), axes_domain) From 915519a714d1d96990e1f19333d6035a0e9d4d7b Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 14:51:49 -0400 Subject: [PATCH 07/15] Reconstruct factorization factors with unmatricize directly, split aux-space inference Drops the `unmatricize_codomain`/`unmatricize_domain` helpers (both the generic definitions and the TensorKit-backend overrides). Each factorization now reconstructs its factors by calling `unmatricize` directly, reading the freshly created bond axis off the returned factor: the factor's last axis on a codomain factor, its first axis on a domain factor, on every backend. This matches how the SVD spectrum factor was already reconstructed inline, so the module uses one primitive throughout. Also splits the TensorKit `allocate_project` into a short dispatcher and an `infer_aux_space` helper. The dispatcher handles the passthrough and rank check, and `infer_aux_space` holds the sector scan that works out the space of a trailing auxiliary leg so the projected operator is symmetry-allowed. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 93 +++++++++++++++++--------------- src/factorizations.jl | 39 ++++++-------- 2 files changed, 67 insertions(+), 65 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 97ceaa8..fe2d7d2 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -145,55 +145,64 @@ end # defines for an `AbstractTensorMap`, so no dedicated method is needed here. # ============================= allocate_project (aux-leg derivation) ===================== -# `allocate_project` for `TensorMap` spaces. An optional trailing surplus axis in `raw` (an -# auxiliary leg appended as the last domain axis) has its space derived here so the result is -# symmetry-allowed; without one this is plain `similar_map`. Candidates are the operator content -# `codomain ⊗ conj(domain)`, scanned in canonical (sorted) sector order — a `GradedSpace` sorts -# its sectors and the dense layout follows, so the aux slices must appear in that order. The -# derived aux may span several sectors (a direct-sum, MPO-style virtual leg). +# `allocate_project` for `TensorMap` spaces. With no surplus axis this is plain `similar_map`; +# a single trailing surplus axis in `raw` is an auxiliary leg whose space is derived (see +# `infer_aux_space`) and appended as the last domain axis so the result is symmetry-allowed. function TensorAlgebra.allocate_project( raw::AbstractArray, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} ) where {S <: ElementarySpace} nphys = length(codomain_axes) + length(domain_axes) - if ndims(raw) > nphys - ndims(raw) == nphys + 1 || throw( - ArgumentError( - "`project`: expected at most one trailing auxiliary axis beyond the \ - $nphys given axes, got a rank-$(ndims(raw)) input" - ) + ndims(raw) <= nphys && + return TensorAlgebra.similar_map(raw, codomain_axes, domain_axes) + ndims(raw) == nphys + 1 || throw( + ArgumentError( + "`project`: expected at most one trailing auxiliary axis beyond the $nphys \ + given axes, got a rank-$(ndims(raw)) input" ) - auxdim = size(raw, nphys + 1) - content = fuse(codomain_axes..., dual.(domain_axes)...) - # A slice keeps the aux axis (width `dim(s)`), so its rank matches the candidate - # axes exactly and `tryproject` allocates, fills, and round-trip-verifies without - # re-entering the derivation branch. - function slice_is_covariant(r, s) - slice = selectdim(raw, nphys + 1, r) - return !isnothing( - TensorAlgebra.tryproject(slice, codomain_axes, (domain_axes..., S(s => 1))) - ) - end - seccounts = Pair{TensorKit.sectortype(S), Int}[] - pos = 1 - for s in sectors(content) - d = dim(S(s => 1)) - m = 0 - while pos + d - 1 <= auxdim && slice_is_covariant(pos:(pos + d - 1), s) - m += 1 - pos += d - end - m > 0 && push!(seccounts, s => m) - end - pos == auxdim + 1 || throw( - ArgumentError( - "`project`: could not derive a covariant auxiliary space for the surplus axis \ - of dimension $auxdim; the aux slices must be ordered by the canonical (sorted) \ - sector order of the derived space" - ) + ) + aux = infer_aux_space(raw, codomain_axes, domain_axes) + return TensorAlgebra.similar_map(raw, codomain_axes, (domain_axes..., aux)) +end + +# The space of `raw`'s trailing auxiliary axis, derived so the projected result is +# symmetry-allowed. Candidates are the operator content `codomain ⊗ conj(domain)`, scanned in +# canonical (sorted) sector order — a `GradedSpace` sorts its sectors and the dense layout +# follows, so the aux slices must appear in that order. The result may span several sectors (a +# direct-sum, MPO-style virtual leg). +function infer_aux_space( + raw, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} + ) where {S <: ElementarySpace} + naux = length(codomain_axes) + length(domain_axes) + 1 + auxdim = size(raw, naux) + content = fuse(codomain_axes..., dual.(domain_axes)...) + # A slice keeps the aux axis (width `dim(s)`), so its rank matches the candidate axes + # exactly and `tryproject` allocates, fills, and round-trip-verifies without re-entering + # the derivation branch. + function slice_is_covariant(r, s) + slice = selectdim(raw, naux, r) + return !isnothing( + TensorAlgebra.tryproject(slice, codomain_axes, (domain_axes..., S(s => 1))) ) - domain_axes = (domain_axes..., S(seccounts...)) end - return TensorAlgebra.similar_map(raw, codomain_axes, domain_axes) + seccounts = Pair{TensorKit.sectortype(S), Int}[] + pos = 1 + for s in sectors(content) + d = dim(S(s => 1)) + m = 0 + while pos + d - 1 <= auxdim && slice_is_covariant(pos:(pos + d - 1), s) + m += 1 + pos += d + end + m > 0 && push!(seccounts, s => m) + end + pos == auxdim + 1 || throw( + ArgumentError( + "`project`: could not derive a covariant auxiliary space for the surplus axis of \ + dimension $auxdim; the aux slices must be ordered by the canonical (sorted) sector \ + order of the derived space" + ) + ) + return S(seccounts...) end # ================================ bipermutedimsopadd! ===================================== diff --git a/src/factorizations.jl b/src/factorizations.jl index ea19715..ef17ba8 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -1,19 +1,12 @@ using LinearAlgebra: LinearAlgebra using MatrixAlgebraKit: MatrixAlgebraKit -# Reconstruct one side of a factorization factor, keeping its other (bond) axis as is: a -# factor carries the codomain (or domain) group of the input on one side and the new bond -# on the other, and only that group needs reconstructing. The bond is the factor's last -# axis (codomain factor) or first axis (domain factor) on every backend: a fusing backend -# returns a rank-2 factor `[fused group, bond]`, a non-fusing backend (e.g. a `TensorMap`) -# returns the group's original legs plus the bond. So `unmatricize` reconstructs both with -# no per-backend override; the bond axis is dualized to stay codomain-facing. -function unmatricize_codomain(style::FusionStyle, X, axes_codomain) - return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)) -end -function unmatricize_domain(style::FusionStyle, Y, axes_domain) - return unmatricize(style, Y, (axes(Y, 1),), axes_domain) -end +# Each factorization reconstructs its factors with `unmatricize`, reading the freshly created +# bond axis off the factor itself: it is the factor's last axis on a codomain factor +# (`[group…, bond]`) and its first axis on a domain factor (`[bond, group…]`), on every backend +# (a fusing backend returns a rank-2 factor, a `TensorMap` keeps the group's original legs). The +# bond is dualized to codomain-facing form (`conj`, a no-op on a dense axis) when it lands on the +# domain side of the reconstruction, matching the `unmatricize`/`similar_map` axis convention. # Two-output factorizations: the first factor `X` has the codomain axes plus a trailing # rank axis, the second factor `Y` has a leading rank axis plus the domain axes. @@ -26,8 +19,8 @@ for f in ( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize_codomain(style, X, axes_codomain), - unmatricize_domain(style, Y, axes_domain) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), + unmatricize(style, Y, (axes(Y, 1),), axes_domain) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -225,9 +218,9 @@ for f in (:svd_compact, :svd_full, :svd_trunc) A_mat = matricize(style, A, ndims_codomain) U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize_codomain(style, U, axes_codomain), + return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)), - unmatricize_domain(style, Vᴴ, axes_domain) + unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -243,7 +236,7 @@ for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) A_mat = matricize(style, A, ndims_codomain) D, V = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return D, unmatricize_codomain(style, V, axes_codomain) + return D, unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),)) end function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -421,7 +414,7 @@ function left_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize_codomain(style, N, axes_codomain) + return unmatricize(style, N, axes_codomain, (conj(axes(N, ndims(N))),)) end function left_null!!(A, ndims_codomain::Val; kwargs...) return left_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -458,7 +451,7 @@ function right_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) _, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize_domain(style, Nᴴ, axes_domain) + return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end function right_null!!(A, ndims_codomain::Val; kwargs...) return right_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -514,7 +507,7 @@ function gram_eigh_full!!( A_mat = matricize(style, A, ndims_codomain) X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize_codomain(style, X, axes_codomain) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)) end function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(FusionStyle(A), A, ndims_codomain; kwargs...) @@ -575,8 +568,8 @@ function gram_eigh_full_with_pinv!!( A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) - return unmatricize_codomain(style, X, axes_codomain), - unmatricize_domain(style, Y, axes_codomain) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), + unmatricize(style, Y, (axes(Y, 1),), axes_codomain) end function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(FusionStyle(A), A, ndims_codomain; kwargs...) From b76744003c15e697a2993b4ca14b99e8e0e04047 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 15:16:23 -0400 Subject: [PATCH 08/15] Overload the diagonal power per backend rather than on diagview shape `pow_diag_safe` passes the whole matrix to `_pow_diag!` for each backend to overload, instead of passing `MAK.diagview(D)` and dispatching on what it returns. A `TensorMap`'s `diagview` is a `SectorVector` for a `DiagonalTensorMap` and a per-sector dict otherwise, which forced the shared code to carry a dict-recursion method. The TensorKit extension now clamps the diagonal per block, where each block's `diagview` is a plain vector view, and the generic path keeps the single-vector case for dense and graded diagonals. A shared helper holds the below-tolerance clamp and PSD-precondition rule for one entry so both backends use the same definition. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 15 ++++++++++++- src/MatrixAlgebra.jl | 36 +++++++++++++++++--------------- 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index fe2d7d2..4c8da86 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -1,9 +1,10 @@ module TensorAlgebraTensorKitExt +using MatrixAlgebraKit: diagview using Random: AbstractRNG using TensorAlgebra: TensorAlgebra using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, - TensorMapWithStorage, codomain, dim, domain, dual, fuse, numind, permute, + TensorMapWithStorage, blocks, codomain, dim, domain, dual, fuse, numind, permute, project_symmetric!, sectors, space, spacetype, zerovector!, ← using TensorOperations: TensorOperations as TO @@ -295,4 +296,16 @@ function Base.copy(::Base.Broadcast.Broadcasted{TensorMapStyle}) ) end +# ==================================== pow_diag_safe ====================================== +# `MAK.diagview` of a `TensorMap` is shape-shifting (a `SectorVector` for a `DiagonalTensorMap`, +# a per-sector dict otherwise), so clamp the diagonal per block instead: each block's `diagview` +# is a plain vector view regardless of the map's type. +function TensorAlgebra.MatrixAlgebra._pow_diag!(D::AbstractTensorMap, p, tol) + for (_, b) in blocks(D) + σ = diagview(b) + copyto!(σ, map(d -> TensorAlgebra.MatrixAlgebra._clamped_pow(d, p, tol), σ)) + end + return D +end + end diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index 17a58f8..41b169f 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -33,11 +33,10 @@ Negative `d` above `tol` cause `d^p` to error for fractional `p` (e.g. `p = 1//2`) and pass through for integer `p`, so the operation itself enforces the PSD precondition per-power. Errors if `isdiag(D)` is `false`. -The implementation writes the clamped powers back through `MAK.diagview` -onto a `copy` of `D`, so the result has the input's type and structure, and -types extending `diagview` (e.g. graded or block diagonal, a `TensorMap`) -automatically extend [`sqrt_diag_safe`](@ref), [`invsqrt_diag_safe`](@ref), -and the [`powh_safe`](@ref) family. +The clamped powers are written onto a `copy` of `D`, so the result keeps the +input's type and structure. This drives [`sqrt_diag_safe`](@ref), +[`invsqrt_diag_safe`](@ref), and the [`powh_safe`](@ref) family, and works for +any diagonal-structured backend (dense, graded, or a `TensorMap`). ## Keyword arguments @@ -54,21 +53,24 @@ function pow_diag_safe( ) tol = max(atol, rtol * norm(D, Inf)) Dp = copy(D) - _pow_diag!(MAK.diagview(Dp), p, tol) + _pow_diag!(Dp, p, tol) return Dp end -# `copyto!` rather than `.=`: block-structured diagonal views (e.g. a graded fused -# vector) have a blockwise `map` and `copyto!` but no broadcast support. -function _pow_diag!(σ, p, tol) - copyto!(σ, map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ)) - return σ -end -# A backend's `diagview` may be a dict of per-block diagonal views (e.g. a `TensorMap`, -# keyed by sector) rather than a single vector view. -function _pow_diag!(σ::AbstractDict, p, tol) - foreach(v -> _pow_diag!(v, p, tol), values(σ)) - return σ +# Clamp-then-power one diagonal entry: entries below `tol` go to zero, and a negative entry +# above `tol` lets `real(d)^p` error for fractional `p`, enforcing the PSD precondition +# per-power. Branching also skips the power for clamped entries. +_clamped_pow(d, p, tol) = abs(d) < tol ? zero(d) : real(d)^p + +# Clamp the powers in place on `D`'s diagonal, returning `D` with its type and structure +# intact. This generic path reads the diagonal with `MAK.diagview`, which for a dense or +# graded diagonal is a single mappable vector; a backend whose diagonal is stored otherwise +# (e.g. a `TensorMap`, per sector) overloads `_pow_diag!` directly. `copyto!` rather than +# `.=`: a graded fused diagonal vector has a blockwise `map` and `copyto!` but no broadcast. +function _pow_diag!(D, p, tol) + σ = MAK.diagview(D) + copyto!(σ, map(d -> _clamped_pow(d, p, tol), σ)) + return D end """ From 1ed7ef4bc4910716dec64b573310c4da4db1ea5a Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 15:31:27 -0400 Subject: [PATCH 09/15] Update the diagonal in place with map! to avoid an allocation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_pow_diag!` updates the diagonal with `map!(f, σ, σ)` instead of `copyto!(σ, map(f, σ))`, which allocated a temporary vector on every call. The generic path targets a dense, in-place-mappable diagonal view and the TensorKit extension maps over each block's plain vector view. A backend whose diagonal view is not in-place-mappable overloads `_pow_diag!`, so the generic path no longer needs the allocating map-and-copy form. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 2 +- src/MatrixAlgebra.jl | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 4c8da86..45f3a9a 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -303,7 +303,7 @@ end function TensorAlgebra.MatrixAlgebra._pow_diag!(D::AbstractTensorMap, p, tol) for (_, b) in blocks(D) σ = diagview(b) - copyto!(σ, map(d -> TensorAlgebra.MatrixAlgebra._clamped_pow(d, p, tol), σ)) + map!(d -> TensorAlgebra.MatrixAlgebra._clamped_pow(d, p, tol), σ, σ) end return D end diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index 41b169f..654f6c9 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -63,13 +63,12 @@ end _clamped_pow(d, p, tol) = abs(d) < tol ? zero(d) : real(d)^p # Clamp the powers in place on `D`'s diagonal, returning `D` with its type and structure -# intact. This generic path reads the diagonal with `MAK.diagview`, which for a dense or -# graded diagonal is a single mappable vector; a backend whose diagonal is stored otherwise -# (e.g. a `TensorMap`, per sector) overloads `_pow_diag!` directly. `copyto!` rather than -# `.=`: a graded fused diagonal vector has a blockwise `map` and `copyto!` but no broadcast. +# intact. The generic path maps over `MAK.diagview(D)`, a single in-place-mappable vector +# (dense); a backend whose diagonal view is not (e.g. a `TensorMap`, stored per sector) +# overloads `_pow_diag!` directly. function _pow_diag!(D, p, tol) σ = MAK.diagview(D) - copyto!(σ, map(d -> _clamped_pow(d, p, tol), σ)) + map!(d -> _clamped_pow(d, p, tol), σ, σ) return D end From 26779eb0a8b46513117d68a47a00535395aa1144 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 15:41:12 -0400 Subject: [PATCH 10/15] Clarify aux-axis names and note the slice-probe cost in aux inference Renames the two locals in `infer_aux_space` to `aux_dim` (the surplus axis's position) and `aux_length` (its size), and folds a note into the comment that the slice-by-slice `tryproject` probe builds one `TensorMap` per candidate column. A faster derivation would read the covariant sectors from the fusion-tree block structure instead. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 45f3a9a..8f317c2 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -173,14 +173,17 @@ end function infer_aux_space( raw, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} ) where {S <: ElementarySpace} - naux = length(codomain_axes) + length(domain_axes) + 1 - auxdim = size(raw, naux) + aux_dim = length(codomain_axes) + length(domain_axes) + 1 + aux_length = size(raw, aux_dim) content = fuse(codomain_axes..., dual.(domain_axes)...) - # A slice keeps the aux axis (width `dim(s)`), so its rank matches the candidate axes - # exactly and `tryproject` allocates, fills, and round-trip-verifies without re-entering - # the derivation branch. + # Probe the surplus axis slice by slice: a slice keeps the aux axis (width `dim(s)`), so its + # rank matches the candidate axes exactly and `tryproject` allocates, fills, and round-trip- + # verifies without re-entering the derivation branch. This builds one `TensorMap` per candidate + # column, which is fine for operator-sized inputs but not cheap. A faster derivation would read + # the covariant sectors from the fusion-tree block structure (as the `TensorMap` constructor + # does) instead of constructing a projection per slice. function slice_is_covariant(r, s) - slice = selectdim(raw, naux, r) + slice = selectdim(raw, aux_dim, r) return !isnothing( TensorAlgebra.tryproject(slice, codomain_axes, (domain_axes..., S(s => 1))) ) @@ -190,16 +193,16 @@ function infer_aux_space( for s in sectors(content) d = dim(S(s => 1)) m = 0 - while pos + d - 1 <= auxdim && slice_is_covariant(pos:(pos + d - 1), s) + while pos + d - 1 <= aux_length && slice_is_covariant(pos:(pos + d - 1), s) m += 1 pos += d end m > 0 && push!(seccounts, s => m) end - pos == auxdim + 1 || throw( + pos == aux_length + 1 || throw( ArgumentError( "`project`: could not derive a covariant auxiliary space for the surplus axis of \ - dimension $auxdim; the aux slices must be ordered by the canonical (sorted) sector \ + length $aux_length; the aux slices must be ordered by the canonical (sorted) sector \ order of the derived space" ) ) From fbccdb8c950930c7238dba61c2367432cdad4cd3 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 15:50:42 -0400 Subject: [PATCH 11/15] Correct the aux-inference perf comment The note in `infer_aux_space` claimed a faster derivation would read sectors "as the `TensorMap` constructor does", but `project_symmetric!` is given the tensor's space and only sorts dense data into its already-known blocks rather than discovering sectors. Rewords the note to say only that reading each column's sector from the fused content's block structure could avoid the per-slice projection. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 8f317c2..b8e7a6f 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -179,9 +179,8 @@ function infer_aux_space( # Probe the surplus axis slice by slice: a slice keeps the aux axis (width `dim(s)`), so its # rank matches the candidate axes exactly and `tryproject` allocates, fills, and round-trip- # verifies without re-entering the derivation branch. This builds one `TensorMap` per candidate - # column, which is fine for operator-sized inputs but not cheap. A faster derivation would read - # the covariant sectors from the fusion-tree block structure (as the `TensorMap` constructor - # does) instead of constructing a projection per slice. + # column, which is fine for operator-sized inputs but not cheap. Reading each column's sector + # directly from the block structure of the fused content could avoid the per-slice projection. function slice_is_covariant(r, s) slice = selectdim(raw, aux_dim, r) return !isnothing( From bc4963db01ba807cf74d4abb5f090c2d9218d9b2 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 16:30:18 -0400 Subject: [PATCH 12/15] Reject a misshapen project input instead of reinterpreting it The projection verbs reshaped `raw`/`src` to the destination size with a bare `reshape`, which silently reinterprets same-length data of a genuinely different shape (a 2x3 input into a 3x2 destination, say). Adds an internal `check_project_shape` that throws unless the two sizes agree ignoring trailing length-1 axes, mirroring `Base.size(A, d)` for `d > ndims(A)`, and guards the `reshape` in the generic `projectto!` and `is_projected` and in the TensorKit `projectto!` with it. Trailing length-1 axes may still be added or dropped. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 3 ++- src/projectto.jl | 35 ++++++++++++++++++++++++-------- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index b8e7a6f..aecb817 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -137,8 +137,9 @@ end # this makes `project(dense, codomain_axes, domain_axes)` build a `TensorMap` from a dense matrix. # `project_symmetric!` requires a matching dense size, so reshape `src` to `size(dest)` first (a # no-op when the ranks already match); this lets a lower-rank `src` omit trailing length-1 axes, -# matching the generic `projectto!`. +# matching the generic `projectto!`, and rejects a genuine shape mismatch. function TensorAlgebra.projectto!(dest::AbstractTensorMap, src::AbstractArray) + TensorAlgebra.check_project_shape(size(src), TensorAlgebra.size(dest)) return project_symmetric!(dest, reshape(src, TensorAlgebra.size(dest))) end diff --git a/src/projectto.jl b/src/projectto.jl index 45769c4..4207576 100644 --- a/src/projectto.jl +++ b/src/projectto.jl @@ -1,15 +1,31 @@ +# Throw unless `sz1` and `sz2` are equal ignoring trailing length-1 axes: an axis beyond one +# size's rank counts as length 1, mirroring `Base.size(A, d)` for `d > ndims(A)`. Guards a +# `reshape` against silently reinterpreting same-length data of a genuinely different shape, while +# still allowing trailing length-1 axes to be added or dropped (the projection verbs' "omit +# trailing length-1 axes" convention). +function check_project_shape(sz1::Dims, sz2::Dims) + all(i -> get(sz1, i, 1) == get(sz2, i, 1), 1:max(length(sz1), length(sz2))) || throw( + DimensionMismatch("sizes $sz1 and $sz2 differ beyond trailing length-1 axes") + ) + return nothing +end + """ projectto!(dest, src) -> dest Project `src` into the restricted space of `dest` without checking which -components may have been projected out. Defaults to `copyto!`, which copies -by linear index, so a lower-rank `src` may omit trailing length-1 axes (e.g. -an auxiliary flux-canceling leg a codomain/domain split introduces on a -symmetric state). A size-strict backend overloads this to reshape `src` to -`size(dest)` for the same effect. This is the in-place fill primitive that -[`unchecked_project`](@ref) allocates a destination for. +components may have been projected out. The default reshapes `src` to +`size(dest)` up to trailing length-1 axes (so a lower-rank `src` may omit +them, e.g. an auxiliary flux-canceling leg a codomain/domain split introduces +on a symmetric state) and `copyto!`s, throwing on a genuine shape mismatch +rather than reinterpreting the data. A backend whose arrays are not +`copyto!`-compatible with a dense array overloads this. This is the in-place +fill primitive that [`unchecked_project`](@ref) allocates a destination for. """ -projectto!(dest, src) = copyto!(dest, src) +function projectto!(dest, src) + check_project_shape(size(src), size(dest)) + return copyto!(dest, reshape(src, size(dest))) +end """ allocate_project(raw, codomain_axes, domain_axes) -> dest @@ -57,13 +73,14 @@ negligible component of `src`. Keyword arguments are forwarded to `isapprox`. Together with [`unchecked_project`](@ref) this is the backend customization point ([`project`](@ref) and [`tryproject`](@ref) derive from the two). The -generic method reshapes `src` to `size(dest)` (so a lower-rank `src` that -omits trailing length-1 axes lines up) and compares against +generic method reshapes `src` to `size(dest)` up to trailing length-1 axes +(so a lower-rank `src` that omits them lines up) and compares against `convert(Array, dest)`, so a backend whose arrays are not elementwise comparable to a dense array (opaque block storage, a `TensorMap`) only needs that conversion. """ function is_projected(dest, src; kwargs...) + check_project_shape(size(src), size(dest)) return isapprox(reshape(src, size(dest)), convert(Array, dest); kwargs...) end From cf3588c4442635a8ea400584f17282f4068ddd48 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 16:47:20 -0400 Subject: [PATCH 13/15] Handle an empty codomain in the TensorMap projection allocation `allocate_project` for `TensorMap` spaces only dispatched on a non-empty codomain, so an empty-codomain projection (a functional over the domain) skipped the auxiliary-leg derivation and fell through to the generic allocation. Splits it into two entries, codomain-led and domain-led, that read the elementary space type from whichever side is non-empty and share an `allocate_project_tensormap` builder, the same two-entry split `similar_map` uses, and threads that space type into `infer_aux_space`. Also renames the shape guard `check_project_shape` to `check_project_size` since it takes sizes. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/TensorAlgebraTensorKitExt.jl | 26 ++++++++++++++++++++------ src/projectto.jl | 6 +++--- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index aecb817..52781c9 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -139,7 +139,7 @@ end # no-op when the ranks already match); this lets a lower-rank `src` omit trailing length-1 axes, # matching the generic `projectto!`, and rejects a genuine shape mismatch. function TensorAlgebra.projectto!(dest::AbstractTensorMap, src::AbstractArray) - TensorAlgebra.check_project_shape(size(src), TensorAlgebra.size(dest)) + TensorAlgebra.check_project_size(size(src), TensorAlgebra.size(dest)) return project_symmetric!(dest, reshape(src, TensorAlgebra.size(dest))) end @@ -147,12 +147,26 @@ end # defines for an `AbstractTensorMap`, so no dedicated method is needed here. # ============================= allocate_project (aux-leg derivation) ===================== -# `allocate_project` for `TensorMap` spaces. With no surplus axis this is plain `similar_map`; -# a single trailing surplus axis in `raw` is an auxiliary leg whose space is derived (see -# `infer_aux_space`) and appended as the last domain axis so the result is symmetry-allowed. +# `allocate_project` for `TensorMap` spaces routes both the codomain-led and the (empty-codomain) +# domain-led cases to `allocate_project_tensormap`, reading the elementary space type `S` from +# whichever side is non-empty, the same two-entry split `similar_map` uses. function TensorAlgebra.allocate_project( raw::AbstractArray, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} ) where {S <: ElementarySpace} + return allocate_project_tensormap(raw, S, codomain_axes, domain_axes) +end +function TensorAlgebra.allocate_project( + raw::AbstractArray, codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}} + ) where {S <: ElementarySpace} + return allocate_project_tensormap(raw, S, codomain_axes, domain_axes) +end + +# With no surplus axis this is plain `similar_map`; a single trailing surplus axis in `raw` is an +# auxiliary leg whose space is derived (see `infer_aux_space`) and appended as the last domain axis +# so the result is symmetry-allowed. +function allocate_project_tensormap( + raw, ::Type{S}, codomain_axes, domain_axes + ) where {S <: ElementarySpace} nphys = length(codomain_axes) + length(domain_axes) ndims(raw) <= nphys && return TensorAlgebra.similar_map(raw, codomain_axes, domain_axes) @@ -162,7 +176,7 @@ function TensorAlgebra.allocate_project( given axes, got a rank-$(ndims(raw)) input" ) ) - aux = infer_aux_space(raw, codomain_axes, domain_axes) + aux = infer_aux_space(raw, S, codomain_axes, domain_axes) return TensorAlgebra.similar_map(raw, codomain_axes, (domain_axes..., aux)) end @@ -172,7 +186,7 @@ end # follows, so the aux slices must appear in that order. The result may span several sectors (a # direct-sum, MPO-style virtual leg). function infer_aux_space( - raw, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} + raw, ::Type{S}, codomain_axes, domain_axes ) where {S <: ElementarySpace} aux_dim = length(codomain_axes) + length(domain_axes) + 1 aux_length = size(raw, aux_dim) diff --git a/src/projectto.jl b/src/projectto.jl index 4207576..759e395 100644 --- a/src/projectto.jl +++ b/src/projectto.jl @@ -3,7 +3,7 @@ # `reshape` against silently reinterpreting same-length data of a genuinely different shape, while # still allowing trailing length-1 axes to be added or dropped (the projection verbs' "omit # trailing length-1 axes" convention). -function check_project_shape(sz1::Dims, sz2::Dims) +function check_project_size(sz1::Dims, sz2::Dims) all(i -> get(sz1, i, 1) == get(sz2, i, 1), 1:max(length(sz1), length(sz2))) || throw( DimensionMismatch("sizes $sz1 and $sz2 differ beyond trailing length-1 axes") ) @@ -23,7 +23,7 @@ rather than reinterpreting the data. A backend whose arrays are not fill primitive that [`unchecked_project`](@ref) allocates a destination for. """ function projectto!(dest, src) - check_project_shape(size(src), size(dest)) + check_project_size(size(src), size(dest)) return copyto!(dest, reshape(src, size(dest))) end @@ -80,7 +80,7 @@ comparable to a dense array (opaque block storage, a `TensorMap`) only needs that conversion. """ function is_projected(dest, src; kwargs...) - check_project_shape(size(src), size(dest)) + check_project_size(size(src), size(dest)) return isapprox(reshape(src, size(dest)), convert(Array, dest); kwargs...) end From fc029ec92f47a9ac6284fa278ddddd88f3f884bd Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 20:19:07 -0400 Subject: [PATCH 14/15] Replace ConjArray with a backend-agnostic ConjBroadcasted node Eliminate `TensorAlgebra.conjugate` and the `ConjArray` wrapper in favor of `ConjBroadcasted`, a lazy conjugation node that is a `LinearBroadcasted` rather than an `AbstractArray`. `ConjArray` could only wrap an `AbstractArray`, so a `TensorMap` had no lazy conjugation. `ConjBroadcasted` folds into `bipermutedimsopadd!` like the other linear-broadcast nodes and works on any backend, so `conj` now lowers to it and a nested `conj` cancels. The `conjed` name is dropped for now. It will return as a public interface for lazy conjugated contraction such as `X * conjed(X)`, backed by a proper conjugation wrapper. --- ext/TensorAlgebraTensorKitExt.jl | 12 --- src/TensorAlgebra.jl | 3 +- src/conjarray.jl | 61 --------------- src/interface.jl | 10 --- src/linearbroadcasted.jl | 77 ++++++++++++++++--- ...t_conjarray.jl => test_conjbroadcasted.jl} | 37 ++++----- test/test_exports.jl | 2 +- test/test_linearbroadcasted.jl | 6 +- test/test_permutedimsadd.jl | 12 +-- 9 files changed, 97 insertions(+), 123 deletions(-) delete mode 100644 src/conjarray.jl rename test/{test_conjarray.jl => test_conjbroadcasted.jl} (51%) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 52781c9..37c39ee 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -21,18 +21,6 @@ TensorAlgebra.axes(t::AbstractTensorMap) = ntuple(i -> space(t, i), numind(t)) TensorAlgebra.size(t::AbstractTensorMap, i::Int) = dim(space(t, i)) TensorAlgebra.size(t::AbstractTensorMap) = ntuple(i -> dim(space(t, i)), numind(t)) -# TensorKit deliberately defines no `Base.conj` for tensors; its native conjugation is -# `adjoint`, which also swaps the codomain and domain. `TensorAlgebra.conjugate` -# (conjugated elements on dualized spaces, leg order unchanged) is the adjoint followed by -# the block swap returning each leg to its original flat position — the inverse of -# TensorKit's `adjointtensorindex` convention (`i <= numout ? numin + i : i - numout`), -# the same lowering its TensorOperations interface uses for `conj` arguments. -function TensorAlgebra.conjugate(t::AbstractTensorMap) - m, n = TensorKit.numout(t), TensorKit.numin(t) - p = ntuple(i -> i <= m ? n + i : i - m, m + n) - return permute(adjoint(t), (p[1:m], p[(m + 1):(m + n)])) -end - # `t[]` on a rank-0 `TensorMap` requires a trivial sector type; `TensorKit.scalar` is the # general spelling. TensorAlgebra.scalar(t::AbstractTensorMap) = TensorKit.scalar(t) diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index ce3d708..9785f9f 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -9,7 +9,7 @@ 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, conjed, ConjArray" + "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, permutedims, permutedims!, to_range, zero!, scale!, permuteddims, PermutedDims" ) ) end @@ -19,7 +19,6 @@ include("inplace.jl") include("MatrixAlgebra.jl") include("bituple.jl") include("permutedimsadd.jl") -include("conjarray.jl") include("matricize.jl") include("diagonal.jl") include("to_range.jl") diff --git a/src/conjarray.jl b/src/conjarray.jl deleted file mode 100644 index 0757558..0000000 --- a/src/conjarray.jl +++ /dev/null @@ -1,61 +0,0 @@ -# Lazy conjugated-array wrapper, mirroring `Adjoint`/`adjoint`. `conjed(a)` is the -# friendly lazy companion to the eager `conj(a)`. The op primitives absorb a `ConjArray` -# by unwrapping to the parent and folding `conj` into `op` (see `bipermutedimsopadd!` -# below), so conjugation flows through broadcasting, contraction, and matricize uniformly. - -""" - ConjArray(a::AbstractArray) - -Lazy conjugate of `a`: an `AbstractArray` whose axes are the conjugated parent axes and -whose elements are the conjugated parent elements. Constructed by [`conjed`](@ref). The op -primitives absorb it by unwrapping to the parent and folding `conj` into their `op`. -""" -struct ConjArray{T, N, P <: AbstractArray{T, N}} <: AbstractArray{T, N} - parent::P -end - -""" - conjed(a::AbstractArray) - -Lazy conjugate of `a`, returning a `ConjArray`. The lazy companion to the eager `conj(a)`, -mirroring `adjoint`/`Adjoint`. `conjed(conjed(a))` returns `a`. -""" -conjed(a::AbstractArray) = ConjArray(a) -conjed(a::ConjArray) = parent(a) - -Base.parent(a::ConjArray) = a.parent - -# Conjugating an axis is identity for plain integer ranges and dualizes graded axes (where -# `conj` is overloaded as `dual` downstream). -Base.axes(a::ConjArray) = map(conj, axes(parent(a))) -Base.size(a::ConjArray) = size(parent(a)) - -Base.IndexStyle(::Type{<:ConjArray{<:Any, <:Any, P}}) where {P} = IndexStyle(P) -Base.getindex(a::ConjArray, I::Int...) = conj(parent(a)[I...]) - -# Absorb a `ConjArray` source by unwrapping to the parent and folding `conj` into `op` -# (`identity -> conj`, `conj -> identity`). Mirrors the `PermutedDimsArray` absorption, -# which instead folds its permutation into `perm`. -function bipermutedimsopadd!( - dest::AbstractArray, op, src::ConjArray, - perm_codomain, perm_domain, - α::Number, β::Number - ) - return bipermutedimsopadd!( - dest, _compose_op(op, conj), parent(src), - perm_codomain, perm_domain, - α, β - ) -end - -# Materialize eagerly through the op primitive, so the conjugation (and any parent-specific -# behavior such as a graded fermion sign) is applied by the same `op = conj` path the -# primitives use everywhere else. `copy(conjed(a))` for a plain `Array` is `conj(a)`. -Base.copy(a::ConjArray) = add!(similar(parent(a), eltype(a), axes(a)), a, true, false) -Base.copyto!(dest::AbstractArray, src::ConjArray) = add!(dest, src, true, false) -Base.Broadcast.materialize(a::ConjArray) = copy(a) - -# Broadcast like the parent (e.g. preserve a graded style), not the default array style. -function Base.Broadcast.BroadcastStyle(::Type{<:ConjArray{<:Any, <:Any, P}}) where {P} - return Base.Broadcast.BroadcastStyle(P) -end diff --git a/src/interface.jl b/src/interface.jl index 932eff5..660e2be 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -14,16 +14,6 @@ function size end size(a) = Base.size(a) size(a, i::Int) = Base.size(a, i) -# Eager conjugation in the same vocabulary: conjugated elements on conjugated (dualized) -# axes, leg order unchanged. Backends whose native conjugation is `adjoint`-shaped (such as -# a `TensorMap`, where the adjoint also swaps codomain and domain) overload this to restore -# the original leg order. Named `conjugate` (the eager companion of the lazy `conjed`) -# rather than `conj`: `Base.conj` is passed around as an `op` value inside this package -# (`op === conj` guards, `::typeof(conj)` dispatch), so a module-level `conj` binding -# would silently change what those mean. -function conjugate end -conjugate(a) = Base.conj(a) - # 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). function scalar end diff --git a/src/linearbroadcasted.jl b/src/linearbroadcasted.jl index 125419f..aaafdd0 100644 --- a/src/linearbroadcasted.jl +++ b/src/linearbroadcasted.jl @@ -65,10 +65,55 @@ Base.ndims(a::ScaledBroadcasted) = ndims(unscaled(a)) operation(::ScaledBroadcasted) = * arguments(a::ScaledBroadcasted) = (coeff(a), unscaled(a)) -# Conjugation is represented by the `ConjArray` lazy wrapper (a real `AbstractArray` with -# conjugated axes), not a dedicated `LinearBroadcasted` node. The lowering below produces -# `ConjArray` leaves, and the op primitives absorb them by folding `conj` into `op` (see -# `conjarray.jl`). +# --- ConjBroadcasted ---------------------------------------------------------- + +""" + ConjBroadcasted(parent) + +Lazy conjugate node in the linear-combination broadcast fold, the [`LinearBroadcasted`](@ref) +counterpart of `ScaledBroadcasted`/`AddBroadcasted`. Holds `parent` (any array-like operand or +backend tensor, not necessarily an `AbstractArray`) and presents its axes conjugated (dualized); +the op primitives absorb it by unwrapping to the parent and folding `conj` into their `op` (see +the `bipermutedimsopadd!` method below). Produced internally by the `conj` lowering of a +broadcast (`linearbroadcasted(conj, a)`), not a user-facing lazy-conjugate wrapper. Because it +is not an `AbstractArray`, it works for non-array backends (e.g. a `TensorMap`) as well as dense +arrays. +""" +struct ConjBroadcasted{P} <: LinearBroadcasted + parent::P +end + +Base.parent(a::ConjBroadcasted) = a.parent + +# Conjugating an axis is identity for plain integer ranges and dualizes graded axes (where +# `conj` is overloaded as `dual` downstream). +Base.axes(a::ConjBroadcasted) = map(conj, axes(parent(a))) +Base.eltype(a::ConjBroadcasted) = eltype(parent(a)) +Base.ndims(a::ConjBroadcasted) = ndims(parent(a)) + +operation(::ConjBroadcasted) = conj +arguments(a::ConjBroadcasted) = (parent(a),) + +function BC.BroadcastStyle(::Type{<:ConjBroadcasted{P}}) where {P} + return BC.BroadcastStyle(P) +end + +# Absorb a `ConjBroadcasted` source by unwrapping to the parent and folding `conj` into `op` +# (`identity -> conj`, `conj -> identity`). Mirrors the `PermutedDims` absorption, which +# instead folds its permutation into `perm`. `dest` is unrestricted so a non-`AbstractArray` +# backend (e.g. a `TensorMap`) receives it and applies `conj` through its own op-aware +# `bipermutedimsopadd!`. +function bipermutedimsopadd!( + dest, op, src::ConjBroadcasted, + perm_codomain, perm_domain, + α::Number, β::Number + ) + return bipermutedimsopadd!( + dest, _compose_op(op, conj), parent(src), + perm_codomain, perm_domain, + α, β + ) +end # --- AddBroadcasted ----------------------------------------------------------- @@ -79,7 +124,18 @@ end addends(a::AddBroadcasted) = a.args -Base.axes(a::AddBroadcasted) = BC.combine_axes(addends(a)...) +# All addends of a linear combination share axes (a linear combination does not broadcast +# differing shapes), so combine by verifying equality through `axes` (TensorAlgebra's, which +# works for a non-`AbstractArray` backend like a `TensorMap`) rather than Base's `combine_axes`, +# which would call `Base.axes`/`Base.size` on the operands. A mismatch (e.g. a half-conjugated +# `conj.(a) .- b`, whose dualized and non-dualized axes differ) throws here. +function Base.axes(a::AddBroadcasted) + axs = map(axes, addends(a)) + ax = first(axs) + all(x -> x == ax, axs) || + throw(DimensionMismatch("linear-combination operands have mismatched axes: $axs")) + return ax +end Base.eltype(a::AddBroadcasted) = Base.promote_op(+, eltype.(addends(a))...) Base.ndims(a::AddBroadcasted) = ndims(first(addends(a))) @@ -195,7 +251,7 @@ Analogous to `Base.Broadcast.broadcasted(f, args...)`. ```julia linearbroadcasted(*, 2.0, a) # ScaledBroadcasted(2.0, a) -linearbroadcasted(conj, a) # ConjArray(a) +linearbroadcasted(conj, a) # ConjBroadcasted(a) linearbroadcasted(+, a, b) # AddBroadcasted(a, b) ``` """ @@ -208,10 +264,11 @@ linearbroadcasted(::typeof(*), a, α::Number) = ScaledBroadcasted(α, a) function linearbroadcasted(::typeof(*), α::Number, a::ScaledBroadcasted) return ScaledBroadcasted(α * coeff(a), unscaled(a)) end -# Conjugation lowers to the `ConjArray` lazy wrapper. A scaled `ConjArray` (e.g. -# `conj.(a) ./ β`) is handled by the generic `AbstractArray` scaling method above, since -# `ConjArray <: AbstractArray`. -linearbroadcasted(::typeof(conj), a) = conjed(a) +# Conjugation lowers to the `ConjBroadcasted` node; a nested conjugate cancels (`conj∘conj = +# identity`). A scaled conjugate (e.g. `conj.(a) ./ β`) is handled by the generic scaling method +# above, whose operand slot is untyped, so it wraps the `ConjBroadcasted` in a `ScaledBroadcasted`. +linearbroadcasted(::typeof(conj), a) = ConjBroadcasted(a) +linearbroadcasted(::typeof(conj), a::ConjBroadcasted) = parent(a) function linearbroadcasted(::typeof(conj), a::ScaledBroadcasted) return ScaledBroadcasted(conj(coeff(a)), linearbroadcasted(conj, unscaled(a))) end diff --git a/test/test_conjarray.jl b/test/test_conjbroadcasted.jl similarity index 51% rename from test/test_conjarray.jl rename to test/test_conjbroadcasted.jl index 6b1be16..c473a3d 100644 --- a/test/test_conjarray.jl +++ b/test/test_conjbroadcasted.jl @@ -1,40 +1,41 @@ -using TensorAlgebra: TensorAlgebra as TA, ConjArray, bipermutedimsopadd!, conjed +using TensorAlgebra: TensorAlgebra as TA, ConjBroadcasted, PermutedDims, + bipermutedimsopadd!, linearbroadcasted using Test: @test, @testset -@testset "ConjArray (plain arrays)" begin +@testset "ConjBroadcasted (plain arrays)" begin a = randn(ComplexF64, 3, 4) - c = conjed(a) - @test c isa ConjArray + c = ConjBroadcasted(a) @test parent(c) === a - @test conjed(c) === a # involution unwraps, no nesting @test eltype(c) === eltype(a) @test ndims(c) == 2 - @test size(c) == size(a) @test axes(c) == axes(a) # plain axes are unchanged by conj - # Indexing and materialization match eager conj of the parent. - @test c[2, 3] == conj(a[2, 3]) - @test collect(c) == conj(a) + # The `conj` lowering produces a `ConjBroadcasted`, and a nested conjugate cancels. + @test linearbroadcasted(conj, a) ≡ c + @test linearbroadcasted(conj, c) === a + + # Materialization matches eager conj of the parent (via the `LinearBroadcasted` protocol; + # `ConjBroadcasted` is not an `AbstractArray`, so it is not indexed or `collect`ed). + @test copy(c) == conj(a) # Real eltype: conj is a value no-op, but the wrapper still wraps. r = randn(Float64, 2, 2) - cr = conjed(r) - @test cr isa ConjArray - @test collect(cr) == r + cr = ConjBroadcasted(r) + @test copy(cr) == r @test axes(cr) == axes(r) end -@testset "ConjArray composes with PermutedDimsArray" begin +@testset "ConjBroadcasted composes with permutation wrappers" begin a = randn(ComplexF64, 2, 3, 4, 5) w = (3, 1, 4, 2) - # ConjArray outside, PermutedDimsArray inside, and vice versa. - c_out = ConjArray(PermutedDimsArray(a, w)) - c_in = PermutedDimsArray(ConjArray(a), w) - @test collect(c_out) == conj(permutedims(a, w)) + # conj outside permute wraps a Base `PermutedDimsArray`; permute outside conj uses the + # generic `PermutedDims` node, since `ConjBroadcasted` is not an `AbstractArray` and so + # cannot be a `PermutedDimsArray` parent. + c_out = ConjBroadcasted(PermutedDimsArray(a, w)) + c_in = PermutedDims(ConjBroadcasted(a), w) @test copy(c_out) ≈ conj(permutedims(a, w)) - @test collect(c_in) == permutedims(conj(a), w) # Both unwrap through bipermutedimsopadd! (the nested wrappers fold into op and perm). for (pc, pd) in (((1, 2, 3, 4), ()), ((2, 4), (1, 3))) diff --git a/test/test_exports.jl b/test/test_exports.jl index cafc339..c2f2f6f 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -40,7 +40,7 @@ using Test: @test, @testset [ :biperm, :bipartition, :contractopadd!, :label_type, :matricizeopperm, :permutedims, :permutedims!, :to_range, :zero!, :scale!, :permuteddims, - :PermutedDims, :conjed, :ConjArray, + :PermutedDims, ] ) end diff --git a/test/test_linearbroadcasted.jl b/test/test_linearbroadcasted.jl index 5535f9a..fa444c7 100644 --- a/test/test_linearbroadcasted.jl +++ b/test/test_linearbroadcasted.jl @@ -13,7 +13,7 @@ using Test: @test, @test_throws, @testset @test copy(x) ≈ 2a x = linearbroadcasted(conj, a) - @test x ≡ TA.ConjArray(a) + @test x ≡ TA.ConjBroadcasted(a) @test copy(x) ≈ conj(a) x = linearbroadcasted(+, a, b) @@ -64,7 +64,7 @@ using Test: @test, @test_throws, @testset # Conjugation of scaled x = linearbroadcasted(conj, linearbroadcasted(*, 2im, a)) - @test x ≡ TA.ScaledBroadcasted(-2im, TA.ConjArray(a)) + @test x ≡ TA.ScaledBroadcasted(-2im, TA.ConjBroadcasted(a)) # Double conjugation cancels @test linearbroadcasted(conj, linearbroadcasted(conj, a)) ≡ a @@ -160,7 +160,7 @@ using Test: @test, @test_throws, @testset TA.add!(dest, linearbroadcasted(+, a, b), true, false) @test dest ≈ a + b - # add! with a ConjArray + # add! with a ConjBroadcasted dest = zeros(ComplexF64, 3, 3) TA.add!(dest, linearbroadcasted(conj, a), true, false) @test dest ≈ conj(a) diff --git a/test/test_permutedimsadd.jl b/test/test_permutedimsadd.jl index 07c9ba6..c897591 100644 --- a/test/test_permutedimsadd.jl +++ b/test/test_permutedimsadd.jl @@ -1,7 +1,7 @@ using Adapt: adapt using JLArrays: JLArray -using TensorAlgebra: TensorAlgebra, ConjArray, PermutedDims, add!, bipermutedimsopadd!, - conjed, permuteddims, permutedimsadd!, permutedimsopadd! +using TensorAlgebra: TensorAlgebra, ConjBroadcasted, PermutedDims, add!, + bipermutedimsopadd!, permuteddims, permutedimsadd!, permutedimsopadd! using Test: @test, @testset # A non-`AbstractArray` operand, to check that `permuteddims` falls back to `PermutedDims`. @@ -104,14 +104,14 @@ end end end end - @testset "bipermutedimsopadd! unwraps ConjArray src (arraytype=$arrayt)" for arrayt in + @testset "bipermutedimsopadd! unwraps ConjBroadcasted src (arraytype=$arrayt)" for arrayt in ( Array, JLArray, ) dev = adapt(arrayt) parent = dev(randn(ComplexF64, 2, 3, 4, 5)) - src = ConjArray(parent) + src = ConjBroadcasted(parent) for (pc, pd) in (((1, 2, 3, 4), ()), ((2, 4), (1, 3)), ((3, 1), (2, 4))) perm = (pc..., pd...) # `op = identity` composes with the wrapper's `conj`: result is the conjugated, @@ -131,7 +131,7 @@ end end end end - @testset "add!(b, conjed(a)) matches eager conj (arraytype=$arrayt)" for arrayt in + @testset "add!(b, ConjBroadcasted(a)) matches eager conj (arraytype=$arrayt)" for arrayt in ( Array, JLArray, @@ -143,7 +143,7 @@ end b = dev(randn(ComplexF64, 2, 3, 4)) b_lazy = copy(b) b_eager = copy(b) - add!(b_lazy, conjed(a), α, β) + add!(b_lazy, ConjBroadcasted(a), α, β) add!(b_eager, conj(a), α, β) @test b_lazy ≈ b_eager end From 71933924c7a28e54d6b592be166e6834ee058158 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 6 Jul 2026 20:20:16 -0400 Subject: [PATCH 15/15] Add a public pow_diag_safe! backend hook for the diagonal power Split `pow_diag_safe` into a public trio so a backend has a friendly overload point instead of the internal `_pow_diag!`. The in-place `pow_diag_safe!(Dp, D, p, tol)` is the method a backend overloads, with the keyword-tolerance `pow_diag_safe(D, p; atol, rtol)` and a positional-`tol` form on top. The generic kernel maps over `MAK.diagview`, and the TensorKit extension clamps per block. --- ext/TensorAlgebraTensorKitExt.jl | 15 ++++++--- src/MatrixAlgebra.jl | 52 +++++++++++++++++++++----------- test/test_exports.jl | 1 + test/test_matrixalgebra.jl | 9 ++++++ 4 files changed, 54 insertions(+), 23 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 37c39ee..16d10e9 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -305,12 +305,17 @@ end # `MAK.diagview` of a `TensorMap` is shape-shifting (a `SectorVector` for a `DiagonalTensorMap`, # a per-sector dict otherwise), so clamp the diagonal per block instead: each block's `diagview` # is a plain vector view regardless of the map's type. -function TensorAlgebra.MatrixAlgebra._pow_diag!(D::AbstractTensorMap, p, tol) - for (_, b) in blocks(D) - σ = diagview(b) - map!(d -> TensorAlgebra.MatrixAlgebra._clamped_pow(d, p, tol), σ, σ) +function TensorAlgebra.MatrixAlgebra.pow_diag_safe!( + Dp::AbstractTensorMap, D::AbstractTensorMap, p, tol + ) + for ((_, bp), (_, b)) in zip(blocks(Dp), blocks(D)) + map!( + d -> TensorAlgebra.MatrixAlgebra._clamped_pow(d, p, tol), + diagview(bp), + diagview(b) + ) end - return D + return Dp end end diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index 654f6c9..82aeece 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -5,6 +5,7 @@ export gram_eigh_full, invsqrt_diag_safe, invsqrth_safe, pow_diag_safe, + pow_diag_safe!, powh_safe, sqrt_diag_safe, sqrth_safe, @@ -24,19 +25,22 @@ end """ pow_diag_safe(D, p; atol=0, rtol=eps(real(eltype(D)))^(3//4)) -> D^p + pow_diag_safe(D, p, tol) -> D^p Raise a diagonal-structured matrix `D` to the power `p`. Diagonal entries `d` with `abs(d) < tol` are clamped to zero before exponentiation, where `tol = max(atol, rtol * norm(D, Inf))` (the largest-magnitude entry, which -is the largest-magnitude diagonal entry for a diagonal-structured matrix). -Negative `d` above `tol` cause `d^p` to error for fractional `p` (e.g. -`p = 1//2`) and pass through for integer `p`, so the operation itself -enforces the PSD precondition per-power. Errors if `isdiag(D)` is `false`. - -The clamped powers are written onto a `copy` of `D`, so the result keeps the -input's type and structure. This drives [`sqrt_diag_safe`](@ref), -[`invsqrt_diag_safe`](@ref), and the [`powh_safe`](@ref) family, and works for -any diagonal-structured backend (dense, graded, or a `TensorMap`). +is the largest-magnitude diagonal entry for a diagonal-structured matrix). The +three-argument form takes `tol` directly. Negative `d` above `tol` cause `d^p` +to error for fractional `p` (e.g. `p = 1//2`) and pass through for integer `p`, +so the operation itself enforces the PSD precondition per-power. Errors if +`isdiag(D)` is `false`. + +The clamped powers are written onto a `copy` of `D` via [`pow_diag_safe!`](@ref), +so the result keeps the input's type and structure. This drives +[`sqrt_diag_safe`](@ref), [`invsqrt_diag_safe`](@ref), and the [`powh_safe`](@ref) +family, and works for any diagonal-structured backend (dense, graded, or a +`TensorMap`). ## Keyword arguments @@ -52,8 +56,12 @@ function pow_diag_safe( ArgumentError("pow_diag_safe requires a diagonal-structured matrix") ) tol = max(atol, rtol * norm(D, Inf)) + return pow_diag_safe(D, p, tol) +end + +function pow_diag_safe(D, p, tol) Dp = copy(D) - _pow_diag!(Dp, p, tol) + pow_diag_safe!(Dp, D, p, tol) return Dp end @@ -62,14 +70,22 @@ end # per-power. Branching also skips the power for clamped entries. _clamped_pow(d, p, tol) = abs(d) < tol ? zero(d) : real(d)^p -# Clamp the powers in place on `D`'s diagonal, returning `D` with its type and structure -# intact. The generic path maps over `MAK.diagview(D)`, a single in-place-mappable vector -# (dense); a backend whose diagonal view is not (e.g. a `TensorMap`, stored per sector) -# overloads `_pow_diag!` directly. -function _pow_diag!(D, p, tol) - σ = MAK.diagview(D) - map!(d -> _clamped_pow(d, p, tol), σ, σ) - return D +""" + pow_diag_safe!(Dp, D, p, tol) -> Dp + +In-place kernel behind [`pow_diag_safe`](@ref): write the clamped powers of +`D`'s diagonal onto `Dp`, where entries `d` with `abs(d) < tol` become zero and +the rest become `d^p`. `Dp` must be diagonal-structured with the same structure +as `D` (in the allocating [`pow_diag_safe`](@ref) path it is `copy(D)`). + +This is the backend overload point. The generic method maps over `MAK.diagview`, +valid when the diagonal view is a single in-place-mappable vector (dense). A +backend whose diagonal is stored per sector (a `TensorMap`, or a graded matrix) +overloads this method to clamp each reduced block. +""" +function pow_diag_safe!(Dp, D, p, tol) + map!(d -> _clamped_pow(d, p, tol), MAK.diagview(Dp), MAK.diagview(D)) + return Dp end """ diff --git a/test/test_exports.jl b/test/test_exports.jl index c2f2f6f..67b2046 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -53,6 +53,7 @@ using Test: @test, @testset :invsqrt_diag_safe, :invsqrth_safe, :pow_diag_safe, + :pow_diag_safe!, :powh_safe, :sqrt_diag_safe, :sqrth_invsqrth_safe, diff --git a/test/test_matrixalgebra.jl b/test/test_matrixalgebra.jl index 68a61e3..af75ebd 100644 --- a/test/test_matrixalgebra.jl +++ b/test/test_matrixalgebra.jl @@ -216,5 +216,14 @@ elts = (Float32, Float64, ComplexF32, ComplexF64) # `pow_diag_safe` rejects non-diagonal inputs. @test_throws ArgumentError MatrixAlgebra.pow_diag_safe(A, 1 // 2) + + # Positional-`tol` form and the in-place `pow_diag_safe!` kernel agree with + # the keyword form. Entries below `tol` clamp to zero. + Dc = Diagonal(real(elt)[4, 9, 0]) + @test MatrixAlgebra.pow_diag_safe(Dc, 1 // 2, 1.0e-8) == + Diagonal(real(elt)[2, 3, 0]) + Dp = copy(Dc) + @test MatrixAlgebra.pow_diag_safe!(Dp, Dc, 1 // 2, 1.0e-8) === Dp + @test Dp == Diagonal(real(elt)[2, 3, 0]) end end