Make project strict and add project_aux for aux derivation

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.17.11"
+version = "0.18.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

docs/Project.toml

@@ -11,4 +11,4 @@ path = ".."
 Documenter = "1.8.1"
 ITensorFormatter = "0.2.27"
 Literate = "2.20.1"
-TensorAlgebra = "0.17"
+TensorAlgebra = "0.18"

examples/Project.toml

@@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 path = ".."
 
 [compat]
-TensorAlgebra = "0.17"
+TensorAlgebra = "0.18"

ext/TensorAlgebraTensorKitExt.jl

@@ -162,46 +162,28 @@ function TensorAlgebra.unproject(t::AbstractTensorMap, ::Val{K}) where {K}
     return convert(Array, t)
 end
 
-# =============================  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(
+# =============================  infer_aux_space (project_aux derivation)  ===================
+# `infer_aux_space` for `TensorMap` spaces routes both the codomain-led and the (empty-codomain)
+# domain-led cases to `infer_aux_space_tensormap`, reading the elementary space type `S` from
+# whichever side is non-empty, the same two-entry split `similar_map` uses. `project`'s allocation
+# stays generic (strict `similar_map`); only the `project_aux` derivation is `TensorMap`-specific.
+function TensorAlgebra.infer_aux_space(
         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)
+    return infer_aux_space_tensormap(raw, S, codomain_axes, domain_axes)
 end
-function TensorAlgebra.allocate_project(
+function TensorAlgebra.infer_aux_space(
         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"
-        )
-    )
-    aux = infer_aux_space(raw, S, codomain_axes, domain_axes)
-    return TensorAlgebra.similar_map(raw, codomain_axes, (domain_axes..., aux))
+    return infer_aux_space_tensormap(raw, S, codomain_axes, domain_axes)
 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(
+function infer_aux_space_tensormap(
         raw, ::Type{S}, codomain_axes, domain_axes
     ) where {S <: ElementarySpace}
     aux_dim = length(codomain_axes) + length(domain_axes) + 1

src/projectto.jl

@@ -32,26 +32,24 @@ end
 
 Allocate the destination that projecting `raw` onto
 `codomain_axes`/`domain_axes` fills. This is a backend customization point
-(with [`projectto!`](@ref) and [`is_projected`](@ref)): the allocation may
-depend on the data, since a 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) has its space taken from `raw` itself on a dense
-backend and derived from the sector structure on a symmetric one.
+(with [`projectto!`](@ref) and [`is_projected`](@ref)); the default is plain
+`similar_map(raw, codomain_axes, domain_axes)`.
 
-The generic method keeps that trailing surplus axis (its `raw` axis appended
-to the domain), so the result's rank matches `raw`'s; with no surplus it is
-plain `similar_map(raw, codomain_axes, domain_axes)`.
+`project` projects into exactly the given axes, so `raw` must not have more
+axes than they account for. To append a derived flux-carrying auxiliary axis
+for a charge-shifting operator or a non-invariant state, use
+[`project_aux`](@ref) instead.
 """
 function allocate_project(raw, codomain_axes, domain_axes)
     nphys = length(codomain_axes) + length(domain_axes)
-    ndims(raw) <= nphys && return similar_map(raw, codomain_axes, domain_axes)
-    ndims(raw) == nphys + 1 || throw(
+    ndims(raw) <= nphys || throw(
         ArgumentError(
-            "`project`: expected at most one trailing auxiliary axis beyond the $nphys \
-            given axes, got a rank-$(ndims(raw)) input"
+            "`project` projects into exactly the given axes and does not derive an auxiliary \
+            axis; got a rank-$(ndims(raw)) input for $nphys given axes. Use `project_aux` to \
+            append a derived flux-carrying leg, or pass the axis explicitly."
         )
     )
-    return similar_map(raw, codomain_axes, (domain_axes..., axes(raw, nphys + 1)))
+    return similar_map(raw, codomain_axes, domain_axes)
 end
 
 """
@@ -71,8 +69,7 @@ domain.
 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.
+# The flat all-codomain (state) form: a list of `axes` with an empty domain.
 unchecked_project(raw, axes) = unchecked_project(raw, axes, ())
 
 # The codomain rank a destination reports when no split is given: its full rank by default (no
@@ -140,14 +137,11 @@ 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.
 
