fix slicing with BlockVector

Author: ogauthe

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

@@ -112,14 +112,14 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test blocksize(m) == (3, 3)
     @test a == splitdims(m, (d1, d2), (d1, d2))
   end
+
   @testset "dual axes" begin
     r = gradedrange([U1(0) => 2, U1(1) => 2])
     for ax in ((r, r), (dual(r), r), (r, dual(r)), (dual(r), dual(r)))
       a = BlockSparseArray{elt}(ax...)
       @views for b in [Block(1, 1), Block(2, 2)]
         a[b] = randn(elt, size(a[b]))
       end
-      # TODO: Define and use `isdual` here.
       for dim in 1:ndims(a)
         @test typeof(ax[dim]) === typeof(axes(a, dim))
         @test isdual(ax[dim]) == isdual(axes(a, dim))
@@ -176,7 +176,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test a[I, :] isa AbstractBlockArray
       @test a[:, I] isa AbstractBlockArray
       @test size(a[I, I]) == (1, 1)
-      @test !GradedAxes.isdual(axes(a[I, I], 1))
+      @test !isdual(axes(a[I, I], 1))
     end
 
     @testset "GradedUnitRange" begin
@@ -190,14 +190,19 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test Array(b) == 2 * Array(a)
       for i in 1:2
         @test axes(b, i) isa GradedUnitRange
-        @test_broken axes(a[:, :], i) isa GradedUnitRange
+        @test axes(a[:, :], i) isa GradedUnitRange
       end
 
       I = [Block(1)[1:1]]
-      @test_broken a[I, :] isa AbstractBlockArray
-      @test_broken a[:, I] isa AbstractBlockArray
+      @test a[I, :] isa AbstractBlockArray
+      @test axes(a[I, :], 1) isa GradedOneTo
+      @test axes(a[I, :], 2) isa GradedUnitRange
+
+      @test a[:, I] isa AbstractBlockArray
+      @test axes(a[:, I], 2) isa GradedOneTo
+      @test axes(a[:, I], 1) isa GradedUnitRange
       @test size(a[I, I]) == (1, 1)
-      @test_broken GradedAxes.isdual(axes(a[I, I], 1))
+      @test !isdual(axes(a[I, I], 1))
     end
 
     # Test case when all axes are dual.
@@ -212,13 +217,18 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test Array(b) == 2 * Array(a)
       for i in 1:2
         @test axes(b, i) isa GradedUnitRangeDual
-        @test_broken axes(a[:, :], i) isa GradedUnitRangeDual
+        @test axes(a[:, :], i) isa GradedUnitRangeDual
       end
       I = [Block(1)[1:1]]
-      @test_broken a[I, :] isa AbstractBlockArray
-      @test_broken a[:, I] isa AbstractBlockArray
+      @test a[I, :] isa AbstractBlockArray
+      @test a[:, I] isa AbstractBlockArray
       @test size(a[I, I]) == (1, 1)
-      @test_broken GradedAxes.isdual(axes(a[I, I], 1))
+      @test isdual(axes(a[I, :], 2))
+      @test isdual(axes(a[:, I], 1))
+      @test_broken isdual(axes(a[I, :], 1))
+      @test_broken isdual(axes(a[:, I], 2))
+      @test_broken isdual(axes(a[I, I], 1))
+      @test_broken isdual(axes(a[I, I], 2))
     end
 
     @testset "dual GradedUnitRange" begin
@@ -232,14 +242,19 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test Array(b) == 2 * Array(a)
       for i in 1:2
         @test axes(b, i) isa GradedUnitRangeDual
-        @test_broken axes(a[:, :], i) isa GradedUnitRangeDual
+        @test axes(a[:, :], i) isa GradedUnitRangeDual
       end
 
       I = [Block(1)[1:1]]
-      @test_broken a[I, :] isa AbstractBlockArray
-      @test_broken a[:, I] isa AbstractBlockArray
+      @test a[I, :] isa AbstractBlockArray
+      @test a[:, I] isa AbstractBlockArray
       @test size(a[I, I]) == (1, 1)
-      @test_broken GradedAxes.isdual(axes(a[I, I], 1))
+      @test isdual(axes(a[I, :], 2))
+      @test isdual(axes(a[:, I], 1))
+      @test_broken isdual(axes(a[I, :], 1))
+      @test_broken isdual(axes(a[:, I], 2))
+      @test_broken isdual(axes(a[I, I], 1))
+      @test_broken isdual(axes(a[I, I], 2))
     end
 
     @testset "dual BlockedUnitRange" begin    # self dual
@@ -261,7 +276,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test a[I, :] isa BlockSparseArray
       @test a[:, I] isa BlockSparseArray
       @test size(a[I, I]) == (1, 1)
