Introduce SectorVector (#48)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.4.14"
+version = "0.4.15"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

src/sectorunitrange.jl

@@ -1,5 +1,25 @@
 # This files defines SectorUnitRange, a unit range associated with a sector and an arrow
 
+struct SectorVector{T,Sector,Values<:AbstractVector{T}} <: AbstractVector{T}
+  sector::Sector
+  values::Values
+  isdual::Bool
+end
+
+sector(sv::SectorVector) = sv.sector
+ungrade(sv::SectorVector) = sv.values
+isdual(sv::SectorVector) = sv.isdual
+
+function sectorvector(s, v::AbstractVector, b::Bool=false)
+  return SectorVector(to_sector(s), v, b)
+end
+Base.length(sv::SectorVector) = length(sv.values)
+Base.size(sv::SectorVector) = (length(sv),)
+function Base.axes(sv::SectorVector)
+  (sectorrange(sector(sv), only(axes(ungrade(sv))), isdual(sv)),)
+end
+Base.getindex(sv::SectorVector, i::Integer) = ungrade(sv)[i]
+
 # =====================================  Definition  =======================================
 
 # This implementation contains the "full range"
@@ -51,21 +71,34 @@ Base.iterate(sr::SectorUnitRange) = iterate(ungrade(sr))
 Base.iterate(sr::SectorUnitRange, i::Integer) = iterate(ungrade(sr), i)
 
 Base.length(sr::SectorUnitRange) = length(ungrade(sr))
+Base.size(sr::SectorUnitRange) = (length(sr),)
+function Base.axes(sr::SectorUnitRange)
+  (sectorrange(sector(sr), only(axes(ungrade(sr))), isdual(sr)),)
+end
 
 Base.last(sr::SectorUnitRange) = last(ungrade(sr))
 
 # slicing
 Base.getindex(sr::SectorUnitRange, i::Integer) = ungrade(sr)[i]
 
-function Base.getindex(sr::SectorUnitRange, r::AbstractUnitRange{T}) where {T<:Integer}
-  return sr[SymmetryStyle(sr), r]
+function Base.getindex(sr::SectorUnitRange, I::AbstractVector{<:Integer})
+  return sr[SymmetryStyle(sr), I]
 end
-function Base.getindex(sr::SectorUnitRange, ::NotAbelianStyle, r::AbstractUnitRange)
+function Base.getindex(sr::SectorUnitRange, I::AbstractUnitRange{<:Integer})
+  return sr[SymmetryStyle(sr), I]
+end
+function Base.getindex(sr::SectorUnitRange, ::NotAbelianStyle, r::AbstractVector{<:Integer})
   return ungrade(sr)[r]
 end
-function Base.getindex(sr::SectorUnitRange, ::AbelianStyle, r::AbstractUnitRange)
+function Base.getindex(sr::SectorUnitRange, ::AbelianStyle, r::AbstractUnitRange{<:Integer})
   return sectorrange(sector(sr), ungrade(sr)[ungrade(r)], isdual(sr))
 end
+function Base.getindex(sr::SectorUnitRange, ::AbelianStyle, r::AbstractVector{<:Integer})
+  return sectorvector(sector(sr), ungrade(sr)[ungrade(r)], isdual(sr))
+end
+function Base.getindex(sr::SectorUnitRange, I::AbstractVector{Bool})
+  return sr[to_indices(sr, (I,))...]
+end
 
 # TODO replace (:,x) indexing with kronecker(:, x)
 Base.getindex(sr::SectorUnitRange, t::Tuple{Colon,<:Integer}) = sr[(:, last(t):last(t))]
@@ -114,3 +147,7 @@ sector_type(::Type{<:SectorUnitRange{T,Sector}}) where {T,Sector} = Sector
 # TBD error for non-integer?
 sector_multiplicity(sr::SectorUnitRange) = length(sr) ÷ length(sector(sr))
 sector_multiplicities(sr::SectorUnitRange) = [sector_multiplicity(sr)]  # TBD remove?
+
+function Base.similar(A::Type{<:AbstractArray}, ax::Tuple{SectorOneTo,Vararg{SectorOneTo}})
+  return similar(A, ungrade.(ax))
+end

test/test_sectorunitrange.jl

@@ -1,5 +1,3 @@
-using Test: @test, @test_throws, @testset
-
 using BlockArrays:
   Block,
   BlockBoundsError,
@@ -11,12 +9,12 @@ using BlockArrays:
   blockisequal,
   blocks,
   findblock
-using TestExtras: @constinferred
-
 using GradedArrays:
   U1,
   SU,
+  SectorOneTo,
   SectorUnitRange,
+  SectorVector,
   dual,
   flip,
   isdual,
@@ -29,6 +27,8 @@ using GradedArrays:
   sectors,
   space_isequal,
   ungrade
+using Test: @test, @test_throws, @testset
+using TestExtras: @constinferred
 
 @testset "SectorUnitRange" begin
   sr = sectorrange(SU((1, 0)), 2)
@@ -49,7 +49,8 @@ using GradedArrays:
   @test eltype(sr) === Int
   @test step(sr) == 1
   @test eachindex(sr) == Base.oneto(6)
-  @test only(axes(sr)) isa Base.OneTo
+  @test only(axes(sr)) isa SectorOneTo
+  @test sector(only(axes(sr))) == sector(sr)
   @test only(axes(sr)) == 1:6
   @test iterate(sr) == (1, 1)
   for i in 1:5
@@ -148,6 +149,11 @@ using GradedArrays:
   @test (@constinferred getindex(srab, 2:2)) isa SectorUnitRange
   @test space_isequal(srab[2:2], sectorrange(U1(1), 2:2))
   @test space_isequal(dual(srab)[2:2], sectorrange(U1(1), 2:2, true))
+  @test srab[[1, 3]] isa SectorVector{Int}
+  @test sector(srab[[1, 3]]) == sector(srab)
+  @test ungrade(srab[[1, 3]]) == [1, 3]
+  @test length(srab[[1, 3]]) == 2
+  @test space_isequal(only(axes(srab[[1, 3]])), sectorrange(U1(1), 2))
 
   # Slice sector range with sector range
   sr1 = sectorrange(U1(1), 4)

test/test_tensor_product.jl

@@ -18,30 +18,46 @@ using TensorProducts: ⊗, OneToOne, tensor_product
 using Test: @test, @testset
 using TestExtras: @constinferred
 
-GradedArrays.SymmetryStyle(::Type{<:String}) = NotAbelianStyle()
-GradedArrays.tensor_product(s1::String, s2::String) = gradedrange([s1 * s2 => 1])
+struct NotAbelianString
+  str::String
+end
+Base.:*(s1::NotAbelianString, s2::NotAbelianString) = NotAbelianString(s1.str * s2.str)
+GradedArrays.SymmetryStyle(::Type{<:NotAbelianString}) = NotAbelianStyle()
+function GradedArrays.tensor_product(s1::NotAbelianString, s2::NotAbelianString)
+  gradedrange([s1 * s2 => 1])
+end
+Base.length(s::NotAbelianString) = length(s.str)
 
 @testset "unmerged_tensor_product" begin
   @test unmerged_tensor_product() isa OneToOne
   @test unmerged_tensor_product(OneToOne(), OneToOne()) isa OneToOne
   @test unmerged_tensor_product(1:1, 1:1) == 1:1
   @test sectormergesort(1:1) isa UnitRange
 
-  a = gradedrange(["x" => 2, "y" => 3])
+  a = gradedrange([NotAbelianString("x") => 2, NotAbelianString("y") => 3])
   @test space_isequal(unmerged_tensor_product(a), a)
 
   b = unmerged_tensor_product(a, a)
   @test b isa GradedOneTo
   @test length(b) == 50
   @test blocklength(b) == 4
   @test blocklengths(b) == [8, 12, 12, 18]
-  @test space_isequal(b, gradedrange(["xx" => 4, "yx" => 6, "xy" => 6, "yy" => 9]))
+  @test space_isequal(
+    b,
+    gradedrange([
+      NotAbelianString("xx") => 4,
+      NotAbelianString("yx") => 6,
+      NotAbelianString("xy") => 6,
+      NotAbelianString("yy") => 9,
+    ]),
+  )
 
   c = unmerged_tensor_product(a, a, a)
   @test c isa GradedOneTo
   @test length(c) == 375
   @test blocklength(c) == 8
-  @test sectors(c) == ["xxx", "yxx", "xyx", "yyx", "xxy", "yxy", "xyy", "yyy"]
+  @test sectors(c) ==
+    NotAbelianString.(["xxx", "yxx", "xyx", "yyx", "xxy", "yxy", "xyy", "yyy"])
 
   a = gradedrange([U1(1) => 1, U1(2) => 3, U1(1) => 1])
   @test space_isequal(