Matrix.ipp

#pragma once
 
// LAPACK LU routines, used for in-place matrix inversion.
// Complex arguments are passed as opaque pointers (all arithmetic happens on
// the Fortran side); the buffers are std::complex<double>, which matches
// LAPACK's storage layout.
extern "C" {
using lapack_complex = long double;
void dgetrf_(int *m, int *n, double *a, int *lda, int *ipiv, int *info);
void dgetri_(int *n, double *a, int *lda, int *ipiv, double *work, int *lwork,
             int *info);
void zgetrf_(int *m, int *n, lapack_complex *a, int *lda, int *ipiv, int *info);
void zgetri_(int *n, lapack_complex *a, int *lda, int *ipiv,
             lapack_complex *work, int *lwork, int *info);
}
 
namespace LinAlg {
 
//==============================================================================
// Returns the determinant. Uses GSL; via LU decomposition. Only works for
// double/complex<double>
template <typename T>
T Matrix<T>::determinant() const {
  static_assert(std::is_same_v<T, double> ||
                  std::is_same_v<T, std::complex<double>>,
                "Determinant only works for double");
 
  assert(rows() == cols() && "Determinant only defined for square matrix");
  // Make a copy, since this is destructive. (Performs LU decomp)
  auto LU = *this; // will become LU decomposed version
  int sLU = 0;
  auto gsl_view = LU.as_gsl_view();
  gsl_permutation *permutn = gsl_permutation_alloc(rows());
  if constexpr (std::is_same_v<T, double>) {
    gsl_linalg_LU_decomp(&gsl_view.matrix, permutn, &sLU);
    gsl_permutation_free(permutn);
    return gsl_linalg_LU_det(&gsl_view.matrix, sLU);
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    gsl_linalg_complex_LU_decomp(&gsl_view.matrix, permutn, &sLU);
    gsl_permutation_free(permutn);
    const auto gsl_cmplx = gsl_linalg_complex_LU_det(&gsl_view.matrix, sLU);
    // Can probably avoid this copy? doesn't really matter.
    return {GSL_REAL(gsl_cmplx), GSL_IMAG(gsl_cmplx)};
  }
}
 
//==============================================================================
// Inverts the matrix, in place. Uses LAPACK (getrf/getri); LU decomposition
// with partial pivoting. Only works for double/complex<double>.
// Row-major data is seen by LAPACK (column-major) as the transpose;
// (A^T)^-1 = (A^-1)^T, so the result read back row-major is A^-1, and no
// transposes are required.
template <typename T>
Matrix<T> &Matrix<T>::invert_in_place() {
  static_assert(
    std::is_same_v<T, double> || std::is_same_v<T, std::complex<double>>,
    "invert only works for Matrix<double> or Matrix<complex<double>>");
 
  assert(rows() == cols() && "Inverse only defined for square matrix");
  int dim = int(m_rows);
  int info = 0;
  std::vector<int> ipiv(m_rows);
  if constexpr (std::is_same_v<T, double>) {
    dgetrf_(&dim, &dim, data(), &dim, ipiv.data(), &info);
    assert(info == 0 && "LU factorisation failed (singular matrix?)");
    // workspace query, then invert
    int lwork = -1;
    double wkopt = 0.0;
    dgetri_(&dim, data(), &dim, ipiv.data(), &wkopt, &lwork, &info);
    lwork = std::max(1, int(wkopt));
    std::vector<double> work(static_cast<std::size_t>(lwork));
    dgetri_(&dim, data(), &dim, ipiv.data(), work.data(), &lwork, &info);
    assert(info == 0 && "LU inversion failed (singular matrix?)");
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    auto *a_ptr = reinterpret_cast<lapack_complex *>(data());
    zgetrf_(&dim, &dim, a_ptr, &dim, ipiv.data(), &info);
    assert(info == 0 && "LU factorisation failed (singular matrix?)");
    // workspace query, then invert
    int lwork = -1;
    std::complex<double> wkopt{0.0, 0.0};
    zgetri_(&dim, a_ptr, &dim, ipiv.data(),
            reinterpret_cast<lapack_complex *>(&wkopt), &lwork, &info);
    lwork = std::max(1, int(wkopt.real()));
    std::vector<std::complex<double>> work(static_cast<std::size_t>(lwork));
    zgetri_(&dim, a_ptr, &dim, ipiv.data(),
            reinterpret_cast<lapack_complex *>(work.data()), &lwork, &info);
    assert(info == 0 && "LU inversion failed (singular matrix?)");
  }
  return *this;
}
 
//==============================================================================
template <typename T>
Matrix<T> Matrix<T>::transpose() const {
  Matrix<T> Tr(m_cols, m_rows);
  if constexpr (std::is_same_v<T, double>) {
    auto Tr_gsl = Tr.as_gsl_view();
    const auto this_gsl = as_gsl_view();
    gsl_matrix_transpose_memcpy(&Tr_gsl.matrix, &this_gsl.matrix);
  } else if constexpr (std::is_same_v<T, float>) {
    auto Tr_gsl = Tr.as_gsl_view();
    const auto this_gsl = as_gsl_view();
    gsl_matrix_float_transpose_memcpy(&Tr_gsl.matrix, &this_gsl.matrix);
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    auto Tr_gsl = Tr.as_gsl_view();
    const auto this_gsl = as_gsl_view();
    gsl_matrix_complex_transpose_memcpy(&Tr_gsl.matrix, &this_gsl.matrix);
  } else if constexpr (std::is_same_v<T, std::complex<float>>) {
    auto Tr_gsl = Tr.as_gsl_view();
    const auto this_gsl = as_gsl_view();
    gsl_matrix_complex_float_transpose_memcpy(&Tr_gsl.matrix, &this_gsl.matrix);
  } else {
    // backup, works for any type
    for (auto i = 0ul; i < Tr.rows(); ++i) {
      for (auto j = 0ul; j < Tr.cols(); ++j) {
        Tr[i][j] = (*this)[j][i];
      }
    }
  }
  return Tr;
}
 
//==============================================================================
// Constructs a diagonal unit matrix (identity)
template <typename T>
Matrix<T> &Matrix<T>::make_identity() {
  assert(m_rows == m_cols && "Can only call make_identity() for square matrix");
  for (auto i = 0ul; i < m_rows; ++i) {
    for (auto j = 0ul; j < m_cols; ++j) {
      at(i, j) = i == j ? T(1) : T(0);
    }
  }
  return *this;
}
// Sets all elements to zero
template <typename T>
Matrix<T> &Matrix<T>::zero() {
  for (std::size_t i = 0; i < size(); ++i) {
    m_data[i] = T(0);
  }
  return *this;
}
 
//==============================================================================
template <typename T>
Matrix<T> Matrix<T>::conj() const {
  static_assert(is_complex_v<T>, "conj() only available for complex Matrix");
  std::vector<T> conj_data;
  conj_data.reserve(m_data.size());
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    conj_data.push_back(std::conj(m_data[i]));
  }
  return Matrix<T>{m_rows, m_cols, std::move(conj_data)};
}
 
