[Sectors] Non-abelian fusion

NDTensors/src/imports.jl

@@ -33,9 +33,9 @@ for lib in [
   :MetalExtensions,
   :BroadcastMapConversion,
   :RankFactorization,
-  :Sectors,
   :LabelledNumbers,
   :GradedAxes,
+  :Sectors,
   :TensorAlgebra,
   :SparseArrayInterface,
   :SparseArrayDOKs,

NDTensors/src/lib/GradedAxes/src/GradedAxes.jl

@@ -1,8 +1,7 @@
 module GradedAxes
 include("blockedunitrange.jl")
 include("gradedunitrange.jl")
-include("fusion.jl")
 include("dual.jl")
 include("unitrangedual.jl")
-include("../ext/GradedAxesSectorsExt/src/GradedAxesSectorsExt.jl")
+include("fusion.jl")
 end

NDTensors/src/lib/GradedAxes/src/dual.jl

@@ -5,3 +5,5 @@ using NDTensors.LabelledNumbers:
 label_dual(x) = label_dual(LabelledStyle(x), x)
 label_dual(::NotLabelled, x) = x
 label_dual(::IsLabelled, x) = labelled(unlabel(x), dual(label(x)))
+
+flip(g::AbstractGradedUnitRange) = dual(gradedrange(label_dual.(blocklengths(g))))

NDTensors/src/lib/GradedAxes/src/fusion.jl

@@ -6,6 +6,10 @@ OneToOne() = OneToOne{Bool}()
 Base.first(a::OneToOne) = one(eltype(a))
 Base.last(a::OneToOne) = one(eltype(a))
 
+gradedisequal(::AbstractUnitRange, ::OneToOne) = false
+gradedisequal(::OneToOne, ::AbstractUnitRange) = false
+gradedisequal(::OneToOne, ::OneToOne) = true
+
 # https://github.com/ITensor/ITensors.jl/blob/v0.3.57/NDTensors/src/lib/GradedAxes/src/tensor_product.jl
 # https://en.wikipedia.org/wiki/Tensor_product
 # https://github.com/KeitaNakamura/Tensorial.jl
@@ -18,26 +22,39 @@ function tensor_product(
   return foldl(tensor_product, (a1, a2, a3, a_rest...))
 end
 
-function tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)
+function tensor_product(::AbstractUnitRange, ::AbstractUnitRange)
   return error("Not implemented yet.")
 end
 
 function tensor_product(a1::Base.OneTo, a2::Base.OneTo)
   return Base.OneTo(length(a1) * length(a2))
 end
 
-function tensor_product(a1::OneToOne, a2::AbstractUnitRange)
+function tensor_product(::OneToOne, a2::AbstractUnitRange)
   return a2
 end
 
-function tensor_product(a1::AbstractUnitRange, a2::OneToOne)
+function tensor_product(a1::AbstractUnitRange, ::OneToOne)
   return a1
 end
 
-function tensor_product(a1::OneToOne, a2::OneToOne)
+function tensor_product(::OneToOne, ::OneToOne)
   return OneToOne()
 end
 
+# Handle dual. Always return a non-dual GradedUnitRange.
+function tensor_product(a1::AbstractUnitRange, a2::UnitRangeDual)
+  return tensor_product(a1, flip(a2))
+end
+
+function tensor_product(a1::UnitRangeDual, a2::AbstractUnitRange)
+  return tensor_product(flip(a1), a2)
+end
+
+function tensor_product(a1::UnitRangeDual, a2::UnitRangeDual)
+  return tensor_product(flip(a1), flip(a2))
+end
+
 function fuse_labels(x, y)
   return error(
     "`fuse_labels` not implemented for object of type `$(typeof(x))` and `$(typeof(y))`."
@@ -53,21 +70,28 @@ function fuse_blocklengths(x::LabelledInteger, y::LabelledInteger)
   return labelled(unlabel(x) * unlabel(y), fuse_labels(label(x), label(y)))
 end
 
+flatten_maybe_nested(v::Vector{<:Integer}) = v
+flatten_maybe_nested(v::Vector{<:AbstractGradedUnitRange}) = reduce(vcat, blocklengths.(v))
+
 using BlockArrays: blockedrange, blocks
 function tensor_product(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)
-  blocklengths = map(vec(collect(Iterators.product(blocks(a1), blocks(a2))))) do x
-    return mapreduce(length, fuse_blocklengths, x)
-  end
+  maybe_nested = map(
+    it -> mapreduce(length, fuse_blocklengths, it),
+    Iterators.flatten((Iterators.product(blocks(a1), blocks(a2)),)),
+  )
+  blocklengths = flatten_maybe_nested(maybe_nested)
   return blockedrange(blocklengths)
 end
 
 function blocksortperm(a::AbstractBlockedUnitRange)
-  # TODO: Figure out how to deal with dual sectors.
-  # TODO: `rev=isdual(a)`  may not be correct for symmetries beyond `U(1)`.
-  ## return Block.(sortperm(nondual_sectors(a); rev=isdual(a)))
   return Block.(sortperm(blocklabels(a)))
 end
 
+# convention: sort UnitRangeDual according to nondual blocks
+function blocksortperm(a::UnitRangeDual)
+  return Block.(sortperm(blocklabels(nondual(a))))
+end
+
 using BlockArrays: Block, BlockVector
 using SplitApplyCombine: groupcount
 # Get the permutation for sorting, then group by common elements.
@@ -83,24 +107,38 @@ end
 # Get the permutation for sorting, then group by common elements.
 # groupsortperm([2, 1, 2, 3]) == [[2], [1, 3], [4]]
 function blockmergesortperm(a::AbstractBlockedUnitRange)
-  # If it is dual, reverse the sorting so the sectors
-  # end up sorted in the same way whether or not the space
-  # is dual.
-  # TODO: Figure out how to deal with dual sectors.
-  # TODO: `rev=isdual(a)`  may not be correct for symmetries beyond `U(1)`.
-  ## return Block.(groupsortperm(nondual_sectors(a); rev=isdual(a)))
   return Block.(groupsortperm(blocklabels(a)))
 end
 
 # Used by `TensorAlgebra.splitdims` in `BlockSparseArraysGradedAxesExt`.
 invblockperm(a::Vector{<:Block{1}}) = Block.(invperm(Int.(a)))
 
-# Used by `TensorAlgebra.fusedims` in `BlockSparseArraysGradedAxesExt`.
-function blockmergesortperm(a::GradedUnitRange)
-  # If it is dual, reverse the sorting so the sectors
-  # end up sorted in the same way whether or not the space
-  # is dual.
-  # TODO: Figure out how to deal with dual sectors.
-  # TODO: `rev=isdual(a)`  may not be correct for symmetries beyond `U(1)`.
-  return Block.(groupsortperm(blocklabels(a)))
+function blockmergesortperm(a::UnitRangeDual)
+  return Block.(groupsortperm(blocklabels(nondual(a))))
+end
+
+function blockmergesort(g::AbstractGradedUnitRange)
+  glabels = blocklabels(g)
+  gblocklengths = blocklengths(g)
+  new_blocklengths = map(
+    la -> labelled(sum(gblocklengths[findall(==(la), glabels)]; init=0), la),
+    sort(unique(glabels)),
+  )
+  return GradedAxes.gradedrange(new_blocklengths)
+end
+
+blockmergesort(g::UnitRangeDual) = dual(blockmergesort(flip(g)))
+blockmergesort(g::OneToOne) = g
+
+# fusion_product produces a sorted, non-dual GradedUnitRange
+function fusion_product(g1, g2)
+  return blockmergesort(tensor_product(g1, g2))
+end
+
+fusion_product(g::AbstractUnitRange) = blockmergesort(g)
+fusion_product(g::UnitRangeDual) = fusion_product(flip(g))
+
+# recursive fusion_product. Simpler than reduce + fix type stability issues with reduce
+function fusion_product(g1, g2, g3...)
+  return fusion_product(fusion_product(g1, g2), g3...)
 end

NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl

@@ -12,7 +12,7 @@ using BlockArrays:
   blockedrange,
   BlockIndexRange,
   blockfirsts,
-  blocklasts,
+  blockisequal,
   blocklength,
   blocklengths,
   findblock,
@@ -35,9 +35,9 @@ function Base.OrdinalRange{T,T}(a::GradedOneTo{<:LabelledInteger{T}}) where {T}
   return unlabel_blocks(a)
 end
 
-# TODO: See if this is needed.
-function Base.AbstractUnitRange{T}(a::GradedOneTo{<:LabelledInteger{T}}) where {T}
-  return unlabel_blocks(a)
+# == is just a range comparison that ignores labels. Need dedicated function to check equality.
+function gradedisequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
+  return blockisequal(a1, a2) && (blocklabels(a1) == blocklabels(a2))
 end
 
 # TODO: Use `TypeParameterAccessors`.

NDTensors/src/lib/GradedAxes/src/unitrangedual.jl

@@ -6,7 +6,10 @@ UnitRangeDual(a::AbstractUnitRange) = UnitRangeDual{eltype(a),typeof(a)}(a)
 dual(a::AbstractUnitRange) = UnitRangeDual(a)
 nondual(a::UnitRangeDual) = a.nondual_unitrange
 dual(a::UnitRangeDual) = nondual(a)
+flip(a::UnitRangeDual) = dual(flip(nondual(a)))
 nondual(a::AbstractUnitRange) = a
+isdual(::AbstractGradedUnitRange) = false
+isdual(::UnitRangeDual) = true
 ## TODO: Define this to instantiate a dual unit range.
 ## materialize_dual(a::UnitRangeDual) = materialize_dual(nondual(a))
 
@@ -16,6 +19,16 @@ Base.step(a::UnitRangeDual) = label_dual(step(nondual(a)))
 
 Base.view(a::UnitRangeDual, index::Block{1}) = a[index]
 
+function Base.show(io::IO, a::UnitRangeDual)
+  return print(io, UnitRangeDual, "(", blocklasts(a), ")")
+end
+
+function Base.show(io::IO, mimetype::MIME"text/plain", a::UnitRangeDual)
+  return Base.invoke(
+    show, Tuple{typeof(io),MIME"text/plain",AbstractArray}, io, mimetype, a
+  )
+end
+
 function Base.getindex(a::UnitRangeDual, indices::AbstractUnitRange{<:Integer})
   return dual(getindex(nondual(a), indices))
 end
@@ -73,6 +86,14 @@ BlockArrays.blockaxes(a::UnitRangeDual) = blockaxes(nondual(a))
 BlockArrays.blockfirsts(a::UnitRangeDual) = label_dual.(blockfirsts(nondual(a)))
 BlockArrays.blocklasts(a::UnitRangeDual) = label_dual.(blocklasts(nondual(a)))
 BlockArrays.findblock(a::UnitRangeDual, index::Integer) = findblock(nondual(a), index)
+
+blocklabels(a::UnitRangeDual) = dual.(blocklabels(nondual(a)))
+
+gradedisequal(::UnitRangeDual, ::AbstractGradedUnitRange) = false
+gradedisequal(::AbstractGradedUnitRange, ::UnitRangeDual) = false
+function gradedisequal(a1::UnitRangeDual, a2::UnitRangeDual)
+  return gradedisequal(nondual(a1), nondual(a2))
+end
 function BlockArrays.combine_blockaxes(a1::UnitRangeDual, a2::UnitRangeDual)
   return dual(combine_blockaxes(dual(a1), dual(a2)))
 end

NDTensors/src/lib/GradedAxes/test/runtests.jl

@@ -2,7 +2,7 @@
 using Test: @testset
 @testset "GradedAxes" begin
   include("test_basics.jl")
-  include("test_tensor_product.jl")
   include("test_dual.jl")
+  include("test_tensor_product.jl")
 end
 end

NDTensors/src/lib/GradedAxes/test/test_basics.jl

@@ -10,7 +10,8 @@ using BlockArrays:
   blocklengths,
   blocks
 using NDTensors.BlockSparseArrays: BlockSparseVector
-using NDTensors.GradedAxes: GradedOneTo, GradedUnitRange, blocklabels, gradedrange
+using NDTensors.GradedAxes:
+  GradedOneTo, GradedUnitRange, blocklabels, gradedisequal, gradedrange
 using NDTensors.LabelledNumbers: LabelledUnitRange, islabelled, label, labelled, unlabel
 using Test: @test, @test_broken, @testset
 @testset "GradedAxes basics" begin
@@ -41,6 +42,7 @@ using Test: @test, @test_broken, @testset
       @test label(x) == "y"
     end
     @test isnothing(iterate(a, labelled(5, "y")))
+    @test gradedisequal(a, a)
     @test length(a) == 5
     @test step(a) == 1
     @test !islabelled(step(a))

NDTensors/src/lib/GradedAxes/test/test_dual.jl

@@ -1,19 +1,40 @@
 @eval module $(gensym())
-using BlockArrays: Block, blockaxes, blockfirsts, blocklasts, blocks, findblock
-using NDTensors.GradedAxes: GradedAxes, UnitRangeDual, dual, gradedrange, nondual
+using BlockArrays:
+  Block, blockaxes, blockfirsts, blocklasts, blocklength, blocklengths, blocks, findblock
+using NDTensors.GradedAxes:
+  GradedAxes,
+  UnitRangeDual,
+  blocklabels,
+  blockmergesortperm,
+  blocksortperm,
+  dual,
+  flip,
+  gradedisequal,
+  gradedrange,
+  isdual,
+  nondual
 using NDTensors.LabelledNumbers: LabelledInteger, label, labelled
 using Test: @test, @test_broken, @testset
 struct U1
   n::Int
 end
 GradedAxes.dual(c::U1) = U1(-c.n)
+Base.isless(c1::U1, c2::U1) = c1.n < c2.n
 @testset "dual" begin
   a = gradedrange([U1(0) => 2, U1(1) => 3])
   ad = dual(a)
   @test eltype(ad) == LabelledInteger{Int,U1}
-  @test dual(ad) == a
-  @test nondual(ad) == a
-  @test nondual(a) == a
+
+  @test gradedisequal(dual(ad), a)
+  @test gradedisequal(nondual(ad), a)
+  @test gradedisequal(nondual(a), a)
+  @test gradedisequal(ad, ad)
+  @test !gradedisequal(a, ad)
+  @test !gradedisequal(ad, a)
+
+  @test isdual(ad)
+  @test !isdual(a)
+
   @test blockfirsts(ad) == [labelled(1, U1(0)), labelled(3, U1(-1))]
   @test blocklasts(ad) == [labelled(2, U1(0)), labelled(5, U1(-1))]
   @test findblock(ad, 4) == Block(2)
@@ -34,5 +55,36 @@ GradedAxes.dual(c::U1) = U1(-c.n)
   @test label(ad[[Block(2), Block(1)]][Block(1)]) == U1(-1)
   @test ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)] == 3:4
   @test label(ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)]) == U1(-1)
