lu.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/getrf.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 <typename LP = F_layout, MemoryMatrix A>
    requires(mem::have_host_compatible_addr_space<A> and blas_lapack::has_F_layout<A>)
  auto get_lu_matrices(A const &a) {
    using output_t = matrix<get_value_t<A>, LP>;
 
    // copy the first k columns from A to L and U
    auto const [m, n] = a.shape();
    auto const k      = std::min(m, n);
    auto L            = output_t::zeros(m, k);
    auto U            = output_t::zeros(k, n);
    for (int i = 0; i < k; ++i) {
      L(i, i)               = get_value_t<A>{1};
      L(range(i + 1, m), i) = a(range(i + 1, m), i);
      U(range(i + 1), i)    = a(range(i + 1), i);
    }
 
    // in case of n > m, copy the remaining columns to U
    for (int i = k; i < n; ++i) U(range::all, i) = a(range::all, i);
 
    return std::make_tuple(L, U);
  }

  template <typename LP = F_layout, blas_lapack::BlasArray<2> A>
    requires(mem::have_host_compatible_addr_space<A> and blas_lapack::has_F_layout<A>)
  auto lu_in_place(A &&a, bool allow_singular = false) { // NOLINT (temporary views are allowed here)
    // pivot indices vector
    auto ipiv = vector<int>{};
 
    // call getrf
    int info = lapack::getrf(a, ipiv);
    if (info < 0) {
      NDA_RUNTIME_ERROR << "Error in nda::linalg::lu_in_place: getrf failed with invalid argument (info = " << info << ")";
    } else if (info > 0 and not allow_singular) {
      NDA_RUNTIME_ERROR << "Error in nda::linalg::lu_in_place: Matrix is singular, U(" << info << "," << info << ") is exactly zero";
    }
 
    // extract sigma, L, U from the output of getrf
    auto sigma  = get_permutation_vector(ipiv, a.extent(0));
    auto [L, U] = get_lu_matrices<LP>(a);
 
    return std::make_tuple(sigma, L, U);
  }

  template <Matrix A>
    requires(mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>>)
  auto lu(A const &a, bool allow_singular = false) {
    auto a_copy = matrix<get_value_t<A>, F_layout>(a);
    if constexpr (blas_lapack::has_F_layout<A>) {
      return lu_in_place(a_copy, allow_singular);
    } else {
      return lu_in_place<C_layout>(a_copy, allow_singular);
    }
  }

 
} // namespace nda::linalg