Bump version

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.5.3"
+version = "0.6.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -39,7 +39,7 @@ GPUArraysCore = "0.1.0, 0.2"
 LinearAlgebra = "1.10"
 MacroTools = "0.5.13"
 MapBroadcast = "0.1.5"
-MatrixAlgebraKit = "0.1.2"
+MatrixAlgebraKit = "0.2"
 SparseArraysBase = "0.5"
 SplitApplyCombine = "1.2.3"
 TensorAlgebra = "0.3.2"

docs/Project.toml

@@ -6,6 +6,6 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.5"
+BlockSparseArrays = "0.6"
 Documenter = "1"
 Literate = "2"

examples/Project.toml

@@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.5"
+BlockSparseArrays = "0.6"
 Test = "1"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -290,13 +290,6 @@ function blockrange(axis::AbstractUnitRange, r::Int)
   return error("Slicing with integer values isn't supported.")
 end
 
-function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
-  for b in r
-    @assert b ∈ blockaxes(axis, 1)
-  end
-  return r
-end
-
 # This handles changing the blocking, for example:
 # a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
 # I = blockedrange([4, 4])
@@ -315,13 +308,20 @@ end
 # I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
 # I = BlockVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
 # a[I, I]
-function blockrange(axis::BlockedOneTo{<:Integer}, r::AbstractBlockVector{<:Block{1}})
+function blockrange(axis::AbstractUnitRange, r::AbstractBlockVector{<:Block{1}})
   for b in r
     @assert b ∈ blockaxes(axis, 1)
   end
   return only(blockaxes(r))
 end
 
+function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
+  for b in r
+    @assert b ∈ blockaxes(axis, 1)
+  end
+  return r
+end
+
 using BlockArrays: BlockSlice
 function blockrange(axis::AbstractUnitRange, r::BlockSlice)
   return blockrange(axis, r.block)

src/BlockArraysExtensions/blockedunitrange.jl

@@ -9,11 +9,32 @@ using BlockArrays:
   BlockSlice,
   BlockVector,
   block,
+  blockedrange,
   blockindex,
+  blocklengths,
   findblock,
   findblockindex,
   mortar
 
+# Get the axes of each block of a block array.
+function eachblockaxes(a::AbstractArray)
+  return map(axes, blocks(a))
+end
+
+axis(a::AbstractVector) = axes(a, 1)
+
+# Get the axis of each block of a blocked unit
+# range.
+function eachblockaxis(a::AbstractVector)
+  return map(axis, blocks(a))
+end
+
+# Take a collection of axes and mortar them
+# into a single blocked axis.
+function mortar_axis(axs)
+  return blockedrange(length.(axs))
+end
+
 # Custom `BlockedUnitRange` constructor that takes a unit range
 # and a set of block lengths, similar to `BlockArray(::AbstractArray, blocklengths...)`.
 function blockedunitrange(a::AbstractUnitRange, blocklengths)

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -445,7 +445,7 @@ end
 
 to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block)
 to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks)
-to_blocks_indices(I::Base.Slice{<:BlockedOneTo}) = Base.OneTo(blocklength(I.indices))
+to_blocks_indices(I::Base.Slice) = Base.OneTo(blocklength(I.indices))
 
 @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks(
   a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}}

src/blocksparsearrayinterface/map.jl

@@ -1,6 +1,52 @@
+using BlockArrays: BlockRange, blockisequal
 using DerivableInterfaces: @interface, AbstractArrayInterface, interface
 using GPUArraysCore: @allowscalar
 
+# Check if the block structures are the same.
+function same_block_structure(as::AbstractArray...)
+  isempty(as) && return true
+  return all(
+    ntuple(ndims(first(as))) do dim
+      ax = map(Base.Fix2(axes, dim), as)
+      return blockisequal(ax...)
+    end,
+  )
+end
+
+# Find the common stored blocks, assuming the block structures are the same.
+function union_eachblockstoredindex(as::AbstractArray...)
+  return ∪(map(eachblockstoredindex, as)...)
+end
+
+function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...)
+  # TODO: This assumes element types are numbers, generalize this logic.
+  f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest))
+  Is = if f_preserves_zeros
+    union_eachblockstoredindex(a_dest, a_srcs...)
+  else
+    BlockRange(a_dest)
+  end
+  for I in Is
+    # TODO: Use:
+    # block_dest = @view a_dest[I]
+    # or:
+    # block_dest = @view! a_dest[I]
+    block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(I))...]
+    # TODO: Use:
+    # block_srcs = map(a_src -> @view(a_src[I]), a_srcs)
+    block_srcs = map(a_srcs) do a_src
+      return blocks_maybe_single(a_src)[Int.(Tuple(I))...]
+    end
+    # TODO: Use `map!!` to handle immutable blocks.
+    map!(f, block_dest, block_srcs...)
+    # Replace the entire block, handles initializing new blocks
+    # or if blocks are immutable.
+    # TODO: Use `a_dest[I] = block_dest`.
+    blocks(a_dest)[Int.(Tuple(I))...] = block_dest
+  end
+  return a_dest
+end
+
 # TODO: Rewrite this so that it takes the blocking structure
 # made by combining the blocking of the axes (i.e. the blocking that
 # is used to determine `union_stored_blocked_cartesianindices(...)`).
