precompute.h

// Copyright (c) 2020 Matthew Scroggs
// FEniCS Project
// SPDX-License-Identifier:    MIT
 
#pragma once
 
#include "math.h"
#include "mdspan.hpp"
#include <concepts>
#include <span>
#include <tuple>
#include <type_traits>
#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 block_size = 1)
{
  for (std::size_t i = 0; i < perm.size(); ++i)
  {
    for (std::size_t b = 0; b < block_size; ++b)
    {
      std::swap(data[block_size * (offset + i) + b],
                data[block_size * (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 block_size = 1)
{
  for (std::size_t i = 0; i < perm.size(); ++i)
  {
    for (std::size_t b = 0; b < block_size; ++b)
    {
      std::swap(data[block_size * emap[i] + b],
                data[block_size * emap[perm[i]] + b]);
    }
  }
}
 
template <typename E>
void apply_permutation_to_transpose(std::span<const std::size_t> perm,
                                    std::span<E> data, std::size_t offset = 0,
                                    std::size_t block_size = 1)
{
  const std::size_t dim = perm.size();
  const std::size_t data_size
      = (data.size() + (dim < block_size ? block_size - dim : 0)) / block_size;
  for (std::size_t b = 0; b < block_size; ++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,
    MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        const T, MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>
        M,
    std::span<E> data, std::size_t offset = 0, std::size_t block_size = 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, block_size);
  for (std::size_t b = 0; b < block_size; ++b)
  {
    for (std::size_t i = 0; i < dim; ++i)
    {
      for (std::size_t j = i + 1; j < dim; ++j)
      {
        data[block_size * (offset + i) + b]
            += static_cast<U>(M(i, j)) * data[block_size * (offset + j) + b];
      }
    }
    for (std::size_t i = 1; i <= dim; ++i)
    {
      data[block_size * (offset + dim - i) + b]
          *= static_cast<U>(M(dim - i, dim - i));
      for (std::size_t j = 0; j < dim - i; ++j)
      {
        data[block_size * (offset + dim - i) + b]
            += static_cast<U>(M(dim - i, j))
               * data[block_size * (offset + j) + b];
      }
    }
  }
}
 
template <typename T, typename E>
void apply_matrix_to_transpose(
    std::span<const std::size_t> v_size_t,
    MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        const T, MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>
        M,
    std::span<E> data, std::size_t offset = 0, std::size_t block_size = 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 < block_size ? block_size - dim : 0)) / block_size;
  apply_permutation_to_transpose(v_size_t, data, offset, block_size);
  for (std::size_t b = 0; b < block_size; ++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