diff --git a/Project.toml b/Project.toml index d2840c0..0412584 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorBase" uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" -version = "0.10.7" +version = "0.10.8" authors = ["ITensor developers and contributors"] [workspace] @@ -54,7 +54,7 @@ OMEinsumContractionOrders = "1" OrderedCollections = "1.6" Random = "1.10" SimpleTraits = "0.9.4" -TensorAlgebra = "0.16" +TensorAlgebra = "0.16.5" TensorKit = "0.17" TensorOperations = "5.3.1" TermInterface = "2" diff --git a/src/abstractnamedtensor.jl b/src/abstractnamedtensor.jl index abc89fc..7043938 100644 --- a/src/abstractnamedtensor.jl +++ b/src/abstractnamedtensor.jl @@ -944,7 +944,32 @@ function aligndims(a::AbstractNamedTensor, dims) "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." ) ) - return nameddims(permutedims(unnamed(a), perm), new_dimnames) + return nameddims(TensorAlgebra.permutedims(unnamed(a), perm), new_dimnames) +end + +""" + aligndims(a::AbstractNamedTensor, codomain, domain) + +Reorder the dimensions of `a` into `(codomain..., domain...)`, matched by name, and forward +the codomain/domain split to the underlying storage. Like the two-argument form, the result +has the same data and dimension names as `a`, and a `NameMismatch` is thrown if +`(codomain..., domain...)` is not a permutation of `a`'s dimension names. A storage backend +that supports a bipartition (such as a TensorKit `TensorMap`) uses it, while a dense backend +stores the result flat. +""" +function aligndims(a::AbstractNamedTensor, codomain, domain) + new_dimnames = (name.(codomain)..., name.(domain)...) + perm = Tuple(getperm(dimnames(a), new_dimnames)) + isperm(perm) || throw( + NameMismatch( + "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." + ) + ) + perm_codomain = perm[1:length(codomain)] + perm_domain = perm[(length(codomain) + 1):end] + return nameddims( + TensorAlgebra.permutedims(unnamed(a), perm_codomain, perm_domain), new_dimnames + ) end """ @@ -1051,6 +1076,83 @@ for f in [:zeros, :ones], dimtype in [:NamedInteger, :NamedUnitRange] Base.$f(dim1::$dimtype, dims::Vararg{$dimtype}) = $f((dim1, dims...)) end end +# Map-shaped construction takes a codomain and a domain index tuple and forwards the split +# down to `TensorAlgebra`'s map constructors: a `TensorMap` backend stores it as a +# `codomain ← domain` map, a dense backend stores flat. Following the `similar_map` +# convention, the domain is conjugated in the flattened/outward view (a `TensorMap` stores +# its domain dual, the dense fallback conjugates it), so a domain index appears as its dual, +# and the result is named with the codomain names followed by the domain names. The +# `rand`/`randn`/`zeros` two-tuple forms (`randn((i,), (j,))`) forward to these. +# +# Each constructor is a shared `*_nameddims` builder (strip the names, call the map hook on the +# raw axes, reattach the names) plus two forwarding methods: one for a nonempty codomain and one +# for an empty codomain with a nonempty domain. The two-way split (rather than a single +# `Tuple{Vararg{NamedUnitRange}}` on both sides) reads the index type from whichever side is +# nonempty and keeps the empty-codomain case from re-dispatching to the same named overload once +# `unnamed` has stripped the names. An all-empty `((), ())` has no map meaning and is left to +# error rather than recurse. +for f in [:rand, :randn] + f_map = Symbol(f, :_map) + f_nameddims = Symbol(f, :_nameddims) + @eval function $f_nameddims(rng::AbstractRNG, elt::Type{<:Number}, codomain, domain) + a = TensorAlgebra.$f_map(rng, elt, unnamed.(codomain), unnamed.(domain)) + return a[Name.(name.((codomain..., domain...)))...] + end + for (codomain_type, domain_type) in [ + ( + :(Tuple{NamedUnitRange, Vararg{NamedUnitRange}}), + :(Tuple{Vararg{NamedUnitRange}}), + ), + (:(Tuple{}), :(Tuple{NamedUnitRange, Vararg{NamedUnitRange}})), + ] + @eval begin + function TensorAlgebra.$f_map( + rng::AbstractRNG, elt::Type{<:Number}, + codomain::$codomain_type, domain::$domain_type + ) + return $f_nameddims(rng, elt, codomain, domain) + end + function Base.$f( + rng::AbstractRNG, elt::Type{<:Number}, + codomain::$codomain_type, domain::$domain_type + ) + return TensorAlgebra.$f_map(rng, elt, codomain, domain) + end + function Base.$f( + elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type + ) + return Base.$f(Random.default_rng(), elt, codomain, domain) + end + function Base.$f(codomain::$codomain_type, domain::$domain_type) + return Base.$f(default_eltype(), codomain, domain) + end + end + end +end +function zeros_nameddims(elt::Type{<:Number}, codomain, domain) + a = TensorAlgebra.zeros_map(elt, unnamed.(codomain), unnamed.(domain)) + return a[Name.(name.((codomain..., domain...)))...] +end +for (codomain_type, domain_type) in [ + (:(Tuple{NamedUnitRange, Vararg{NamedUnitRange}}), :(Tuple{Vararg{NamedUnitRange}})), + (:(Tuple{}), :(Tuple{NamedUnitRange, Vararg{NamedUnitRange}})), + ] + @eval begin + function TensorAlgebra.zeros_map( + elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type + ) + return zeros_nameddims(elt, codomain, domain) + end + function Base.zeros( + elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type + ) + return TensorAlgebra.zeros_map(elt, codomain, domain) + end + function Base.zeros(codomain::$codomain_type, domain::$domain_type) + return Base.zeros(default_eltype(), codomain, domain) + end + end +end for dimtype in [:NamedInteger, :NamedUnitRange] @eval begin function Base.fill(value, ax::Tuple{$dimtype, Vararg{$dimtype}}) diff --git a/test/test_nameddims_basics.jl b/test/test_nameddims_basics.jl index 4d96ce9..f442c3c 100644 --- a/test/test_nameddims_basics.jl +++ b/test/test_nameddims_basics.jl @@ -257,6 +257,31 @@ end na′ = aligneddims(na, (j, i)) @test unnamed(na′) isa PermutedDimsArray{elt} @test a == permutedims(unnamed(na′), (2, 1)) + # The map form of `aligndims` takes a codomain and a domain tuple. A dense backend + # ignores the split and stores the reordered result flat, matching the flat form. + na′ = aligndims(na, (j,), (i,)) + @test unnamed(na′) isa Matrix{elt} + @test a == permutedims(unnamed(na′), (2, 1)) + # Two-tuple `randn`/`zeros` take a codomain and a domain index tuple; a dense backend + # ignores the split and stores flat, named by the codomain then the domain. + ci, cj = namedoneto(3, "i"), namedoneto(4, "j") + nab = randn(elt, (ci,), (cj,)) + @test unnamed(nab) isa Matrix{elt} + @test dimnames(nab) == [name(ci), name(cj)] + @test unnamed(zeros(elt, (ci,), (cj,))) == zeros(elt, 3, 4) + # An empty codomain lands every index in the domain, the mirror of an empty domain. A + # dense backend ignores the split, so these match the flat forms. + na′ = aligndims(na, (), (j, i)) + @test unnamed(na′) isa Matrix{elt} + @test a == permutedims(unnamed(na′), (2, 1)) + nbra = randn(elt, (), (cj,)) + @test unnamed(nbra) isa Vector{elt} + @test dimnames(nbra) == [name(cj)] + @test unnamed(zeros(elt, (), (cj,))) == zeros(elt, 4) + # An all-empty split has no map meaning, so it errors rather than recursing or + # silently building a scalar. + @test_throws MethodError randn(elt, (), ()) + @test_throws MethodError zeros(elt, (), ()) na = nameddims(randn(elt, 2, 3), (:i, :j)) nb = nameddims(randn(elt, 3, 2), (:j, :i)) diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index c1be719..24efcd0 100644 --- a/test/test_tensorkitext.jl +++ b/test/test_tensorkitext.jl @@ -1,8 +1,8 @@ -using ITensorBase: ITensorBase, Index, dimnames, name, prime, unnamed +using ITensorBase: ITensorBase, Index, aligndims, dimnames, name, prime, unnamed using LinearAlgebra: norm using MatrixAlgebraKit: qr_compact, svd_compact using StableRNGs: StableRNG -using TensorAlgebra: checked_project, project +using TensorAlgebra: TensorAlgebra, checked_project, project using TensorKit: TensorKit, @tensor, AbstractTensorMap, SU2Irrep, U1Irrep, Vect, dim, dual, scalar, space, ←, ⊗ using Test: @test, @test_throws, @testset @@ -68,6 +68,49 @@ using Test: @test, @test_throws, @testset @test sca((u * s * v) * w) ≈ sca(a3 * w) q, r = qr_compact(a3, (i,), (j, k)) @test sca((q * r) * w) ≈ sca(a3 * w) + + # Map-shaped construction forwards the codomain/domain split to the `TensorMap`: + # `randn((i,), (j,))` stores a `Vi ← Vj` map rather than flattening all-codomain. + # Following the `similar_map` convention the domain appears dualized in the outward + # view, the same as a flat graded array would have axes `(Vi, dual(Vj))`. + m = randn(rng, elt, (i,), (j,)) + @test unnamed(m) isa AbstractTensorMap + @test space(unnamed(m)) == (Vi ← Vj) + @test space(unnamed(m), 1) == Vi + @test space(unnamed(m), 2) == dual(Vj) + @test unnamed(rand(rng, elt, (i,), (j,))) isa AbstractTensorMap + @test norm(unnamed(zeros(elt, (i,), (j,)))) == 0 + # The friendly forms agree with the underlying `TensorAlgebra` map hooks. + @test space(unnamed(TensorAlgebra.randn_map(elt, (i, j), (k,)))) == + space(unnamed(randn(elt, (i, j), (k,)))) + # An empty codomain builds an all-domain `TensorMap`, the mirror of an empty domain. The + # space type is read from the domain, since the empty codomain carries none. An all-empty + # split has no map meaning and errors rather than recursing. + cd = randn(rng, elt, (), (j,)) + @test unnamed(cd) isa AbstractTensorMap + @test space(unnamed(cd)) == (one(Vj) ← Vj) + @test dimnames(cd) == [name(j)] + @test space(unnamed(zeros(elt, (), (j,)))) == (one(Vj) ← Vj) + @test_throws MethodError randn(rng, elt, (), ()) + + # `aligndims` reorders a `TensorMap`-backed tensor. The flat form gives an all-codomain + # result and the map form re-expresses the requested codomain/domain split, both + # carrying each index with its arrow to the new position. + mf = aligndims(m, (j, i)) + @test dimnames(mf) == [name(j), name(i)] + @test space(unnamed(mf), 1) == dual(Vj) + @test space(unnamed(mf), 2) == Vi + md = aligndims(m, (j,), (i,)) + @test dimnames(md) == [name(j), name(i)] + @test space(unnamed(md)) == (dual(Vj) ← dual(Vi)) + @test space(unnamed(md), 1) == dual(Vj) + @test space(unnamed(md), 2) == Vi + # An empty codomain moves both indices into the domain, preserving the outward axes. + me = aligndims(m, (), (i, j)) + @test dimnames(me) == [name(i), name(j)] + @test space(unnamed(me)) == (one(Vi) ← (dual(Vi) ⊗ Vj)) + @test space(unnamed(me), 1) == Vi + @test space(unnamed(me), 2) == dual(Vj) end # `project` builds a `TensorMap`-backed operator/state from a dense basis matrix: the index