[BlockSparseArrays] Fix eachindex(::BlockSparseArray) with dual axes

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.15"
+version = "0.3.16"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

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

@@ -1,8 +1,15 @@
 module BlockSparseArraysGradedAxesExt
-using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange
-using ..BlockSparseArrays: BlockSparseArrays, block_merge
+using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange, blocks
+using ..BlockSparseArrays:
+  BlockSparseArrays, AbstractBlockSparseArray, BlockSparseArray, block_merge
 using ...GradedAxes:
-  GradedUnitRange, OneToOne, blockmergesortperm, blocksortperm, invblockperm, tensor_product
+  GradedUnitRange,
+  OneToOne,
+  blockmergesortperm,
+  blocksortperm,
+  invblockperm,
+  nondual,
+  tensor_product
 using ...TensorAlgebra:
   TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims
 
@@ -45,4 +52,28 @@ function TensorAlgebra.splitdims(
   a_blockpermed = a[blockperms...]
   return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...)
 end
+
+# This is a temporary fix for `eachindex` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.eachindex(a::AbstractBlockSparseArray)
+  return CartesianIndices(nondual.(axes(a)))
+end
+
+# This is a temporary fix for `show` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.show(io::IO, mime::MIME"text/plain", a::BlockSparseArray; kwargs...)
+  a_nondual = BlockSparseArray(blocks(a), nondual.(axes(a)))
+  println(io, "typeof(axes) = ", typeof(axes(a)), "\n")
+  println(
+    io,
+    "Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.\n",
+  )
+  return invoke(
+    show, Tuple{IO,MIME"text/plain",AbstractArray}, io, mime, a_nondual; kwargs...
+  )
+end
 end

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

@@ -3,9 +3,8 @@ using Compat: Returns
 using Test: @test, @testset, @test_broken
 using BlockArrays: Block, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
-using NDTensors.GradedAxes: GradedUnitRange, gradedrange
+using NDTensors.GradedAxes: GradedAxes, GradedUnitRange, dual, gradedrange
 using NDTensors.LabelledNumbers: label
-using NDTensors.Sectors: U1
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
 using Random: randn!
@@ -16,6 +15,14 @@ function blockdiagonal!(f, a::AbstractArray)
   end
   return a
 end
+
+struct U1
+  n::Int
+end
+GradedAxes.dual(c::U1) = U1(-c.n)
+GradedAxes.fuse_labels(c1::U1, c2::U1) = U1(c1.n + c2.n)
+Base.isless(c1::U1, c2::U1) = isless(c1.n, c2.n)
+
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "BlockSparseArraysGradedAxesExt (eltype=$elt)" for elt in elts
   @testset "map" begin
@@ -70,5 +77,18 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test_broken 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])
+    a = BlockSparseArray{elt}(dual(r), r)
+    a[Block(1, 1)] = randn(size(a[Block(1, 1)]))
+    a[Block(2, 2)] = randn(size(a[Block(2, 2)]))
+    a_dense = Array(a)
+    @test eachindex(a) == CartesianIndices(size(a))
+    for I in eachindex(a)
+      @test a[I] == a_dense[I]
+    end
+
+    @test isnothing(show(devnull, MIME("text/plain"), a))
+  end
 end
 end