[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