solve.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 "../basic_array.hpp"
#include "../blas/tools.hpp"
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../lapack/getrf.hpp"
#include "../lapack/getrs.hpp"
#include "../layout/policies.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../mem/policies.hpp"
#include "../traits.hpp"
 
#include <type_traits>
 
namespace nda::linalg {


  template <MemoryMatrix A, MemoryArray B>
    requires(have_same_value_type_v<A, B> and nda::mem::have_compatible_addr_space<A, B> and is_blas_lapack_v<get_value_t<A>>)
  void solve_in_place(A &&a, B &&b) { // NOLINT (temporary views are allowed here)
    static_assert(get_rank<B> == 1 || get_rank<B> == 2, "Error in nda::linalg::solve_in_place: Right hand side must have rank 1 or 2");
    static_assert(nda::blas::has_F_layout<B>, "Error in nda::linalg::solve_in_place: Right hand side must have Fortran layout");
 
    constexpr auto addr_space = nda::mem::common_addr_space<A, B>;
 
    // check the dimensions of the input/output arrays/views
    EXPECTS_WITH_MESSAGE(a.shape()[0] == a.shape()[1], "Error in nda::linalg::solve_in_place: Matrix A is not square");
    EXPECTS_WITH_MESSAGE(a.shape()[0] == b.extent(0), "Error in nda::linalg::solve_in_place: Dimension mismatch between matrix A and B");
 
    // pivot indices vector
    auto ipiv = vector<int, heap<addr_space>>(a.extent(0));
 
    // call lapack getrf to compute LU factorization
    int info = lapack::getrf(a, ipiv);
    if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::solve_in_place: getrf returned a non-zero value: info = " << info;
 
    // call lapack getrs to solve using the LU factorization
    info = lapack::getrs(a, b, ipiv);
    if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::solve_in_place: getrs returned a non-zero value: info = " << info;
  }

  template <Matrix A, Array B>
    requires(have_same_value_type_v<A, B> and nda::mem::have_compatible_addr_space<A, B> and is_blas_lapack_v<get_value_t<A>>)
  auto solve(A const &a, B const &b) { // NOLINT (temporary views are allowed here)
    // copy A and preserve its layout
    using a_layout_policy = nda::detail::layout_to_policy<typename std::remove_cvref_t<A>::layout_t>::type;
    auto a_copy           = matrix<get_value_t<A>, a_layout_policy, heap<nda::mem::common_addr_space<A, B>>>(a);
 
    // copy B and enforce Fortran layout for the matrix case
    using vector_t = vector<get_value_t<A>, heap<nda::mem::common_addr_space<A, B>>>;
    using matrix_t = matrix<get_value_t<A>, F_layout, heap<nda::mem::common_addr_space<A, B>>>;
    using b_type   = std::conditional_t<get_rank<B> == 1, vector_t, matrix_t>;
    auto b_copy    = b_type(b);
 
    // call solve_in_place with the copies
    solve_in_place(a_copy, b_copy);
    return b_copy;
  }

 
} // namespace nda::linalg