diff --git a/Project.toml b/Project.toml index 5a019d19..826dd156 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.6.6" +version = "0.6.7" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/abstractblocksparsearray/linearalgebra.jl b/src/abstractblocksparsearray/linearalgebra.jl index 144ea475..8121477a 100644 --- a/src/abstractblocksparsearray/linearalgebra.jl +++ b/src/abstractblocksparsearray/linearalgebra.jl @@ -1,4 +1,4 @@ -using LinearAlgebra: Adjoint, Transpose +using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, norm, tr # Like: https://github.com/JuliaLang/julia/blob/v1.11.1/stdlib/LinearAlgebra/src/transpose.jl#L184 # but also takes the dual of the axes. @@ -16,3 +16,19 @@ function Base.copy(a::Transpose{T,<:AbstractBlockSparseMatrix{T}}) where {T} a_dest .= a return a_dest end + +function LinearAlgebra.norm(a::AnyAbstractBlockSparseArray, p::Real=2) + nrmᵖ = float(norm(zero(eltype(a)))) + for I in eachblockstoredindex(a) + nrmᵖ += norm(@view(a[I]), p)^p + end + return nrmᵖ^(1/p) +end + +function LinearAlgebra.tr(a::AnyAbstractBlockSparseMatrix) + tr_a = zero(eltype(a)) + for I in eachstoredblockdiagindex(a) + tr_a += tr(@view(a[I])) + end + return tr_a +end diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 0403b789..84c24ed8 100644 --- a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -338,14 +338,17 @@ end # TODO: Implement this in a more generic way using a smarter `copyto!`, # which is ultimately what `Array{T,N}(::AbstractArray{<:Any,N})` calls. # These are defined for now to avoid scalar indexing issues when there -# are blocks on GPU. +# are blocks on GPU, and also work with exotic block types like +# KroneckerArrays. function Base.Array{T,N}(a::AnyAbstractBlockSparseArray{<:Any,N}) where {T,N} - # First make it dense, then move to CPU. - # Directly copying to CPU causes some issues with - # scalar indexing on GPU which we have to investigate. - a_dest = similartype(blocktype(a), T)(undef, size(a)) - a_dest .= a - return Array{T,N}(a_dest) + a_dest = zeros(T, size(a)) + for I in eachblockstoredindex(a) + # TODO: Use: `I′ = CartesianIndices(axes(a))[I]`, unfortunately this + # outputs `Matrix{CartesianIndex}` instead of `CartesianIndices`. + I′ = CartesianIndices(ntuple(dim -> axes(a, dim)[Tuple(I)[dim]], ndims(a))) + a_dest[I′] = Array{T,N}(@view(a[I])) + end + return a_dest end function Base.Array{T}(a::AnyAbstractBlockSparseArray) where {T} return Array{T,ndims(a)}(a) diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 06eaa6f6..0015de0b 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -36,6 +36,16 @@ function eachblockstoredindex(a::AbstractArray) return Block.(Tuple.(eachstoredindex(blocks(a)))) end +using DiagonalArrays: diagindices +# Block version of `DiagonalArrays.diagindices`. +function blockdiagindices(a::AbstractArray) + return Block.(Tuple.(diagindices(blocks(a)))) +end + +function eachstoredblockdiagindex(a::AbstractArray) + return eachblockstoredindex(a) ∩ blockdiagindices(a) +end + # Like `BlockArrays.eachblock` but only iterating # over stored blocks. function eachstoredblock(a::AbstractArray) diff --git a/test/test_basics.jl b/test/test_basics.jl index 82e86ea7..b0ee9a58 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -22,17 +22,19 @@ using BlockSparseArrays: BlockSparseMatrix, BlockSparseVector, BlockView, + blockdiagindices, blockreshape, blockstoredlength, blockstype, blocktype, eachblockstoredindex, eachstoredblock, + eachstoredblockdiagindex, sparsemortar, view! using GPUArraysCore: @allowscalar using JLArrays: JLArray, JLMatrix -using LinearAlgebra: Adjoint, Transpose, dot, norm +using LinearAlgebra: Adjoint, Transpose, dot, norm, tr using SparseArraysBase: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, storedlength using Test: @test, @test_broken, @test_throws, @testset, @inferred using TestExtras: @constinferred @@ -217,10 +219,26 @@ arrayts = (Array, JLArray) a[Block(1, 2)] = randn(elt, 2, 3) @test issetequal(eachstoredblock(a), [a[Block(2, 1)], a[Block(1, 2)]]) @test issetequal(eachblockstoredindex(a), [Block(2, 1), Block(1, 2)]) + @test issetequal(blockdiagindices(a), [Block(1, 1), Block(2, 2)]) + @test isempty(eachstoredblockdiagindex(a)) + @test norm(a) ≈ norm(Array(a)) + for p in 1:3 + @test norm(a, p) ≈ norm(Array(a), p) + end + @test tr(a) ≈ tr(Array(a)) a[3, 3] = NaN @test isnan(norm(a)) + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + @test issetequal(eachstoredblockdiagindex(a), [Block(1, 1)]) + @test norm(a) ≈ norm(Array(a)) + for p in 1:3 + @test norm(a, p) ≈ norm(Array(a), p) + end + @test tr(a) ≈ tr(Array(a)) + # Empty constructor for a in (dev(BlockSparseArray{elt}(undef)),) @test size(a) == ()