added column pivoted QR solve to math::linalg::non_distributed

Author: kshitij-05

src/TiledArray/math/linalg/non-distributed/qr.h

@@ -34,6 +34,21 @@ auto householder_qr(const ArrayV& V, TiledRange q_trange = TiledRange(),
   }
 }
 
+template <typename ArrayA, typename ArrayB, typename T = ArrayB::numeric_type>
+auto qr_solve(const ArrayA& A, const ArrayB& B,
+              const TiledArray::detail::real_t<T> cond = 1e8,
+              TiledRange x_trange = TiledRange()) {
+  (void)detail::array_traits<ArrayB>{};
+  auto& world = B.world();
+  auto A_eig = detail::make_matrix(A);
+  auto B_eig = detail::make_matrix(B);
+  TA_LAPACK_ON_RANK_ZERO(qr_solve, world, A_eig, B_eig, cond);
+  world.gop.broadcast_serializable(B_eig, 0);
+  if (x_trange.rank() == 0) x_trange = B.trange();
+  auto X = eigen_to_array<ArrayB>(world, x_trange, B_eig);
+  return X;
+}
+
 }  // namespace TiledArray::math::linalg::non_distributed
 
 #endif

src/TiledArray/math/linalg/rank-local.cpp

@@ -112,6 +112,22 @@ void cholesky_lsolve(Op transpose, Matrix<T>& A, Matrix<T>& X) {
   TA_LAPACK(trtrs, uplo, transpose, diag, n, nrhs, a, lda, b, ldb);
 }
 
+template <typename T>
+void qr_solve(Matrix<T>& A, Matrix<T>& B,
+              const TiledArray::detail::real_t<T> cond) {
+  integer m = A.rows();
+  integer n = A.cols();
+  integer nrhs = B.cols();
+  T* a = A.data();
+  integer lda = A.rows();
+  T* b = B.data();
+  integer ldb = B.rows();
+  std::vector<integer> jpiv(n);
+  const TiledArray::detail::real_t<T> rcond = 1 / cond;
+  integer rank = -1;
+  TA_LAPACK(gelsy, m, n, nrhs, a, lda, b, ldb, jpiv.data(), rcond, &rank);
+}
+
 template <typename T>
 void heig(Matrix<T>& A, std::vector<TiledArray::detail::real_t<T>>& W) {
   auto jobz = lapack::Job::Vec;
@@ -250,7 +266,7 @@ void householder_qr(Matrix<T>& V, Matrix<T>& R) {
     lapack::orgqr(m, n, k, v, ldv, tau.data());
 }
 
-#define TA_LAPACK_EXPLICIT(MATRIX, VECTOR)                         \
+#define TA_LAPACK_EXPLICIT(MATRIX, VECTOR, DOUBLE)                 \
   template void cholesky(MATRIX&);                                 \
   template void cholesky_linv(MATRIX&);                            \
   template void cholesky_solve(MATRIX&, MATRIX&);                  \
@@ -261,11 +277,12 @@ void householder_qr(Matrix<T>& V, Matrix<T>& R) {
   template void lu_solve(MATRIX&, MATRIX&);                        \
   template void lu_inv(MATRIX&);                                   \
   template void householder_qr<true>(MATRIX&, MATRIX&);            \
-  template void householder_qr<false>(MATRIX&, MATRIX&);
+  template void householder_qr<false>(MATRIX&, MATRIX&);           \
+  template void qr_solve(MATRIX&, MATRIX&, DOUBLE)
 
-TA_LAPACK_EXPLICIT(Matrix<double>, std::vector<double>);
-TA_LAPACK_EXPLICIT(Matrix<float>, std::vector<float>);
-TA_LAPACK_EXPLICIT(Matrix<std::complex<double>>, std::vector<double>);
-TA_LAPACK_EXPLICIT(Matrix<std::complex<float>>, std::vector<float>);
+TA_LAPACK_EXPLICIT(Matrix<double>, std::vector<double>, double );
+TA_LAPACK_EXPLICIT(Matrix<float>, std::vector<float>, float);
+TA_LAPACK_EXPLICIT(Matrix<std::complex<double>>, std::vector<double>, double);
+TA_LAPACK_EXPLICIT(Matrix<std::complex<float>>, std::vector<float>, float);
 
 }  // namespace TiledArray::math::linalg::rank_local

src/TiledArray/math/linalg/rank-local.h

@@ -41,6 +41,10 @@ void cholesky_solve(Matrix<T> &A, Matrix<T> &X);
 template <typename T>
 void cholesky_lsolve(Op transpose, Matrix<T> &A, Matrix<T> &X);
 
+template <typename T>
+void qr_solve(Matrix<T> &A, Matrix<T> &B,
+              const TiledArray::detail::real_t<T> cond = 1e8);
+
 template <typename T>
 void heig(Matrix<T> &A, std::vector<TiledArray::detail::real_t<T>> &W);