Loading...
Searching...
No Matches
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
basix::math
Mathematical functions.
Definition math.h:51
basix::math::is_singular
bool is_singular(md::mdspan< const T, md::dextents< std::size_t, 2 > > A)
Check if A is a singular matrix.
Definition math.h:230
basix::math::cross
std::array< typename U::value_type, 3 > cross(const U &u, const V &v)
Compute the cross product u x v.
Definition math.h:111
basix::math::transpose_lu
std::vector< std::size_t > transpose_lu(std::pair< std::vector< T >, std::array< std::size_t, 2 > > &A)
Compute the LU decomposition of the transpose of a square matrix A.
Definition math.h:271
basix::math::orthogonalise
void orthogonalise(md::mdspan< T, md::dextents< std::size_t, 2 > > wcoeffs, std::size_t start=0)
Orthogonalise the rows of a matrix (in place).
Definition math.h:365
basix::math::solve
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)
Solve A X = B.
Definition math.h:187
basix::math::dot
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)
Compute C = alpha A * B + beta C.
Definition math.h:306
basix::math::eigh
std::pair< std::vector< T >, std::vector< T > > eigh(std::span< const T > A, std::size_t n)
Compute the eigenvalues and eigenvectors of a square Hermitian matrix A.
Definition math.h:127
basix::math::eye
std::vector< T > eye(std::size_t n)
Build an identity matrix.
Definition math.h:351
basix::math::outer
std::pair< std::vector< typename U::value_type >, std::array< std::size_t, 2 > > outer(const U &u, const V &v)
Compute the outer product of vectors u and v.
Definition math.h:97