Solvers.ipp

//==============================================================================
// Implementations:
 
extern "C" {
 
using complex_double = long double;
// https://gcc.gnu.org/onlinedocs/gcc-4.5.4/gfortran/ISO_005fC_005fBINDING.html
// on apple M1, long double is just double
// nb: it actually doesn't matter, since all operations happen on fortran end
 
void dsyevd_(char *, char *, int *, double *, int *, double *, double *, int *,
             int *, int *, int *);
void zheev_(char *, char *, int *, complex_double *, int *, double *,
            complex_double *, int *, double *, int *);
void dsygv_(int *, char *, char *, int *, double *, int *, double *, int *,
            double *, double *, int *, int *);
void zhegv_(int *, char *, char *, int *, complex_double *, int *,
            complex_double *, int *, double *, complex_double *, int *,
            double *, int *);
 
void dsyevr_(char *jobz, char *range, char *uplo, int *n, double *a, int *lda,
             double *vl, double *vu, int *il, int *iu, double *abstol, int *m,
             double *w, double *z, int *ldz, int *isuppz, double *work,
             int *lwork, int *iwork, int *liwork, int *info);
 
void dgetrf_(int *m, int *n, double *a, int *lda, int *ipiv, int *info);
void dgetrs_(char *trans, int *n, int *nrhs, double *a, int *lda, int *ipiv,
             double *b, int *ldb, int *info);
void zgetrf_(int *m, int *n, complex_double *a, int *lda, int *ipiv, int *info);
void zgetrs_(char *trans, int *n, int *nrhs, complex_double *a, int *lda,
             int *ipiv, complex_double *b, int *ldb, int *info);
}
 
