undef_blocks constructor for SparseArrayDOK, more sparsemortar constructors

Author: mtfishman

src/BlockSparseArrays.jl

@@ -39,7 +39,6 @@ include("abstractblocksparsearray/cat.jl")
 include("abstractblocksparsearray/adapt.jl")
 
 # functions specifically for BlockSparseArray
-include("blocksparsearray/defaults.jl")
 include("blocksparsearray/blocksparsearray.jl")
 include("blocksparsearray/blockdiagonalarray.jl")
 

src/blocksparsearray/blockdiagonalarray.jl

@@ -12,9 +12,7 @@ const BlockSparseDiagonal{T,A<:AbstractBlockSparseVector{T}} = Diagonal{T,A}
 end
 
 function BlockDiagonal(blocks::AbstractVector{<:AbstractMatrix})
-  return sparsemortar(
-    Diagonal(blocks), (blockedrange(size.(blocks, 1)), blockedrange(size.(blocks, 2)))
-  )
+  return sparsemortar(Diagonal(blocks), size.(blocks, 1), size.(blocks, 2))
 end
 
 function DiagonalArrays.diagonal(S::BlockSparseVector)

src/blocksparsearray/blocksparsearray.jl

@@ -1,8 +1,64 @@
-using BlockArrays: BlockArrays, Block, BlockedUnitRange, blockedrange, blocklength
+using BlockArrays:
+  BlockArrays,
+  Block,
+  BlockedUnitRange,
+  UndefBlocksInitializer,
+  blockedrange,
+  blocklength,
+  undef_blocks
 using DerivableInterfaces: @interface
 using Dictionaries: Dictionary
 using SparseArraysBase: SparseArrayDOK
 
+"""
+    SparseArrayDOK{T}(::UndefBlocksInitializer, axes)
+    SparseArrayDOK{T,N}(::UndefBlocksInitializer, axes)
+
+Construct the block structure of an undefined BlockSparseArray that will have
+blocked axes `axes`.
+"""
+function SparseArraysBase.SparseArrayDOK{T,N}(
+  ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef, blocklength.(ax), GetUnstoredBlock(ax))
+end
+function SparseArraysBase.SparseArrayDOK{T,N}(
+  ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer},N}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, ax)
+end
+function SparseArraysBase.SparseArrayDOK{T,N}(
+  ::UndefBlocksInitializer, dims::Tuple{Vararg{AbstractVector{<:Integer},N}}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, blockedrange.(dims))
+end
+function SparseArraysBase.SparseArrayDOK{T,N}(
+  ::UndefBlocksInitializer, dims::Vararg{AbstractVector{<:Integer},N}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, blockedrange.(dims))
+end
+
+function SparseArraysBase.SparseArrayDOK{T}(
+  ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, ax)
+end
+function SparseArraysBase.SparseArrayDOK{T}(
+  ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer},N}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, ax)
+end
+function SparseArraysBase.SparseArrayDOK{T}(
+  ::UndefBlocksInitializer, dims::Tuple{Vararg{AbstractVector{<:Integer},N}}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, blockedrange.(dims))
+end
+function SparseArraysBase.SparseArrayDOK{T}(
+  ::UndefBlocksInitializer, dims::Vararg{AbstractVector{<:Integer},N}
+) where {T,N}
+  return SparseArrayDOK{T,N}(undef_blocks, blockedrange.(dims))
+end
+
 function _BlockSparseArray end
 
 struct BlockSparseArray{
@@ -46,12 +102,22 @@ function sparsemortar(
 ) where {T,N}
   return _BlockSparseArray(blocks, axes)
 end
-
-function BlockArrays.mortar(
-  blocks::SparseArrayDOK{<:AbstractArray{T,N},N},
-  axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}},
+function sparsemortar(
+  blocks::AbstractArray{<:AbstractArray{T,N},N},
+  axes::Vararg{AbstractUnitRange{<:Integer},N},
 ) where {T,N}
-  return _BlockSparseArray(blocks, axes)
+  return sparsemortar(blocks, axes)
+end
+function sparsemortar(
+  blocks::AbstractArray{<:AbstractArray{T,N},N},
+  dims::Tuple{Vararg{AbstractVector{<:Integer},N}},
+) where {T,N}
+  return sparsemortar(blocks, blockedrange.(dims))
+end
+function sparsemortar(
+  blocks::AbstractArray{<:AbstractArray{T,N},N}, dims::Vararg{AbstractVector{<:Integer},N}
+) where {T,N}
+  return sparsemortar(blocks, dims)
 end
 
 @doc """
@@ -64,20 +130,19 @@ of block lengths in each dimension or a list of blocked ranges representing the
 """ BlockSparseArray
 
 function BlockSparseArray{T,N,A}(
-  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange,N}}
+  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}
 ) where {T,N,A<:AbstractArray{T,N}}
-  blocks = default_blocks(A, axes)
-  return _BlockSparseArray(blocks, axes)
+  return _BlockSparseArray(SparseArrayDOK{A}(undef_blocks, axes), axes)
 end
 
 function BlockSparseArray{T,N,A}(
-  ::UndefInitializer, axes::Vararg{AbstractUnitRange,N}
+  ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer},N}
 ) where {T,N,A<:AbstractArray{T,N}}
   return BlockSparseArray{T,N,A}(undef, axes)
 end
 
 function BlockSparseArray{T,N,A}(
-  ::UndefInitializer, dims::Tuple{Vararg{Vector{Int},N}}
+  ::UndefInitializer, dims::Tuple{Vararg{AbstractVector{<:Integer},N}}
 ) where {T,N,A<:AbstractArray{T,N}}
   return BlockSparseArray{T,N,A}(undef, blockedrange.(dims))
 end