template <typename T>
Matrix<T> &Matrix<T>::conj_in_place() {
  static_assert(is_complex_v<T>, "conj() only available for complex Matrix");
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    m_data[i] = std::conj(m_data[i]);
  }
  return *this;
}
//------------------------------------------------------------------------------
template <typename T>
auto Matrix<T>::real() const {
  static_assert(is_complex_v<T>, "real() only available for complex Matrix");
  std::vector<typename T::value_type> real_data;
  real_data.reserve(m_data.size());
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    real_data.push_back(std::real(m_data[i]));
  }
  return Matrix<typename T::value_type>{m_rows, m_cols, std::move(real_data)};
}
//------------------------------------------------------------------------------
template <typename T>
auto Matrix<T>::imag() const {
  static_assert(is_complex_v<T>, "imag() only available for complex Matrix");
  std::vector<typename T::value_type> imag_data;
  imag_data.reserve(m_data.size());
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    imag_data.push_back(std::imag(m_data[i]));
  }
  return Matrix<typename T::value_type>{m_rows, m_cols, std::move(imag_data)};
}
//------------------------------------------------------------------------------
template <typename T>
auto Matrix<T>::complex() const {
  static_assert(!is_complex_v<T>, "complex() only available for real Matrix");
  // use move constructor to avoid default Matrix construction
  std::vector<std::complex<T>> new_data;
  new_data.reserve(m_data.size());
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    new_data.push_back(m_data[i]);
  }
  return Matrix<std::complex<T>>{m_rows, m_cols, std::move(new_data)};
}
 
