Add flux-canceling randn(sector, inds) constructors

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.13.5"
+version = "0.13.6"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -25,12 +25,15 @@ WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
+TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 
 [extensions]
 ITensorBaseAdaptExt = "Adapt"
+ITensorBaseGradedArraysExt = ["GradedArrays", "TensorKitSectors"]
 ITensorBaseMooncakeExt = "Mooncake"
 ITensorBaseOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 ITensorBaseTensorKitExt = "TensorKit"
@@ -42,14 +45,16 @@ Adapt = "4.1.1"
 ArrayLayouts = "1.11"
 Combinatorics = "1"
 ConstructionBase = "1.6"
+GradedArrays = "0.14.2"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
 OMEinsumContractionOrders = "1.3"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.17.8"
+TensorAlgebra = "0.17.10"
 TensorKit = "0.17"
+TensorKitSectors = "0.3.9"
 TermInterface = "2"
 TupleTools = "1.6"
 UUIDs = "1.10"

ext/ITensorBaseGradedArraysExt.jl

@@ -0,0 +1,79 @@
+module ITensorBaseGradedArraysExt
+
+using GradedArrays: SectorRange
+using ITensorBase: ITensorBase, Index
+using Random: AbstractRNG, default_rng
+using TensorKitSectors: Sector
+
+const NamedUnitRange = ITensorBase.NamedUnitRange
+
+# Flux-canceling constructors at the `Index` level: mint a fresh auxiliary `Index` carrying
+# sector `c` (auto-generated unique name) and forward to the map-shaped split constructor,
+# appending the aux to the (implicitly dualized) domain dangling last, so the physical axes
+# fuse to `c`. Mirrors the GradedArrays backend method. The sector may be a bare
+# `TensorKitSectors.Sector` or a `SectorRange`. This lives in an extension because ITensorBase
+# does not depend on the sector types.
+for S in (Sector, SectorRange)
+    for f in (:rand, :randn)
+        @eval begin
+            function Base.$f(
+                    rng::AbstractRNG, elt::Type{<:Number}, c::$S,
+                    codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(rng, elt, codomain, (domain..., Index([c => 1])))
+            end
+            function Base.$f(
+                    rng::AbstractRNG, c::$S,
+                    codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(rng, ITensorBase.default_eltype(), c, codomain, domain)
+            end
+            function Base.$f(
+                    elt::Type{<:Number}, c::$S,
+                    codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(default_rng(), elt, c, codomain, domain)
+            end
+            function Base.$f(
+                    c::$S, codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(
+                    default_rng(),
+                    ITensorBase.default_eltype(),
+                    c,
+                    codomain,
+                    domain
+                )
+            end
+        end
+    end
+    for f in (:zeros, :ones)
+        @eval begin
+            function Base.$f(
+                    elt::Type{<:Number}, c::$S,
+                    codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(elt, codomain, (domain..., Index([c => 1])))
+            end
+            function Base.$f(
+                    c::$S, codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+                    domain::Tuple{Vararg{NamedUnitRange}} = ()
+                )
+                return Base.$f(ITensorBase.default_eltype(), c, codomain, domain)
+            end
+        end
+    end
+    @eval function Base.fill(
+            value, c::$S, codomain::Tuple{NamedUnitRange, Vararg{NamedUnitRange}},
+            domain::Tuple{Vararg{NamedUnitRange}} = ()
+        )
+        return Base.fill(value, codomain, (domain..., Index([c => 1])))
+    end
+end
+
+end

src/abstractnamedtensor.jl

@@ -1379,6 +1379,12 @@ for f in [:rand, :randn]
                 )
                 return $f(rng, elt, (dim1, dims...))
             end
+            Base.$f(rng::AbstractRNG, dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(
+                rng, default_eltype(), dims
+            )
+            Base.$f(rng::AbstractRNG, dim1::$dimtype, dims::Vararg{$dimtype}) = $f(
+                rng, (dim1, dims...)
+            )
             Base.$f(elt::Type{<:Number}, dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(
                 Random.default_rng(), elt, dims
             )
@@ -1390,21 +1396,17 @@ for f in [:rand, :randn]
         end
     end
 end
-# `zeros` routes through `TensorAlgebra.zeros_map` (all-codomain map, trivial domain), which
-# dispatches on the axis type to build a dense `Array`, a block-sparse array, or a TensorKit
-# `TensorMap`. `ones` has no map hook (an all-ones symmetric tensor is not well-defined), so it
-# stays on `Base.ones`, which already accepts axes.
-for f in [:zeros, :ones], dimtype in [:NamedInteger, :NamedUnitRange]
-    parent = if f === :zeros
-        :(TensorAlgebra.zeros_map(elt, unnamed.(ax), ()))
-    else
-        :(Base.ones(elt, unnamed.(ax)))
-    end
+# `zeros`/`ones` route through `TensorAlgebra.zeros_map`/`ones_map` (all-codomain map, trivial
+# domain), which dispatch on the axis type to build a dense `Array`, a block-sparse array, or a
+# TensorKit `TensorMap`.
+for (f, f_map) in [(:zeros, :zeros_map), (:ones, :ones_map)],
+        dimtype in [:NamedInteger, :NamedUnitRange]
+
     @eval begin
         function Base.$f(
                 elt::Type{<:Number}, ax::Tuple{$dimtype, Vararg{$dimtype}}
             )
-            a = $parent
+            a = TensorAlgebra.$f_map(elt, unnamed.(ax), ())
             return a[Name.(name.(ax))...]
         end
         function Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype})
@@ -1456,6 +1458,11 @@ for f in [:rand, :randn]
                 )
                 return TensorAlgebra.$f_map(rng, elt, codomain, domain)
             end
+            function Base.$f(
+                    rng::AbstractRNG, codomain::$codomain_type, domain::$domain_type
+                )
+                return Base.$f(rng, default_eltype(), codomain, domain)
+            end
             function Base.$f(
                     elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type
                 )
@@ -1467,27 +1474,56 @@ for f in [:rand, :randn]
         end
     end
 end
-function zeros_nameddims(elt::Type{<:Number}, codomain, domain)
-    a = TensorAlgebra.zeros_map(elt, unnamed.(codomain), unnamed.(domain))
+for (f, f_map) in [(:zeros, :zeros_map), (:ones, :ones_map)]
+    f_nameddims = Symbol(f, :_nameddims)
+    @eval function $f_nameddims(elt::Type{<:Number}, codomain, domain)
+        a = TensorAlgebra.$f_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.$f_map(
+                    elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type
+                )
+                return $f_nameddims(elt, codomain, domain)
+            end
+            function Base.$f(
+                    elt::Type{<:Number}, codomain::$codomain_type, domain::$domain_type
+                )
+                return TensorAlgebra.$f_map(elt, codomain, domain)
+            end
+            function Base.$f(codomain::$codomain_type, domain::$domain_type)
+                return Base.$f(default_eltype(), codomain, domain)
+            end
+        end
+    end
+end
+# `fill` takes the fill value first, so it does not fit the eltype-leading forms above; it gets
+# the same map-shaped split via `fill_map`.
+function fill_nameddims(value, codomain, domain)
+    a = TensorAlgebra.fill_map(value, 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
+        function TensorAlgebra.fill_map(
+                value,
+                codomain::$codomain_type,
+                domain::$domain_type
             )
-            return TensorAlgebra.zeros_map(elt, codomain, domain)
+            return fill_nameddims(value, codomain, domain)
         end
-        function Base.zeros(codomain::$codomain_type, domain::$domain_type)
-            return Base.zeros(default_eltype(), codomain, domain)
+        function Base.fill(value, codomain::$codomain_type, domain::$domain_type)
+            return TensorAlgebra.fill_map(value, codomain, domain)
         end
     end
 end

test/Project.toml

@@ -3,6 +3,7 @@ AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 ITensorPkgSkeleton = "3d388ab1-018a-49f4-ae50-18094d5f71ea"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
@@ -16,6 +17,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
+TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
@@ -30,6 +32,7 @@ AbstractTrees = "0.4.5"
 Adapt = "4"
 Aqua = "0.8.9"
 Combinatorics = "1"
+GradedArrays = "0.14.2"
 ITensorBase = "0.13"
 ITensorPkgSkeleton = "0.3.42"
 JLArrays = "0.2, 0.3"
@@ -43,6 +46,7 @@ StableRNGs = "1"
 Suppressor = "0.2"
 TensorAlgebra = "0.17.5"
 TensorKit = "0.17"
+TensorKitSectors = "0.3.9"
 TermInterface = "2"
 Test = "1.10"
 UUIDs = "1.10"

test/test_gradedarraysext.jl

@@ -0,0 +1,65 @@
+using GradedArrays: U1, sectors
+using ITensorBase: ITensorBase, Index, inds, space
+using StableRNGs: StableRNG
+using TensorAlgebra: TensorAlgebra, isdual
+using TensorKitSectors: FermionNumber
+using Test: @test, @testset
+
+# The flux-canceling constructor mints an auxiliary `Index` carrying the requested charge and
+# appends it to the domain, so an `ITensor` over graded (block-sparse) indices can be built
+# with a nonzero total flux directly. Covers an abelian (U₁) and a fermionic sector.
+@testset "GradedArraysExt flux-canceling constructor (eltype = $elt)" for elt in
+    (
+        Float64,
+        ComplexF64,
+    )
+    rng = StableRNG(1234)
+
+    i = Index([U1(0) => 1, U1(1) => 2]; tags = "i")
+    j = Index([U1(0) => 2, U1(1) => 1]; tags = "j")
+
+    # Flat form: all physical legs in the codomain, the aux the sole domain leg.
+    a = randn(rng, elt, U1(1), (i, j))
+    @test length(inds(a)) == 3
+    @test i in inds(a)
+    @test j in inds(a)
+    aux = only(setdiff(collect(inds(a)), [i, j]))
+    @test length(aux) == 1                     # multiplicity-1 aux leg
+    @test isdual(aux)                          # dualized, in the domain
+    @test only(sectors(space(aux))) == U1(1)   # carries the requested flux
+    @test eltype(a) == elt
+
+    # Map form: the aux is appended after the given domain leg.
+    b = randn(rng, elt, U1(1), (i,), (j,))
+    @test length(inds(b)) == 3
+    auxb = only(setdiff(collect(inds(b)), [i, j]))
+    @test isdual(auxb) && length(auxb) == 1 && only(sectors(space(auxb))) == U1(1)
+
+    # The rng-first flux forms (default eltype) accept both flat and split axes.
+    @test length(inds(randn(rng, U1(1), (i, j)))) == 3
+    @test length(inds(randn(rng, U1(1), (i,), (j,)))) == 3
+
+    # A bare `TensorKitSectors.Sector` (fermionic) works as the flux.
+    s = [
+        Index([FermionNumber(0) => 2, FermionNumber(1) => 2]; tags = "s" => "$n") for
+            n in 1:4
+    ]
+    t = randn(rng, elt, FermionNumber(2), (s[1], s[2], s[3], s[4]))
+    @test length(inds(t)) == 5
+    auxt = only(setdiff(collect(inds(t)), s))
+    @test isdual(auxt) && length(auxt) == 1 &&
+        only(sectors(space(auxt))) == FermionNumber(2)
+
+    # `zeros`/`ones`/`fill` mirror `randn` (`fill` takes the value first). Each carries the
+    # flux on an aux leg the same way.
+    z = zeros(U1(1), (i, j))
+    @test length(inds(z)) == 3
+    @test eltype(zeros(elt, U1(1), (i, j))) == elt
+    @test iszero(z)
+    o = ones(elt, U1(1), (i, j))
+    @test length(inds(o)) == 3
+    @test only(sectors(space(only(setdiff(collect(inds(o)), [i, j]))))) == U1(1)
+    fl = fill(elt(2), U1(1), (i,), (j,))
+    @test length(inds(fl)) == 3
+    @test eltype(fl) == elt
+end

test/test_nameddims_basics.jl

@@ -4,6 +4,7 @@ using ITensorBase: @names, AbstractNamedTensor, Name, NameMismatch, NamedDimsCar
     dimnametype, dims, inds, isnamed, mapinds, name, named, nameddims, namedoneto, product,
     replacedimnames, replaceinds, setdimnames, unname, unnamed, unnamedtype
 using LinearAlgebra: LinearAlgebra
+using Random: default_rng
 using TensorAlgebra: datatype
 using Test: @test, @test_throws, @testset
 using VectorInterface: scalartype
@@ -287,6 +288,12 @@ end
         @test unnamed(nab) isa Matrix{elt}
         @test dimnames(nab) == [name(ci), name(cj)]
         @test unnamed(zeros(elt, (ci,), (cj,))) == zeros(elt, 3, 4)
+        # The rng-first forms (mirroring `Base.randn(rng, dims...)`, default eltype) accept flat
+        # axes as varargs or a tuple, and the split codomain/domain.
+        rng = default_rng()
+        @test dimnames(randn(rng, ci, cj)) == [name(ci), name(cj)]
+        @test dimnames(randn(rng, (ci, cj))) == [name(ci), name(cj)]
+        @test dimnames(randn(rng, (ci,), (cj,))) == [name(ci), name(cj)]
         # 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))