Towards more general truncation and slicing

Author: mtfishman

Project.toml

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

src/BlockArraysExtensions/blockedunitrange.jl

@@ -314,6 +314,8 @@ end
 # `Base.getindex(a::Block, b...)`.
 _getindex(a::Block{N}, b::Vararg{Any,N}) where {N} = GenericBlockIndex(a, b)
 _getindex(a::Block{N}, b::Vararg{Integer,N}) where {N} = a[b...]
+_getindex(a::Block{N}, b::Vararg{AbstractUnitRange{<:Integer},N}) where {N} = a[b...]
+_getindex(a::Block{N}, b::Vararg{AbstractVector,N}) where {N} = BlockIndexVector(a, b)
 # Fix ambiguity.
 _getindex(a::Block{0}) = a[]
 
@@ -372,7 +374,11 @@ function blockedunitrange_getindices(
   a::AbstractBlockedUnitRange,
   indices::BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},
 )
-  return mortar(map(b -> a[b], blocks(indices)))
+  blks = map(b -> a[b], blocks(indices))
+  # Preserve any extra structure in the axes, like a
+  # Kronecker structure, symmetry sectors, etc.
+  ax = mortar_axis(map(b -> axis(a[b]), blocks(indices)))
+  return mortar(blks, (ax,))
 end
 
 # This is a specialization of `BlockArrays.unblock`:

src/abstractblocksparsearray/linearalgebra.jl

@@ -1,4 +1,4 @@
-using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, norm, tr
+using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, diag, 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.
@@ -33,6 +33,24 @@ function LinearAlgebra.tr(a::AnyAbstractBlockSparseMatrix)
   return tr_a
 end
 
+# TODO: Define in DiagonalArrays.jl.
+function diagaxis(a::AbstractArray)
+  LinearAlgebra.checksquare(a)
+  return axes(a, 1)
+end
+function LinearAlgebra.diag(a::AnyAbstractBlockSparseMatrix)
+  # TODO: Add `checkblocksquare` to also check it is square blockwise.
+  LinearAlgebra.checksquare(a)
+  diagaxes = map(blockdiagindices(a)) do I
+    return diagaxis(@view(a[I]))
+  end
+  r = blockrange(diagaxes)
+  stored_blocks = Dict((
+    Tuple(I)[1] => diag(@view!(a[I])) for I in eachstoredblockdiagindex(a)
+  ))
+  return blocksparse(stored_blocks, (r,))
+end
+
 # TODO: Define `SparseArraysBase.isdiag`, define as
 # `isdiag(blocks(a))`.
 function blockisdiag(a::AbstractArray)

src/factorizations/truncation.jl

@@ -1,14 +1,11 @@
-using MatrixAlgebraKit: TruncationStrategy, diagview, eig_trunc!, eigh_trunc!, svd_trunc!
-
-function MatrixAlgebraKit.diagview(A::BlockSparseMatrix{T,Diagonal{T,Vector{T}}}) where {T}
-  D = BlockSparseVector{T}(undef, axes(A, 1))
-  for I in eachblockstoredindex(A)
-    if ==(Int.(Tuple(I))...)
-      D[Tuple(I)[1]] = diagview(A[I])
-    end
-  end
-  return D
-end
+using MatrixAlgebraKit:
+  TruncationStrategy,
+  diagview,
+  eig_trunc!,
+  eigh_trunc!,
+  findtruncated,
+  svd_trunc!,
+  truncate!
 
 """
     BlockPermutedDiagonalTruncationStrategy(strategy::TruncationStrategy)
@@ -27,7 +24,7 @@ function MatrixAlgebraKit.truncate!(
   strategy::TruncationStrategy,
 )
   # TODO assert blockdiagonal
-  return MatrixAlgebraKit.truncate!(
+  return truncate!(
     svd_trunc!, (U, S, Vᴴ), BlockPermutedDiagonalTruncationStrategy(strategy)
   )
 end
@@ -38,9 +35,7 @@ for f in [:eig_trunc!, :eigh_trunc!]
       (D, V)::NTuple{2,AbstractBlockSparseMatrix},
       strategy::TruncationStrategy,
     )
-      return MatrixAlgebraKit.truncate!(
-        $f, (D, V), BlockPermutedDiagonalTruncationStrategy(strategy)
-      )
+      return truncate!($f, (D, V), BlockPermutedDiagonalTruncationStrategy(strategy))
     end
   end
 end
@@ -50,18 +45,30 @@ end
 function MatrixAlgebraKit.findtruncated(
   values::AbstractVector, strategy::BlockPermutedDiagonalTruncationStrategy
 )
-  ind = MatrixAlgebraKit.findtruncated(values, strategy.strategy)
+  ind = findtruncated(Vector(values), strategy.strategy)
   indexmask = falses(length(values))
   indexmask[ind] .= true
-  return indexmask
+  return to_truncated_indices(values, indexmask)
+end
+
+# Allow customizing the indices output by `findtruncated`
+# based on the type of `values`, for example to preserve
+# a block or Kronecker structure.
+to_truncated_indices(values, I) = I
+function to_truncated_indices(values::AbstractBlockVector, I::AbstractVector{Bool})
+  I′ = BlockedVector(I, blocklengths(axis(values)))
+  blocks = map(BlockRange(values)) do b
+    return _getindex(b, to_truncated_indices(values[b], I′[b]))
+  end
+  return blocks
 end
 
 function MatrixAlgebraKit.truncate!(
   ::typeof(svd_trunc!),
   (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix},
   strategy::BlockPermutedDiagonalTruncationStrategy,
 )
-  I = MatrixAlgebraKit.findtruncated(diagview(S), strategy)
+  I = MatrixAlgebraKit.findtruncated(diag(S), strategy)
   return (U[:, I], S[I, I], Vᴴ[I, :])
 end
 for f in [:eig_trunc!, :eigh_trunc!]
@@ -71,7 +78,7 @@ for f in [:eig_trunc!, :eigh_trunc!]
       (D, V)::NTuple{2,AbstractBlockSparseMatrix},
       strategy::BlockPermutedDiagonalTruncationStrategy,
     )
-      I = MatrixAlgebraKit.findtruncated(diagview(D), strategy)
+      I = MatrixAlgebraKit.findtruncated(diag(D), strategy)
       return (D[I, I], V[:, I])
     end
   end