Preserve the codomain/domain split in map construction and aligndims

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 <support@itensor.org> 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"

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}})

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))

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