+  @test blocksortperm(a) == [Block(1), Block(2)]
+  @test blocksortperm(ad) == [Block(1), Block(2)]
+  @test blocklength(blockmergesortperm(a)) == 2
+  @test blocklength(blockmergesortperm(ad)) == 2
+  @test blockmergesortperm(a) == [Block(1), Block(2)]
+  @test blockmergesortperm(ad) == [Block(1), Block(2)]
+end
+
+@testset "flip" begin
+  a = gradedrange([U1(0) => 2, U1(1) => 3])
+  ad = dual(a)
+  @test gradedisequal(flip(a), dual(gradedrange([U1(0) => 2, U1(-1) => 3])))
+  @test gradedisequal(flip(ad), gradedrange([U1(0) => 2, U1(-1) => 3]))
+
+  @test blocklabels(a) == [U1(0), U1(1)]
+  @test blocklabels(dual(a)) == [U1(0), U1(-1)]
+  @test blocklabels(flip(a)) == [U1(0), U1(1)]
+  @test blocklabels(flip(dual(a))) == [U1(0), U1(-1)]
+  @test blocklabels(dual(flip(a))) == [U1(0), U1(-1)]
+
+  @test blocklengths(a) == [2, 3]
+  @test blocklengths(dual(a)) == [2, 3]
+  @test blocklengths(flip(a)) == [2, 3]
+  @test blocklengths(flip(dual(a))) == [2, 3]
+  @test blocklengths(dual(flip(a))) == [2, 3]
+
+  @test !isdual(a)
+  @test isdual(dual(a))
+  @test isdual(flip(a))
+  @test !isdual(flip(dual(a)))
+  @test !isdual(dual(flip(a)))
 end
 end