[GradedAxes] Replace GradedAxes with GradedAxesNext (#1355)

Author: mtfishman

NDTensors/src/imports.jl

@@ -37,7 +37,6 @@ for lib in [
   :RankFactorization,
   :Sectors,
   :LabelledNumbers,
-  :GradedAxesNext,
   :GradedAxes,
   :TensorAlgebra,
   :SparseArrayInterface,

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/src/BlockSparseArraysGradedAxesExt.jl

@@ -2,25 +2,18 @@ module BlockSparseArraysGradedAxesExt
 using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange
 using ..BlockSparseArrays: BlockSparseArrays, block_merge
 using ...GradedAxes:
-  AbstractGradedUnitRange,
-  OneToOne,
-  blockmergesortperm,
-  blocksortperm,
-  invblockperm,
-  tensor_product
+  GradedUnitRange, OneToOne, blockmergesortperm, blocksortperm, invblockperm, tensor_product
 using ...TensorAlgebra:
   TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims
 
 # TODO: Make a `ReduceWhile` library.
 include("reducewhile.jl")
 
-TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion()
+TensorAlgebra.FusionStyle(::GradedUnitRange) = SectorFusion()
 
 # TODO: Need to implement this! Will require implementing
 # `block_merge(a::AbstractUnitRange, blockmerger::BlockedUnitRange)`.
-function BlockSparseArrays.block_merge(
-  a::AbstractGradedUnitRange, blockmerger::BlockedUnitRange
-)
+function BlockSparseArrays.block_merge(a::GradedUnitRange, blockmerger::BlockedUnitRange)
   return a
 end
 

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -22,6 +22,7 @@ end
 
 # TODO: Implement as `copy(@view a[I...])`, which is then implemented
 # through `ArrayLayouts.sub_materialize`.
+using ..SparseArrayInterface: set_getindex_zero_function
 function blocksparse_getindex(
   a::AbstractArray{<:Any,N}, I::Vararg{AbstractVector{<:Block{1}},N}
 ) where {N}
@@ -30,8 +31,9 @@ function blocksparse_getindex(
   CI = map(i -> Int.(i), I)
   subblocks_a = blocks_a[CI...]
   subaxes = ntuple(ndims(a)) do i
-    return axes(a, i)[I[i]]
+    return only(axes(axes(a, i)[I[i]]))
   end
+  subblocks_a = set_getindex_zero_function(subblocks_a, BlockZero(subaxes))
   return typeof(a)(subblocks_a, subaxes)
 end
 

NDTensors/src/lib/GradedAxes/Project.toml

@@ -1,2 +1,2 @@
 [deps]
-BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"

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

@@ -1,8 +1,9 @@
 module GradedAxesSectorsExt
 using ..GradedAxes: GradedAxes
-using ...Sectors: Sectors, AbstractCategory, ⊗, dual
+using ...Sectors: Sectors, AbstractCategory, ⊗ # , dual
 
-GradedAxes.fuse(c1::AbstractCategory, c2::AbstractCategory) = only(c1 ⊗ c2)
+GradedAxes.fuse_labels(c1::AbstractCategory, c2::AbstractCategory) = only(c1 ⊗ c2)
 
-GradedAxes.dual(c::AbstractCategory) = dual(c)
+# TODO: Decide the fate of `dual`.
+## GradedAxes.dual(c::AbstractCategory) = dual(c)
 end

NDTensors/src/lib/GradedAxes/ext/GradedAxesSectorsExt/src/GradedAxesSectorsExt.jl

@@ -1,14 +1,14 @@
 @eval module $(gensym())
-using NDTensors.GradedAxes: dual, fuse
+using NDTensors.GradedAxes: dual, fuse_labels
 using NDTensors.Sectors: U1, Z
 using Test: @test, @testset
 
 @testset "GradedAxesSectorsExt" begin
-  @test fuse(U1(1), U1(2)) == U1(3)
+  @test fuse_labels(U1(1), U1(2)) == U1(3)
   @test dual(U1(2)) == U1(-2)
 
-  @test fuse(Z{2}(1), Z{2}(1)) == Z{2}(0)
-  @test fuse(Z{2}(0), Z{2}(1)) == Z{2}(1)
+  @test fuse_labels(Z{2}(1), Z{2}(1)) == Z{2}(0)
+  @test fuse_labels(Z{2}(0), Z{2}(1)) == Z{2}(1)
   @test dual(Z{2}(1)) == Z{2}(1)
   @test dual(Z{2}(0)) == Z{2}(0)
 end

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

@@ -1,7 +1,5 @@
 module GradedAxes
-include("groupsortperm.jl")
-include("tensor_product.jl")
-include("abstractgradedunitrange.jl")
 include("gradedunitrange.jl")
+include("fusion.jl")
 include("../ext/GradedAxesSectorsExt/src/GradedAxesSectorsExt.jl")
 end

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

@@ -0,0 +1,110 @@
+using BlockArrays: BlockedUnitRange
+
+# TODO: Decide what to do about `dual`. Should there just
+# be a version in `Sectors`?
+## function dual end
+
+# 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))
+
+# 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
+function tensor_product(
+  a1::AbstractUnitRange,
+  a2::AbstractUnitRange,
+  a3::AbstractUnitRange,
+  a_rest::Vararg{AbstractUnitRange},
+)
+  return foldl(tensor_product, (a1, a2, a3, a_rest...))
+end
+
+function tensor_product(a1::AbstractUnitRange, a2::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)
+  return a2
+end
+
+function tensor_product(a1::AbstractUnitRange, a2::OneToOne)
+  return a1
+end
+
+function tensor_product(a1::OneToOne, a2::OneToOne)
+  return OneToOne()
+end
+
+function fuse_labels(x, y)
+  return error(
+    "`fuse_labels` not implemented for object of type `$(typeof(x))` and `$(typeof(y))`."
+  )
+end
+
+function fuse_blocklengths(x::Integer, y::Integer)
+  return 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)))
+end
+
+using BlockArrays: blockedrange, blocks
+function tensor_product(a1::BlockedUnitRange, a2::BlockedUnitRange)
+  blocklengths = map(vec(collect(Iterators.product(blocks(a1), blocks(a2))))) do x
+    return mapreduce(length, fuse_blocklengths, x)
+  end
+  return blockedrange(blocklengths)
+end
+
+function blocksortperm(a::BlockedUnitRange)
+  # 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
+
+using BlockArrays: Block, BlockVector
+using SplitApplyCombine: groupcount
+# Get the permutation for sorting, then group by common elements.
+# groupsortperm([2, 1, 2, 3]) == [[2], [1, 3], [4]]
+function groupsortperm(v; kwargs...)
+  perm = sortperm(v; kwargs...)
+  v_sorted = @view v[perm]
+  group_lengths = collect(groupcount(identity, v_sorted))
+  return BlockVector(perm, group_lengths)
+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::BlockedUnitRange)
+  # 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)))
+end