[NDTensors] Introduce LabelledNumbers and GradedAxesNext

NDTensors/src/imports.jl

@@ -36,6 +36,8 @@ for lib in [
   :BroadcastMapConversion,
   :RankFactorization,
   :Sectors,
+  :LabelledNumbers,
+  :GradedAxesNext,
   :GradedAxes,
   :TensorAlgebra,
   :SparseArrayInterface,

NDTensors/src/lib/GradedAxesNext/.JuliaFormatter.toml

@@ -0,0 +1,2 @@
+style = "blue"
+indent = 2

NDTensors/src/lib/GradedAxesNext/src/GradedAxesNext.jl

@@ -0,0 +1,3 @@
+module GradedAxesNext
+include("gradedunitrange.jl")
+end

NDTensors/src/lib/GradedAxesNext/src/gradedunitrange.jl

@@ -0,0 +1,245 @@
+using BlockArrays:
+  BlockArrays,
+  Block,
+  BlockedUnitRange,
+  BlockRange,
+  BlockVector,
+  blockedrange,
+  BlockIndexRange,
+  blockfirsts,
+  blocklasts,
+  blocklengths,
+  findblock,
+  findblockindex,
+  mortar
+using ..LabelledNumbers: LabelledNumbers, LabelledInteger, label, labelled, unlabel
+
+# Custom `BlockedUnitRange` constructor that takes a unit range
+# and a set of block lengths, similar to `BlockArray(::AbstractArray, blocklengths...)`.
+function blockedunitrange(a::AbstractUnitRange, blocklengths)
+  blocklengths_shifted = copy(blocklengths)
+  blocklengths_shifted[1] += (first(a) - 1)
+  blocklasts = cumsum(blocklengths_shifted)
+  return BlockArrays._BlockedUnitRange(first(a), blocklasts)
+end
+
+# Circumvents issue in `findblock` that assumes the `BlockedUnitRange`
+# starts at 1.
+# TODO: Raise an issue with `BlockArrays`.
+function blockedunitrange_findblock(a::BlockedUnitRange, index::Integer)
+  @boundscheck index in 1:length(a) || throw(BoundsError(a, index))
+  return @inbounds findblock(a, index + first(a) - 1)
+end
+
+# Circumvents issue in `findblockindex` that assumes the `BlockedUnitRange`
+# starts at 1.
+# TODO: Raise an issue with `BlockArrays`.
+function blockedunitrange_findblockindex(a::BlockedUnitRange, index::Integer)
+  @boundscheck index in 1:length(a) || throw(BoundsError())
+  return @inbounds findblockindex(a, index + first(a) - 1)
+end
+
+const GradedUnitRange{BlockLasts<:Vector{<:LabelledInteger}} = BlockedUnitRange{BlockLasts}
+
+function gradedrange(lblocklengths::AbstractVector{<:LabelledInteger})
+  brange = blockedrange(unlabel.(lblocklengths))
+  lblocklasts = labelled.(blocklasts(brange), label.(lblocklengths))
+  # TODO: `first` is forced to be `Int` in `BlockArrays.BlockedUnitRange`,
+  # so this doesn't do anything right now. Make a PR to generalize it.
+  firstlength = first(lblocklengths)
+  lfirst = oneunit(firstlength)
+  return BlockArrays._BlockedUnitRange(lfirst, lblocklasts)
+end
+
+Base.last(a::GradedUnitRange) = isempty(a.lasts) ? first(a) - 1 : last(a.lasts)
+
+function gradedrange(lblocklengths::AbstractVector{<:Pair{<:Any,<:Integer}})
+  return gradedrange(labelled.(last.(lblocklengths), first.(lblocklengths)))
+end
+
+function labelled_blocks(a::BlockedUnitRange, labels)
+  return BlockArrays._BlockedUnitRange(a.first, labelled.(a.lasts, labels))
+end
+
+function BlockArrays.findblock(a::GradedUnitRange, index::Integer)
+  return blockedunitrange_findblock(unlabel_blocks(a), index)
+end
+
+function blockedunitrange_findblock(a::GradedUnitRange, index::Integer)
+  return blockedunitrange_findblock(unlabel_blocks(a), index)
+end
+
+function blockedunitrange_findblockindex(a::GradedUnitRange, index::Integer)
+  return blockedunitrange_findblockindex(unlabel_blocks(a), index)
+end
+
+function BlockArrays.findblockindex(a::GradedUnitRange, index::Integer)
+  return blockedunitrange_findblockindex(unlabel_blocks(a), index)
+end
+
+## Block label interface
+
+# Internal function
+function get_label(a::BlockedUnitRange, index::Block{1})
+  return label(blocklasts(a)[Int(index)])
+end
+
+# Internal function
+function get_label(a::BlockedUnitRange, index::Integer)
+  return get_label(a, blockedunitrange_findblock(a, index))
+end
+
+function blocklabels(a::BlockVector)
+  return map(BlockRange(a)) do block
+    return label(@view(a[block]))
+  end
+end
+
+function blocklabels(a::BlockedUnitRange)
+  # Using `a.lasts` here since that is what is stored
+  # inside of `BlockedUnitRange`, maybe change that.
+  # For example, it could be something like:
+  #
+  # map(BlockRange(a)) do block
+  #   return label(@view(a[block]))
+  # end
+  #
+  return label.(a.lasts)
+end
+
+# TODO: This relies on internals of `BlockArrays`, maybe redesign
+# to try to avoid that.
+# TODO: Define `set_grades`, `set_sector_labels`, `set_labels`.
+function unlabel_blocks(a::BlockedUnitRange)
+  return BlockArrays._BlockedUnitRange(a.first, unlabel.(a.lasts))
+end
+
+## BlockedUnitRage interface
+
+function Base.axes(ga::GradedUnitRange)
+  return map(axes(unlabel_blocks(ga))) do a
+    return labelled_blocks(a, blocklabels(ga))
+  end
+end
+
+function BlockArrays.blockfirsts(a::GradedUnitRange)
+  return labelled.(blockfirsts(unlabel_blocks(a)), blocklabels(a))
+end
+
+function BlockArrays.blocklasts(a::GradedUnitRange)
+  return labelled.(blocklasts(unlabel_blocks(a)), blocklabels(a))
+end
+
+function BlockArrays.blocklengths(a::GradedUnitRange)
+  return labelled.(blocklengths(unlabel_blocks(a)), blocklabels(a))
+end
+
+function Base.first(a::GradedUnitRange)
+  return labelled(first(unlabel_blocks(a)), label(a[Block(1)]))
+end
+
+function firstblockindices(a::GradedUnitRange)
+  return labelled.(firstblockindices(unlabel_blocks(a)), blocklabels(a))
+end
+
+function blockedunitrange_getindex(a::GradedUnitRange, index)
+  # This uses `blocklasts` since that is what is stored
+  # in `BlockedUnitRange`, maybe abstract that away.
+  return labelled(unlabel_blocks(a)[index], get_label(a, index))
+end
+
+# Like `a[indices]` but preserves block structure.
+using BlockArrays: block, blockindex
+function blockedunitrange_getindices(
+  a::BlockedUnitRange, indices::AbstractUnitRange{<:Integer}
+)
+  first_blockindex = blockedunitrange_findblockindex(a, first(indices))
+  last_blockindex = blockedunitrange_findblockindex(a, last(indices))
+  first_block = block(first_blockindex)
+  last_block = block(last_blockindex)
+  blocklengths = if first_block == last_block
+    [length(indices)]
+  else
+    map(first_block:last_block) do block
+      if block == first_block
+        return length(a[first_block]) - blockindex(first_blockindex) + 1
+      end
+      if block == last_block
+        return blockindex(last_blockindex)
+      end
+      return length(a[block])
+    end
+  end
+  return blockedunitrange(indices .+ (first(a) - 1), blocklengths)
+end
+
+function blockedunitrange_getindices(a::BlockedUnitRange, indices::BlockIndexRange)
+  return a[block(indices)][only(indices.indices)]
+end
+
+function blockedunitrange_getindices(a::BlockedUnitRange, indices::Vector{<:Integer})
+  return map(index -> a[index], indices)
+end
+
+function blockedunitrange_getindices(
+  a::BlockedUnitRange, indices::Vector{<:Union{Block{1},BlockIndexRange{1}}}
+)
+  return mortar(map(index -> a[index], indices))
+end
+
+function blockedunitrange_getindices(a::BlockedUnitRange, indices)
+  return error("Not implemented.")
+end
+
+# The blocks of the corresponding slice.
+_blocks(a::AbstractUnitRange, indices) = error("Not implemented")
+function _blocks(a::AbstractUnitRange, indices::AbstractUnitRange)
+  return findblock(a, first(indices)):findblock(a, last(indices))
+end
+function _blocks(a::AbstractUnitRange, indices::BlockRange)
+  return indices
+end
+
+# The block labels of the corresponding slice.
+function blocklabels(a::AbstractUnitRange, indices)
+  return map(_blocks(a, indices)) do block
+    return label(a[block])
+  end
+end
+
+function blockedunitrange_getindices(
+  ga::GradedUnitRange, indices::AbstractUnitRange{<:Integer}
+)
+  a_indices = blockedunitrange_getindices(unlabel_blocks(ga), indices)
+  return labelled_blocks(a_indices, blocklabels(ga, indices))
+end
+
+function blockedunitrange_getindices(ga::GradedUnitRange, indices::BlockRange)
+  return labelled_blocks(unlabel_blocks(ga)[indices], blocklabels(ga, indices))
+end
+
+function Base.getindex(a::GradedUnitRange, index::Integer)
+  return blockedunitrange_getindex(a, index)
+end
+
+function Base.getindex(a::GradedUnitRange, index::Block{1})
+  return blockedunitrange_getindex(a, index)
+end
+
+function Base.getindex(a::GradedUnitRange, indices::BlockIndexRange)
+  return blockedunitrange_getindices(a, indices)
+end
+
+function Base.getindex(
+  a::GradedUnitRange, indices::BlockRange{1,<:Tuple{AbstractUnitRange{Int}}}
+)
+  return blockedunitrange_getindices(a, indices)
+end
+
+function Base.getindex(a::GradedUnitRange, indices)
+  return blockedunitrange_getindices(a, indices)
+end
+
+function Base.getindex(a::GradedUnitRange, indices::AbstractUnitRange{<:Integer})
+  return blockedunitrange_getindices(a, indices)
+end

NDTensors/src/lib/GradedAxesNext/test/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

NDTensors/src/lib/GradedAxesNext/test/runtests.jl

@@ -0,0 +1,107 @@
+@eval module $(gensym())
+using BlockArrays:
+  Block, BlockVector, blockedrange, blockfirsts, blocklasts, blocklength, blocklengths
+using NDTensors.GradedAxesNext: GradedUnitRange, blocklabels, gradedrange
+using NDTensors.LabelledNumbers: LabelledUnitRange, label, unlabel
+using Test: @test, @test_broken, @testset
+@testset "GradedAxes" begin
+  a = gradedrange(["x" => 2, "y" => 3])
+  @test a isa GradedUnitRange
+  @test length(a) == 5
+  @test a[Block(2)] == 3:5
+  @test label(a[Block(2)]) == "y"
+  @test a[Block(2)] isa LabelledUnitRange
+  @test a[4] == 4
+  @test label(a[4]) == "y"
+  @test unlabel(a[4]) == 4
+  @test blocklengths(a) == [2, 3]
+  @test blocklabels(a) == ["x", "y"]
+  @test label.(blocklengths(a)) == ["x", "y"]
+  @test blockfirsts(a) == [1, 3]
+  @test label.(blockfirsts(a)) == ["x", "y"]
+  @test first(a) == 1
+  @test label(first(a)) == "x"
+  @test blocklasts(a) == [2, 5]
+  @test label.(blocklasts(a)) == ["x", "y"]
+  @test last(a) == 5
+  @test label(last(a)) == "y"
+  @test a[Block(2)] == 3:5
+  @test label(a[Block(2)]) == "y"
+  @test length(a[Block(2)]) == 3
+  @test blocklengths(only(axes(a))) == blocklengths(a)
+  @test blocklabels(only(axes(a))) == blocklabels(a)
+
+  # Slicing operations
+  x = gradedrange(["x" => 2, "y" => 3])
+  a = x[2:4]
+  @test a isa GradedUnitRange
+  @test length(a) == 3
+  @test blocklength(a) == 2
+  @test a[Block(1)] == 2:2
+  @test label(a[Block(1)]) == "x"
+  @test a[Block(2)] == 3:4
+  @test label(a[Block(2)]) == "y"
+  @test isone(first(only(axes(a))))
+  @test length(only(axes(a))) == length(a)
+  @test blocklengths(only(axes(a))) == blocklengths(a)
+
+  x = gradedrange(["x" => 2, "y" => 3])
+  a = x[3:4]
+  @test a isa GradedUnitRange
+  @test length(a) == 2
+  @test blocklength(a) == 1
+  @test a[Block(1)] == 3:4
+  @test label(a[Block(1)]) == "y"
+
+  x = gradedrange(["x" => 2, "y" => 3])
+  a = x[2:4][1:2]
+  @test a isa GradedUnitRange
+  @test length(a) == 2
+  @test blocklength(a) == 2
+  @test a[Block(1)] == 2:2
+  @test label(a[Block(1)]) == "x"
+  @test a[Block(2)] == 3:3
+  @test label(a[Block(2)]) == "y"
+
+  x = gradedrange(["x" => 2, "y" => 3])
+  a = x[Block(2)[2:3]]
+  @test a isa LabelledUnitRange
+  @test length(a) == 2
+  @test a == 4:5
+  @test label(a) == "y"
+
+  x = gradedrange(["x" => 2, "y" => 3, "z" => 4])
+  a = x[Block(2):Block(3)]
+  @test a isa GradedUnitRange
+  @test length(a) == 7
+  @test blocklength(a) == 2
+  @test blocklengths(a) == [3, 4]
+  @test blocklabels(a) == ["y", "z"]
+  @test a[Block(1)] == 3:5
+  @test a[Block(2)] == 6:9
+
+  x = gradedrange(["x" => 2, "y" => 3, "z" => 4])
+  a = x[[Block(3), Block(2)]]
+  @test a isa BlockVector
+  @test length(a) == 7
+  @test blocklength(a) == 2
+  # TODO: `BlockArrays` doesn't define `blocklengths`
+  # for `BlockVector`, should it?
+  @test_broken blocklengths(a) == [4, 3]
+  @test blocklabels(a) == ["z", "y"]
+  @test a[Block(1)] == 6:9
+  @test a[Block(2)] == 3:5
+
+  x = gradedrange(["x" => 2, "y" => 3, "z" => 4])
+  a = x[[Block(3)[2:3], Block(2)[2:3]]]
+  @test a isa BlockVector
+  @test length(a) == 4
+  @test blocklength(a) == 2
+  # TODO: `BlockArrays` doesn't define `blocklengths`
+  # for `BlockVector`, should it?
+  @test_broken blocklengths(a) == [2, 2]
+  @test blocklabels(a) == ["z", "y"]
+  @test a[Block(1)] == 7:8
+  @test a[Block(2)] == 4:5
+end
+end

NDTensors/src/lib/LabelledNumbers/.JuliaFormatter.toml

@@ -0,0 +1,2 @@
+style = "blue"
+indent = 2

NDTensors/src/lib/LabelledNumbers/src/LabelledNumbers.jl

@@ -0,0 +1,7 @@
+module LabelledNumbers
+include("labelled_interface.jl")
+include("labellednumber.jl")
+include("labelledinteger.jl")
+include("labelledarray.jl")
+include("labelledunitrange.jl")
+end