Add matrixU and matrixL in SparseLU

Author: Lucas-Haubert

include/nanoeigenpy/decompositions/sparse/sparse-lu.hpp

@@ -10,20 +10,74 @@ namespace nanoeigenpy {
 namespace nb = nanobind;
 using namespace nb::literals;
 
+template <typename LTypeOrUType, typename MatrixOrVector>
+static void solveInPlace(const LTypeOrUType &self,
+                         Eigen::Ref<MatrixOrVector> mat_vec) {
+  self.solveInPlace(mat_vec);
+}
+
+template <typename MappedSupernodalType>
+void exposeMatrixL(nb::module_ m) {
+  using LType = Eigen::SparseLUMatrixLReturnType<MappedSupernodalType>;
+  using Scalar = typename MappedSupernodalType::Scalar;
+  using VectorXs = Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Eigen::ColMajor>;
+  using MatrixXs =
+      Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;
+
+  if (check_registration_alias<LType>(m)) {
+    return;
+  }
+
+  nb::class_<LType>(m, "SparseLUMatrixLReturnType")
+      .def("rows", &LType::rows)
+      .def("cols", &LType::cols)
+      .def("solveInPlace", &solveInPlace<LType, MatrixXs>, "X"_a)
+      .def("solveInPlace", &solveInPlace<LType, VectorXs>, "x"_a);
+}
+
+template <typename MatrixLType, typename MatrixUType>
+void exposeMatrixU(nb::module_ m) {
+  using UType = Eigen::SparseLUMatrixUReturnType<MatrixLType, MatrixUType>;
+  using Scalar = typename MatrixLType::Scalar;
+  using VectorXs = Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Eigen::ColMajor>;
+  using MatrixXs =
+      Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;
+
+  if (check_registration_alias<UType>(m)) {
+    return;
+  }
+
+  nb::class_<UType>(m, "SparseLUMatrixUReturnType")
+      .def("rows", &UType::rows)
+      .def("cols", &UType::cols)
+      .def("solveInPlace", &solveInPlace<UType, MatrixXs>, "X"_a)
+      .def("solveInPlace", &solveInPlace<UType, VectorXs>, "x"_a);
+}
+
 template <typename _MatrixType, typename _Ordering = Eigen::COLAMDOrdering<
                                     typename _MatrixType::StorageIndex>>
 void exposeSparseLU(nb::module_ m, const char *name) {
   using MatrixType = _MatrixType;
   using Solver = Eigen::SparseLU<MatrixType>;
+  using Scalar = typename MatrixType::Scalar;
   using RealScalar = typename MatrixType::RealScalar;
+  using StorageIndex = typename MatrixType::StorageIndex;
   using SparseLUTransposeViewTrue = Eigen::SparseLUTransposeView<true, Solver>;
   using SparseLUTransposeViewFalse =
       Eigen::SparseLUTransposeView<false, Solver>;
+  using SCMatrix = typename Solver::SCMatrix;
+  using MappedSparseMatrix =
+      typename Eigen::MappedSparseMatrix<Scalar, Eigen::ColMajor, StorageIndex>;
+  using LType = Eigen::SparseLUMatrixLReturnType<SCMatrix>;
+  using UType = Eigen::SparseLUMatrixUReturnType<SCMatrix, MappedSparseMatrix>;
 
   if (check_registration_alias<Solver>(m)) {
     return;
   }
 
+  exposeMatrixL<SCMatrix>(m);
+  exposeMatrixU<SCMatrix, MappedSparseMatrix>(m);
+
   nb::class_<SparseLUTransposeViewFalse>(m, "SparseLUTransposeView")
       .def(SparseSolverBaseVisitor())
       .def("setIsInitialized", &SparseLUTransposeViewFalse::setIsInitialized)
@@ -104,6 +158,13 @@ void exposeSparseLU(nb::module_ m, const char *name) {
           },
           "Returns an expression of the adjoint of the factored matrix.")
 
+      .def(
+          "matrixL", [](const Solver &self) -> LType { return self.matrixL(); },
+          "Returns an expression of the matrix L.")
+      .def(
+          "matrixU", [](const Solver &self) -> UType { return self.matrixU(); },
+          "Returns an expression of the matrix U.")
+
       .def("rows", &Solver::rows, "Returns the number of rows of the matrix.")
       .def("cols", &Solver::cols, "Returns the number of cols of the matrix.")
 

src/module.cpp

@@ -27,14 +27,22 @@ using FullPivHhJacobiSVD =
 using HhJacobiSVD = JacobiSVD<Matrix, HouseholderQRPreconditioner>;
 using NoPrecondJacobiSVD = JacobiSVD<Matrix, NoQRPreconditioner>;
 
+using SparseQR = Eigen::SparseQR<SparseMatrix, Eigen::COLAMDOrdering<int>>;
+using SparseLU = Eigen::SparseLU<SparseMatrix>;
+using SCMatrix = typename SparseLU::SCMatrix;
+using StorageIndex = typename Matrix::StorageIndex;
+using MappedSparseMatrix =
+    typename Eigen::MappedSparseMatrix<Scalar, Options, StorageIndex>;
+
 NB_MAKE_OPAQUE(ColPivHhJacobiSVD)
 NB_MAKE_OPAQUE(FullPivHhJacobiSVD)
 NB_MAKE_OPAQUE(HhJacobiSVD)
 NB_MAKE_OPAQUE(NoPrecondJacobiSVD)
 
-using SparseQR = Eigen::SparseQR<SparseMatrix, Eigen::COLAMDOrdering<int>>;
 NB_MAKE_OPAQUE(Eigen::SparseQRMatrixQReturnType<SparseQR>)
 NB_MAKE_OPAQUE(Eigen::SparseQRMatrixQTransposeReturnType<SparseQR>)
+NB_MAKE_OPAQUE(Eigen::SparseLUMatrixLReturnType<SCMatrix>)
+NB_MAKE_OPAQUE(Eigen::SparseLUMatrixUReturnType<SCMatrix, MappedSparseMatrix>)
 
 NB_MAKE_OPAQUE(Eigen::LLT<Eigen::MatrixXd>)
 NB_MAKE_OPAQUE(Eigen::LDLT<Eigen::MatrixXd>)

tests/test_sparse_lu.py

@@ -35,3 +35,29 @@
 assert isinstance(X_est, spa.csc_matrix)
 assert nanoeigenpy.is_approx(X_est.toarray(), X_sparse.toarray())
 assert nanoeigenpy.is_approx(A.dot(X_est.toarray()), B_sparse.toarray())
+
+assert splu.nnzL() > 0
+assert splu.nnzU() > 0
+
+L = splu.matrixL()
+U = splu.matrixU()
+
+assert L.rows() == dim
+assert L.cols() == dim
+assert U.rows() == dim
+assert U.cols() == dim
+
+x_true = rng.random(dim)
+b_true = A.dot(x_true)
+P_rows_indices = splu.rowsPermutation().indices()
+P_cols_indices = splu.colsPermutation().indices()
+
+b_permuted = b_true[P_rows_indices]
+z = b_permuted.copy()
+L.solveInPlace(z)
+y = z.copy()
+U.solveInPlace(y)
+x_reconstructed = np.zeros(dim)
+x_reconstructed[P_cols_indices] = y
+
+assert nanoeigenpy.is_approx(x_reconstructed, x_true, 1e-6)