@@ -86,48 +151,51 @@ end
 function BlockSparseArray{T,0,A}(
   ::UndefInitializer, axes::Tuple{}
 ) where {T,A<:AbstractArray{T,0}}
-  blocks = default_blocks(A, axes)
-  return _BlockSparseArray(blocks, axes)
+  return _BlockSparseArray(SparseArrayDOK{A}(undef_blocks, axes), axes)
 end
 
 function BlockSparseArray{T,N,A}(
-  ::UndefInitializer, dims::Vararg{Vector{Int},N}
+  ::UndefInitializer, dims::Vararg{AbstractVector{<:Integer},N}
 ) where {T,N,A<:AbstractArray{T,N}}
   return BlockSparseArray{T,N,A}(undef, dims)
 end
 
 function BlockSparseArray{T,N}(
-  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange,N}}
+  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}
 ) where {T,N}
-  return BlockSparseArray{T,N,default_arraytype(T, axes)}(undef, axes)
+  return BlockSparseArray{T,N,Array{T,N}}(undef, axes)
 end
 
 function BlockSparseArray{T,N}(
-  ::UndefInitializer, axes::Vararg{AbstractUnitRange,N}
+  ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer},N}
 ) where {T,N}
   return BlockSparseArray{T,N}(undef, axes)
 end
 
 function BlockSparseArray{T,0}(::UndefInitializer, axes::Tuple{}) where {T}
-  return BlockSparseArray{T,0,default_arraytype(T, axes)}(undef, axes)
+  return BlockSparseArray{T,0,Array{T,0}}(undef, axes)
 end
 
 function BlockSparseArray{T,N}(
-  ::UndefInitializer, dims::Tuple{Vararg{Vector{Int},N}}
+  ::UndefInitializer, dims::Tuple{Vararg{AbstractVector{<:Integer},N}}
 ) where {T,N}
   return BlockSparseArray{T,N}(undef, blockedrange.(dims))
 end
 
-function BlockSparseArray{T,N}(::UndefInitializer, dims::Vararg{Vector{Int},N}) where {T,N}
+function BlockSparseArray{T,N}(
+  ::UndefInitializer, dims::Vararg{AbstractVector{<:Integer},N}
+) where {T,N}
   return BlockSparseArray{T,N}(undef, dims)
 end
 
-function BlockSparseArray{T}(::UndefInitializer, dims::Tuple{Vararg{Vector{Int}}}) where {T}
+function BlockSparseArray{T}(
+  ::UndefInitializer, dims::Tuple{Vararg{AbstractVector{<:Integer}}}
+) where {T}
   return BlockSparseArray{T,length(dims)}(undef, dims)
 end
 
 function BlockSparseArray{T}(
-  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange}}
+  ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}
 ) where {T}
   return BlockSparseArray{T,length(axes)}(undef, axes)
 end
@@ -136,11 +204,15 @@ function BlockSparseArray{T}(::UndefInitializer, axes::Tuple{}) where {T}
   return BlockSparseArray{T,length(axes)}(undef, axes)
 end
 
-function BlockSparseArray{T}(::UndefInitializer, dims::Vararg{Vector{Int}}) where {T}
+function BlockSparseArray{T}(
+  ::UndefInitializer, dims::Vararg{AbstractVector{<:Integer}}
+) where {T}
   return BlockSparseArray{T}(undef, dims)
 end
 
-function BlockSparseArray{T}(::UndefInitializer, axes::Vararg{AbstractUnitRange}) where {T}
+function BlockSparseArray{T}(
+  ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer}}
+) where {T}
   return BlockSparseArray{T}(undef, axes)
 end
 

src/blocksparsearray/defaults.jl

@@ -17,7 +17,8 @@ using BlockArrays:
   blocklengths,
   blocksize,
   blocksizes,
-  mortar
+  mortar,
+  undef_blocks
 using BlockSparseArrays:
   @view!,
   BlockSparseArray,
@@ -158,6 +159,24 @@ arrayts = (Array, JLArray)
     @test (@constinferred blocktype(a)) <: SubArray{elt,2,arrayt{elt,2}}
     # TODO: This is difficult to determine just from type information.
     @test_broken blocktype(typeof(a)) <: SubArray{elt,2,arrayt{elt,2}}
+
+    # sparsemortar
+    for ax in (
+      ([2, 3], [2, 3]),
+      (([2, 3], [2, 3]),),
+      blockedrange.(([2, 3], [2, 3])),
+      (blockedrange.(([2, 3], [2, 3])),),
+    )
+      blocks = SparseArrayDOK{arrayt{elt,2}}(undef_blocks, ax...)
+      blocks[2, 1] = arrayt(randn(elt, 3, 2))
+      blocks[1, 2] = arrayt(randn(elt, 2, 3))
+      a = sparsemortar(blocks, ax...)
+      @test a isa BlockSparseArray{elt,2,arrayt{elt,2}}
+      @test iszero(a[Block(1, 1)])
+      @test a[Block(2, 1)] == blocks[2, 1]
+      @test a[Block(1, 2)] == blocks[1, 2]
+      @test iszero(a[Block(2, 2)])
+    end
   end
   @testset "Basics" begin
     a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))