-      @test !GradedAxes.isdual(axes(a[I, I], 1))
+      @test !isdual(axes(a[I, I], 1))
     end
 
     # Test case when all axes are dual from taking the adjoint.
@@ -281,8 +296,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       end
 
       I = [Block(1)[1:1]]
-      @test_broken size(b[I, :]) == (1, 4)
-      @test_broken size(b[:, I]) == (4, 1)
+      @test size(b[I, :]) == (1, 4)
+      @test size(b[:, I]) == (4, 1)
       @test size(b[I, I]) == (1, 1)
     end
   end

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

@@ -222,6 +222,13 @@ function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::Vector
   return map(index -> a[index], indices)
 end
 
+function blockedunitrange_getindices(
+  a::AbstractGradedUnitRange,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+
 function blockedunitrange_getindices(a::AbstractGradedUnitRange, index)
   return labelled(unlabel_blocks(a)[index], get_label(a, index))
 end
@@ -345,6 +352,9 @@ function BlockArrays.combine_blockaxes(a::GradedUnitRange, b::GradedUnitRange)
   return GradedUnitRange(new_first, new_blocklasts)
 end
 
+# preserve axes in SubArray
+Base.axes(S::Base.Slice{<:AbstractGradedUnitRange}) = (S.indices,)
+
 # Version of length that checks that all blocks have the same label
 # and returns a labelled length with that label.
 function labelled_length(a::AbstractBlockVector{<:Integer})

NDTensors/src/lib/GradedAxes/test/test_basics.jl

@@ -10,7 +10,8 @@ using BlockArrays:
   blocklength,
   blocklengths,
   blocks,
-  combine_blockaxes
+  combine_blockaxes,
+  mortar
 using NDTensors.GradedAxes: GradedOneTo, GradedUnitRange, OneToOne, blocklabels, gradedrange
 using NDTensors.LabelledNumbers:
   LabelledUnitRange, islabelled, label, labelled, labelled_isequal, unlabel
@@ -97,6 +98,7 @@ end
     @test blocklabels(only(axes(a))) == blocklabels(a)
 
     @test combine_blockaxes(a, a) isa GradedOneTo
+    @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
   end
 
   # Slicing operations
@@ -145,6 +147,7 @@ end
   @test length(ax) == length(a)
   @test blocklengths(ax) == blocklengths(a)
   @test blocklabels(ax) == blocklabels(a)
+  @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
 
   x = gradedrange(["x" => 2, "y" => 3])
   a = x[2:4][1:2]
@@ -227,5 +230,17 @@ end
   # once `blocklengths(::BlockVector)` is defined.
   @test blocklengths(ax) == [2, 2]
   @test blocklabels(ax) == blocklabels(a)
+
+  x = gradedrange(["x" => 2, "y" => 3])
+  I = mortar([Block(1)[1:1]])
+  a = x[I]
+  @test length(a) == 1
+  @test label(first(a)) == "x"
+
+  x = gradedrange(["x" => 2, "y" => 3])[1:5]
+  I = mortar([Block(1)[1:1]])
+  a = x[I]
+  @test length(a) == 1
+  @test label(first(a)) == "x"
 end
 end

NDTensors/src/lib/GradedAxes/test/test_dual.jl

@@ -10,7 +10,8 @@ using BlockArrays:
   blocklength,
   blocklengths,
   blocks,
-  findblock
+  findblock,
+  mortar
 using NDTensors.GradedAxes:
   AbstractGradedUnitRange,
   GradedAxes,
@@ -26,7 +27,7 @@ using NDTensors.GradedAxes:
   isdual,
   nondual
 using NDTensors.LabelledNumbers: LabelledInteger, label, labelled, labelled_isequal
-using Test: @test, @testset
+using Test: @test, @test_broken, @testset
 struct U1
   n::Int
 end
@@ -75,6 +76,7 @@ end
 
     @test isdual(ad)
     @test !isdual(a)
+    @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
 
     @test blockfirsts(ad) == [labelled(1, U1(0)), labelled(3, U1(-1))]
     @test blocklasts(ad) == [labelled(2, U1(0)), labelled(5, U1(-1))]
@@ -106,6 +108,12 @@ end
     @test blocklength(blockmergesortperm(ad)) == 2
     @test blockmergesortperm(a) == [Block(1), Block(2)]
     @test blockmergesortperm(ad) == [Block(1), Block(2)]
+
+    I = mortar([Block(2)[1:1]])
+    g = ad[I]
+    @test length(g) == 1
+    @test label(first(g)) == U1(-1)
+    @test_broken isdual(g[Block(1)])
   end
 end