Restore GradedUnitRange slicing

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.6.17"
+version = "0.6.18"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/GradedArrays.jl

@@ -15,7 +15,6 @@ export dual, flip, gradedrange, isdual,
 
 # imports
 # -------
-import KroneckerArrays: ×
 using BlockArrays: BlockArrays, Block, blocksize
 using BlockSparseArrays: @view!, AbstractBlockSparseArray, BlockOneTo, blockrange
 using KroneckerArrays

src/gradedarray.jl

@@ -38,16 +38,16 @@ isdual(g::GradedUnitRange) = isdual(first(eachblockaxis(g)))  # crash for empty.
 
 flux(a::AbstractBlockedUnitRange, I::Block{1}) = flux(a[I])
 
-function ×(g1::GradedOneTo, g2::GradedOneTo)
+function KroneckerArrays.:×(g1::GradedOneTo, g2::GradedOneTo)
     v = vec([a × b for a in eachblockaxis(g1), b in eachblockaxis(g2)])
     return mortar_axis(v)
 end
 
-×(g::GradedUnitRange, a::AbstractUnitRange) = ×(g, to_gradedrange(a))
-×(a::AbstractUnitRange, g::GradedUnitRange) = ×(to_gradedrange(a), g)
-×(g::GradedUnitRange, a::SectorRange) = ×(g, to_gradedrange(a))
-×(a::SectorRange, g::GradedUnitRange) = ×(to_gradedrange(a), g)
-function ×(g1::GradedUnitRange, g2::GradedUnitRange)
+KroneckerArrays.:×(g::GradedUnitRange, a::AbstractUnitRange) = ×(g, to_gradedrange(a))
+KroneckerArrays.:×(a::AbstractUnitRange, g::GradedUnitRange) = ×(to_gradedrange(a), g)
+KroneckerArrays.:×(g::GradedUnitRange, a::SectorRange) = ×(g, to_gradedrange(a))
+KroneckerArrays.:×(a::SectorRange, g::GradedUnitRange) = ×(to_gradedrange(a), g)
+function KroneckerArrays.:×(g1::GradedUnitRange, g2::GradedUnitRange)
     v = vec([a × b for a in eachblockaxis(g1), b in eachblockaxis(g2)])
     return mortar_axis(v)
 end
@@ -74,6 +74,40 @@ function BlockSparseArrays.mortar_axis(geachblockaxis::AbstractVector{<:SectorUn
     return blockrange(geachblockaxis)
 end
 
+# Keep graded labels when slicing by block vectors.
+function BlockSparseArrays.blockedunitrange_getindices(
+        g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}}
+    )
+    gblocks = map(index -> g[index], indices)
+    new_multiplicities = sector_multiplicity.(gblocks)
+    new_axis = mortar_axis(sectorrange.(sector.(gblocks), Base.OneTo.(new_multiplicities)))
+    return BlockArrays.mortar(gblocks, (new_axis,))
+end
+
+function BlockSparseArrays.blockedunitrange_getindices(
+        g::GradedUnitRange, indices::BlockArrays.AbstractBlockVector{<:BlockArrays.Block{1}}
+    )
+    blks = map(bs -> BlockArrays.mortar(map(b -> g[b], bs)), BlockArrays.blocks(indices))
+    new_sectors = map(bs -> sectors(g)[Int.(bs)], BlockArrays.blocks(indices))
+    @assert all(allequal.(new_sectors))
+    new_multiplicities = map(BlockArrays.blocks(indices)) do bs
+        return sum(b -> sector_multiplicity(g[b]), bs; init = 0)
+    end
+    new_axis = mortar_axis(
+        sectorrange.(first.(new_sectors), Base.OneTo.(new_multiplicities))
+    )
+    return BlockArrays.mortar(blks, (new_axis,))
+end
+function Base.getindex(g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}})
+    return BlockSparseArrays.blockedunitrange_getindices(g, indices)
+end
+function Base.getindex(
+        g::GradedUnitRange,
+        indices::BlockArrays.AbstractBlockVector{<:BlockArrays.Block{1}}
+    )
+    return BlockSparseArrays.blockedunitrange_getindices(g, indices)
+end
+
 to_gradedrange(g::GradedUnitRange) = g
 
 function Base.show(io::IO, g::GradedUnitRange)

src/sectorarray.jl

