Also fix factorizations

Author: lkdvos

src/factorizations/truncation.jl

@@ -1,4 +1,6 @@
 using MatrixAlgebraKit:
+  MatrixAlgebraKit,
+  TruncatedAlgorithm,
   TruncationStrategy,
   diagview,
   eig_trunc!,
@@ -8,42 +10,58 @@ using MatrixAlgebraKit:
   truncate!
 
 """
-    BlockPermutedDiagonalTruncationStrategy(strategy::TruncationStrategy)
+    BlockDiagonalTruncationStrategy(strategy::TruncationStrategy)
 
 A wrapper for `TruncationStrategy` that implements the wrapped strategy on a block-by-block
-basis, which is possible if the input matrix is a block-diagonal matrix or a block permuted
-block-diagonal matrix.
+basis, which is possible if the input matrix is a block-diagonal matrix.
 """
-struct BlockPermutedDiagonalTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy
+struct BlockDiagonalTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy
   strategy::T
 end
 
-function MatrixAlgebraKit.truncate!(
-  ::typeof(svd_trunc!),
-  (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix},
-  strategy::TruncationStrategy,
+function BlockDiagonalTruncationStrategy(alg::BlockPermutedDiagonalAlgorithm)
+  return BlockDiagonalTruncationStrategy(alg.strategy)
+end
+
+function MatrixAlgebraKit.svd_trunc!(
+  A::AbstractBlockSparseMatrix,
+  out,
+  alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm},
 )
-  # TODO assert blockdiagonal
-  return truncate!(
-    svd_trunc!, (U, S, Vᴴ), BlockPermutedDiagonalTruncationStrategy(strategy)
-  )
+  Ad, (invrowperm, invcolperm) = blockdiagonalize(A)
+  blockalg = BlockDiagonalAlgorithm(alg.alg)
+  blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc)
+  Ud, S, Vᴴd = svd_trunc!(Ad, TruncatedAlgorithm(blockalg, blockstrategy))
+
+  U = transform_rows(Ud, invrowperm)
+  Vᴴ = transform_cols(Vᴴd, invcolperm)
+
+  return U, S, Vᴴ
 end
+
 for f in [:eig_trunc!, :eigh_trunc!]
   @eval begin
-    function MatrixAlgebraKit.truncate!(
-      ::typeof($f),
-      (D, V)::NTuple{2,AbstractBlockSparseMatrix},
-      strategy::TruncationStrategy,
+    function MatrixAlgebraKit.$f(
+      A::AbstractBlockSparseMatrix,
+      out,
+      alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm},
     )
-      return truncate!($f, (D, V), BlockPermutedDiagonalTruncationStrategy(strategy))
+      Ad, (invrowperm, invcolperm) = blockdiagonalize(A)
+      blockalg = BlockDiagonalAlgorithm(alg.alg)
+      blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc)
+      Dd, Vd = $f(Ad, TruncatedAlgorithm(blockalg, blockstrategy))
+
+      D = transform_rows(Dd, invrowperm)
+      V = transform_cols(Vd, invcolperm)
+      return D, V
     end
   end
 end
 
 # cannot use regular slicing here: I want to slice without altering blockstructure
 # solution: use boolean indexing and slice the mask, effectively cheaply inverting the map
 function MatrixAlgebraKit.findtruncated(
-  values::AbstractVector, strategy::BlockPermutedDiagonalTruncationStrategy
+  values::AbstractVector, strategy::BlockDiagonalTruncationStrategy
 )
   ind = findtruncated(Vector(values), strategy.strategy)
   indexmask = falses(length(values))
@@ -66,7 +84,7 @@ end
 function MatrixAlgebraKit.truncate!(
   ::typeof(svd_trunc!),
   (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix},
-  strategy::BlockPermutedDiagonalTruncationStrategy,
+  strategy::BlockDiagonalTruncationStrategy,
 )
   I = findtruncated(diag(S), strategy)
   return (U[:, I], S[I, I], Vᴴ[I, :])
@@ -76,7 +94,7 @@ for f in [:eig_trunc!, :eigh_trunc!]
     function MatrixAlgebraKit.truncate!(
       ::typeof($f),
       (D, V)::NTuple{2,AbstractBlockSparseMatrix},
-      strategy::BlockPermutedDiagonalTruncationStrategy,
+      strategy::BlockDiagonalTruncationStrategy,
     )
       I = findtruncated(diag(D), strategy)
       return (D[I, I], V[:, I])