precompute.h

// Copyright (c) 2020 Matthew Scroggs
// FEniCS Project
// SPDX-License-Identifier:    MIT
 
#pragma once
 
#include "math.h"
#include "mdspan.hpp"
#include <array>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <span>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

namespace basix::precompute
{
namespace impl
{
template <typename T, typename = void>
struct scalar_value_type
{
  typedef T value_type;
};
template <typename T>
struct scalar_value_type<T, std::void_t<typename T::value_type>>
{
  typedef typename T::value_type value_type;
};
template <typename T>
using scalar_value_type_t = typename scalar_value_type<T>::value_type;
} // namespace impl

void prepare_permutation(std::span<std::size_t> perm);

template <typename E>
void apply_permutation(std::span<const std::size_t> perm, std::span<E> data,
                       std::size_t offset = 0, std::size_t n = 1)
{
  for (std::size_t i = 0; i < perm.size(); ++i)
    for (std::size_t b = 0; b < n; ++b)
      std::swap(data[n * (offset + i) + b], data[n * (offset + perm[i]) + b]);
}

template <typename E>
void apply_permutation_mapped(std::span<const std::size_t> perm,
                              std::span<E> data, std::span<const int> emap,
                              std::size_t n = 1)
{
  for (std::size_t i = 0; i < perm.size(); ++i)
    for (std::size_t b = 0; b < n; ++b)
      std::swap(data[n * emap[i] + b], data[n * emap[perm[i]] + b]);
}

template <typename E>
void apply_inv_permutation_right(std::span<const std::size_t> perm,
                                 std::span<E> data, std::size_t offset = 0,
                                 std::size_t n = 1)
{
  const std::size_t dim = perm.size();
  const std::size_t data_size = (data.size() + (dim < n ? n - dim : 0)) / n;
  for (std::size_t b = 0; b < n; ++b)
  {
    for (std::size_t i = 0; i < dim; ++i)
    {
      std::swap(data[data_size * b + offset + i],
                data[data_size * b + offset + perm[i]]);
    }
  }
}

template <std::floating_point T>
std::vector<std::size_t>
prepare_matrix(std::pair<std::vector<T>, std::array<std::size_t, 2>>& A)
{
  return math::transpose_lu<T>(A);
}

template <typename T, typename E>
void apply_matrix(std::span<const std::size_t> v_size_t,
                  md::mdspan<const T, md::dextents<std::size_t, 2>> M,
                  std::span<E> data, std::size_t offset = 0, std::size_t n = 1)
{
  using U = typename impl::scalar_value_type_t<E>;
 
  const std::size_t dim = v_size_t.size();
  apply_permutation(v_size_t, data, offset, n);
 
  // data has n contiguous (unit-stride) values per dof. Looping over b
  // innermost vectorises each elimination step, but for small n (common
  // in practice, e.g. 1-3) the loop-entry cost per (i, j) pair outweighs
  // that benefit, so fall back to the original b-outermost ordering.
  // Threshold chosen from measurements showing the crossover is n=12-16.
  if (n <= 12)
  {
    for (std::size_t b = 0; b < n; ++b)
    {
      for (std::size_t i = 0; i < dim; ++i)
      {
        for (std::size_t j = i + 1; j < dim; ++j)
        {
          data[n * (offset + i) + b]
              += static_cast<U>(M(i, j)) * data[n * (offset + j) + b];
        }
      }
 
      for (std::size_t i = 1; i <= dim; ++i)
      {
        data[n * (offset + dim - i) + b] *= static_cast<U>(M(dim - i, dim - i));
        for (std::size_t j = 0; j < dim - i; ++j)
        {
          data[n * (offset + dim - i) + b]
              += static_cast<U>(M(dim - i, j)) * data[n * (offset + j) + b];
        }
      }
    }
  }
  else
  {
    for (std::size_t i = 0; i < dim; ++i)
    {
      for (std::size_t j = i + 1; j < dim; ++j)
      {
        const U Mij = static_cast<U>(M(i, j));
        for (std::size_t b = 0; b < n; ++b)
        {
          data[n * (offset + i) + b] += Mij * data[n * (offset + j) + b];
        }
      }
    }
 
    for (std::size_t i = 1; i <= dim; ++i)
    {
      const U Mdiag = static_cast<U>(M(dim - i, dim - i));
      for (std::size_t b = 0; b < n; ++b)
        data[n * (offset + dim - i) + b] *= Mdiag;
 
      for (std::size_t j = 0; j < dim - i; ++j)
      {
        const U Mij = static_cast<U>(M(dim - i, j));
        for (std::size_t b = 0; b < n; ++b)
        {
          data[n * (offset + dim - i) + b] += Mij * data[n * (offset + j) + b];
        }
      }
    }
  }
}

template <typename T, typename E>
void apply_transpose_matrix_right(
    std::span<const std::size_t> v_size_t,
    md::mdspan<const T, md::dextents<std::size_t, 2>> M, std::span<E> data,
    std::size_t offset = 0, std::size_t n = 1)
{
  using U = typename impl::scalar_value_type_t<E>;
 
  const std::size_t dim = v_size_t.size();
  const std::size_t data_size = (data.size() + (dim < n ? n - dim : 0)) / n;
  apply_inv_permutation_right(v_size_t, data, offset, n);
  for (std::size_t b = 0; b < n; ++b)
  {
    for (std::size_t i = 0; i < dim; ++i)
    {
      for (std::size_t j = i + 1; j < dim; ++j)
      {
        data[data_size * b + offset + i]
            += static_cast<U>(M(i, j)) * data[data_size * b + offset + j];
      }
    }
    for (std::size_t i = 1; i <= dim; ++i)
    {
      data[data_size * b + offset + dim - i]
          *= static_cast<U>(M(dim - i, dim - i));
      for (std::size_t j = 0; j < dim - i; ++j)
      {
        data[data_size * b + offset + dim - i]
            += static_cast<U>(M(dim - i, j)) * data[data_size * b + offset + j];
      }
    }
  }
}
 
} // namespace basix::precompute