diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index b4e21b6198..7e09c978b3 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.10" +version = "0.3.11" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl index 2d22efd277..22f9350605 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl @@ -1,5 +1,6 @@ using ArrayLayouts: LayoutArray using BlockArrays: blockisequal +using LinearAlgebra: Adjoint, Transpose using ..SparseArrayInterface: SparseArrayInterface, SparseArrayStyle, @@ -73,6 +74,20 @@ function Base.copyto!(a_dest::LayoutArray, a_src::BlockSparseArrayLike) return a_dest end +function Base.copyto!( + a_dest::AbstractMatrix, a_src::Transpose{T,<:AbstractBlockSparseMatrix{T}} +) where {T} + sparse_copyto!(a_dest, a_src) + return a_dest +end + +function Base.copyto!( + a_dest::AbstractMatrix, a_src::Adjoint{T,<:AbstractBlockSparseMatrix{T}} +) where {T} + sparse_copyto!(a_dest, a_src) + return a_dest +end + function Base.permutedims!(a_dest, a_src::BlockSparseArrayLike, perm) sparse_permutedims!(a_dest, a_src, perm) return a_dest diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index dc40010526..784b4e27c8 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -3,8 +3,8 @@ using SplitApplyCombine: groupcount using Adapt: Adapt, WrappedArray -const WrappedAbstractBlockSparseArray{T,N,A} = WrappedArray{ - T,N,<:AbstractBlockSparseArray,<:AbstractBlockSparseArray{T,N} +const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{ + T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N} } # TODO: Rename `AnyBlockSparseArray`. diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 5fe63e9804..56e6e7c96b 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -8,6 +8,7 @@ using BlockArrays: blocks, blocklengths, findblockindex +using LinearAlgebra: Adjoint, Transpose using ..SparseArrayInterface: perm, iperm, nstored ## using MappedArrays: mappedarray @@ -86,10 +87,15 @@ end # BlockArrays -using ..SparseArrayInterface: SparseArrayInterface, AbstractSparseArray +using ..SparseArrayInterface: + SparseArrayInterface, AbstractSparseArray, AbstractSparseMatrix -# Represents the array of arrays of a `SubArray` -# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`. +_perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P +_getindices(t::Tuple, indices) = map(i -> t[i], indices) +_getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), indices)) + +# Represents the array of arrays of a `PermutedDimsArray` +# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `PermutedDimsArray`. struct SparsePermutedDimsArrayBlocks{T,N,Array<:PermutedDimsArray{T,N}} <: AbstractSparseArray{T,N} array::Array @@ -97,24 +103,80 @@ end function blocksparse_blocks(a::PermutedDimsArray) return SparsePermutedDimsArrayBlocks(a) end -_perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P -_getindices(t::Tuple, indices) = map(i -> t[i], indices) -_getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), indices)) -function SparseArrayInterface.stored_indices(a::SparsePermutedDimsArrayBlocks) - return map(I -> _getindices(I, _perm(a.array)), stored_indices(blocks(parent(a.array)))) -end function Base.size(a::SparsePermutedDimsArrayBlocks) return _getindices(size(blocks(parent(a.array))), _perm(a.array)) end -function Base.getindex(a::SparsePermutedDimsArrayBlocks, index::Vararg{Int}) +function Base.getindex( + a::SparsePermutedDimsArrayBlocks{<:Any,N}, index::Vararg{Int,N} +) where {N} return PermutedDimsArray( blocks(parent(a.array))[_getindices(index, _perm(a.array))...], _perm(a.array) ) end +function SparseArrayInterface.stored_indices(a::SparsePermutedDimsArrayBlocks) + return map(I -> _getindices(I, _perm(a.array)), stored_indices(blocks(parent(a.array)))) +end +# TODO: Either make this the generic interface or define +# `SparseArrayInterface.sparse_storage`, which is used +# to defined this. +SparseArrayInterface.nstored(a::SparsePermutedDimsArrayBlocks) = length(stored_indices(a)) function SparseArrayInterface.sparse_storage(a::SparsePermutedDimsArrayBlocks) return error("Not implemented") end +reverse_index(index) = reverse(index) +reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index))) + +# Represents the array of arrays of a `Transpose` +# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Transpose`. +struct SparseTransposeBlocks{T,Array<:Transpose{T}} <: AbstractSparseMatrix{T} + array::Array +end +function blocksparse_blocks(a::Transpose) + return SparseTransposeBlocks(a) +end +function Base.size(a::SparseTransposeBlocks) + return reverse(size(blocks(parent(a.array)))) +end +function Base.getindex(a::SparseTransposeBlocks, index::Vararg{Int,2}) + return transpose(blocks(parent(a.array))[reverse(index)...]) +end +function SparseArrayInterface.stored_indices(a::SparseTransposeBlocks) + return map(reverse_index, stored_indices(blocks(parent(a.array)))) +end +# TODO: Either make this the generic interface or define +# `SparseArrayInterface.sparse_storage`, which is used +# to defined this. +SparseArrayInterface.nstored(a::SparseTransposeBlocks) = length(stored_indices(a)) +function SparseArrayInterface.sparse_storage(a::SparseTransposeBlocks) + return error("Not implemented") +end + +# Represents the array of arrays of a `Adjoint` +# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Adjoint`. +struct SparseAdjointBlocks{T,Array<:Adjoint{T}} <: AbstractSparseMatrix{T} + array::Array +end +function blocksparse_blocks(a::Adjoint) + return SparseAdjointBlocks(a) +end +function Base.size(a::SparseAdjointBlocks) + return reverse(size(blocks(parent(a.array)))) +end +function Base.getindex(a::SparseAdjointBlocks, index::Vararg{Int,2}) + return blocks(parent(a.array))[reverse(index)...]' +end +function SparseArrayInterface.stored_indices(a::SparseAdjointBlocks) + return map(reverse_index, stored_indices(blocks(parent(a.array)))) +end +# TODO: Either make this the generic interface or define +# `SparseArrayInterface.sparse_storage`, which is used +# to defined this. +SparseArrayInterface.nstored(a::SparseAdjointBlocks) = length(stored_indices(a)) +function SparseArrayInterface.sparse_storage(a::SparseAdjointBlocks) + return error("Not implemented") +end + # TODO: Move to `BlockArraysExtensions`. # This takes a range of indices `indices` of array `a` # and maps it to the range of indices within block `block`. @@ -167,9 +229,6 @@ end function Base.size(a::SparseSubArrayBlocks) return length.(axes(a)) end -function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks) - return stored_indices(view(blocks(parent(a.array)), axes(a)...)) -end function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) where {N} return a[Tuple(I)...] end @@ -192,6 +251,13 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe # TODO: Implement this properly. return true end +function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks) + return stored_indices(view(blocks(parent(a.array)), axes(a)...)) +end +# TODO: Either make this the generic interface or define +# `SparseArrayInterface.sparse_storage`, which is used +# to defined this. +SparseArrayInterface.nstored(a::SparseSubArrayBlocks) = length(stored_indices(a)) function SparseArrayInterface.sparse_storage(a::SparseSubArrayBlocks) return error("Not implemented") end diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index c9d982a6cc..62cc23d9db 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -48,6 +48,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(a) == 2 @test nstored(a) == 2 * 4 + 3 * 3 + 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)]))) b = similar(a, complex(elt)) @test eltype(b) == complex(eltype(a)) @test iszero(b) @@ -56,37 +59,58 @@ include("TestBlockSparseArraysUtils.jl") @test size(b) == size(a) @test blocksize(b) == blocksize(a) + 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)]))) b = copy(a) b[1, 1] = 11 @test b[1, 1] == 11 @test a[1, 1] ≠ 11 + 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)]))) b = copy(a) b .*= 2 @test b ≈ 2a + 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)]))) b = copy(a) b ./= 2 @test b ≈ a / 2 + 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)]))) b = 2 * a @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) b = (2 + 3im) * a @test Array(b) ≈ (2 + 3im) * Array(a) @test eltype(b) == complex(elt) @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) b = a + a @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) x = BlockSparseArray{elt}(undef, ([3, 4], [2, 3])) x[Block(1, 2)] = randn(elt, size(@view(x[Block(1, 2)]))) x[Block(2, 1)] = randn(elt, size(@view(x[Block(2, 1)]))) @@ -96,12 +120,18 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) b = permutedims(a, (2, 1)) @test Array(b) ≈ permutedims(Array(a), (2, 1)) @test eltype(b) == elt @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) b = map(x -> 2x, a) @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @@ -110,6 +140,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 2 @test nstored(b) == 2 * 4 + 3 * 3 + 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)]))) b = a[[Block(2), Block(1)], [Block(2), Block(1)]] @test b[Block(1, 1)] == a[Block(2, 2)] @test b[Block(1, 2)] == a[Block(2, 1)] @@ -120,6 +153,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == nstored(a) @test block_nstored(b) == 2 + 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)]))) b = a[Block(1):Block(2), Block(1):Block(2)] @test b == a @test size(b) == size(a) @@ -127,6 +163,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == nstored(a) @test block_nstored(b) == 2 + 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)]))) b = a[Block(1):Block(1), Block(1):Block(2)] @test b == Array(a)[1:2, 1:end] @test b[Block(1, 1)] == a[Block(1, 1)] @@ -136,6 +175,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == nstored(a[Block(1, 2)]) @test block_nstored(b) == 1 + 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)]))) b = a[2:4, 2:4] @test b == Array(a)[2:4, 2:4] @test size(b) == (3, 3) @@ -143,6 +185,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 1 * 1 + 2 * 2 @test block_nstored(b) == 2 + 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)]))) b = a[Block(2, 1)[1:2, 2:3]] @test b == Array(a)[3:4, 2:3] @test size(b) == (2, 2) @@ -150,27 +195,63 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 2 @test block_nstored(b) == 1 + 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)]))) + b = PermutedDimsArray(a, (2, 1)) + @test block_nstored(b) == 2 + @test Array(b) == permutedims(Array(a), (2, 1)) + c = 2 * b + @test block_nstored(c) == 2 + @test Array(c) == 2 * permutedims(Array(a), (2, 1)) + + 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)]))) + b = a' + @test block_nstored(b) == 2 + @test Array(b) == Array(a)' + c = 2 * b + @test block_nstored(c) == 2 + @test Array(c) == 2 * Array(a)' + + 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)]))) + b = transpose(a) + @test block_nstored(b) == 2 + @test Array(b) == transpose(Array(a)) + c = 2 * b + @test block_nstored(c) == 2 + @test Array(c) == 2 * transpose(Array(a)) + ## Broken, need to fix. + 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)]))) @test_broken a[Block(1), Block(1):Block(2)] # 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)]))) b = a[Block(2):Block(2), Block(1):Block(2)] @test_broken block_nstored(b) == 1 @test_broken b == Array(a)[3:5, 1:end] - b = a' - @test_broken block_nstored(b) == 2 - - b = transpose(a) - @test_broken block_nstored(b) == 2 - + 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)]))) b = copy(a) x = randn(size(@view(a[Block(2, 2)]))) b[Block(2), Block(2)] = x @test_broken b[Block(2, 2)] == x # Doesnt' set the block + 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)]))) b = copy(a) b[Block(1, 1)] .= 1 @test_broken b[1, 1] == trues(size(@view(b[1, 1])))