SparseQR: Add matrixQ method

Author: Lucas-Haubert

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

@@ -10,21 +10,82 @@ namespace nanoeigenpy {
 namespace nb = nanobind;
 using namespace nb::literals;
 
+template <typename SparseQRType>
+void exposeMatrixQ(nb::module_ m) {
+  using Scalar = typename SparseQRType::Scalar;
+  using QType = Eigen::SparseQRMatrixQReturnType<SparseQRType>;
+  using QTransposeType =
+      Eigen::SparseQRMatrixQTransposeReturnType<SparseQRType>;
+  using VectorXd = Eigen::VectorXd;
+  using MatrixXd = Eigen::MatrixXd;
+  using QRMatrixType = typename SparseQRType::QRMatrixType;
+
+  if (!check_registration_alias<QTransposeType>(m)) {
+    nb::class_<QTransposeType>(m, "SparseQRMatrixQTransposeReturnType")
+        .def(nb::init<const SparseQRType&>(), "qr"_a)
+
+        .def(
+            "__matmul__",
+            [](QTransposeType& self, const MatrixXd& other) -> MatrixXd {
+              return MatrixXd(self * other);
+            },
+            "other"_a)
+
+        .def(
+            "__matmul__",
+            [](QTransposeType& self, const VectorXd& other) -> VectorXd {
+              return VectorXd(self * other);
+            },
+            "other"_a);
+  }
+
+  if (!check_registration_alias<QType>(m)) {
+    nb::class_<QType>(m, "SparseQRMatrixQReturnType")
+        .def(nb::init<const SparseQRType&>(), "qr"_a)
+
+        .def_prop_ro("rows", &QType::rows)
+        .def_prop_ro("cols", &QType::cols)
+
+        .def(
+            "__matmul__",
+            [](QType& self, const MatrixXd& other) -> MatrixXd {
+              return MatrixXd(self * other);
+            },
+            "other"_a)
+
+        .def(
+            "__matmul__",
+            [](QType& self, const VectorXd& other) -> VectorXd {
+              return VectorXd(self * other);
+            },
+            "other"_a)
+
+        .def("adjoint",
+             [](const QType& self) -> QTransposeType { return self.adjoint(); })
+
+        .def("transpose", [](const QType& self) -> QTransposeType {
+          return self.transpose();
+        });
+  }
+}
+
 template <typename _MatrixType, typename _Ordering = Eigen::COLAMDOrdering<
                                     typename _MatrixType::StorageIndex>>
-void exposeSparseQR(nb::module_ m, const char *name) {
+void exposeSparseQR(nb::module_ m, const char* name) {
   using MatrixType = _MatrixType;
   using Ordering = _Ordering;
   using Solver = Eigen::SparseQR<MatrixType, Ordering>;
   using Scalar = typename MatrixType::Scalar;
   using RealScalar = typename MatrixType::RealScalar;
-  using QRMatrixType = Eigen::SparseMatrix<Scalar, Eigen::ColMajor,
-                                           typename MatrixType::StorageIndex>;
+  using QRMatrixType = typename Solver::QRMatrixType;
+  using QType = Eigen::SparseQRMatrixQReturnType<Solver>;
 
   if (check_registration_alias<Solver>(m)) {
     return;
   }
 
+  exposeMatrixQ<Solver>(m);
+
   nb::class_<Solver>(
       m, name,
       "Sparse left-looking QR factorization with numerical column pivoting. "
@@ -51,7 +112,7 @@ void exposeSparseQR(nb::module_ m, const char *name) {
       "factor of full rank.")
 
       .def(nb::init<>(), "Default constructor.")
-      .def(nb::init<const MatrixType &>(), "matrix"_a,
+      .def(nb::init<const MatrixType&>(), "matrix"_a,
            "Constructs a LU factorization from a given matrix.")
 
       .def(SparseSolverBaseVisitor())
@@ -71,9 +132,13 @@ void exposeSparseQR(nb::module_ m, const char *name) {
            "The input matrix should be in compressed mode "
            "(see SparseMatrix::makeCompressed()).")
 
+      .def(
+          "matrixQ", [](const Solver& self) -> QType { return self.matrixQ(); },
+          "Returns an expression of the matrix Q as products of sparse "
+          "Householder reflectors.")
       .def(
           "matrixR",
-          [](const Solver &self) -> const QRMatrixType & {
+          [](const Solver& self) -> const QRMatrixType& {
             return self.matrixR();
           },
           "Returns a const reference to the \b sparse upper triangular matrix "
@@ -99,7 +164,7 @@ void exposeSparseQR(nb::module_ m, const char *name) {
 
       .def(
           "setPivotThreshold",
-          [](Solver &self, const RealScalar &thresh) -> void {
+          [](Solver& self, const RealScalar& thresh) -> void {
             return self.setPivotThreshold(thresh);
           },
           "Set the threshold used for a diagonal entry to be an acceptable "

src/module.cpp

@@ -32,6 +32,10 @@ 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::LLT<Eigen::MatrixXd>)
 NB_MAKE_OPAQUE(Eigen::LDLT<Eigen::MatrixXd>)
 NB_MAKE_OPAQUE(Eigen::FullPivLU<Eigen::MatrixXd>)

tests/test_sparse_qr.py

@@ -35,3 +35,39 @@
 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())
+
+Q = spqr.matrixQ()
+R = spqr.matrixR()
+P = spqr.colsPermutation()
+
+assert spqr.matrixQ().rows == dim
+assert spqr.matrixQ().cols == dim
+assert R.shape[0] == dim
+assert R.shape[1] == dim
+assert P.indices().size == dim
+
+test_vec = rng.random(dim)
+test_matrix = rng.random((dim, 20))
+
+Qv = Q @ test_vec
+QM = Q @ test_matrix
+Qt = Q.transpose()
+QtV = Qt @ test_vec
+QtM = Qt @ test_matrix
+
+assert Qv.shape == (dim,)
+assert QM.shape == (dim, 20)
+assert QtV.shape == (dim,)
+assert QtM.shape == (dim, 20)
+
+Qa_real_mat = Q.adjoint()
+QaV = Qa_real_mat @ test_vec
+assert nanoeigenpy.is_approx(QtV, QaV)
+
+A_dense = A.toarray()
+P_indices = np.array([P.indices()[i] for i in range(dim)])
+A_permuted = A_dense[:, P_indices]
+
+QtAP = Qt @ A_permuted
+R_dense = spqr.matrixR().toarray()
+assert nanoeigenpy.is_approx(QtAP, R_dense)