namespace LinAlg {
 
//==============================================================================
// Solves matrix equation Ax=b for x, for known square matrix A and vector b.
template <typename T>
Vector<T> solve_Axeqb(Matrix<T> Am, const Vector<T> &b) {
  static_assert(std::is_same_v<T, double> ||
                  std::is_same_v<T, std::complex<double>>,
                "solve_Axeqb only available for Matrix<double> or "
                "Matrix<complex<double>>");
  assert(Am.rows() == b.size());
 
  // LAPACK solves in place, so start from a copy of b
  Vector<T> x = b;
 
  int dim = int(Am.rows());
  int nrhs = 1;
  int info = 0;
  std::vector<int> ipiv(static_cast<std::size_t>(dim));
 
  // Matrix is row-major, so LAPACK sees the transpose of Am. Factorise that,
  // then solve with trans='T', which gives Am*x=b
  char trans = 'T';
 
  if constexpr (std::is_same_v<T, double>) {
    dgetrf_(&dim, &dim, Am.data(), &dim, ipiv.data(), &info);
    dgetrs_(&trans, &dim, &nrhs, Am.data(), &dim, ipiv.data(), x.data(), &dim,
            &info);
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    complex_double *Am_ptr = reinterpret_cast<complex_double *>(Am.data());
    complex_double *x_ptr = reinterpret_cast<complex_double *>(x.data());
    zgetrf_(&dim, &dim, Am_ptr, &dim, ipiv.data(), &info);
    zgetrs_(&trans, &dim, &nrhs, Am_ptr, &dim, ipiv.data(), x_ptr, &dim, &info);
  }
 
  return x;
}
 
//============================================================================*
// Solves Av = ev for eigenvalues e and eigenvectors v of symmetric/Hermetian
// matrix A. Returns [e,v], where v(i,j) is the jth element of the ith
// eigenvector corresponding to ith eigenvalue, e(i). e is always real.
template <typename T>
std::pair<Vector<double>, Matrix<T>> symmhEigensystem(Matrix<T> A) {
  assert(A.rows() == A.cols());
 
  // E. values of Hermetian complex matrix are _real_
  auto eigen_vv = std::make_pair(Vector<double>(A.rows()), std::move(A));
  auto &[e_values, e_vectors] = eigen_vv;
 
  int dim = (int)e_vectors.rows();
  char jobz{'V'};
  char uplo{'U'};
  int info = 0;
 
  if constexpr (std::is_same_v<T, double>) {
 
    // Workspace query
    int lwork = -1, liwork = -1;
    double work_query = 0.0;
    int iwork_query = 0;
    dsyevd_(&jobz, &uplo, &dim, e_vectors.data(), &dim, e_values.data(),
            &work_query, &lwork, &iwork_query, &liwork, &info);
    lwork = static_cast<int>(work_query);
    liwork = iwork_query;
    std::vector<double> work(std::size_t(std::max(1, lwork)));
    std::vector<int> iwork(std::size_t(std::max(1, liwork)));
    dsyevd_(&jobz, &uplo, &dim, e_vectors.data(), &dim, e_values.data(),
            work.data(), &lwork, iwork.data(), &liwork, &info);
 
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    std::vector<T> work(3 * e_vectors.rows());
    int workspace_size = (int)work.size();
    // static_assert(sizeof(complex_double) == 16); // LAPACK assumption
    // nb: this isn't actually required
 
    // For complex double (Hermetian) matrix
    std::vector<double> Rwork(3 * e_vectors.rows() - 2ul);
    complex_double *evec_ptr =
      reinterpret_cast<complex_double *>(e_vectors.data());
    complex_double *work_ptr = reinterpret_cast<complex_double *>(work.data());
    zheev_(&jobz, &uplo, &dim, evec_ptr, &dim, e_values.data(), work_ptr,
           &workspace_size, Rwork.data(), &info);
  }
 
  if (info != 0) {
    std::cerr << "\nError 91: symmhEigensystem " << info << " " << dim << " "
              << std::endl;
    if (info < 0) {
      std::cerr << "The " << -info << "-th argument had an illegal value\n";
    } else {
      std::cerr << "The algorithm failed to converge; " << info
                << " off-diagonal elements of an intermediate "
                   "tridiagonal form did not converge to zero.\n";
    }
    std::cerr << info << " " << A.size() << "\n";
  }
 
  return eigen_vv;
}
 
//============================================================================*
template <typename T>
std::pair<Vector<double>, Matrix<T>> symmhEigensystem(Matrix<T> A, int number) {
  assert(A.rows() == A.cols());
  assert(number < (int)A.rows());
  static_assert(std::is_same_v<T, double>,
                "This overload uses dsyevr_ (real symmetric). "
                "For complex Hermitian use zheevr/cheevr.");
 
  // E. values of Hermetian complex matrix are _real_
  auto eigen_vv = std::make_pair(Vector<double>(A.rows()), Matrix<T>(A.rows()));
  auto &[e_values, e_vectors] = eigen_vv;
 
  int dim = (int)A.rows();
  char jobz{'V'};
  char uplo{'U'};
  char range{'I'};
  double vl{0.0};
  double vu{0.0};
  int il{1};
  int iu{number};
  double abstol{1.0e-12};
  int m{0};
  int info{0};
 
  std::vector<int> isuppz(std::size_t(2 * std::max(1, number)));
 
  // Workspace query
  int lwork = -1, liwork = -1;
  double work_query = 0.0;
  int iwork_query = 0;
  dsyevr_(&jobz, &range, &uplo, &dim, A.data(), &dim, &vl, &vu, &il, &iu,
          &abstol, &m, e_values.data(), e_vectors.data(), &dim, isuppz.data(),
          &work_query, &lwork, &iwork_query, &liwork, &info);
 
  if (info != 0) {
    std::cerr << "\nError: dsyevr workspace query info=" << info << "\n";
  }
 
  lwork = static_cast<int>(work_query);
  liwork = iwork_query;
  std::vector<double> work(std::size_t(std::max(1, lwork)));
  std::vector<int> iwork(std::size_t(std::max(1, liwork)));
 
  dsyevr_(&jobz, &range, &uplo, &dim, A.data(), &dim, &vl, &vu, &il, &iu,
          &abstol, &m, e_values.data(), e_vectors.data(), &dim, isuppz.data(),
          work.data(), &lwork, iwork.data(), &liwork, &info);
 
  if (info != 0) {
    std::cerr << "\nError: dsyevr (range=I) symmhEigensystem info=" << info
              << " dim=" << dim << "\n";
  }
 
  return eigen_vv;
}
 
//============================================================================*
template <typename T>
std::tuple<int, Vector<double>, Matrix<T>> symmhEigensystem(Matrix<T> A,
                                                            double all_below) {
  static_assert(std::is_same_v<T, double>,
                "This overload uses dsyevr (real symmetric). "
                "For complex Hermitian use zheevr/cheevr.");
 
  assert(A.rows() == A.cols());
  int n = static_cast<int>(A.rows());
 
  // Return containers (same idea as your code: values size n, vectors n x n)
  auto eigen_vv =
    std::make_tuple(int{0}, Vector<double>(A.rows()), Matrix<T>(A.rows()));
  auto &[num_solutions, e_values, e_vectors] = eigen_vv;
 
  char jobz{'V'};  // eigenvectors
  char range{'V'}; // value range
  char uplo{'U'};
 
  // dsyevr uses (vl,vu]; pick vl very negative so we get everything <= all_below
  // Avoid +/-inf: many Fortran ABIs dislike it.
  double vl = -std::numeric_limits<double>::max();
  double vu = all_below;
 
  int il{0}, iu{0};       // unused for RANGE='V'
  double abstol{1.0e-12}; // your choice; could also use DLAMCH("S")
  num_solutions = 0;
  int info{0};
 
  // Required by dsyevr
  std::vector<int> isuppz(std::size_t(2 * std::max(1, n)));
 
  // Workspace query
  int lwork{-1}, liwork{-1};
  double work_query{0.0};
  int iwork_query{0};
 
  dsyevr_(&jobz, &range, &uplo, &n, A.data(), &n, &vl, &vu, &il, &iu, &abstol,
          &num_solutions, e_values.data(), e_vectors.data(), &n, isuppz.data(),
          &work_query, &lwork, &iwork_query, &liwork, &info);
 
  if (info != 0) {
    std::cerr << "\nError: symmhEigensystem dsyevr(work query) info=" << info
              << " n=" << n << "\n";
    return eigen_vv;
  }
 
  lwork = static_cast<int>(work_query);
  liwork = iwork_query;
  std::vector<double> work(std::size_t(std::max(1, lwork)));
  std::vector<int> iwork(std::size_t(std::max(1, liwork)));
 
  // Actual call
  dsyevr_(&jobz, &range, &uplo, &n, A.data(), &n, &vl, &vu, &il, &iu, &abstol,
          &num_solutions, e_values.data(), e_vectors.data(), &n, isuppz.data(),
          work.data(), &lwork, iwork.data(), &liwork, &info);
 
  if (info != 0) {
    std::cerr << "\nError: symmhEigensystem dsyevr info=" << info << " n=" << n
              << " m=" << num_solutions << "\n";
  }
 
  return eigen_vv;
}
 
//==============================================================================
// Solves Av = eBv for eigenvalues e and eigenvectors v of symmetric/Hermetian
// matrix pair A,B. Returns [e,v], where v(i,j) is the jth element of the ith
// eigenvector corresponding to ith eigenvalue, e(i). e is always real.
template <typename T>
std::pair<Vector<double>, Matrix<T>> symmhEigensystem(Matrix<T> A,
                                                      Matrix<T> B) {
  assert(A.rows() == A.cols());
  assert(B.rows() == B.cols());
  assert(A.rows() == B.rows());
 
  // E. values of Hermetian complex matrix are _real_
  auto eigen_vv = std::make_pair(Vector<double>(A.rows()), std::move(A));
  auto &[e_values, e_vectors] = eigen_vv;
 
  int itype = 1;
  int dim = (int)e_vectors.rows();
  char jobz{'V'};
  char uplo{'U'};
  int info = 0;
 
  std::vector<T> work(3 * e_vectors.rows());
  int workspace_size = (int)work.size();
 
  if constexpr (std::is_same_v<T, double>) {
 
    // For double (symmetric real) matrix
    dsygv_(&itype, &jobz, &uplo, &dim, e_vectors.data(), &dim, B.data(), &dim,
           e_values.data(), work.data(), &workspace_size, &info);
 
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    // static_assert(sizeof(complex_double) == 16); // LAPACK assumption
    // nb: this isn't actually required
 
    // For complex double (Hermetian) matrix
    std::vector<double> Rwork(3 * e_vectors.rows() - 2ul);
    complex_double *evec_ptr =
      reinterpret_cast<complex_double *>(e_vectors.data());
    complex_double *Bmat_ptr = reinterpret_cast<complex_double *>(B.data());
    complex_double *work_ptr = reinterpret_cast<complex_double *>(work.data());
    zhegv_(&itype, &jobz, &uplo, &dim, evec_ptr, &dim, Bmat_ptr, &dim,
           e_values.data(), work_ptr, &workspace_size, Rwork.data(), &info);
    e_vectors.conj_in_place(); // ?? why?
  }
 
  if (info != 0) {
    std::cerr << "\nError 198: symmhEigensystem " << info << " " << dim << " "
              << B.size() << std::endl;
    if (info < 0) {
      std::cerr << "The " << -info << "-th argument had an illegal value\n";
    } else if (info <= dim) {
      std::cerr << "DSYEV failed to converge; " << info
                << " off-diagonal elements of an intermediate tridiagonal form "
                   "did not converge to zero;";
    } else {
      std::cerr << "The leading minor of order " << info
                << " of B is not positive definite. "
                   "The factorization of B could not be completed and no "
                   "eigenvalues or eigenvectors were computed.";
    }
  }
 
  return eigen_vv;
}
 
//==============================================================================
// Solves Av = ev for eigenvalues e and eigenvectors v of non-symmetric real
// matrix A. Returns [e,v], where v(i,j) is the jth element of the ith
// eigenvector corresponding to ith eigenvalue, e(i). A must be real, while e
// and v are complex.
template <typename T>
std::pair<Vector<std::complex<double>>, Matrix<std::complex<double>>>
genEigensystem(Matrix<T> A, bool sort) {
  assert(A.rows() == A.cols());
  static_assert(std::is_same_v<T, double>,
                "genEigensystem only for Real matrix");
 
  // Solves Av = ev for eigenvalues e and eigenvectors v
  const auto dim = A.rows();
  auto eigen_vv = std::make_pair(Vector<std::complex<double>>(dim),
                                 Matrix<std::complex<double>>(dim));
  auto &[e_values, e_vectors] = eigen_vv;
 
  // GSL: This function computes eigenvalues and right eigenvectors of the
  // n-by-n real nonsymmetric matrix A.
  auto A_gsl = A.as_gsl_view();
  auto val_gsl = e_values.as_gsl_view();
  auto vec_gsl = e_vectors.as_gsl_view();
  gsl_eigen_nonsymmv_workspace *work = gsl_eigen_nonsymmv_alloc(dim);
  gsl_eigen_nonsymmv(&A_gsl.matrix, &val_gsl.vector, &vec_gsl.matrix, work);
  gsl_eigen_nonsymmv_free(work);
  if (sort)
    gsl_eigen_nonsymmv_sort(&val_gsl.vector, &vec_gsl.matrix,
                            GSL_EIGEN_SORT_ABS_ASC);
 
  // Eigensystems using GSL. NOTE: e-vectors are stored in COLUMNS (not rows) of
  // matrix! Therefore, we transpose the matrix (duuumb)
  auto tmp = e_vectors.transpose();
  e_vectors = std::move(tmp);
 
  return eigen_vv;
}
 
} // namespace LinAlg