[ENH] define SectorUnitRange

src/GradedUnitRanges.jl

@@ -10,4 +10,7 @@ include("gradedunitrangedual.jl")
 include("onetoone.jl")
 include("fusion.jl")
 
+include("sectorunitrange.jl")
+include("newgradedunitrange.jl")
+
 end

src/newgradedunitrange.jl

@@ -0,0 +1,109 @@
+
+using BlockArrays: BlockArrays
+
+struct NewGradedUnitRange{T,BlockLasts,Sector} <: AbstractGradedUnitRange{T,BlockLasts}
+  nondual_labels::Vector{Sector}
+  multiplicity_range::BlockedOneTo{T,BlockLasts}  # TBD offset != 1?
+  isdual::Bool
+
+  function NewGradedUnitRange(glabels, grange, gisdual)
+    return new{eltype(grange),Vector{eltype(grange)},eltype(glabels)}(
+      glabels, grange, gisdual
+    )
+  end
+end
+
+#
+# Accessors
+#
+nondual_labels(g::NewGradedUnitRange) = g.nondual_labels
+multiplicity_range(g::NewGradedUnitRange) = g.multiplicity_range
+isdual(g::NewGradedUnitRange) = g.isdual
+
+#
+# GradedUnitRanges interface
+#
+function GradedUnitRanges.sector_type(
+  ::Type{<:NewGradedUnitRange{T,BlockLasts,Sector}}
+) where {T,BlockLasts,Sector}
+  return Sector
+end
+
+function blocklabels(g::NewGradedUnitRange)
+  return isdual(g) ? dual.(nondual_labels(g)) : nondual_labels(g)
+end
+
+function dual(g::NewGradedUnitRange)
+  return NewGradedUnitRange(nondual_labels(g), multiplicity_range(g), !isdual(g))
+end
+
+function flip(g::NewGradedUnitRange)
+  return NewGradedUnitRange(dual.(nondual_labels(g)), multiplicity_range(g), !isdual(g))
+end
+
+function space_isequal(g1::NewGradedUnitRange, g2::NewGradedUnitRange)
+  return nondual_labels(g1) == nondual_labels(g2) &&
+         blockisequal(multiplicity_range(g1), multiplicity_range(g2)) &&
+         isdual(g1) == isdual(g2)
+end
+
+#
+# Base interface
+#
+Base.length(g::NewGradedUnitRange) = sum(length.(blocks(g)))
+
+function Base.show(io::IO, g::NewGradedUnitRange)
+  return print(io, nameof(typeof(g)), blocklabels(g), multiplicity_range(g))
+end
+
+function Base.show(io::IO, ::MIME"text/plain", g::NewGradedUnitRange)
+  print(io, typeof(g))
+  if isdual(g)
+    print(io, " dual")
+  end
+  println()
+  return print(io, join(repr.(blocks(g)), '\n'))
+end
+
+#
+# BlockArrays interface
+#
+function BlockArrays.mortar(v::Vector{<:SectorUnitRange})
+  glabels = nondual_sector.(v)
+  grange = blockedrange(length.(multiplicity_range.(v)))
+  gisdual = isdual(first(v))
+  # TODO add checks
+  return NewGradedUnitRange(glabels, grange, gisdual)
+end
+
+function BlockArrays.blocks(g::NewGradedUnitRange)
+  return sectorunitrange.(
+    nondual_labels(g), blocks(multiplicity_range(g)), isdual(g), blockfirsts(g)
+  )
+end
+
+BlockArrays.blocklengths(g::NewGradedUnitRange) = length.(blocks(g))
+
+function BlockArrays.blockfirsts(g::NewGradedUnitRange)
+  return vcat([1], blocklasts(g)[begin:(end - 1)])
+end
+
+function BlockArrays.blocklasts(g::NewGradedUnitRange)
+  return cumsum(length.(nondual_labels(g)) .* blocklengths(multiplicity_range(g)))
+end
+
+#
+# slicing
+#
+Base.getindex(g::NewGradedUnitRange, b::Block{1}) = blocks(g)[Int(b)]
+
+function Base.getindex(
+  g::NewGradedUnitRange, br::BlockRange{1,R}
+) where {R<:Tuple{AbstractUnitRange{Int64}}}  # TODO remove ambiguities & use more generic def
+  return mortar(blocks(g)[Int.(br)])
+end
+
+# TODO replace Tuple with kronecker
+function Base.getindex(g::NewGradedUnitRange, bx::Tuple{<:Block{1},<:Any})
+  return blocks(g)[Int(first(bx))][(:, last(bx))]
+end

