Better definitions of norm and tr (#127)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.6.6"
+version = "0.6.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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

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)

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)

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) == ()