math.h

// Copyright (C) 2021-2024 Igor Baratta and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "mdspan.hpp"
#include "types.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
extern "C"
{
  void ssyevd_(char* jobz, char* uplo, int* n, float* a, int* lda, float* w,
               float* work, int* lwork, int* iwork, int* liwork, int* info);
  void dsyevd_(char* jobz, char* uplo, int* n, double* a, int* lda, double* w,
               double* work, int* lwork, int* iwork, int* liwork, int* info);
 
  void sgesv_(int* N, int* NRHS, float* A, int* LDA, int* IPIV, float* B,
              int* LDB, int* INFO);
  void dgesv_(int* N, int* NRHS, double* A, int* LDA, int* IPIV, double* B,
              int* LDB, int* INFO);
 
  void sgemm_(char* transa, char* transb, int* m, int* n, int* k, float* alpha,
              float* a, int* lda, float* b, int* ldb, float* beta, float* c,
              int* ldc);
  void dgemm_(char* transa, char* transb, int* m, int* n, int* k, double* alpha,
              double* a, int* lda, double* b, int* ldb, double* beta, double* c,
              int* ldc);
 
  int sgetrf_(const int* m, const int* n, float* a, const int* lda, int* lpiv,
              int* info);
  int dgetrf_(const int* m, const int* n, double* a, const int* lda, int* lpiv,
              int* info);
}

