precompute.h

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