//==============================================================================
template <typename T>
Matrix<T> &Matrix<T>::operator+=(const Matrix<T> &rhs) {
  assert(rows() == rhs.rows() && cols() == rhs.cols() &&
         "Matrices must have same dimensions for addition");
// this->m_data += rhs.m_data;
#pragma omp parallel for
  for (auto i = 0ul; i < m_data.size(); ++i) {
    m_data[i] += rhs.m_data[i];
  }
  return *this;
}
template <typename T>
Matrix<T> &Matrix<T>::operator-=(const Matrix<T> &rhs) {
  assert(rows() == rhs.rows() && cols() == rhs.cols() &&
         "Matrices must have same dimensions for subtraction");
// using namespace qip::overloads;
// this->m_data -= rhs.m_data;
#pragma omp parallel for
  for (auto i = 0ul; i < m_data.size(); ++i) {
    m_data[i] -= rhs.m_data[i];
  }
  return *this;
}
template <typename T>
Matrix<T> &Matrix<T>::operator*=(const T x) {
// using namespace qip::overloads;
// this->m_data *= x;
#pragma omp parallel for
  for (auto i = 0ul; i < m_data.size(); ++i) {
    m_data[i] *= x;
  }
  return *this;
}
template <typename T>
Matrix<T> &Matrix<T>::operator/=(const T x) {
// using namespace qip::overloads;
// this->m_data /= x;
#pragma omp parallel for
  for (auto i = 0ul; i < m_data.size(); ++i) {
    m_data[i] /= x;
  }
  return *this;
}
 
//==============================================================================
// Matrix<T> += T : T assumed to be *Identity!
template <typename T>
Matrix<T> &Matrix<T>::operator+=(T aI) {
  // Adds 'a' to diagonal elements (Assume a*Ident)
  assert(m_rows == m_cols && "Can only call M+a for square matrix");
  for (auto i = 0ul; i < m_rows; ++i) {
    at(i, i) += aI;
  }
  return *this;
}
// Matrix<T> -= T : T assumed to be *Identity!
template <typename T>
Matrix<T> &Matrix<T>::operator-=(T aI) {
  // Adds 'a' to diagonal elements (Assume a*Ident)
  assert(m_rows == m_cols && "Can only call M-a for square matrix");
  for (auto i = 0ul; i < m_rows; ++i) {
    at(i, i) -= aI;
  }
  return *this;
}
 
//==============================================================================
template <typename T>
Matrix<T> &Matrix<T>::mult_elements_by(const Matrix<T> &a) {
  assert(rows() == a.rows() && cols() == a.cols() &&
         "Matrices must have same dimensions for mult_elements_by");
  for (auto i = 0ul; i < m_data.size(); ++i) {
    m_data[i] *= a.m_data[i];
  }
  return *this;
}
 
//==============================================================================
template <typename T>
[[nodiscard]] Matrix<T> operator*(const Matrix<T> &a, const Matrix<T> &b) {
  // https://www.gnu.org/software/gsl/doc/html/blas.html
  assert(a.cols() == b.rows() &&
         "Matrices a and b must have correct dimension for multiplication");
  Matrix<T> product(a.rows(), b.cols());
 
  GEMM(a, b, &product);
 
  return product;
}
 
