diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 17a14ceb90..407445f8d7 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.3.23" +version = "0.3.24" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl index 43b8afab24..57f2eb2846 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -159,8 +159,8 @@ function blockrange(axis::AbstractUnitRange, r::UnitRange) end function blockrange(axis::AbstractUnitRange, r::Int) - error("Slicing with integer values isn't supported.") - return findblock(axis, r) + ## return findblock(axis, r) + return error("Slicing with integer values isn't supported.") end function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}}) @@ -187,6 +187,24 @@ function blockrange(axis::AbstractUnitRange, r::BlockIndexRange) return Block(r):Block(r) end +function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:BlockIndexRange{1}}) + return error("Slicing not implemented for range of type `$(typeof(r))`.") +end + +function blockrange( + axis::AbstractUnitRange, + r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, +) + return map(b -> Block(b), blocks(r)) +end + +# This handles slicing with `:`/`Colon()`. +function blockrange(axis::AbstractUnitRange, r::Base.Slice) + # TODO: Maybe use `BlockRange`, but that doesn't output + # the same thing. + return only(blockaxes(axis)) +end + function blockrange(axis::AbstractUnitRange, r) return error("Slicing not implemented for range of type `$(typeof(r))`.") end @@ -228,6 +246,22 @@ function blockindices(a::AbstractUnitRange, b::Block, r::BlockIndices) return blockindices(a, b, r.blocks) end +function blockindices( + a::AbstractUnitRange, + b::Block, + r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, +) + # TODO: Change to iterate over `BlockRange(r)` + # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404 + # is fixed. + for bl in blocks(r) + if b == Block(bl) + return only(bl.indices) + end + end + return error("Block not found.") +end + function cartesianindices(a::AbstractArray, b::Block) return cartesianindices(axes(a), b) end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 2a5350f3ec..5499c957eb 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -1,5 +1,6 @@ using Adapt: Adapt, WrappedArray -using BlockArrays: BlockArrays, BlockedUnitRange, BlockRange, blockedrange, unblock +using BlockArrays: + BlockArrays, BlockedUnitRange, BlockIndexRange, BlockRange, blockedrange, mortar, unblock using SplitApplyCombine: groupcount const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{ @@ -46,6 +47,15 @@ function Base.to_indices( return blocksparse_to_indices(a, I) end +# Handle case of indexing with `[Block(1)[1:2], Block(2)[1:2]]` +# by converting it to a `BlockVector` with +# `mortar([Block(1)[1:2], Block(2)[1:2]])`. +function Base.to_indices( + a::BlockSparseArrayLike, inds, I::Tuple{AbstractVector{<:BlockIndexRange{1}},Vararg{Any}} +) + return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) +end + # Fixes ambiguity error with BlockArrays. function Base.to_indices(a::BlockSparseArrayLike, I::Tuple{BlockRange{1},Vararg{Any}}) return blocksparse_to_indices(a, I) @@ -126,14 +136,25 @@ function Base.getindex( return blocksparse_getindex(a, block...) end -# TODO: Define `issasigned(a, ::Block{N})`. +# TODO: Define `blocksparse_isassigned`. function Base.isassigned( a::BlockSparseArrayLike{<:Any,N}, index::Vararg{Block{1},N} ) where {N} - # TODO: Define `blocksparse_isassigned`. return isassigned(blocks(a), Int.(index)...) end +function Base.isassigned(a::BlockSparseArrayLike{<:Any,N}, index::Block{N}) where {N} + return isassigned(a, Tuple(index)...) +end + +# TODO: Define `blocksparse_isassigned`. +function Base.isassigned( + a::BlockSparseArrayLike{<:Any,N}, index::Vararg{BlockIndex{1},N} +) where {N} + b = block.(index) + return isassigned(a, b...) && isassigned(@view(a[b...]), blockindex.(index)...) +end + function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::BlockIndex{N}) where {N} blocksparse_setindex!(a, value, I) return a diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index 7c4f998e65..00ff6b70df 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -1,6 +1,14 @@ @eval module $(gensym()) using BlockArrays: - Block, BlockRange, BlockedUnitRange, BlockVector, blockedrange, blocklength, blocksize + Block, + BlockRange, + BlockedUnitRange, + BlockVector, + blockedrange, + blocklength, + blocklengths, + blocksize, + mortar using LinearAlgebra: mul! using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored, block_reshape using NDTensors.SparseArrayInterface: nstored @@ -331,13 +339,69 @@ include("TestBlockSparseArraysUtils.jl") @test b[4, 4] == 44 end - # This is outputting only zero blocks. a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)]))) - a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[Block(2):Block(2), Block(1):Block(2)] @test block_nstored(b) == 1 @test b == Array(a)[3:5, 1:end] + + a = BlockSparseArray{elt}(undef, ([2, 3, 4], [2, 3, 4])) + # TODO: Define `block_diagindices`. + @views for b in [Block(1, 1), Block(2, 2), Block(3, 3)] + a[b] = randn(elt, size(a[b])) + end + for (I1, I2) in ( + (mortar([Block(2)[2:3], Block(3)[1:3]]), mortar([Block(2)[2:3], Block(3)[2:3]])), + ([Block(2)[2:3], Block(3)[1:3]], [Block(2)[2:3], Block(3)[2:3]]), + ) + for b in (a[I1, I2], @view(a[I1, I2])) + # TODO: Rename `block_stored_length`. + @test block_nstored(b) == 2 + @test b[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] + @test b[Block(2, 2)] == a[Block(3, 3)[1:3, 2:3]] + end + end + + a = BlockSparseArray{elt}(undef, ([3, 3], [3, 3])) + # TODO: Define `block_diagindices`. + @views for b in [Block(1, 1), Block(2, 2)] + a[b] = randn(elt, size(a[b])) + end + I = mortar([Block(1)[1:2], Block(2)[1:2]]) + b = a[:, I] + @test b[Block(1, 1)] == a[Block(1, 1)][:, 1:2] + @test b[Block(2, 1)] == a[Block(2, 1)][:, 1:2] + @test b[Block(1, 2)] == a[Block(1, 2)][:, 1:2] + @test b[Block(2, 2)] == a[Block(2, 2)][:, 1:2] + @test blocklengths.(axes(b)) == ([3, 3], [2, 2]) + # TODO: Rename `block_stored_length`. + @test blocksize(b) == (2, 2) + @test block_nstored(b) == 2 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + @test isassigned(a, 1, 1) + @test isassigned(a, 5, 7) + @test !isassigned(a, 0, 1) + @test !isassigned(a, 5, 8) + @test isassigned(a, Block(1), Block(1)) + @test isassigned(a, Block(2), Block(2)) + @test !isassigned(a, Block(1), Block(0)) + @test !isassigned(a, Block(3), Block(2)) + @test isassigned(a, Block(1, 1)) + @test isassigned(a, Block(2, 2)) + @test !isassigned(a, Block(1, 0)) + @test !isassigned(a, Block(3, 2)) + @test isassigned(a, Block(1)[1], Block(1)[1]) + @test isassigned(a, Block(2)[3], Block(2)[4]) + @test !isassigned(a, Block(1)[0], Block(1)[1]) + @test !isassigned(a, Block(2)[3], Block(2)[5]) + @test !isassigned(a, Block(1)[1], Block(0)[1]) + @test !isassigned(a, Block(3)[3], Block(2)[4]) end @testset "LinearAlgebra" begin a1 = BlockSparseArray{elt}([2, 3], [2, 3]) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl index 6f755dd405..6ab1f77a88 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl @@ -160,16 +160,11 @@ function sparse_setindex!(a::AbstractArray, value, I::NotStoredIndex) end # isassigned -function sparse_isassigned(a::AbstractArray, I::Integer...) - return sparse_isassigned(a, CartesianIndex(I)) +function sparse_isassigned(a::AbstractArray{<:Any,N}, I::CartesianIndex{N}) where {N} + return sparse_isassigned(a, Tuple(I)...) end -sparse_isassigned(a::AbstractArray, I::NotStoredIndex) = true -sparse_isassigned(a::AbstractArray, I::StoredIndex) = sparse_isassigned(a, StorageIndex(I)) -function sparse_isassigned(a::AbstractArray, I::StorageIndex) - return isassigned(sparse_storage(a), index(I)) -end -function sparse_isassigned(a::AbstractArray, I::CartesianIndex) - return sparse_isassigned(a, storage_index(a, I)) +function sparse_isassigned(a::AbstractArray{<:Any,N}, I::Vararg{Integer,N}) where {N} + return all(dim -> I[dim] ∈ axes(a, dim), 1:ndims(a)) end # A set of indices into the storage of the sparse array. diff --git a/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl index 1be4d5d8dc..78bab9aecf 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl @@ -21,6 +21,14 @@ using Test: @test, @testset for I in eachindex(a) @test iszero(a) end + for I in CartesianIndices(a) + @test isassigned(a, Tuple(I)...) + @test isassigned(a, I) + end + @test !isassigned(a, 0, 1) + @test !isassigned(a, CartesianIndex(0, 1)) + @test !isassigned(a, 1, 4) + @test !isassigned(a, CartesianIndex(1, 4)) a = SparseArray{elt}(2, 3) fill!(a, 0) @@ -60,6 +68,14 @@ using Test: @test, @testset @test iszero(a[I]) end end + for I in CartesianIndices(a) + @test isassigned(a, Tuple(I)...) + @test isassigned(a, I) + end + @test !isassigned(a, 0, 1) + @test !isassigned(a, CartesianIndex(0, 1)) + @test !isassigned(a, 1, 4) + @test !isassigned(a, CartesianIndex(1, 4)) a = SparseArray{elt}(2, 3) a[1, 2] = 12