diff --git a/Project.toml b/Project.toml index 3cc75bc..3f14182 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 714c771..b49ecb7 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/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` diff --git a/src/projectto.jl b/src/projectto.jl index 7f61a37..15649c6 100644 --- a/src/projectto.jl +++ b/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 diff --git a/test/test_projectto.jl b/test/test_projectto.jl index a582349..4bbbd4e 100644 --- a/test/test_projectto.jl +++ b/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 diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index d6715ee..236ed82 100644 --- a/test/test_tensorkitext.jl +++ b/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