eig.hpp

Go to the documentation of this file.

// Copyright (c) 2024--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 "../blas/tools.hpp"
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../lapack/geev.hpp"
#include "../layout/policies.hpp"
#include "../macros.hpp"
#include "../matrix_functions.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <array>
#include <cmath>
#include <complex>
#include <concepts>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace nda::linalg {


  template <Vector WR, Vector WI>
    requires(mem::have_host_compatible_addr_space<WR, WI> and AnyOf<get_value_t<WR>, float, double> and have_same_value_type_v<WR, WI>)
  auto get_geev_eigenvalues(const WR &wr, const WI &wi) {
    // check the dimensions of the input arrays/views
    auto const n = wr.size();
    EXPECTS(n == wi.size());
 
    // generate eigenvalues
    using fp_t  = get_fp_t<WR>;
    auto lambda = array<std::complex<fp_t>, 1>(n);
    for (long i = 0; i < n; ++i) lambda(i) = std::complex<fp_t>(wr(i), wi(i));
    return lambda;
  }

  template <Vector WI, Matrix VA>
    requires(mem::have_host_compatible_addr_space<WI, VA> and FloatOrDouble<get_value_t<WI>> and have_same_value_type_v<WI, VA>)
  auto unpack_eigenvectors(const WI &wi, const VA &va) {
    using namespace std::complex_literals;
 
    // check the dimensions of the input arrays/views
    auto const n = wi.size();
    EXPECTS(va.shape() == (std::array<long, 2>{n, n}));
 
    // unpack eigenvectors
    using fp_t = get_fp_t<WI>;
    auto X     = matrix<std::complex<fp_t>, F_layout>(n, n);
    long j     = 0;
    while (j < n) {
      if (wi(j) > 0.0) {
        // complex conjugate eigenvalue pair --> we need to unpack the eigenvectors
        for (long i = 0; i < n; ++i) {
          X(i, j)     = std::complex<fp_t>{va(i, j), va(i, j + 1)};
          X(i, j + 1) = std::complex<fp_t>{va(i, j), -va(i, j + 1)};
        }
        j += 2;
      } else {
        // real eigenvalue --> eigenvector is purely real
        X(range::all, j) = va(range::all, j);
        ++j;
      }
    }
 
    return X;
  }

  template <Vector AR, Vector AI, Vector B>
    requires(mem::have_host_compatible_addr_space<AR, AI, B> and FloatOrDouble<get_value_t<AR>> and have_same_value_type_v<AR, AI, B>)
  auto get_ggev_eigenvalues(const AR &alphar, const AI &alphai, const B &beta) {
    // check the dimensions of the input arrays/views
    auto const n = alphar.size();
    EXPECTS(n == alphai.size());
    EXPECTS(n == beta.size());
 
    // generate eigenvalues: lambda_j = (alphar_j + i * alphai_j) / beta_j
    using fp_t  = get_fp_t<AR>;
    auto lambda = array<std::complex<fp_t>, 1>(n);
    for (long i = 0; i < n; ++i) lambda(i) = std::complex<fp_t>(alphar(i), alphai(i)) / std::complex<fp_t>(beta(i));
    return lambda;
  }

  template <Vector A, Vector B>
    requires(mem::have_host_compatible_addr_space<A, B> and AnyOf<get_value_t<A>, std::complex<float>, std::complex<double>>
             and have_same_value_type_v<A, B>)
  auto get_ggev_eigenvalues(const A &alpha, const B &beta) {
    // check the dimensions of the input arrays/views
    auto const n = alpha.size();
    EXPECTS(n == beta.size());
 
    // generate eigenvalues: lambda_j = alpha_j / beta_j
    using val_t = get_value_t<A>;
    auto lambda = array<val_t, 1>(n);
    for (long i = 0; i < n; ++i) lambda(i) = alpha(i) / beta(i);
    return lambda;
  }
 
  namespace detail {
 
    // Implementation for complex matrices - straightforward call to geev.
    template <blas_lapack::BlasArrayCplx<2> A>
    auto eig_impl(A &&a, char jobvl, char jobvr) { // NOLINT (temporary views are allowed here)
      using arr_t  = array<get_value_t<A>, 1>;
      using mat_t  = matrix<get_value_t<A>, F_layout>;
      auto const n = a.extent(0);
 
      // early return if the matrix is empty
      if (a.empty()) return std::make_tuple(arr_t{}, mat_t{}, mat_t{});
 
      // allocate outputs
      auto lambda = arr_t(n);
      auto U      = (jobvl == 'V') ? mat_t(n, n) : mat_t();
      auto V      = (jobvr == 'V') ? mat_t(n, n) : mat_t();
 
      // make the call to geev
      int info = lapack::geev(a, lambda, U, V, jobvl, jobvr);
      if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::detail::eig_impl: geev routine failed: info = " << info;
 
      return std::make_tuple(std::move(lambda), std::move(U), std::move(V));
    }
 
    // Implementation for real matrices.
    template <blas_lapack::BlasArrayReal<2> A>
    auto eig_impl(A &&a, char jobvl, char jobvr) { // NOLINT (temporary views are allowed here)
      using fp_t   = get_fp_t<A>;
      using arr_t  = array<std::complex<fp_t>, 1>;
      using mat_t  = matrix<std::complex<fp_t>, F_layout>;
      auto const n = a.extent(0);
 
      // early return if the matrix is empty
      if (a.empty()) return std::make_tuple(arr_t{}, mat_t{}, mat_t{});
 
      // allocate outputs
      auto wr = array<fp_t, 1>(n);
      auto wi = array<fp_t, 1>(n);
      auto vl = (jobvl == 'V') ? matrix<fp_t, F_layout>(n, n) : matrix<fp_t, F_layout>{};
      auto vr = (jobvr == 'V') ? matrix<fp_t, F_layout>(n, n) : matrix<fp_t, F_layout>{};
 
      // make the call to geev
      int info = lapack::geev(a, wr, wi, vl, vr, jobvl, jobvr);
      if (info != 0) NDA_RUNTIME_ERROR << "Error in nda::linalg::detail::eig_impl: geev routine failed: info = " << info;
 
      // get eigenvalues and eigenvectors from geev output
      auto lambda = get_geev_eigenvalues(wr, wi);
      auto U      = (jobvl == 'V') ? unpack_eigenvectors(wi, vl) : mat_t{};
      auto V      = (jobvr == 'V') ? unpack_eigenvectors(wi, vr) : mat_t{};
 
      return std::make_tuple(std::move(lambda), std::move(U), std::move(V));
    }
 
  } // namespace detail

  template <blas_lapack::BlasArray<2> A>
    requires(mem::have_host_compatible_addr_space<A> and blas_lapack::has_F_layout<A>)
  auto eig_in_place(A &&a) {
    auto [lambda, U, V] = detail::eig_impl(std::forward<A>(a), 'N', 'V');
    return std::make_pair(std::move(lambda), std::move(V));
  }

  template <blas_lapack::BlasArray<2> A>
    requires(mem::have_host_compatible_addr_space<A> and blas_lapack::has_F_layout<A>)
  auto eigvals_in_place(A &&a) {
    auto [lambda, U, V] = detail::eig_impl(std::forward<A>(a), 'N', 'N');
    return lambda;
  }

  template <Matrix A>
    requires(mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>>)
  auto eig(A const &a) {
    auto m_copy = matrix<get_value_t<A>, F_layout>(a);
    return eig_in_place(m_copy);
  }

  template <Matrix A>
    requires(mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>>)
  auto eigvals(A const &a) {
    auto m_copy = matrix<get_value_t<A>, F_layout>(a);
    return eigvals_in_place(m_copy);
  }

 
} // namespace nda::linalg