[BlockSparseArrays] Redesign block views again (#1513)

* [BlockSparseArrays] Redesign block views again

* [NDTensors] Bump to v0.3.38

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
-version = "0.3.37"
+version = "0.3.38"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -396,6 +396,33 @@ function blocked_cartesianindices(axes::Tuple, subaxes::Tuple, blocks)
  end
 end
 
+# Represents a view of a block of a blocked array.
+struct BlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
+ array::Array
+ block::Tuple{Vararg{Block{1,Int},N}}
+end
+function Base.axes(a::BlockView)
+ # TODO: Try to avoid conversion to `Base.OneTo{Int}`, or just convert
+ # the element type to `Int` with `Int.(...)`.
+ # When the axes of `a.array` are `GradedOneTo`, the block is `LabelledUnitRange`,
+ # which has element type `LabelledInteger`. That causes conversion problems
+ # in some generic Base Julia code, for example when printing `BlockView`.
+ return ntuple(ndims(a)) do dim
+ return Base.OneTo{Int}(only(axes(axes(a.array, dim)[a.block[dim]])))
+ end
+end
+function Base.size(a::BlockView)
+ return length.(axes(a))
+end
+function Base.getindex(a::BlockView{<:Any,N}, index::Vararg{Int,N}) where {N}
+ return blocks(a.array)[Int.(a.block)...][index...]
+end
+function Base.setindex!(a::BlockView{<:Any,N}, value, index::Vararg{Int,N}) where {N}
+ blocks(a.array)[Int.(a.block)...] = blocks(a.array)[Int.(a.block)...]
+ blocks(a.array)[Int.(a.block)...][index...] = value
+ return a
+end
+
 function view!(a::BlockSparseArray{<:Any,N}, index::Block{N}) where {N}
  return view!(a, Tuple(index)...)
 end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -42,3 +42,18 @@ function Base.setindex!(
  blocksparse_setindex!(a, value, I...)
  return a
 end
+
+function Base.setindex!(
+ a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Block{1},N}
+) where {N}
+ blocksize = ntuple(dim -> length(axes(a, dim)[I[dim]]), N)
+ if size(value) ≠ blocksize
+ throw(
+ DimensionMismatch(
+ "Trying to set block $(Block(Int.(I)...)), which has a size $blocksize, with data of size $(size(value)).",
+ ),
+ )
+ end
+ blocks(a)[Int.(I)...] = value
+ return a
+end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl

@@ -1,4 +1,4 @@
-using BlockArrays: Block, BlockSlices
+using BlockArrays: BlockArrays, Block, BlockSlices, viewblock
 
 function blocksparse_view(a, I...)
  return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...)
@@ -22,3 +22,19 @@ function Base.view(
 )
  return blocksparse_view(a, I)
 end
+
+# Specialized code for getting the view of a block.
+function BlockArrays.viewblock(
+ a::AbstractBlockSparseArray{<:Any,N}, block::Block{N}
+) where {N}
+ return viewblock(a, Tuple(block)...)
+end
+function BlockArrays.viewblock(
+ a::AbstractBlockSparseArray{<:Any,N}, block::Vararg{Block{1},N}
+) where {N}
+ I = CartesianIndex(Int.(block))
+ if I ∈ stored_indices(blocks(a))
+ return blocks(a)[I]
+ end
+ return BlockView(a, block)
+end

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

@@ -17,7 +17,7 @@ using BlockArrays:
 using Compat: @compat
 using LinearAlgebra: mul!
 using NDTensors.BlockSparseArrays:
- @view!, BlockSparseArray, block_nstored, block_reshape, view!
+ @view!, BlockSparseArray, BlockView, block_nstored, block_reshape, view!
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: contract
 using Test: @test, @test_broken, @test_throws, @testset
@@ -362,10 +362,10 @@ include("TestBlockSparseArraysUtils.jl")
  b = @view a[Block(2, 2)]
  @test size(b) == (3, 4)
  for i in parentindices(b)
- @test i isa BlockSlice{<:Block{1}}
+ @test i isa Base.OneTo{Int}
  end
- @test parentindices(b)[1] == BlockSlice(Block(2), 3:5)
- @test parentindices(b)[2] == BlockSlice(Block(2), 4:7)
+ @test parentindices(b)[1] == 1:3
+ @test parentindices(b)[2] == 1:4
 
  a = BlockSparseArray{elt}([2, 3], [3, 4])
  b = @view a[Block(2, 2)[1:2, 2:2]]
@@ -392,9 +392,9 @@ include("TestBlockSparseArraysUtils.jl")
 
  a = BlockSparseArray{elt}([2, 3], [3, 4])
  b = @views a[Block(2, 2)][1:2, 2:3]
- @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
+ @test b isa SubArray{<:Any,<:Any,<:BlockView}
  for i in parentindices(b)
- @test i isa BlockSlice{<:BlockIndexRange{1}}
+ @test i isa UnitRange{Int}
  end
  x = randn(elt, 2, 2)
  b .= x

NDTensors/src/lib/GradedAxes/src/unitrangedual.jl

@@ -68,6 +68,21 @@ end
 using NDTensors.LabelledNumbers: LabelledNumbers, label
 LabelledNumbers.label(a::UnitRangeDual) = dual(label(nondual(a)))
 
+using NDTensors.LabelledNumbers: LabelledUnitRange
+# The Base version of `length(::AbstractUnitRange)` drops the label.
+function Base.length(a::UnitRangeDual{<:Any,<:LabelledUnitRange})
+ return dual(length(nondual(a)))
+end
+function Base.iterate(a::UnitRangeDual, i)
+ i == last(a) && return nothing
+ return dual.(iterate(nondual(a), i))
+end
+# TODO: Is this a good definition?
+Base.unitrange(a::UnitRangeDual{<:Any,<:AbstractUnitRange}) = a
+
+using NDTensors.LabelledNumbers: LabelledInteger, label, labelled, unlabel
+dual(i::LabelledInteger) = labelled(unlabel(i), dual(label(i)))
+
 using BlockArrays: BlockArrays, blockaxes, blocklasts, combine_blockaxes, findblock
 BlockArrays.blockaxes(a::UnitRangeDual) = blockaxes(nondual(a))
 BlockArrays.blockfirsts(a::UnitRangeDual) = label_dual.(blockfirsts(nondual(a)))

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

@@ -28,6 +28,9 @@ function Base.OrdinalRange{T,T}(a::LabelledUnitRange) where {T<:Integer}
  return OrdinalRange{T,T}(unlabel(a))
 end
 
+# TODO: Is this a good definition?
+Base.unitrange(a::LabelledUnitRange) = a
+
 for f in [:first, :getindex, :last, :length, :step]
  @eval Base.$f(a::LabelledUnitRange, args...) = labelled($f(unlabel(a), args...), label(a))
 end