@@ -16,6 +62,10 @@ using GPUArraysCore: @allowscalar
     @interface interface map_zero_dim!(f, a_dest, a_srcs...)
     return a_dest
   end
+  if same_block_structure(a_dest, a_srcs...)
+    map_blockwise!(f, a_dest, a_srcs...)
+    return a_dest
+  end
   # TODO: This assumes element types are numbers, generalize this logic.
   f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest))
   a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs)

src/factorizations/svd.jl

@@ -12,20 +12,37 @@ struct BlockPermutedDiagonalAlgorithm{A<:MatrixAlgebraKit.AbstractAlgorithm} <:
   alg::A
 end
 
-# TODO: this is a hardcoded for now to get around this function not being defined in the
-# type domain
-function MatrixAlgebraKit.default_svd_algorithm(A::AbstractBlockSparseMatrix; kwargs...)
+function default_blocksparse_svd_algorithm(f, A; kwargs...)
   blocktype(A) <: StridedMatrix{<:LinearAlgebra.BLAS.BlasFloat} ||
     error("unsupported type: $(blocktype(A))")
+  # TODO: this is a hardcoded for now to get around this function not being defined in the
+  # type domain
+  # alg = MatrixAlgebraKit.default_algorithm(f, blocktype(A); kwargs...)
   alg = MatrixAlgebraKit.LAPACK_DivideAndConquer(; kwargs...)
   return BlockPermutedDiagonalAlgorithm(alg)
 end
 
