qr.hpp

Go to the documentation of this file.

// Copyright (c) 2019--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "./utils.hpp"
#include "../basic_array.hpp"
#include "../basic_functions.hpp"
#include "../blas/tools.hpp"
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../lapack/geqp3.hpp"
#include "../lapack/orgqr.hpp"
#include "../lapack/ungqr.hpp"
#include "../layout/policies.hpp"
#include "../layout/range.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <algorithm>
#include <tuple>
#include <type_traits>
 
namespace nda::linalg {


  template <blas_lapack::BlasArray<2> A, blas_lapack::BlasArrayFor<A, 1> TAU>
    requires(mem::have_host_compatible_addr_space<A, TAU>)
  auto get_qr_matrices(A const &a, TAU const &tau, bool complete = false) {
    auto const [m, n] = a.shape();
    auto const min_mn = std::min(m, n);
    auto const k      = (complete ? m : min_mn);
    auto Q            = matrix<get_value_t<A>, F_layout>::zeros(m, k);
    auto R            = matrix<get_value_t<A>, F_layout>::zeros(k, n);
 
    // compute Q matrix
    Q(range::all, range(min_mn)) = a(range::all, range(min_mn));
    int info{};
    if constexpr (is_complex_v<get_value_t<A>>) {
      info = lapack::ungqr(Q, tau);
    } else {
      info = lapack::orgqr(Q, tau);
    }
    if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::get_qr_matrices: orgqr/ungqr returned a non-zero value: info = " << info;
 
    // extract R matrix
    for (int i = 0; i < min_mn; ++i) R(range(i + 1), i) = a(range(i + 1), i);
    for (int i = min_mn; i < n; ++i) R(range::all, i) = a(range::all, i);
 
    return std::make_tuple(Q, R);
  }

  template <blas_lapack::BlasArray<2> A>
    requires(mem::have_host_compatible_addr_space<A> and blas_lapack::has_F_layout<A>)
  auto qr_in_place(A &&a, bool complete = false) { // NOLINT (temporary views are allowed here)
    auto const [m, n] = a.shape();
 
    // permutation and tau vector
    auto jpvt = vector<int>::zeros(n);
    auto tau  = vector<get_value_t<A>>(std::min(m, n));
 
    // call lapack geqp3
    int info = lapack::geqp3(a, jpvt, tau);
    if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::qr_in_place: geqp3 returned a non-zero value: info = " << info;
 
    // extract Q and R from the output
    auto [Q, R] = get_qr_matrices(a, tau, complete);
 
    // shift permutation vector to zero-based indexing
    jpvt -= 1;
 
    return std::make_tuple(jpvt, Q, R);
  }

  template <Matrix A>
    requires(mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>>)
  auto qr(A const &a, bool complete = false) {
    auto a_copy = matrix<get_value_t<A>, F_layout>(a);
    if constexpr (blas_lapack::has_F_layout<A>) {
      return qr_in_place(a_copy, complete);
    } else {
      auto [sigma, Q, R] = qr_in_place(a_copy, complete);
      return std::make_tuple(sigma, matrix<get_value_t<A>, C_layout>{Q}, matrix<get_value_t<A>, C_layout>{R});
    }
  }

 
} // namespace nda::linalg