//==============================================================================
template <typename T>
void GEMM(const Matrix<T> &a, const Matrix<T> &b, Matrix<T> *c, bool trans_A,
          bool trans_B) {
  assert(c);
 
  const auto ta = to_cblas_trans(trans_A);
  const auto tb = to_cblas_trans(trans_B);
 
  // Effective dimensions:
  // op(A): (trans_A ? a.cols x a.rows : a.rows x a.cols)
  // op(B): (trans_B ? b.cols x b.rows : b.rows x b.cols)
  const int A_rows = static_cast<int>(trans_A ? a.cols() : a.rows());
  const int A_cols = static_cast<int>(trans_A ? a.rows() : a.cols());
  const int B_rows = static_cast<int>(trans_B ? b.cols() : b.rows());
  const int B_cols = static_cast<int>(trans_B ? b.rows() : b.cols());
 
  // GEMM sizes: C = op(A) * op(B), where
  // M = rows(op(A)), N = cols(op(B)), K = cols(op(A)) = rows(op(B))
  const int M = A_rows;
  const int N = B_cols;
  const int K = A_cols;
 
  assert(A_cols == B_rows && "op(A) cols must equal op(B) rows");
  assert(static_cast<int>(c->rows()) == M && static_cast<int>(c->cols()) == N &&
         "Output matrix c must be sized MxN");
 
  // Row-major leading dimensions:
  // lda = number of columns in A's *storage* (i.e., a.cols()) regardless of trans
  // same for b, c
  const int lda = static_cast<int>(a.cols());
  const int ldb = static_cast<int>(b.cols());
  const int ldc = static_cast<int>(c->cols());
 
  if constexpr (std::is_same_v<T, double>) {
    cblas_dgemm(CblasRowMajor, ta, tb, M, N, K, 1.0, a.data(), lda, b.data(),
                ldb, 0.0, c->data(), ldc);
 
  } else if constexpr (std::is_same_v<T, float>) {
    cblas_sgemm(CblasRowMajor, ta, tb, M, N, K, 1.0f, a.data(), lda, b.data(),
                ldb, 0.0f, c->data(), ldc);
 
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    const std::complex<double> alpha{1.0, 0.0};
    const std::complex<double> beta{0.0, 0.0};
    cblas_zgemm(CblasRowMajor, ta, tb, M, N, K, &alpha, a.data(), lda, b.data(),
                ldb, &beta, c->data(), ldc);
 
  } else if constexpr (std::is_same_v<T, std::complex<float>>) {
    const std::complex<float> alpha{1.0f, 0.0f};
    const std::complex<float> beta{0.0f, 0.0f};
    cblas_cgemm(CblasRowMajor, ta, tb, M, N, K, &alpha, a.data(), lda, b.data(),
                ldb, &beta, c->data(), ldc);
 
  } else {
    static_assert(!sizeof(T), "GEMM: unsupported scalar type");
  }
}
 
//==============================================================================
// M_ab = A_ai B_aj C_ij D_ib E_jb, using BLAS
template <typename T>
void PENTA_GEMM(const Matrix<T> &A, const Matrix<T> &B, const Matrix<T> &C,
                const Matrix<T> &D, const Matrix<T> &E, Matrix<T> *pM) {
  //
  const auto N = A.rows(); // assume all square
  assert(A.cols() == A.rows() && "Must be square");
 
  Matrix<T> X(N, N);
  Matrix<T> Y(N, N);
  auto &M = *pM;
 
  // M_ab = A_ai B_aj C_ij D_ib E_jb
  //      = A_ai B_aj X(i)_jb D_ib
  //      = A_ai Y(i)_ab D_ib
  // X(i)_jb = C_ij * E_j2;
  // Y(i)_aj = B_ij * X(i)_jb
 
  for (std::size_t i = 0; i < N; ++i) {
    for (std::size_t j = 0; j < N; ++j) {
      const auto cij = C[i][j];
      for (std::size_t b = 0; b < N; ++b) {
        X[j][b] = cij * E[j][b];
      }
    }
    GEMM(B, X, &Y);
    for (std::size_t a = 0; a < N; ++a) {
      for (std::size_t b = 0; b < N; ++b) {
        M[a][b] += A[a][i] * Y[a][b] * D[i][b];
      }
    }
  }
}
 