src/sectorunitrange.jl

@@ -0,0 +1,111 @@
+
+using BlockArrays: BlockArrays
+
+struct SectorUnitRange{T,Sector,Range<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
+  nondual_sector::Sector
+  multiplicity_range::Range
+  isdual::Bool
+  offset::T
+
+  function SectorUnitRange(s, r, b, offset)
+    return new{eltype(r),typeof(s),typeof(r)}(s, r, b, offset)
+  end
+end
+
+#
+# Constructors
+#
+sectorunitrange(s, r, b=false, offset=1) = SectorUnitRange(s, r, b, offset)
+
+#
+# accessors
+#
+nondual_sector(sr::SectorUnitRange) = sr.nondual_sector
+multiplicity_range(sr::SectorUnitRange) = sr.multiplicity_range
+isdual(sr::SectorUnitRange) = sr.isdual
+Base.first(sr::SectorUnitRange) = sr.offset
+
+#
+# Base interface
+#
+Base.axes(sr::SectorUnitRange) = Base.oneto(length(sr))
+
+Base.eachindex(sr::SectorUnitRange) = Base.oneto(length(sr))
+
+Base.lastindex(sr::SectorUnitRange) = length(sr)
+
+function Base.length(sr::SectorUnitRange)
+  return length(nondual_sector(sr)) * length(multiplicity_range(sr))
+end  # TBD directly quantum_dimension(nondual_sector(sr))?
+
+Base.iterate(sr::SectorUnitRange) = iterate(first(sr):last(sr))
+Base.iterate(sr::SectorUnitRange, i::Integer) = iterate(first(sr):last(sr), i)
+
+Base.last(sr::SectorUnitRange) = first(sr) + length(sr) - 1
+
+# slicing
+Base.getindex(sr::SectorUnitRange, i::Integer) = range(sr)[i]
+function Base.getindex(sr::SectorUnitRange, r::AbstractUnitRange{T}) where {T<:Integer}
+  return range(sr)[r]
+end
+
+# TODO replace (:,x) indexing with kronecker(:, x)
+Base.getindex(sr::SectorUnitRange, t::Tuple{Colon,<:Integer}) = sr[(:, last(t):last(t))]
+function Base.getindex(sr::SectorUnitRange, t::Tuple{Colon,<:AbstractUnitRange})
+  return sectorunitrange(
+    nondual_sector(sr), multiplicity_range(sr)[last(t)], isdual(sr), first(sr)
+  )
+end
+
+Base.range(sr::SectorUnitRange) = first(sr):last(sr)
+
+function Base.show(io::IO, sr::SectorUnitRange)
+  print(io, nameof(typeof(sr)), " ", first(sr), " .+ ")
+  if isdual(sr)
+    print(io, "dual(", nondual_sector(sr), ")")
+  else
+    print(io, nondual_sector(sr))
+  end
+  return print(io, " => ", multiplicity_range(sr))
+end
+
+#
+# GradedUnitRanges interface
+#
+function blocklabels(sr::SectorUnitRange)
+  return isdual(sr) ? [dual(nondual_sector(sr))] : [nondual_sector(sr)]
+end
+
+function dual(sr::SectorUnitRange)
+  return sectorunitrange(nondual_sector(sr), multiplicity_range(sr), !isdual(sr), first(sr))
+end
+
+function flip(sr::SectorUnitRange)
+  return sectorunitrange(
+    dual(nondual_sector(sr)), multiplicity_range(sr), !isdual(sr), first(sr)
+  )
+end
+
+function map_blocklabels(f, sr::SectorUnitRange)
+  return sectorunitrange(
+    f(nondual_sector(sr)), multiplicity_range(sr), isdual(sr), first(sr)
+  )
+end
+
+sector_type(::Type{<:SectorUnitRange{T,Sector}}) where {T,Sector} = Sector
+
+function space_isequal(sr1::SectorUnitRange, sr2::SectorUnitRange)
+  return nondual_sector(sr1) == nondual_sector(sr2) &&
+         isdual(sr1) == isdual(sr2) &&
+         multiplicity_range(sr1) == multiplicity_range(sr2) &&
+         first(sr1) == first(sr2)
+end
+
+#
+# BlockArrays interface
+#
+BlockArrays.blocks(sr::SectorUnitRange) = [sr]
+
+BlockArrays.blocklength(::SectorUnitRange) = 1
+
+BlockArrays.blocklengths(sr::SectorUnitRange) = [length(sr)]  # TBD return length(multiplicity_range(sr)) ?

test/test_newgradedunitrange.jl

@@ -0,0 +1,81 @@
+using Test: @test, @test_throws, @testset
+
+using BlockArrays: Block, blocklength, blocklengths, blockisequal, blocks, mortar
+
+using GradedUnitRanges:
+  AbstractGradedUnitRange,
+  NewGradedUnitRange,
+  SectorUnitRange,
+  blocklabels,
+  dual,
+  gradedrange,
+  flip,
+  isdual,
+  multiplicity_range,
+  nondual_sector,
+  sector_type,
+  sectorunitrange,
+  space_isequal
+using SymmetrySectors: AbstractSector, SU, quantum_dimension
+
+Base.length(s::AbstractSector) = quantum_dimension(s)
+
+@testset "NewGradedUnitRange" begin
+  sr1 = sectorunitrange(SU((1, 0)), 1:2)
+  g1 = mortar([sr1])
+  @test g1 isa NewGradedUnitRange
+  @test blocklabels(g1) == [SU((1, 0))]
+  @test blockisequal(multiplicity_range(g1), blockedrange([2]))
+  @test !isdual(g1)
+
+  @test length(g1) == 6
+  @test blocklength(g1) == 1
+  @test blocklengths(g1) == [6]
+  @test space_isequal(only(blocks(g1)), sr1)
+
+  sr2 = sectorunitrange(SU((2, 1)), 3:3, false, 6)
+  g2 = mortar([sr1, sr2])
+  @test g2 isa NewGradedUnitRange
+  @test blocklabels(g2) == [SU((1, 0)), SU((2, 1))]
+  @test blockisequal(multiplicity_range(g2), blockedrange([2, 1]))
+  @test !isdual(g2)
+  @test space_isequal(g2, g2)
+  @test !space_isequal(g1, g2)
+
+  @test length(g2) == 14
+  @test blocklength(g2) == 2
+  @test blocklengths(g2) == [6, 8]
+  @test all(space_isequal.(blocks(g2), [sr1, sr2]))
+  @test space_isequal(g2[Block(1)], sr1)
+  @test space_isequal(g2[Block(2)], sr2)
+
+  g2b = dual(g2)
+  @test g2b isa NewGradedUnitRange
+  @test blocklabels(g2b) == [SU((1, 1)), SU((2, 1))]
+  @test blockisequal(multiplicity_range(g2b), blockedrange([2, 1]))
+  @test isdual(g2b)
+
+  g2f = flip(g2)
+  @test g2f isa NewGradedUnitRange
+  @test blocklabels(g2f) == [SU((1, 0)), SU((2, 1))]
+  @test blockisequal(multiplicity_range(g2f), blockedrange([2, 1]))
+  @test isdual(g2f)
+
+  @test blockfirsts(g2) == first.(blocks(g2))
+  @test blocklasts(g2) == [6, 14]
+  @test space_isequal(mortar(blocks(g2)), g2)
+
+  # slicing
+  @test space_isequal(g2[Block.(1:2)], g2)
+  @test space_isequal(g2[Block.(1:2)], g2)
+
+  @test space_isequal(g2[(Block(1), 1)], sectorunitrange(SU((1, 0)), 1:1, false, 1))
+  @test space_isequal(g2[(Block(1), 2)], sectorunitrange(SU((1, 0)), 2:2, false, 1))
+  @test space_isequal(g2[(Block(1), 1:2)], sectorunitrange(SU((1, 0)), 1:2, false, 1))
+  @test_throws BoundsError g2[(Block(1), 3)]
+
+  @test space_isequal(g2[(Block(2), 1)], sectorunitrange(SU((2, 1)), 3:3, false, 6))
+  @test space_isequal(g2[(Block(2), 1:1)], sectorunitrange(SU((2, 1)), 3:3, false, 6))
+  @test_throws BoundsError g2[(Block(2), 2)]
+  @test_throws BoundsError g2[(Block(2), 3:3)]  # misleading?
+end