-When `raw` has one axis more than the given axes account for, that trailing
-surplus axis is an auxiliary leg appended as the last domain axis, so the
-result's shape matches `raw`'s (e.g. a flux-canceling leg for a
-charge-shifting operator). Its space comes from `raw` itself on a dense
-backend and is derived from the sector structure on a symmetric one (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.
+`raw` must not have more axes than `codomain_axes`/`domain_axes` account for:
+`project` projects into exactly the given axes. To append a derived
+flux-carrying auxiliary axis (for a charge-shifting operator or a
+non-invariant state), use [`project_aux`](@ref). The two-argument form takes a
+flat list of `axes` and is equivalent to an empty domain.
 """
 function project(raw, codomain_axes, domain_axes; kwargs...)
     dest = unchecked_project(raw, codomain_axes, domain_axes)
@@ -167,7 +161,7 @@ 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,))
+    @something tryproject(v, (cod,)) project_aux(v, (cod,))
 
 Keyword arguments are forwarded to the `isapprox` tolerance check.
 """
@@ -176,3 +170,79 @@ function tryproject(raw, codomain_axes, domain_axes; kwargs...)
     return is_projected(dest, raw, Val(length(codomain_axes)); kwargs...) ? dest : nothing
 end
 tryproject(raw, axes; kwargs...) = tryproject(raw, axes, (); kwargs...)
+
+"""
+    infer_aux_space(raw, codomain_axes, domain_axes) -> aux
+
+Derive the auxiliary axis the `*_aux` projection verbs append as the last
+domain axis, so the projected result is symmetry-allowed. `raw` carries the
+trailing slice axis whose space is derived. This is the backend customization
+point for that derivation: the generic (dense) method takes the space straight
+from `raw`, while a symmetric backend reads it from the sector structure (a
+graded backend derives per-slice sectors, the `TensorMap` backend scans the
+`codomain ⊗ conj(domain)` content).
+"""
+function infer_aux_space(raw, codomain_axes, domain_axes)
+    return axes(raw, length(codomain_axes) + length(domain_axes) + 1)
+end
+
+# Reshape a physical-rank `raw` up to one trailing slice axis, derive the auxiliary space, and
+# return the `(raw, codomain_axes, domain_axes)` triple to forward to a projection verb, with the
+# aux appended to the domain. A rank beyond one surplus axis is an error. Shared by the three
+# `*_aux` verbs below.
+function project_aux_args(raw, codomain_axes, domain_axes)
+    nphys = length(codomain_axes) + length(domain_axes)
+    nphys <= ndims(raw) <= nphys + 1 || throw(
+        ArgumentError(
+            "`project_aux` expected a rank-$nphys or rank-$(nphys + 1) input for $nphys given \
+            axes, got rank $(ndims(raw))"
+        )
+    )
+    slices = ndims(raw) == nphys ? reshape(raw, (size(raw)..., 1)) : raw
+    aux = infer_aux_space(slices, codomain_axes, domain_axes)
+    return slices, codomain_axes, (domain_axes..., aux)
+end
+
+"""
+    project_aux(raw, codomain_axes, domain_axes; kwargs...) -> dest
+    project_aux(raw, axes; kwargs...) -> dest
+
+Project `raw` and append a derived auxiliary domain axis carrying its flux,
+giving a symmetry-allowed result whose squeezed data is the input (the
+flux-canceling MPO-virtual-leg idiom). Unlike [`project`](@ref), which projects
+into exactly the given axes, `project_aux` derives the extra leg (see
+[`infer_aux_space`](@ref)). `raw` may have the physical rank (a single operator
+or state, its flux on a length-1 leg) or one trailing slice axis (an operator
+multiplet as laid out by `stack`). Like `project`, it verifies that only a
+negligible component is discarded; see [`unchecked_project_aux`](@ref) and
+[`tryproject_aux`](@ref) for the unchecked and nullable siblings.
+"""
+function project_aux(raw, codomain_axes, domain_axes; kwargs...)
+    return project(project_aux_args(raw, codomain_axes, domain_axes)...; kwargs...)
+end
+project_aux(raw, axes; kwargs...) = project_aux(raw, axes, (); kwargs...)
+
+"""
+    unchecked_project_aux(raw, codomain_axes, domain_axes) -> dest
+    unchecked_project_aux(raw, axes) -> dest
+
+The unchecked sibling of [`project_aux`](@ref): derive and append the auxiliary
+axis, then project without verifying which components are discarded.
+"""
+function unchecked_project_aux(raw, codomain_axes, domain_axes)
+    return unchecked_project(project_aux_args(raw, codomain_axes, domain_axes)...)
+end
+unchecked_project_aux(raw, axes) = unchecked_project_aux(raw, axes, ())
+
+"""
+    tryproject_aux(raw, codomain_axes, domain_axes; kwargs...) -> Union{dest, Nothing}
+    tryproject_aux(raw, axes; kwargs...) -> Union{dest, Nothing}
+
+The nullable sibling of [`project_aux`](@ref): derive and append the auxiliary
+axis, returning `nothing` instead of throwing when more than a negligible
+component of `raw` would be discarded.
+"""
+function tryproject_aux(raw, codomain_axes, domain_axes; kwargs...)
+    return tryproject(project_aux_args(raw, codomain_axes, domain_axes)...; kwargs...)
+end
+tryproject_aux(raw, axes; kwargs...) = tryproject_aux(raw, axes, (); kwargs...)

test/Project.toml

@@ -37,7 +37,7 @@ Random = "1.10"
 SafeTestsets = "0.1"
 StableRNGs = "1.0.2"
 Suppressor = "0.2"
-TensorAlgebra = "0.17"
+TensorAlgebra = "0.18"
 TensorKit = "0.17"
 TensorOperations = "5.1.4"
 Test = "1.10"

test/test_projectto.jl

@@ -1,5 +1,6 @@
-using TensorAlgebra: TensorAlgebra, is_projected, project, project!, projectto!, tryproject,
-    unchecked_project, unproject
+using TensorAlgebra: TensorAlgebra, is_projected, project, project!, project_aux,
+    projectto!, tryproject, tryproject_aux, unchecked_project, unchecked_project_aux,
+    unproject
 using Test: @test, @test_throws, @testset
 
 const elts = (Float32, Float64, ComplexF32, ComplexF64)
@@ -110,25 +111,36 @@ end
     @test vec(Msplit) == vec(flat)
 end
 
-@testset "project keeps a trailing surplus axis ($T)" for T in elts
-    # The dual of the padding case: when `raw` carries one axis *more* than the given axes
-    # account for, that trailing surplus axis is an auxiliary leg (e.g. a flux-canceling leg a
-    # codomain/domain split introduces on a symmetric state/operator), and its space is taken
-    # from `raw`. The result keeps the axis rather than reshaping it away, so its rank matches
-    # `raw`'s, matching the symmetric backends. Here the aux leg is the dim-1 leg a caller adds
-    # with `reshape(a, (size(a)..., 1))`.
+@testset "project rejects a surplus axis; project_aux derives one ($T)" for T in elts
+    # `project` projects into exactly the given axes: a surplus axis is an error, not a silent
+    # auxiliary leg. `project_aux` is the deriving entry point that appends that leg (e.g. a
+    # flux-canceling leg a codomain/domain split introduces on a symmetric state/operator). On a
+    # dense backend the aux space is taken from `raw`, and the result keeps the axis rather than
+    # reshaping it away, so its rank matches `raw`'s (matching the symmetric backends).
     raw = randn(T, 2, 3, 1)
+    @test_throws ArgumentError project(raw, (Base.OneTo(2), Base.OneTo(3)))
+    @test_throws ArgumentError project(raw, (Base.OneTo(2),), (Base.OneTo(3),))
 
     # all-codomain (state) form
-    M = project(raw, (Base.OneTo(2), Base.OneTo(3)))
+    M = project_aux(raw, (Base.OneTo(2), Base.OneTo(3)))
     @test size(M) == (2, 3, 1)
     @test M == raw
 
-    # explicit split: the surplus axis lands past the codomain/domain axes given
-    Msplit = project(raw, (Base.OneTo(2),), (Base.OneTo(3),))
+    # explicit split: the aux lands past the codomain/domain axes given
+    Msplit = project_aux(raw, (Base.OneTo(2),), (Base.OneTo(3),))
     @test size(Msplit) == (2, 3, 1)
     @test Msplit == raw
 
-    # more than one surplus axis is rejected
+    # a physical-rank input is reshaped up to a length-1 aux
+    flat = randn(T, 2, 3)
+    Mflat = project_aux(flat, (Base.OneTo(2),), (Base.OneTo(3),))
+    @test size(Mflat) == (2, 3, 1)
+    @test vec(Mflat) == vec(flat)
+
+    # more than one surplus axis is rejected by both
     @test_throws ArgumentError project(randn(T, 2, 3, 1, 1), (Base.OneTo(2), Base.OneTo(3)))
+    @test_throws ArgumentError project_aux(
+        randn(T, 2, 3, 1, 1),
+        (Base.OneTo(2), Base.OneTo(3))
+    )
 end