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..16d10e9 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -1,10 +1,11 @@ module TensorAlgebraTensorKitExt +using MatrixAlgebraKit: diagview using Random: AbstractRNG using TensorAlgebra: TensorAlgebra using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, - TensorMapWithStorage, numind, permute, project_symmetric!, space, spacetype, - zerovector!, ← + TensorMapWithStorage, blocks, codomain, dim, domain, dual, fuse, numind, permute, + project_symmetric!, sectors, space, spacetype, zerovector!, ← using TensorOperations: TensorOperations as TO # ============================ AbstractArray-vocabulary bridge ============================ @@ -16,6 +17,29 @@ 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)) + +# `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; 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 +# 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 +123,92 @@ 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!`, and rejects a genuine shape mismatch. 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.check_project_size(size(src), TensorAlgebra.size(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` 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) + 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" + ) ) - TensorAlgebra.projectto!(dest, src) - isapprox(src, convert(Array, dest); kwargs...) || - throw(InexactError(:checked_projectto!, typeof(dest), src)) - return dest + aux = infer_aux_space(raw, S, 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, ::Type{S}, codomain_axes, domain_axes + ) where {S <: ElementarySpace} + aux_dim = length(codomain_axes) + length(domain_axes) + 1 + aux_length = size(raw, aux_dim) + content = fuse(codomain_axes..., dual.(domain_axes)...) + # 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. 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( + 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 <= aux_length && slice_is_covariant(pos:(pos + d - 1), s) + m += 1 + pos += d + end + m > 0 && push!(seccounts, s => m) + end + pos == aux_length + 1 || throw( + ArgumentError( + "`project`: could not derive a covariant auxiliary space for the surplus axis of \ + length $aux_length; the aux slices must be ordered by the canonical (sorted) sector \ + order of the derived space" + ) + ) + return S(seccounts...) end # ================================ bipermutedimsopadd! ===================================== @@ -148,6 +240,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 @@ -201,4 +301,21 @@ 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_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 Dp +end + end diff --git a/src/MatrixAlgebra.jl b/src/MatrixAlgebra.jl index d3ed5e6..82aeece 100644 --- a/src/MatrixAlgebra.jl +++ b/src/MatrixAlgebra.jl @@ -5,9 +5,11 @@ export gram_eigh_full, invsqrt_diag_safe, invsqrth_safe, pow_diag_safe, + 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 @@ -22,26 +24,30 @@ 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 + pow_diag_safe(D, p, tol) -> 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)))`. -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) -automatically extend [`sqrt_diag_safe`](@ref), [`invsqrt_diag_safe`](@ref), -and the [`powh_safe`](@ref) family. +`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). 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 $(_clamp_kwargs_doc("D")) """ function pow_diag_safe( - D::AbstractMatrix, p; + D, p; atol = zero(real(eltype(D))), rtol = iszero(atol) ? eps(real(eltype(D)))^(3 // 4) : zero(real(eltype(D))) @@ -49,13 +55,41 @@ function pow_diag_safe( 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))))) - return MAK.diagonal(map(d -> abs(d) < tol ? zero(d) : real(d)^p, σ)) + 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_safe!(Dp, D, p, tol) + return Dp +end + +# 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 + +""" + 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 """ - 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)`. @@ -64,10 +98,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 @@ -77,32 +111,36 @@ 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::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, +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 - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. $(_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 = MAK.eigh_full(M; alg) return V * pow_diag_safe(D, p; kwargs...) * V' 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. +Square root of a Hermitian positive semi-definite matrix. Equivalent to `powh_safe(M, 1//2; alg, atol, rtol)`. ## Keyword arguments @@ -111,13 +149,34 @@ 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; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(-1//2) + +Inverse square root of a Hermitian positive semi-definite matrix. +Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. + +## Keyword arguments + + - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`. + +$(_clamp_kwargs_doc("M")) +""" +invsqrth_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) + sqrth_invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2), M^(-1//2) -Inverse square root of an approximately Hermitian positive semi-definite -matrix. Equivalent to `powh_safe(M, -1//2; alg, atol, rtol)`. +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 @@ -125,7 +184,14 @@ 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...) +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 = 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 for (gram, gram_with_pinv, eigh_full) in ( (:gram_eigh_full, :gram_eigh_full_with_pinv, :eigh_full), @@ -133,11 +199,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/TensorAlgebra.jl b/src/TensorAlgebra.jl index 35317cd..9785f9f 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -1,14 +1,15 @@ 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( 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 @@ -18,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/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..ef17ba8 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -1,6 +1,13 @@ using LinearAlgebra: LinearAlgebra using MatrixAlgebraKit: MatrixAlgebraKit +# 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. for f in ( @@ -12,7 +19,7 @@ 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)),)), + 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...) @@ -27,6 +34,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,7 +218,7 @@ 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(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)), unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) end @@ -406,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(style, N, axes_codomain, (conj(axes(N, 2)),)) + 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...) @@ -499,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(style, X, axes_codomain, (conj(axes(X, 2)),)) + 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...) @@ -560,7 +568,7 @@ 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)),)), + 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...) @@ -576,6 +584,120 @@ 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 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 + + - `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 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 + + - `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..660e2be 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -1,11 +1,25 @@ # 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) + +# 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/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/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) 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..759e395 100644 --- a/src/projectto.jl +++ b/src/projectto.jl @@ -1,79 +1,152 @@ +# 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_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") + ) + 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!`. See -[`checked_projectto!`](@ref) for a checked version, and [`project`](@ref) -for the allocating form. +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_size(size(src), size(dest)) + return copyto!(dest, reshape(src, size(dest))) +end """ - 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 + +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`. -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). +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)` 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 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...) + check_project_size(size(src), size(dest)) + 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_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 8828349..67b2046 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, @@ -36,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 @@ -49,8 +53,10 @@ using Test: @test, @testset :invsqrt_diag_safe, :invsqrth_safe, :pow_diag_safe, + :pow_diag_safe!, :powh_safe, :sqrt_diag_safe, + :sqrth_invsqrth_safe, :sqrth_safe, ] @test issetequal(names(TensorAlgebra.MatrixAlgebra), exports) 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_matrixalgebra.jl b/test/test_matrixalgebra.jl index 2ecf63e..af75ebd 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) ≈ @@ -214,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 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 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