precompute.h

// Copyright (c) 2020 Matthew Scroggs
// FEniCS Project
// SPDX-License-Identifier:    MIT
 
#pragma once
 
#include "mdspan.hpp"
#include <span>
#include <tuple>
#include <vector>
 
namespace basix::precompute
{
 
void prepare_permutation(const std::span<std::size_t>& perm);
 
template <typename E>
void apply_permutation(const std::span<const std::size_t>& perm,
                       const std::span<E>& data, std::size_t offset = 0,
                       std::size_t block_size = 1)
{
  for (std::size_t b = 0; b < block_size; ++b)
  {
    for (std::size_t i = 0; i < perm.size(); ++i)
    {
      std::swap(data[block_size * (offset + i) + b],
                data[block_size * (offset + perm[i]) + b]);
    }
  }
}
 
template <typename E>
void apply_permutation_to_transpose(const std::span<const std::size_t>& perm,
                                    const 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]]);
    }
  }
}
 
std::vector<std::size_t>
prepare_matrix(std::pair<std::vector<double>, std::array<std::size_t, 2>>& A);
 
template <typename T, typename E>
void apply_matrix(const std::span<const std::size_t>& v_size_t,
                  const std::experimental::mdspan<
                      const T, std::experimental::dextents<std::size_t, 2>>& M,
                  const std::span<E>& data, std::size_t offset = 0,
                  std::size_t block_size = 1)
{
  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]
            += 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] *= M(dim - i, dim - i);
      for (std::size_t j = 0; j < dim - i; ++j)
      {
        data[block_size * (offset + dim - i) + b]
            += M(dim - i, j) * data[block_size * (offset + j) + b];
      }
    }
  }
}
 
template <typename T, typename E>
void apply_matrix_to_transpose(
    const std::span<const std::size_t>& v_size_t,
    const std::experimental::mdspan<
        const T, std::experimental::dextents<std::size_t, 2>>& M,
    const std::span<E>& data, std::size_t offset = 0,
    std::size_t block_size = 1)
{
  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]
            += 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] *= M(dim - i, dim - i);
      for (std::size_t j = 0; j < dim - i; ++j)
      {
        data[data_size * b + offset + dim - i]
            += M(dim - i, j) * data[data_size * b + offset + j];
      }
    }
  }
}
 
} // namespace basix::precompute