@@ -33,8 +33,8 @@ function sectorrange(
     return sectorrange(to_sector(sector), args...)
 end
 
-×(a::SectorRange, g::AbstractUnitRange) = cartesianrange(a, g)
-×(g::AbstractUnitRange, a::SectorRange) = cartesianrange(a, g)
+KroneckerArrays.:×(a::SectorRange, g::AbstractUnitRange) = cartesianrange(a, g)
+KroneckerArrays.:×(g::AbstractUnitRange, a::SectorRange) = cartesianrange(a, g)
 
 to_gradedrange(g::SectorUnitRange) = mortar_axis([g])
 

src/sectorproduct.jl

@@ -161,38 +161,41 @@ end
 
 # =================================  Cartesian Product  ====================================
 
-const sectorproduct = ×
+const sectorproduct = KroneckerArrays.:×
 
-×(c::SectorRange) = SectorRange(SectorProduct(label(c)))
-×(c1::SectorRange, c2::SectorRange) = SectorRange(×(label(c1), label(c2)))
-×(c1::TKS.Sector, c2::TKS.Sector) = ×(SectorProduct(c1), SectorProduct(c2))
+KroneckerArrays.:×(c::SectorRange) = SectorRange(SectorProduct(label(c)))
+KroneckerArrays.:×(c1::SectorRange, c2::SectorRange) = SectorRange(×(label(c1), label(c2)))
+KroneckerArrays.:×(c1::TKS.Sector, c2::TKS.Sector) = ×(SectorProduct(c1), SectorProduct(c2))
 
-function ×(p1::SectorProduct{<:Tuple}, p2::SectorProduct{<:Tuple})
+function KroneckerArrays.:×(p1::SectorProduct{<:Tuple}, p2::SectorProduct{<:Tuple})
     return SectorProduct(arguments(p1)..., arguments(p2)...)
 end
-function ×(p1::SectorProduct{<:NamedTuple}, p2::SectorProduct{<:NamedTuple})
+function KroneckerArrays.:×(
+        p1::SectorProduct{<:NamedTuple},
+        p2::SectorProduct{<:NamedTuple}
+    )
     isdisjoint(keys(arguments(p1)), keys(arguments(p2))) ||
         throw(ArgumentError("keys of SectorProducts must be distinct"))
     return SectorProduct(merge(arguments(p1), arguments(p2)))
 end
-function ×(a::SectorProduct, b::SectorProduct)
+function KroneckerArrays.:×(a::SectorProduct, b::SectorProduct)
     isempty(arguments(a)) && return b
     isempty(arguments(b)) && return a
     throw(MethodError(×, typeof.((a, b))))
 end
 
-×(nt1::NamedTuple) = to_sector(nt1)
-×(nt1::NamedTuple, nt2::NamedTuple) = ×(to_sector(nt1), to_sector(nt2))
-×(c1::NamedTuple, c2::SectorRange) = ×(to_sector(c1), c2)
-×(c1::SectorRange, c2::NamedTuple) = ×(c1, to_sector(c2))
+KroneckerArrays.:×(nt1::NamedTuple) = to_sector(nt1)
+KroneckerArrays.:×(nt1::NamedTuple, nt2::NamedTuple) = ×(to_sector(nt1), to_sector(nt2))
+KroneckerArrays.:×(c1::NamedTuple, c2::SectorRange) = ×(to_sector(c1), c2)
+KroneckerArrays.:×(c1::SectorRange, c2::NamedTuple) = ×(c1, to_sector(c2))
 
-function ×(pairs::Pair...)
+function KroneckerArrays.:×(pairs::Pair...)
     keys = Symbol.(first.(pairs))
     vals = last.(pairs)
     return ×(NamedTuple{keys}(vals))
 end
 
-function ×(sr1::SectorOneTo, sr2::SectorOneTo)
+function KroneckerArrays.:×(sr1::SectorOneTo, sr2::SectorOneTo)
     isdual(sr1) == isdual(sr2) || throw(ArgumentError("SectorProduct duality must match"))
     sr = sectorrange(
         sector(sr1) × sector(sr2),

src/sectorrange.jl

@@ -1,6 +1,6 @@
 # This file defines the interface for type Sector
 # all fusion categories (Z{2}, SU2, Ising...) are subtypes of Sector
-import TensorKitSectors as TKS
+using TensorKitSectors: TensorKitSectors as TKS
 
 """
     SectorRange(sector::TKS.Sector, isdual::Bool)

src/tensoralgebra.jl

@@ -49,7 +49,7 @@ end
 function TensorAlgebra.tensor_product_axis(
         ::SectorFusion, ::Val{:domain}, r1::SectorUnitRange, r2::SectorUnitRange
     )
-    return dual(r1 ⊗ r2)
+    return flip(r1 ⊗ r2)
 end
 function TensorAlgebra.tensor_product_axis(
         style::BlockReshapeFusion, side::Val{:codomain},
@@ -90,11 +90,20 @@ function TensorAlgebra.matricize(
     ) where {Ncodomain}
     biperm = trivialbiperm(ndims_codomain, Val(ndims(a)))
     ax_codomain, ax_domain = blocks(axes(a)[biperm])
-    ax_codomain = tensor_product(ax_codomain...)
-    ax_domain = flip(tensor_product(ax_domain...))
+    ax_codomain =
+        isempty(ax_codomain) ? trivial(sector_type(a)) : tensor_product(ax_codomain...)
+    ax_domain =
+        isempty(ax_domain) ? trivial(sector_type(a)) : flip(tensor_product(ax_domain...))
     return SectorDelta{eltype(a)}((ax_codomain, ax_domain))
 end
 
+function TensorAlgebra.unmatricize(
+        ::SectorFusion, m::SectorDelta,
+        codomain_axes::Tuple{Vararg{SectorRange}},
+        domain_axes::Tuple{Vararg{SectorRange}}
+    )
+    return SectorDelta{eltype(m)}((codomain_axes..., domain_axes...))
+end
 function TensorAlgebra.unmatricize(
         ::SectorFusion, m::AbstractMatrix,
         codomain_axes::Tuple{Vararg{AbstractUnitRange}},

test/test_gradedunitrange.jl

@@ -214,7 +214,9 @@ using Test: @test, @test_throws, @testset
         v = mortar([[Block(2), Block(2)], [Block(1)]])
         a = g[v]
         @test a isa BlockVector
-        @test_broken only(axes(a)) isa GradedOneTo
+        @test only(axes(a)) isa GradedOneTo
+        @test sector_multiplicities(only(axes(a))) == [6, 2]
+        @test sectors(only(axes(a))) == sectors(g)[[2, 1]]
         @test_broken space_isequal(
             only(axes(a)),
             gradedrange(["y" => 6, "x" => 2]; isdual = isdual(g))

test/test_sectorproduct.jl

@@ -1,8 +1,8 @@
-import TensorKitSectors as TKS
 using BlockArrays: blocklengths
 using GradedArrays: GradedArrays, SU2, SectorProduct, TrivialSector, U1, Z, arguments, dual,
     flip, gradedrange, label, quantum_dimension, sector, sector_type, sectorproduct,
     sectorrange, space_isequal, trivial, ×, ⊗
+using TensorKitSectors: TensorKitSectors as TKS
 using Test: @test, @test_broken, @test_throws, @testset
 using TestExtras: @constinferred
 

test/test_tensoralgebraext.jl

@@ -1,9 +1,11 @@
 using BlockArrays: Block, blocksize
 using BlockSparseArrays: BlockSparseArray, mortar_axis
-using GradedArrays: GradedArray, GradedMatrix, SU2, U1, dual, flip, gradedrange,
-    sector_type, space_isequal, trivial_gradedrange
+using GradedArrays: GradedArray, GradedMatrix, SU2, SectorDelta, U1, dual, flip,
+    gradedrange, isdual, sector, sector_type, sectorrange, space_isequal, trivial,
+    trivial_gradedrange, ⊗
 using Random: randn!
-using TensorAlgebra: contract, matricize, trivial_axis, unmatricize
+using TensorAlgebra:
+    FusionStyle, contract, matricize, tensor_product_axis, trivial_axis, unmatricize
 using Test: @test, @testset
 
 function randn_blockdiagonal(elt::Type, axes::Tuple)
@@ -32,6 +34,37 @@ end
     @test space_isequal(trivial_gradedrange((gN, gS)), gNS)
 end
 
+@testset "SectorDelta domain axis tensor product uses flip" begin
+    r1 = sectorrange(U1(1), 2)
+    r2 = sectorrange(U1(2), 3)
+    rdomain = tensor_product_axis(FusionStyle(SectorDelta), Val(:domain), r1, r2)
+    @test space_isequal(rdomain, flip(r1 ⊗ r2))
+    @test isdual(rdomain)
+    @test sector(rdomain) == flip(U1(3))
+end
+
+@testset "SectorDelta matricize handles empty codomain/domain" begin
+    s1 = U1(1)
+    s2 = U1(2)
+    s3 = U1(-1)
+    d = SectorDelta{Float64}((s1, s2, s3))
+
+    m_left_empty = matricize(d, (), (1, 2, 3))
+    @test axes(m_left_empty, 1) == trivial(sector_type(d))
+    @test axes(m_left_empty, 2) == flip(s1 ⊗ s2 ⊗ s3)
+
+    m_right_empty = matricize(d, (1, 2, 3), ())
+    @test axes(m_right_empty, 1) == s1 ⊗ s2 ⊗ s3
+    @test axes(m_right_empty, 2) == trivial(sector_type(d))
+end
+
+@testset "SectorDelta unmatricize preserves provided axes" begin
+    m = SectorDelta{Float64}((U1(3), flip(U1(3))))
+    @test unmatricize(m, (U1(1), U1(2)), (U1(-2), U1(-1))) ==
+        SectorDelta{Float64}((U1(1), U1(2), U1(-2), U1(-1)))
+    @test unmatricize(m, (U1(1), U1(1)), (U1(-2), U1(-1))) isa SectorDelta
+end
+
 broken = true
 
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})