Add project convenience for symmetry-restricted arrays and TensorMap support

Project.toml

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

ext/TensorAlgebraTensorKitExt.jl

@@ -2,8 +2,9 @@ module TensorAlgebraTensorKitExt
 
 using Random: AbstractRNG
 using TensorAlgebra: TensorAlgebra
-using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, numind,
-    permute, space, spacetype, zerovector!, ←
+using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace,
+    TensorMapWithStorage, numind, permute, project_symmetric!, space, spacetype,
+    zerovector!, ←
 using TensorOperations: TensorOperations as TO
 
 # ============================  AbstractArray-vocabulary bridge  ============================
@@ -27,6 +28,33 @@ function TensorAlgebra.similar_map(
     return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...))
 end
 
+# A plain-array prototype with native (space) axes is the operator/state construction case: `raw`
+# is a dense matrix and the codomain/domain axes are `TensorMap` spaces. Build an uninitialized
+# `TensorMap` whose block storage type `A` follows `raw`'s array type, so the storage stays on the
+# same device (a GPU array's blocks stay on the GPU) while the map structure comes from the spaces.
+# `A` is the return type of `similar(raw, T, ::Int)`, a 1-d vector of `raw`'s family, which is the
+# block storage type. The elementary space type `S` is passed explicitly so the two dispatch
+# entries below can read it from whichever of codomain/domain is non-empty and share one builder,
+# mirroring `_map_homspace` and the map constructors.
+function similar_tensormap(
+        raw::AbstractArray, ::Type{T}, ::Type{S}, codomain_axes, domain_axes
+    ) where {T, S <: ElementarySpace}
+    A = Base.promote_op(similar, typeof(raw), Type{T}, Int)
+    return TensorMapWithStorage{T, A}(undef, _map_homspace(S, codomain_axes, domain_axes))
+end
+function TensorAlgebra.similar_map(
+        raw::AbstractArray, ::Type{T},
+        codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+    ) where {T, S <: ElementarySpace}
+    return similar_tensormap(raw, T, S, codomain_axes, domain_axes)
+end
+function TensorAlgebra.similar_map(
+        raw::AbstractArray, ::Type{T},
+        codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+    ) where {T, S <: ElementarySpace}
+    return similar_tensormap(raw, T, S, codomain_axes, domain_axes)
+end
+
 # ===============================  zeros_map / randn_map / rand_map  ========================
 # A `TensorMap` keeps its codomain and domain as separate `ProductSpace`s rather than a single
 # flattened axis, so build the `codomain ← domain` space directly instead of the dense
@@ -65,6 +93,32 @@ for (f, g) in ((:randn_map, :randn), (:rand_map, :rand))
     end
 end
 
+# =====================================  projectto!  ========================================
+# `projectto!` places dense `src` data into the restricted (symmetric) space of `dest`. A
+# `TensorMap` is not an `AbstractArray`, so the generic `copyto!` default does not apply; delegate
+# 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.
+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
+end
+
 # ================================  bipermutedimsopadd!  =====================================
 # `dest = β * dest + α * permutedims(op.(src), (perm_codomain, perm_domain))`. Delegate to
 # TensorKit's TensorOperations interface: `tensoradd!` realizes the permutation, the `op === conj`

src/projectto.jl

@@ -3,7 +3,8 @@
 
 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.
+[`checked_projectto!`](@ref) for a checked version, and [`project`](@ref)
+for the allocating form.
 """
 projectto!(dest, src) = copyto!(dest, src)
 
@@ -27,8 +28,10 @@ end
     project_map(raw, codomain_axes, domain_axes) -> dest
 
 Allocate a map-shaped array via [`similar_map`](@ref) and project `raw`
-into it with [`projectto!`](@ref). See [`checked_project_map`](@ref) for
-a checked version.
+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.
 """
 function project_map(raw, codomain_axes, domain_axes)
     return projectto!(similar_map(raw, codomain_axes, domain_axes), raw)
@@ -46,3 +49,31 @@ function checked_project_map(raw, codomain_axes, domain_axes; kwargs...)
         similar_map(raw, codomain_axes, domain_axes), raw; kwargs...
     )
 end
