[WIP] Make more operations more agnostic about the block type

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

src/BlockArraysExtensions/blockedunitrange.jl

@@ -33,7 +33,8 @@ end
 # Take a collection of axes and mortar them
 # into a single blocked axis.
 function mortar_axis(axs)
-  return blockedrange(length.(axs))
+  ## return blockedrange(length.(axs))
+  return blockrange(axs)
 end
 
 # Custom `BlockedUnitRange` constructor that takes a unit range

src/abstractblocksparsearray/map.jl

@@ -111,9 +111,13 @@ function Base.isreal(a::AnyAbstractBlockSparseArray)
   return @interface interface(a) isreal(a)
 end
 
+# Helps with specialization.
 function Base.:*(x::Number, a::AnyAbstractBlockSparseArray)
   return map(Base.Fix1(*, x), a)
 end
 function Base.:*(a::AnyAbstractBlockSparseArray, x::Number)
   return map(Base.Fix2(*, x), a)
 end
+function Base.:/(a::AnyAbstractBlockSparseArray, x::Number)
+  return map(Base.Fix2(/, x), a)
+end

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -230,10 +230,16 @@ function Base.similar(
   return similar(arraytype, eltype(arraytype), axes)
 end
 
+# This circumvents some issues with `TypeParameterAccessors.similartype`.
+# TODO: Fix this poperly in `TypeParameterAccessors.jl`.
+function _similartype(arraytype::Type{<:AbstractArray}, elt::Type, axt)
+  return Base.promote_op(similar, arraytype, elt, axt)
+end
+
 function blocksparse_similar(a, elt::Type, axes::Tuple)
-  return BlockSparseArray{elt,length(axes),similartype(blocktype(a), elt, axes)}(
-    undef, axes
-  )
+  block_axt = Tuple{eltype.(eachblockaxis.(axes))...}
+  blockt = _similartype(blocktype(a), Type{elt}, block_axt)
+  return BlockSparseArray{elt,length(axes),blockt}(undef, axes)
 end
 @interface ::AbstractBlockSparseArrayInterface function Base.similar(
   a::AbstractArray, elt::Type, axes::Tuple{Vararg{Int}}

src/blocksparsearray/blocksparsearray.jl

@@ -173,7 +173,9 @@ end
 function BlockSparseArray{T,N}(
   ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}
 ) where {T,N}
-  return BlockSparseArray{T,N,Array{T,N}}(undef, axes)
+  # TODO: Use `similartype` to determine the block type.
+  A = Base.promote_op(similar, Array{T}, Tuple{eltype.(eachblockaxis.(axes))...})
+  return BlockSparseArray{T,N,A}(undef, axes)
 end
 
 function BlockSparseArray{T,N}(
@@ -230,6 +232,20 @@ function BlockSparseArray{T}(
   return BlockSparseArray{T}(undef, axes)
 end
 
+function blocksparsezeros(elt::Type, axes...)
+  return BlockSparseArray{elt}(undef, axes...)
+end
+function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray}
+  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
 

src/factorizations/svd.jl

@@ -20,14 +20,21 @@ function MatrixAlgebraKit.default_svd_algorithm(
   return BlockPermutedDiagonalAlgorithm(alg)
 end
 
+# TODO: Put this in a common location or package,
+# maybe `TypeParameterAccessors.jl`?
+# Also define `imagtype`, `complextype`, etc.
+realtype(a::AbstractArray) = realtype(typeof(a))
+function realtype(A::Type{<:AbstractArray})
+  return Base.promote_op(real, A)
+end
+
+using DiagonalArrays: diagonaltype
 function similar_output(
   ::typeof(svd_compact!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm
 )
   U = similar(A, axes(A, 1), S_axes[1])
   T = real(eltype(A))
-  # TODO: this should be replaced with a more general similar function that can handle setting
-  # the blocktype and element type - something like S = similar(A, BlockType(...))
-  S = BlockSparseMatrix{T,Diagonal{T,Vector{T}}}(undef, S_axes)
+  S = similar(A, BlockType(diagonaltype(realtype(blocktype(A)))), S_axes)
   Vt = similar(A, S_axes[2], axes(A, 2))
   return U, S, Vt
 end
@@ -49,9 +56,9 @@ function MatrixAlgebraKit.initialize_output(
   bcolIs = Int.(last.(Tuple.(bIs)))
   for bI in eachblockstoredindex(A)
     row, col = Int.(Tuple(bI))
-    len = minimum(length, (brows[row], bcols[col]))
-    u_axes[col] = brows[row][Base.OneTo(len)]
-    v_axes[col] = bcols[col][Base.OneTo(len)]
+    b = argmin(length, (brows[row], bcols[col]))
+    u_axes[col] = b
+    v_axes[col] = b
   end
 
   # fill in values for blocks that aren't present, pairing them in order of occurence