qr.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell

 
#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 <MemoryMatrix A, MemoryVector TAU>
    requires(nda::mem::have_host_compatible_addr_space<A, TAU> and have_same_value_type_v<A, TAU> and is_blas_lapack_v<get_value_t<A>>)
  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::qr_in_place: 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 <MemoryMatrix A>
    requires(nda::mem::have_host_compatible_addr_space<A> and nda::blas::has_F_layout<A> and is_blas_lapack_v<get_value_t<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 = nda::vector<int>::zeros(n);
    auto tau  = nda::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::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(nda::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 (nda::blas::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