Convenient constructors like blocksparse and blocksparsezeros (#141)

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.7.6"
+version = "0.7.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/blocksparsearray/blocksparsearray.jl

@@ -230,6 +230,27 @@ function BlockSparseArray{T}(
   return BlockSparseArray{T}(undef, axes)
 end
 
+# Convenient constructors.
+function blocksparsezeros(elt::Type, axes...)
+  return BlockSparseArray{elt}(undef, axes...)
+end
+function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray}
+  # TODO: Use:
+  # ```julia
+  # B = similartype(A, Type{eltype(A)}, Tuple{blockaxistype.(axes)...})
+  # BlockSparseArray{eltype(A),length(axes),B}(undef, axes...)
+  # ```
+  # to make a bit more generic.
+  return BlockSparseArray{eltype(A),ndims(A),A}(undef, axes...)
+end
+function blocksparse(d::Dict{<:Block,<:AbstractArray}, axes...)
+  a = blocksparsezeros(BlockType(valtype(d)), axes...)
+  for I in eachindex(d)
+    a[I] = d[I]
+  end
+  return a
+end
+
 # Base `AbstractArray` interface
 Base.axes(a::BlockSparseArray) = a.axes
 

test/test_basics.jl

@@ -21,9 +21,12 @@ using BlockSparseArrays:
   BlockSparseArray,
   BlockSparseMatrix,
   BlockSparseVector,
+  BlockType,
   BlockView,
   blockdiagindices,
   blockreshape,
+  blocksparse,
+  blocksparsezeros,
   blockstoredlength,
   blockstype,
   blocktype,
@@ -131,6 +134,37 @@ arrayts = (Array, JLArray)
     )
       @test_throws ArgumentError BlockSparseVector{elt}(undef, dims...)
     end
+
+    # Convenient constructors.
+    a = blocksparsezeros(elt, [2, 3], [2, 3])
+    @test iszero(a)
+    @test iszero(blockstoredlength(a))
+    @test a isa BlockSparseMatrix{elt,Matrix{elt}}
+    @test blocktype(a) == Matrix{elt}
+    @test blockstype(a) <: SparseMatrixDOK{Matrix{elt}}
+    @test blocksize(a) == (2, 2)
+    @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)]
+
+    a = blocksparsezeros(BlockType(arrayt{elt,2}), [2, 3], [2, 3])
+    @test iszero(a)
+    @test iszero(blockstoredlength(a))
+    @test a isa BlockSparseMatrix{elt,arrayt{elt,2}}
+    @test blocktype(a) == arrayt{elt,2}
+    @test blockstype(a) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test blocksize(a) == (2, 2)
+    @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)]
+
+    d = Dict(Block(1, 1) => dev(randn(elt, 2, 2)), Block(2, 2) => dev(randn(elt, 3, 3)))
+    a = blocksparse(d, [2, 3], [2, 3])
+    @test !iszero(a)
+    @test a[Block(1, 1)] == d[Block(1, 1)]
+    @test a[Block(2, 2)] == d[Block(2, 2)]
+    @test blockstoredlength(a) == 2
+    @test a isa BlockSparseMatrix{elt,arrayt{elt,2}}
+    @test blocktype(a) == arrayt{elt,2}
+    @test blockstype(a) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test blocksize(a) == (2, 2)
+    @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)]
   end
   @testset "blockstype, blocktype" begin
     a = arrayt(randn(elt, 2, 2))