-# 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(...))
-function _similar_S(A::AbstractBlockSparseMatrix, s_axis)
+function MatrixAlgebraKit.default_algorithm(
+  f::typeof(svd_compact!), A::AbstractBlockSparseMatrix; kwargs...
+)
+  return default_blocksparse_svd_algorithm(f, A; kwargs...)
+end
+function MatrixAlgebraKit.default_algorithm(
+  f::typeof(svd_full!), A::AbstractBlockSparseMatrix; kwargs...
+)
+  return default_blocksparse_svd_algorithm(f, A; kwargs...)
+end
+
+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))
-  return BlockSparseArray{T,2,Diagonal{T,Vector{T}}}(undef, (s_axis, s_axis))
+  # 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)
+  Vt = similar(A, S_axes[2], axes(A, 2))
+  return U, S, Vt
 end
 
 function MatrixAlgebraKit.initialize_output(
@@ -34,33 +51,36 @@ function MatrixAlgebraKit.initialize_output(
   bm, bn = blocksize(A)
   bmn = min(bm, bn)
 
-  brows = blocklengths(axes(A, 1))
-  bcols = blocklengths(axes(A, 2))
-  slengths = Vector{Int}(undef, bmn)
+  brows = eachblockaxis(axes(A, 1))
+  bcols = eachblockaxis(axes(A, 2))
+  u_axes = similar(brows, bmn)
+  v_axes = similar(brows, bmn)
 
   # fill in values for blocks that are present
   bIs = collect(eachblockstoredindex(A))
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
   for bI in eachblockstoredindex(A)
     row, col = Int.(Tuple(bI))
-    nrows = brows[row]
-    ncols = bcols[col]
-    slengths[col] = min(nrows, ncols)
+    len = minimum(length, (brows[row], bcols[col]))
+    u_axes[col] = brows[row][Base.OneTo(len)]
+    v_axes[col] = bcols[col][Base.OneTo(len)]
   end
 
   # fill in values for blocks that aren't present, pairing them in order of occurence
   # this is a convention, which at least gives the expected results for blockdiagonal
   emptyrows = setdiff(1:bm, browIs)
   emptycols = setdiff(1:bn, bcolIs)
   for (row, col) in zip(emptyrows, emptycols)
-    slengths[col] = min(brows[row], bcols[col])
+    len = minimum(length, (brows[row], bcols[col]))
+    u_axes[col] = brows[row][Base.OneTo(len)]
+    v_axes[col] = bcols[col][Base.OneTo(len)]
   end
 
-  s_axis = blockedrange(slengths)
-  U = similar(A, axes(A, 1), s_axis)
-  S = _similar_S(A, s_axis)
-  Vt = similar(A, s_axis, axes(A, 2))
+  u_axis = mortar_axis(u_axes)
+  v_axis = mortar_axis(v_axes)
+  S_axes = (u_axis, v_axis)
+  U, S, Vt = similar_output(svd_compact!, A, S_axes, alg)
 
   # allocate output
   for bI in eachblockstoredindex(A)
@@ -79,40 +99,47 @@ function MatrixAlgebraKit.initialize_output(
   return U, S, Vt
 end
 
+function similar_output(
+  ::typeof(svd_full!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm
+)
+  U = similar(A, axes(A, 1), S_axes[1])
+  T = real(eltype(A))
+  S = similar(A, T, S_axes)
+  Vt = similar(A, S_axes[2], axes(A, 2))
+  return U, S, Vt
+end
+
 function MatrixAlgebraKit.initialize_output(
   ::typeof(svd_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
 )
   bm, bn = blocksize(A)
 
-  brows = blocklengths(axes(A, 1))
-  slengths = copy(brows)
+  brows = eachblockaxis(axes(A, 1))
+  u_axes = similar(brows)
 
   # fill in values for blocks that are present
   bIs = collect(eachblockstoredindex(A))
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
   for bI in eachblockstoredindex(A)
     row, col = Int.(Tuple(bI))
-    nrows = brows[row]
-    slengths[col] = nrows
+    u_axes[col] = brows[row]
   end
 
   # fill in values for blocks that aren't present, pairing them in order of occurence
   # this is a convention, which at least gives the expected results for blockdiagonal
   emptyrows = setdiff(1:bm, browIs)
   emptycols = setdiff(1:bn, bcolIs)
   for (row, col) in zip(emptyrows, emptycols)
-    slengths[col] = brows[row]
+    u_axes[col] = brows[row]
   end
   for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
-    slengths[bn + i] = brows[emptyrows[k]]
+    u_axes[bn + i] = brows[emptyrows[k]]
   end
 
-  s_axis = blockedrange(slengths)
-  U = similar(A, axes(A, 1), s_axis)
-  Tr = real(eltype(A))
-  S = similar(A, Tr, (s_axis, axes(A, 2)))
-  Vt = similar(A, axes(A, 2), axes(A, 2))
+  u_axis = mortar_axis(u_axes)
+  S_axes = (u_axis, axes(A, 2))
+  U, S, Vt = similar_output(svd_full!, A, S_axes, alg)
 
   # allocate output
   for bI in eachblockstoredindex(A)

src/factorizations/truncation.jl

@@ -45,6 +45,22 @@ function MatrixAlgebraKit.findtruncated(
   return indexmask
 end
 
+function similar_truncate(
+  ::typeof(svd_trunc!),
+  (U, S, Vᴴ)::TBlockUSVᴴ,
+  strategy::BlockPermutedDiagonalTruncationStrategy,
+  indexmask=MatrixAlgebraKit.findtruncated(diagview(S), strategy),
+)
+  ax = axes(S, 1)
+  counter = Base.Fix1(count, Base.Fix1(getindex, indexmask))
+  s_lengths = filter!(>(0), map(counter, blocks(ax)))
+  s_axis = blockedrange(s_lengths)
+  Ũ = similar(U, axes(U, 1), s_axis)
+  S̃ = similar(S, s_axis, s_axis)
+  Ṽᴴ = similar(Vᴴ, s_axis, axes(Vᴴ, 2))
+  return Ũ, S̃, Ṽᴴ
+end
+
 function MatrixAlgebraKit.truncate!(
   ::typeof(svd_trunc!),
   (U, S, Vᴴ)::TBlockUSVᴴ,
@@ -54,13 +70,7 @@ function MatrixAlgebraKit.truncate!(
 
   # first determine the block structure of the output to avoid having assumptions on the
   # data structures
-  ax = axes(S, 1)
-  counter = Base.Fix1(count, Base.Fix1(getindex, indexmask))
-  Slengths = filter!(>(0), map(counter, blocks(ax)))
-  Sax = blockedrange(Slengths)
-  Ũ = similar(U, axes(U, 1), Sax)
-  S̃ = similar(S, Sax, Sax)
-  Ṽᴴ = similar(Vᴴ, Sax, axes(Vᴴ, 2))
+  Ũ, S̃, Ṽᴴ = similar_truncate(svd_trunc!, (U, S, Vᴴ), strategy, indexmask)
 
   # then loop over the blocks and assign the data
   # TODO: figure out if we can presort and loop over the blocks -
@@ -70,6 +80,7 @@ function MatrixAlgebraKit.truncate!(
   bI_Vᴴs = collect(eachblockstoredindex(Vᴴ))
 
   I′ = 0 # number of skipped blocks that got fully truncated
+  ax = axes(S, 1)
   for I in 1:blocksize(ax, 1)
     b = ax[Block(I)]
     mask = indexmask[b]

test/Project.toml

@@ -23,12 +23,12 @@ Adapt = "4"
 Aqua = "0.8"
 ArrayLayouts = "1"
 BlockArrays = "1"
-BlockSparseArrays = "0.5"
+BlockSparseArrays = "0.6"
 DiagonalArrays = "0.3"
 GPUArraysCore = "0.2"
 JLArrays = "0.2"
 LinearAlgebra = "1"
-MatrixAlgebraKit = "0.1"
+MatrixAlgebraKit = "0.2"
 Random = "1"
 SafeTestsets = "0.1"
 SparseArraysBase = "0.5"