lu.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/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(nda::mem::have_host_compatible_addr_space<A> and nda::blas::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, 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 lu_in_place(A &&a, bool allow_singular = false) { // NOLINT (temporary views are allowed here)
    // input, output types and static assertions
 
    // pivot indices vector
    auto ipiv = vector<int>{};
 
    // call lapack getrf
    int info = lapack::getrf(a, ipiv);
    if (info < 0) {
      NDA_RUNTIME_ERROR << "Error in nda::lu_in_place: getrf failed with invalid argument (info = " << info << ")";
    } else if (info > 0 and not allow_singular) {
      NDA_RUNTIME_ERROR << "Error in nda::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(nda::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 (nda::blas::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