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 std::same_as<double, get_value_t<WR>> 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
    auto lambda = array<std::complex<double>, 1>(n);
    for (long i = 0; i < n; ++i) lambda(i) = std::complex<double>(wr(i), wi(i));
    return lambda;
  }

  template <Vector WI, Matrix VA>
    requires(mem::have_host_compatible_addr_space<WI, VA> and std::same_as<double, get_value_t<WI>> and have_same_value_type_v<WI, VA>)
  auto get_geev_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
    auto X = matrix<std::complex<double>, 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<double>{va(i, j), va(i, j + 1)};
          X(i, j + 1) = std::complex<double>{va(i, j), -va(i, j + 1)};
        }
        j += 2;
      } else {
        // real eigenvalue --> eigenvector is purely real
        X(nda::range::all, j) = va(nda::range::all, j);
        ++j;
      }
    }
 
    return X;
  }
 
  namespace detail {
 
    // Implementation for complex matrices - straightforward call to geev.
    template <MemoryMatrix A>
      requires(is_complex_v<get_value_t<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 = nda::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 <MemoryMatrix A>
      requires(std::same_as<double, get_value_t<A>>)
    auto eig_impl(A &&a, char jobvl, char jobvr) { // NOLINT (temporary views are allowed here)
      using arr_t  = array<std::complex<double>, 1>;
      using mat_t  = matrix<std::complex<double>, 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<double, 1>(n);
      auto wi = array<double, 1>(n);
      auto vl = (jobvl = 'V') ? matrix<double, F_layout>(n, n) : matrix<double, F_layout>{};
      auto vr = (jobvr = 'V') ? matrix<double, F_layout>(n, n) : matrix<double, F_layout>{};
 
      // make the call to geev
      int info = nda::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') ? get_geev_eigenvectors(wi, vl) : mat_t{};
      auto V      = (jobvr == 'V') ? get_geev_eigenvectors(wi, vr) : mat_t{};
 
      return std::make_tuple(std::move(lambda), std::move(U), std::move(V));
    }
 
  } // namespace detail

  template <MemoryMatrix A>
    requires(nda::mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>> and nda::blas::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 <MemoryMatrix A>
    requires(nda::mem::have_host_compatible_addr_space<A> and is_blas_lapack_v<get_value_t<A>> and nda::blas::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(nda::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 <MemoryMatrix A>
    requires(nda::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