// M_ab = A_ai B_aj C_ij D_ib E_jb
template <typename T, bool PARALLEL>
void PENTA(const Matrix<T> &A, const Matrix<T> &B, const Matrix<T> &C,
           const Matrix<T> &D, const Matrix<T> &E, Matrix<T> *pM) {
  //
  const auto N = A.rows(); // assume all square
  assert(A.cols() == A.rows() && "Must be square");
 
  auto &M = *pM;
 
  // M_ab = A_ai B_aj C_ij D_ib E_jb
  if constexpr (PARALLEL) {
 
#pragma omp parallel for collapse(2)
    for (std::size_t a = 0; a < N; ++a) {
      for (std::size_t b = 0; b < N; ++b) {
        const T *Ba = &B[a][0];
        T Mab = T(0);
        for (std::size_t i = 0; i < N; ++i) {
          const auto AaiDib = A[a][i] * D[i][b];
          const T *Ci = &C[i][0];
          for (std::size_t j = 0; j < N; ++j) {
            Mab += AaiDib * Ba[j] * Ci[j] * E[j][b];
          }
        }
        M[a][b] = Mab;
      }
    }
 
  } else {
 
    for (std::size_t a = 0; a < N; ++a) {
      const T *Ba = &B[a][0];
      for (std::size_t b = 0; b < N; ++b) {
        T Mab = T(0);
        for (std::size_t i = 0; i < N; ++i) {
          const auto AaiDib = A[a][i] * D[i][b];
          const T *Ci = &C[i][0];
          for (std::size_t j = 0; j < N; ++j) {
            Mab += AaiDib * Ba[j] * Ci[j] * E[j][b];
          }
        }
        M[a][b] = Mab;
      }
    }
  }
}
 
//==============================================================================
template <typename T>
auto Matrix<T>::as_gsl_view() {
  if constexpr (std::is_same_v<T, double>) {
    return gsl_matrix_view_array(m_data.data(), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, float>) {
    return gsl_matrix_float_view_array(m_data.data(), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    // reinterpret_cast OK: cppreference.com/w/cpp/numeric/complex
    return gsl_matrix_complex_view_array(
      reinterpret_cast<double *>(m_data.data()), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, std::complex<float>>) {
    return gsl_matrix_complex_float_view_array(
      reinterpret_cast<float *>(m_data.data()), m_rows, m_cols);
  } else {
    assert(false && "as_gsl_view() only available for double/float (or complex "
                    "double/float)");
  }
}
 
template <typename T>
auto Matrix<T>::as_gsl_view() const {
  if constexpr (std::is_same_v<T, double>) {
    return gsl_matrix_const_view_array(m_data.data(), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, float>) {
    return gsl_matrix_float_const_view_array(m_data.data(), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, std::complex<double>>) {
    return gsl_matrix_complex_const_view_array(
      reinterpret_cast<const double *>(m_data.data()), m_rows, m_cols);
  } else if constexpr (std::is_same_v<T, std::complex<float>>) {
    return gsl_matrix_complex_float_const_view_array(
      reinterpret_cast<const float *>(m_data.data()), m_rows, m_cols);
  } else {
    assert(false && "as_gsl_view() only for available double/float (or complex "
                    "double/float)");
  }
}
 
//==============================================================================
template <typename T>
std::ostream &operator<<(std::ostream &os, const Matrix<T> &a) {
  for (auto i = 0ul; i < a.rows(); ++i) {
    for (auto j = 0ul; j < a.cols(); ++j) {
      os << a(i, j) << " ";
    }
    os << "\n";
  }
  os << "\n";
  return os;
}
 
//==============================================================================
//==============================================================================
//==============================================================================
 
//==============================================================================
// Helper for equal()
template <typename T>
constexpr auto myEps() {
  if constexpr (std::is_same_v<T, float> ||
                std::is_same_v<T, std::complex<float>>) {
    return 1.0e-6f;
  } else if constexpr (std::is_same_v<T, double> ||
                       std::is_same_v<T, std::complex<double>>) {
    return 1.0e-12;
  } else {
    return 0;
  }
}
 
// Compares two matrices; returns true iff all elements compare relatively to
// better than eps
template <typename T>
bool equal(const Matrix<T> &lhs, const Matrix<T> &rhs, T eps) {
  if (lhs.rows() != rhs.rows())
    return false;
  if (lhs.cols() != rhs.cols())
    return false;
  for (auto i = 0ul; i < lhs.rows(); ++i) {
    for (auto j = 0ul; j < lhs.cols(); ++j) {
      // need abs on eps in case of complex
      if (std::abs(lhs(i, j) - rhs(i, j)) >
          std::abs(eps * (lhs(i, j) + rhs(i, j))))
        return false;
    }
  }
  return true;
}
 
} // namespace LinAlg