namespace basix::math
{
namespace impl
{
template <std::floating_point T>
void dot_blas(std::span<const T> A, std::array<std::size_t, 2> Ashape,
              std::span<const T> B, std::array<std::size_t, 2> Bshape,
              std::span<T> C, T alpha = 1, T beta = 0)
{
  static_assert(std::is_same_v<T, float> or std::is_same_v<T, double>);
 
  assert(Ashape[1] == Bshape[0]);
  assert(C.size() == Ashape[0] * Bshape[1]);
 
  int M = Ashape[0];
  int N = Bshape[1];
  int K = Ashape[1];
 
  int lda = K;
  int ldb = N;
  int ldc = N;
  char trans = 'N';
  if constexpr (std::is_same_v<T, float>)
  {
    sgemm_(&trans, &trans, &N, &M, &K, &alpha, const_cast<T*>(B.data()), &ldb,
           const_cast<T*>(A.data()), &lda, &beta, C.data(), &ldc);
  }
  else if constexpr (std::is_same_v<T, double>)
  {
    dgemm_(&trans, &trans, &N, &M, &K, &alpha, const_cast<T*>(B.data()), &ldb,
           const_cast<T*>(A.data()), &lda, &beta, C.data(), &ldc);
  }
}
 
} // namespace impl

template <typename U, typename V>
std::pair<std::vector<typename U::value_type>, std::array<std::size_t, 2>>
outer(const U& u, const V& v)
{
  std::vector<typename U::value_type> result(u.size() * v.size());
  for (std::size_t i = 0; i < u.size(); ++i)
    for (std::size_t j = 0; j < v.size(); ++j)
      result[i * v.size() + j] = u[i] * v[j];
  return {std::move(result), {u.size(), v.size()}};
}

template <typename U, typename V>
std::array<typename U::value_type, 3> cross(const U& u, const V& v)
{
  assert(u.size() == 3);
  assert(v.size() == 3);
  return {u[1] * v[2] - u[2] * v[1], u[2] * v[0] - u[0] * v[2],
          u[0] * v[1] - u[1] * v[0]};
}

template <std::floating_point T>
std::pair<std::vector<T>, std::vector<T>> eigh(std::span<const T> A,
                                               std::size_t n)
{
  // Copy A
  std::vector<T> M(A.begin(), A.end());
 
  // Allocate storage for eigenvalues
  std::vector<T> w(n, 0);
 
  int N = n;
  char jobz = 'V'; // Compute eigenvalues and eigenvectors
  char uplo = 'L'; // Lower
  int ldA = n;
  int lwork = -1;
  int liwork = -1;
  int info;
  std::vector<T> work(1);
  std::vector<int> iwork(1);
 
  // Query optimal workspace size
  if constexpr (std::is_same_v<T, float>)
  {
    ssyevd_(&jobz, &uplo, &N, M.data(), &ldA, w.data(), work.data(), &lwork,
            iwork.data(), &liwork, &info);
  }
  else if constexpr (std::is_same_v<T, double>)
  {
    dsyevd_(&jobz, &uplo, &N, M.data(), &ldA, w.data(), work.data(), &lwork,
            iwork.data(), &liwork, &info);
  }
 
  if (info != 0)
    throw std::runtime_error("Could not find workspace size for syevd.");
 
  // Solve eigen problem
  work.resize(work[0]);
  iwork.resize(iwork[0]);
  lwork = work.size();
  liwork = iwork.size();
  if constexpr (std::is_same_v<T, float>)
  {
    ssyevd_(&jobz, &uplo, &N, M.data(), &ldA, w.data(), work.data(), &lwork,
            iwork.data(), &liwork, &info);
  }
  else if constexpr (std::is_same_v<T, double>)
  {
    dsyevd_(&jobz, &uplo, &N, M.data(), &ldA, w.data(), work.data(), &lwork,
            iwork.data(), &liwork, &info);
  }
  if (info != 0)
    throw std::runtime_error("Eigenvalue computation did not converge.");
 
  return {std::move(w), std::move(M)};
}

template <std::floating_point T>
std::vector<T> solve(md::mdspan<const T, md::dextents<std::size_t, 2>> A,
                     md::mdspan<const T, md::dextents<std::size_t, 2>> B)
{
  // Copy A and B to column-major storage
  mdex::mdarray<T, md::dextents<std::size_t, 2>, md::layout_left> _A(
      A.extents()),
      _B(B.extents());
  for (std::size_t i = 0; i < A.extent(0); ++i)
    for (std::size_t j = 0; j < A.extent(1); ++j)
      _A(i, j) = A(i, j);
  for (std::size_t i = 0; i < B.extent(0); ++i)
    for (std::size_t j = 0; j < B.extent(1); ++j)
      _B(i, j) = B(i, j);
 
  int N = _A.extent(0);
  int nrhs = _B.extent(1);
  int lda = _A.extent(0);
  int ldb = _B.extent(0);
 
  // Pivot indices that define the permutation matrix for the LU solver
  std::vector<int> piv(N);
  int info;
  if constexpr (std::is_same_v<T, float>)
    sgesv_(&N, &nrhs, _A.data(), &lda, piv.data(), _B.data(), &ldb, &info);
  else if constexpr (std::is_same_v<T, double>)
    dgesv_(&N, &nrhs, _A.data(), &lda, piv.data(), _B.data(), &ldb, &info);
  if (info != 0)
    throw std::runtime_error("Call to dgesv failed: " + std::to_string(info));
 
  // Copy result to row-major storage
  std::vector<T> rb(_B.extent(0) * _B.extent(1));
  md::mdspan<T, md::dextents<std::size_t, 2>> r(rb.data(), _B.extents());
  for (std::size_t i = 0; i < _B.extent(0); ++i)
    for (std::size_t j = 0; j < _B.extent(1); ++j)
      r(i, j) = _B(i, j);
 
  return rb;
}

template <std::floating_point T>
bool is_singular(md::mdspan<const T, md::dextents<std::size_t, 2>> A)
{
  // Copy to column major matrix
  mdex::mdarray<T, md::dextents<std::size_t, 2>, md::layout_left> _A(
      A.extents());
  for (std::size_t i = 0; i < A.extent(0); ++i)
    for (std::size_t j = 0; j < A.extent(1); ++j)
      _A(i, j) = A(i, j);
 
  std::vector<T> B(A.extent(1), 1);
  int N = _A.extent(0);
  int nrhs = 1;
  int lda = _A.extent(0);
  int ldb = B.size();
 
  // Pivot indices that define the permutation matrix for the LU solver
  std::vector<int> piv(N);
  int info;
  if constexpr (std::is_same_v<T, float>)
    sgesv_(&N, &nrhs, _A.data(), &lda, piv.data(), B.data(), &ldb, &info);
  else if constexpr (std::is_same_v<T, double>)
    dgesv_(&N, &nrhs, _A.data(), &lda, piv.data(), B.data(), &ldb, &info);
 
  if (info < 0)
  {
    throw std::runtime_error("dgesv failed due to invalid value: "
                             + std::to_string(info));
  }
  else if (info > 0)
    return true;
  else
    return false;
}

template <std::floating_point T>
std::vector<std::size_t>
transpose_lu(std::pair<std::vector<T>, std::array<std::size_t, 2>>& A)
{
  std::size_t dim = A.second[0];
  assert(dim == A.second[1]);
  int N = dim;
  int info;
  std::vector<int> lu_perm(dim);
 
  // Comput LU decomposition of M
  if constexpr (std::is_same_v<T, float>)
    sgetrf_(&N, &N, A.first.data(), &N, lu_perm.data(), &info);
  else if constexpr (std::is_same_v<T, double>)
    dgetrf_(&N, &N, A.first.data(), &N, lu_perm.data(), &info);
 
  if (info != 0)
  {
    throw std::runtime_error("LU decomposition failed: "
                             + std::to_string(info));
  }
 
  std::vector<std::size_t> perm(dim);
  for (std::size_t i = 0; i < dim; ++i)
    perm[i] = static_cast<std::size_t>(lu_perm[i] - 1);
 
  return perm;
}

template <typename U, typename V, typename W>
void dot(const U& A, const V& B, W&& C,
         typename std::decay_t<U>::value_type alpha = 1,
         typename std::decay_t<U>::value_type beta = 0)
{
  using T = typename std::decay_t<U>::value_type;
 
  assert(A.extent(1) == B.extent(0));
  assert(C.extent(0) == A.extent(0));
  assert(C.extent(1) == B.extent(1));
  if (A.extent(0) * B.extent(1) * A.extent(1) < 256)
  {
    for (std::size_t i = 0; i < A.extent(0); ++i)
    {
      for (std::size_t j = 0; j < B.extent(1); ++j)
      {
        T C0 = C(i, j);
        C(i, j) = 0;
        T& _C = C(i, j);
        for (std::size_t k = 0; k < A.extent(1); ++k)
          _C += A(i, k) * B(k, j);
        _C = alpha * _C + beta * C0;
      }
    }
  }
  else
  {
    static_assert(std::is_same_v<typename std::decay_t<U>::layout_type,
                                 md::layout_right>);
    static_assert(std::is_same_v<typename std::decay_t<V>::layout_type,
                                 md::layout_right>);
    static_assert(std::is_same_v<typename std::decay_t<W>::layout_type,
                                 md::layout_right>);
    static_assert(std::is_same_v<typename std::decay_t<V>::value_type, T>);
    static_assert(std::is_same_v<typename std::decay_t<W>::value_type, T>);
    impl::dot_blas<T>(
        std::span(A.data_handle(), A.size()), {A.extent(0), A.extent(1)},
        std::span(B.data_handle(), B.size()), {B.extent(0), B.extent(1)},
        std::span(C.data_handle(), C.size()), alpha, beta);
  }
}

template <std::floating_point T>
std::vector<T> eye(std::size_t n)
{
  std::vector<T> I(n * n, 0);
  md::mdspan<T, md::dextents<std::size_t, 2>> Iview(I.data(), n, n);
  for (std::size_t i = 0; i < n; ++i)
    Iview(i, i) = 1;
  return I;
}

template <std::floating_point T>
void orthogonalise(md::mdspan<T, md::dextents<std::size_t, 2>> wcoeffs,
                   std::size_t start = 0)
{
  for (std::size_t i = start; i < wcoeffs.extent(0); ++i)
  {
    T norm = 0;
    for (std::size_t k = 0; k < wcoeffs.extent(1); ++k)
      norm += wcoeffs(i, k) * wcoeffs(i, k);
 
    norm = std::sqrt(norm);
    if (norm < 2 * std::numeric_limits<T>::epsilon())
    {
      throw std::runtime_error("Cannot orthogonalise the rows of a matrix "
                               "with incomplete row rank");
    }
 
    for (std::size_t k = 0; k < wcoeffs.extent(1); ++k)
      wcoeffs(i, k) /= norm;
 
    for (std::size_t j = i + 1; j < wcoeffs.extent(0); ++j)
    {
      T a = 0;
      for (std::size_t k = 0; k < wcoeffs.extent(1); ++k)
        a += wcoeffs(i, k) * wcoeffs(j, k);
      for (std::size_t k = 0; k < wcoeffs.extent(1); ++k)
        wcoeffs(j, k) -= a * wcoeffs(i, k);
    }
  }
}
} // namespace basix::math