[SymmetrySectors] Non-abelian fusion (#1363)

Author: ogauthe

NDTensors/src/imports.jl

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

NDTensors/src/lib/GradedAxes/ext/GradedAxesSectorsExt/Project.toml

@@ -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/ext/GradedAxesSectorsExt/src/GradedAxesSectorsExt.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/ext/GradedAxesSectorsExt/test/Project.toml

@@ -1,10 +1,11 @@
-using BlockArrays: AbstractBlockedUnitRange
+using BlockArrays: AbstractBlockedUnitRange, blocklengths
 
 # Represents the range `1:1` or `Base.OneTo(1)`.
 struct OneToOne{T} <: AbstractUnitRange{T} end
 OneToOne() = OneToOne{Bool}()
 Base.first(a::OneToOne) = one(eltype(a))
 Base.last(a::OneToOne) = one(eltype(a))
+BlockArrays.blockaxes(g::OneToOne) = (Block.(g),)  # BlockArrays default crashes for OneToOne{Bool}
 
 # 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
@@ -18,23 +19,25 @@ function tensor_product(
   return foldl(tensor_product, (a1, a2, a3, a_rest...))
 end
 
+flip_dual(r::AbstractUnitRange) = r
+flip_dual(r::UnitRangeDual) = flip(r)
 function tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)
-  return error("Not implemented yet.")
+  return tensor_product(flip_dual(a1), flip_dual(a2))
 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
 
@@ -45,27 +48,28 @@ function fuse_labels(x, y)
 end
 
 function fuse_blocklengths(x::Integer, y::Integer)
-  return x * y
+  # return blocked unit range to keep non-abelian interface
+  return blockedrange([x * y])
 end
 
 using ..LabelledNumbers: LabelledInteger, label, labelled, unlabel
 function fuse_blocklengths(x::LabelledInteger, y::LabelledInteger)
-  return labelled(unlabel(x) * unlabel(y), fuse_labels(label(x), label(y)))
+  # return blocked unit range to keep non-abelian interface
+  return blockedrange([labelled(x * y, fuse_labels(label(x), label(y)))])
 end
 
 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)
+  nested = map(Iterators.flatten((Iterators.product(blocks(a1), blocks(a2)),))) do it
+    return mapreduce(length, fuse_blocklengths, it)
   end
-  return blockedrange(blocklengths)
+  new_blocklengths = mapreduce(blocklengths, vcat, nested)
+  return blockedrange(new_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)))
+# convention: sort UnitRangeDual according to nondual blocks
+function blocksortperm(a::AbstractUnitRange)
+  return Block.(sortperm(blocklabels(nondual(a))))
 end
 
 using BlockArrays: Block, BlockVector
@@ -82,25 +86,34 @@ end
 # Used by `TensorAlgebra.splitdims` in `BlockSparseArraysGradedAxesExt`.
 # 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)))
+function blockmergesortperm(a::AbstractUnitRange)
+  return Block.(groupsortperm(blocklabels(nondual(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 blockmergesort(g::AbstractGradedUnitRange)
+  glabels = blocklabels(g)
+  gblocklengths = blocklengths(g)
+  new_blocklengths = map(sort(unique(glabels))) do la
+    return labelled(sum(gblocklengths[findall(==(la), glabels)]; init=0), la)
+  end
+  return gradedrange(new_blocklengths)
+end
+
+blockmergesort(g::UnitRangeDual) = flip(blockmergesort(flip(g)))
+blockmergesort(g::AbstractUnitRange) = 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/ext/GradedAxesSectorsExt/test/runtests.jl

@@ -12,13 +12,14 @@ using BlockArrays:
   blockedrange,
   BlockIndexRange,
   blockfirsts,
-  blocklasts,
+  blockisequal,
   blocklength,
   blocklengths,
   findblock,
   findblockindex,
   mortar
 using Compat: allequal
+using FillArrays: Fill
 using ..LabelledNumbers:
   LabelledNumbers, LabelledInteger, LabelledUnitRange, label, labelled, unlabel
 
@@ -37,6 +38,18 @@ function Base.OrdinalRange{T,T}(a::GradedOneTo{<:LabelledInteger{T}}) where {T}
   return unlabel_blocks(a)
 end
 
+# == is just a range comparison that ignores labels. Need dedicated function to check equality.
+struct NoLabel end
+blocklabels(r::AbstractUnitRange) = Fill(NoLabel(), blocklength(r))
+
+function labelled_isequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
+  return blockisequal(a1, a2) && (blocklabels(a1) == blocklabels(a2))
+end
+
+function space_isequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
+  return (isdual(a1) == isdual(a2)) && labelled_isequal(a1, a2)
+end
+
 # This is only needed in certain Julia versions below 1.10
 # (for example Julia 1.6).
 # TODO: Delete this once we drop Julia 1.6 support.

NDTensors/src/lib/GradedAxes/src/GradedAxes.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(::AbstractUnitRange) = 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
@@ -92,6 +105,9 @@ 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)))
+
 function BlockArrays.combine_blockaxes(a1::UnitRangeDual, a2::UnitRangeDual)
   return dual(combine_blockaxes(dual(a1), dual(a2)))
 end

NDTensors/src/lib/GradedAxes/src/dual.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/src/fusion.jl

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

NDTensors/src/lib/GradedAxes/src/gradedunitrange.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,
+  space_isequal,
+  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 space_isequal(dual(ad), a)
+  @test space_isequal(nondual(ad), a)
+  @test space_isequal(nondual(a), a)
+  @test space_isequal(ad, ad)
+  @test !space_isequal(a, ad)
+  @test !space_isequal(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 space_isequal(flip(a), dual(gradedrange([U1(0) => 2, U1(-1) => 3])))
+  @test space_isequal(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