+
+"""
+    project(raw, codomain_axes, domain_axes) -> dest
+    project(raw, axes) -> 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.
+"""
+project(raw, codomain_axes, domain_axes) = project_map(raw, codomain_axes, domain_axes)
+project(raw, axes) = project_map(raw, axes, ())
+
+"""
+    checked_project(raw, codomain_axes, domain_axes; kwargs...) -> dest
+    checked_project(raw, axes; kwargs...) -> dest
+
+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).
+"""
+function checked_project(raw, codomain_axes, domain_axes; kwargs...)
+    return checked_project_map(raw, codomain_axes, domain_axes; kwargs...)
+end
+function checked_project(raw, axes; kwargs...)
+    return checked_project_map(raw, axes, (); kwargs...)
+end

test/test_projectto.jl

@@ -1,5 +1,5 @@
-using TensorAlgebra:
-    TensorAlgebra, checked_project_map, checked_projectto!, project_map, projectto!
+using TensorAlgebra: TensorAlgebra, checked_project, checked_project_map,
+    checked_projectto!, project, project_map, projectto!
 using Test: @test, @test_throws, @testset
 
 const elts = (Float32, Float64, ComplexF32, ComplexF64)
@@ -19,7 +19,7 @@ function TensorAlgebra.projectto!(dest::Rounded, src::AbstractArray)
     return dest
 end
 
-@testset "projectto!/project_map ($T)" for T in elts
+@testset "projectto!/project ($T)" for T in elts
     src = randn(T, 2, 3)
 
     # `projectto!` defaults to `copyto!`.
@@ -50,12 +50,21 @@ end
     raw = randn(T, 2, 3, 2, 3)
     cod = (Base.OneTo(2), Base.OneTo(3))
     dom = (Base.OneTo(2), Base.OneTo(3))
-    M = project_map(raw, cod, dom)
-    @test eltype(M) === T
-    @test size(M) == (2, 3, 2, 3)
-    @test M == raw
+    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
+
+    # `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
 
-    # `checked_project_map` agrees and accepts the same buffer.
-    M2 = checked_project_map(raw, cod, dom)
-    @test M2 == raw
+    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
 end

test/test_tensorkitext.jl

@@ -1,9 +1,11 @@
 using Base.Broadcast: broadcasted
+using LinearAlgebra: norm
 using StableRNGs: StableRNG
-using TensorAlgebra: TensorAlgebra, contract, matricize, rand_map, randn_map, similar_map,
-    tryflattenlinear, unmatricize, zeros_map
-using TensorKit:
-    @tensor, AbstractTensorMap, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗
+using TensorAlgebra: TensorAlgebra, checked_project, contract, matricize, project,
+    project_map, projectto!, rand_map, randn_map, similar_map, tryflattenlinear,
+    unmatricize, zeros_map
+using TensorKit: @tensor, AbstractTensorMap, Rep, SU₂, TensorMap, U₁, fuse, isomorphism,
+    randn, space, storagetype, ←, ⊗
 using Test: @test, @test_throws, @testset
 
 # A shared bond contracts when it sits in one operand's domain and the other's codomain, i.e.
@@ -131,6 +133,43 @@ 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.
+    @testset "project a dense matrix into a TensorMap" begin
+        W = Rep[U₁](0 => 1, 1 => 1)
+        Sz = elt[0.5 0; 0 -0.5]
+        Sx = elt[0 0.5; 0.5 0]
+
+        pz = project(Sz, (W,), (W,))
+        @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)
+
+        # the flat two-argument form builds an all-codomain `TensorMap` (empty domain): only
+        # the trivial-charge component of a dense vector survives the projection
+        pv = project(elt[1, 0], (W,))
+        @test pv isa AbstractTensorMap
+        @test space(pv) == (W ← one(W))
+        @test norm(pv) ≈ 1
+        @test norm(project(elt[0, 1], (W,))) == 0
+    end
+
     # A linear combination of `TensorMap`s flattens to a `LinearBroadcasted` that materializes
     # into a `TensorMap` destination via `copyto!`; a nonlinear broadcast has no linear form.
     @testset "linear-combination broadcast" begin