ZNIrrep for larger N (#31)

* generalize ZNIrrep to handle large N

* update docstrings

* improve modular arithmetic

* test some more edge cases

* fix printing

Author: lkdvos

src/TensorKitSectors.jl

@@ -28,6 +28,9 @@ export FermionParity, FermionNumber, FermionSpin
 export PlanarTrivial, FibonacciAnyon, IsingAnyon
 export IsingBimodule
 
+# accessors
+export charge, modulus
+
 # unicode exports
 # ---------------
 export ⊠, ⊗, ×

src/irreps/znirrep.jl

@@ -1,43 +1,98 @@
 # ZNIrrep: irreps of Z_N are labelled by integers mod N; do we ever want N > 64?
 """
-    struct ZNIrrep{N} <: AbstractIrrep{ℤ{N}}
+    struct ZNIrrep{N, T <: Unsigned} <: AbstractIrrep{ℤ{N}}
     ZNIrrep{N}(n::Integer)
     Irrep[ℤ{N}](n::Integer)
 
-Represents irreps of the group ``ℤ_N`` for some value of `N<64`. (We need 2*(N-1) <= 127 in
-order for a ⊗ b to work correctly.) For `N` equals `2`, `3` or `4`, `ℤ{N}` can be replaced
-by `ℤ₂`, `ℤ₃`, `ℤ₄`. An arbitrary `Integer` `n` can be provided to the constructor, but only
-the value `mod(n, N)` is relevant.
+Represents irreps of the group ``ℤ_N`` for some value of `N`.
+For `N` equals `2`, `3` or `4`, `ℤ{N}` can be replaced by [`ℤ₂`](@ref), [`ℤ₃`](@ref), and [`ℤ₄`](@ref).
+An arbitrary `Integer` `n` can be provided to the constructor, but only the value `mod(n, N)` is relevant.
+
+The type of the stored integer `T` can either be explicitly provided, or will automatically be determined
+to be the smallest unsigned integer type that fits all possible irreps for the given `N`.
+
+See also [`charge`](@ref)` and [`modulus`](@ref) to extract the relevant data.
 
 ## Fields
-- `n::Int8`: the integer label of the irrep, modulo `N`.
+- `n::T`: the integer label of the irrep, modulo `N`.
 """
-struct ZNIrrep{N} <: AbstractIrrep{ℤ{N}}
-    n::Int8
+struct ZNIrrep{N, T <: Unsigned} <: AbstractIrrep{ℤ{N}}
+    n::T
     function ZNIrrep{N}(n::Integer) where {N}
-        @assert N < 64
-        return new{N}(mod(n, N))
+        T = _integer_type(N)
+        return new{N, T}(T(mod(n, N)))
+    end
+    function ZNIrrep{N, T}(n::Integer) where {N, T <: Unsigned}
+        N ≤ typemax(T) + 1 ||
+            throw(TypeError(:ZNIrrep, ZNIrrep{N, T}, ZNIrrep{N, _integer_type(N)}))
+        return new{N, T}(mod(n, N))
     end
 end
-Base.getindex(::IrrepTable, ::Type{ℤ{N}}) where {N} = ZNIrrep{N}
+
+"""
+    modulus(c::ZNIrrep{N}) -> N
+    modulus(::Type{<:ZNIrrep{N}}) -> N
+
+The order of the cyclic group, or the modulus of the charge labels.
+"""
+modulus(c::ZNIrrep) = modulus(typeof(c))
+modulus(::Type{<:ZNIrrep{N}}) where {N} = N
+
+"""
+    charge(c::ZNIrrep) -> Int
+
+The charge label of the irrep `c`.
+"""
+charge(c::ZNIrrep) = Int(c.n)
+
+Base.@assume_effects :foldable function _integer_type(N::Integer)
+    N <= 0 && throw(DomainError(N, "N should be positive"))
+    for T in (UInt8, UInt16, UInt32, UInt64)
+        # T needs to fit a
+        N ≤ (typemax(T) + 1) && return T
+    end
+    throw(DomainError(N, "N is too large"))
+end
+
+Base.getindex(::IrrepTable, ::Type{ℤ{N}}) where {N} = ZNIrrep{N, _integer_type(N)}
 Base.convert(Z::Type{<:ZNIrrep}, n::Real) = Z(n)
-const Z2Irrep = ZNIrrep{2}
-const Z3Irrep = ZNIrrep{3}
-const Z4Irrep = ZNIrrep{4}
-
-unit(::Type{ZNIrrep{N}}) where {N} = ZNIrrep{N}(0)
-dual(c::ZNIrrep{N}) where {N} = ZNIrrep{N}(-c.n)
-⊗(c1::ZNIrrep{N}, c2::ZNIrrep{N}) where {N} = (ZNIrrep{N}(c1.n + c2.n),)
-
-Base.IteratorSize(::Type{SectorValues{ZNIrrep{N}}}) where {N} = HasLength()
-Base.length(::SectorValues{ZNIrrep{N}}) where {N} = N
-function Base.iterate(::SectorValues{ZNIrrep{N}}, i = 0) where {N}
-    return i == N ? nothing : (ZNIrrep{N}(i), i + 1)
+const Z2Irrep = ZNIrrep{2, UInt8}
+const Z3Irrep = ZNIrrep{3, UInt8}
+const Z4Irrep = ZNIrrep{4, UInt8}
+
+unit(::Type{ZNIrrep{N, T}}) where {N, T} = ZNIrrep{N, T}(zero(T))
+# be careful with `-` for unsigned integers!
+dual(c::ZNIrrep{N, T}) where {N, T} = ZNIrrep{N, T}(N - c.n)
+⊗(c1::ZNIrrep{N, T}, c2::ZNIrrep{N, T}) where {N, T} = (ZNIrrep{N, T}(modular_add(c1.n, c2.n, Val(N))),)
+
+Base.IteratorSize(::Type{SectorValues{ZNIrrep{N, T}}}) where {N, T} = HasLength()
+Base.length(::SectorValues{ZNIrrep{N, T}}) where {N, T} = N
+function Base.iterate(::SectorValues{ZNIrrep{N, T}}, i = 0) where {N, T}
+    return i == N ? nothing : (ZNIrrep{N, T}(i), i + 1)
 end
-function Base.getindex(::SectorValues{ZNIrrep{N}}, i::Int) where {N}
-    return 1 <= i <= N ? ZNIrrep{N}(i - 1) : throw(BoundsError(values(ZNIrrep{N}), i))
+function Base.getindex(::SectorValues{ZNIrrep{N, T}}, i::Int) where {N, T}
+    return 1 <= i <= N ? ZNIrrep{N, T}(i - 1) : throw(BoundsError(values(ZNIrrep{N}), i))
 end
-findindex(::SectorValues{ZNIrrep{N}}, c::ZNIrrep{N}) where {N} = c.n + 1
+findindex(::SectorValues{ZNIrrep{N, T}}, c::ZNIrrep{N, T}) where {N, T} = c.n + 1
+
+Base.hash(c::ZNIrrep{N, T}, h::UInt) where {N, T} = hash(c.n, h)
+Base.isless(c1::ZNIrrep{N, T}, c2::ZNIrrep{N, T}) where {N, T} = isless(c1.n, c2.n)
 
-Base.hash(c::ZNIrrep{N}, h::UInt) where {N} = hash(c.n, h)
-Base.isless(c1::ZNIrrep{N}, c2::ZNIrrep{N}) where {N} = isless(c1.n, c2.n)
+# ensure the printing uses `Int`.
+function Base.show(io::IO, c::ZNIrrep)
+    I = typeof(c)
+    print_type = get(io, :typeinfo, nothing) !== I
+    print_type && print(io, type_repr(I), '(')
+    print(io, charge(c))
+    print_type && print(io, ')')
+    return nothing
+end
+
+# compute x + y mod N, requires 0 <= x < N and 0 <= y < N (unchecked!)
+function modular_add(x::T, y::T, ::Val{N}) where {T <: Unsigned, N}
+    Tmax = typemax(T) + 1
+    0 < N ≤ Tmax || throw(DomainError(N, "N is too large"))
+    N ≤ (typemax(T) + 1) ÷ 2 && return x + y
+    r, flag = Base.add_with_overflow(x, y)
+    return ifelse(flag, (Tmax - N) + r, r)
+end

test/runtests.jl

@@ -15,7 +15,7 @@ include("newsectors.jl")
 using .NewSectors
 
 const sectorlist = (
-    Z2Irrep, Z3Irrep, Z4Irrep, U1Irrep,
+    Z2Irrep, Z3Irrep, Z4Irrep, Irrep[ℤ{200}], U1Irrep,
     DNIrrep{3}, DNIrrep{4}, DNIrrep{5}, CU1Irrep,
     SU2Irrep, NewSU2Irrep,
     FibonacciAnyon, IsingAnyon, FermionParity,
@@ -98,6 +98,23 @@ end
     @test repr("text/plain", b ⊗ b) == "Irrep[D₄](1, false) ⊗ Irrep[D₄](1, false):\n (0, false)\n (0, true)\n (2, false)\n (2, true)"
 end
 
+@testset "ZNIrrep edge cases" begin
+    a = Irrep[Cyclic{255}](254)
+    @test typeof(a) == ZNIrrep{255, UInt8}
+    @test charge(dual(a)) == mod(-charge(a), 255)
+    @test charge(only(a ⊗ a)) == mod(charge(a) + charge(a), 255)
+
+    b = Irrep[Cyclic{256}](255)
+    @test typeof(b) == ZNIrrep{256, UInt8}
+    @test charge(dual(b)) == mod(-charge(b), 256)
+    @test charge(only(b ⊗ b)) == mod(charge(b) + charge(b), 256)
+
+    c = Irrep[Cyclic{257}](256)
+    @test typeof(c) == ZNIrrep{257, UInt16}
+    @test charge(dual(c)) == mod(-charge(c), 257)
+    @test charge(only(c ⊗ c)) == mod(charge(c) + charge(c), 257)
+end
+
 include("multifusion.jl")
 
 @testset "Aqua" begin