From 819b9bd559021fac274d4ec3a45282537f4635f3 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 07:13:06 -0500 Subject: [PATCH] `to_blockdiagonal` now returns blocks --- .../src/blocksparsearray/blockdiagonalarray.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blockdiagonalarray.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blockdiagonalarray.jl index ee2d408111..e20eb8b0c7 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blockdiagonalarray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blockdiagonalarray.jl @@ -20,7 +20,7 @@ function try_to_blockdiagonal_perm(A) p = sortperm(J) Jsorted = J[p] allunique(Jsorted) || return nothing - return I[p], Jsorted + return Block.(I[p], Jsorted) end """ @@ -32,8 +32,11 @@ returns nothing, otherwise returns both the blockdiagonal `B` as well as the per function try_to_blockdiagonal(A::AbstractBlockSparseMatrix) perm = try_to_blockdiagonal_perm(A) isnothing(perm) && return perm - I, J = perm - diagblocks = blocks(A)[CartesianIndex.(invperm(I), J)] + I = first.(Tuple.(perm)) + J = last.(Tuple.(perm)) + diagblocks = map(invperm(I), J) do i, j + return A[Block(i, j)] + end return BlockDiagonal(diagblocks), perm end