math.h

// Copyright (C) 2021 Igor Baratta
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "mdspan.hpp"
#include <array>
#include <span>
#include <vector>
 
namespace basix::math
{
 
namespace impl
{
void dot_blas(const std::span<const double>& A,
              std::array<std::size_t, 2> Ashape,
              const std::span<const double>& B,
              std::array<std::size_t, 2> Bshape, const std::span<double>& C);
} // 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]};
}
 
std::pair<std::vector<double>, std::vector<double>>
eigh(const std::span<const double>& A, std::size_t n);
 
std::vector<double>
solve(const std::experimental::mdspan<
          const double, std::experimental::dextents<std::size_t, 2>>& A,
      const std::experimental::mdspan<
          const double, std::experimental::dextents<std::size_t, 2>>& B);
 
bool is_singular(const std::experimental::mdspan<
                 const double, std::experimental::dextents<std::size_t, 2>>& A);
 
std::vector<std::size_t>
transpose_lu(std::pair<std::vector<double>, std::array<std::size_t, 2>>& A);
 
template <typename U, typename V, typename W>
void dot(const U& A, const V& B, W&& C)
{
  assert(A.extent(1) == B.extent(0));
  assert(C.extent(0) == C.extent(0));
  assert(C.extent(1) == B.extent(1));
  if (A.extent(0) * B.extent(1) * A.extent(1) < 4096)
  {
    std::fill_n(C.data_handle(), C.extent(0) * C.extent(1), 0);
    for (std::size_t i = 0; i < A.extent(0); ++i)
      for (std::size_t j = 0; j < B.extent(1); ++j)
        for (std::size_t k = 0; k < A.extent(1); ++k)
          C(i, j) += A(i, k) * B(k, j);
  }
  else
  {
    impl::dot_blas(
        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()));
  }
}
 
std::vector<double> eye(std::size_t n);
 
} // namespace basix::math