PolynomialReconstruction.hpp

#ifndef POLYNOMIAL_RECONSTRUCTION_HPP
#define POLYNOMIAL_RECONSTRUCTION_HPP

#include <algebra/ShrinkMatrixView.hpp>
#include <algebra/ShrinkVectorView.hpp>
#include <algebra/SmallMatrix.hpp>
#include <algebra/SmallVector.hpp>
#include <mesh/MeshTraits.hpp>
#include <scheme/DiscreteFunctionDPkVariant.hpp>

#warning MOVE TO .cpp WHEN FINISHED
#include <algebra/Givens.hpp>
#include <analysis/GaussLegendreQuadratureDescriptor.hpp>
#include <analysis/QuadratureFormula.hpp>
#include <analysis/QuadratureManager.hpp>
#include <geometry/LineTransformation.hpp>
#include <mesh/MeshData.hpp>
#include <mesh/MeshDataManager.hpp>
#include <mesh/StencilManager.hpp>

class DiscreteFunctionDPkVariant;
class DiscreteFunctionVariant;

class PolynomialReconstruction
{
 private:
  class MutableDiscreteFunctionDPkVariant;

  template <MeshConcept MeshType>
  [[nodiscard]] std::vector<std::shared_ptr<const DiscreteFunctionDPkVariant>> _build(
    const size_t degree,
    const std::shared_ptr<const MeshType>& p_mesh,
    const std::vector<std::shared_ptr<const DiscreteFunctionVariant>>& discrete_function_variant_list) const;

 public:
  template <MeshConcept MeshType, typename DataType>
  PUGS_INLINE DiscreteFunctionDPk<MeshType::Dimension, std::remove_const_t<DataType>>
  build(const MeshType& mesh, const DiscreteFunctionP0<DataType> p0_function) const
  {
    using Rd = TinyVector<MeshType::Dimension>;

    const size_t degree    = 1;
    const Stencil& stencil = StencilManager::instance().getStencil(mesh.connectivity());

    auto xj = MeshDataManager::instance().getMeshData(mesh).xj();

    auto cell_is_owned = mesh.connectivity().cellIsOwned();

    const size_t max_stencil_size = [&]() {
      size_t max_size = 0;
      for (CellId cell_id = 0; cell_id < mesh.numberOfCells(); ++cell_id) {
        auto stencil_cell_list = stencil[cell_id];
        if (cell_is_owned[cell_id] and stencil_cell_list.size() > max_size) {
          max_size = stencil_cell_list.size();
        }
      }
      return max_size;
    }();

    Kokkos::Experimental::UniqueToken<Kokkos::DefaultExecutionSpace::execution_space,
                                      Kokkos::Experimental::UniqueTokenScope::Global>
      tokens;

    DiscreteFunctionDPk<MeshType::Dimension, std::remove_const_t<DataType>> dpk{mesh.meshVariant(), degree};

    if constexpr (is_polygonal_mesh_v<MeshType>) {
      if constexpr (std::is_arithmetic_v<DataType>) {
        SmallArray<SmallMatrix<double>> A_pool(Kokkos::DefaultExecutionSpace::concurrency());
        SmallArray<SmallVector<double>> b_pool(Kokkos::DefaultExecutionSpace::concurrency());
        for (size_t i = 0; i < A_pool.size(); ++i) {
          A_pool[i] = SmallMatrix<double>(max_stencil_size, MeshType::Dimension);
          b_pool[i] = SmallVector<double>(max_stencil_size);
        }

        parallel_for(
          mesh.numberOfCells(), PUGS_LAMBDA(const CellId cell_j_id) {
            if (cell_is_owned[cell_j_id]) {
              const int32_t t = tokens.acquire();

              auto stencil_cell_list = stencil[cell_j_id];

              ShrinkVectorView b(b_pool[t], stencil_cell_list.size());

              const double pj = p0_function[cell_j_id];
              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];

                b[i] = p0_function[cell_i_id] - pj;
              }

              ShrinkMatrixView A(A_pool[t], stencil_cell_list.size());

              const Rd& Xj = xj[cell_j_id];

              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const Rd Xi_Xj         = xj[cell_i_id] - Xj;
                for (size_t l = 0; l < MeshType::Dimension; ++l) {
                  A(i, l) = Xi_Xj[l];
                }
              }

              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const double weight    = 1. / l2Norm(xj[cell_i_id] - Xj);
                for (size_t l = 0; l < MeshType::Dimension; ++l) {
                  A(i, l) *= weight;
                }
                b[i] *= weight;
              }

              SmallVector<double> x(MeshType::Dimension);
              Givens::solveInPlace(A, x, b);

              auto dpk_j = dpk.coefficients(cell_j_id);

              dpk_j[0] = pj;

              for (size_t l = 0; l < MeshType::Dimension; ++l) {
                dpk_j[1 + l] = x[l];
              }

              tokens.release(t);
            }
          });
      } else if constexpr (is_tiny_vector_v<DataType>) {
        SmallArray<SmallMatrix<double>> A_pool(Kokkos::DefaultExecutionSpace::concurrency());
        SmallArray<SmallMatrix<double>> B_pool(Kokkos::DefaultExecutionSpace::concurrency());
        for (size_t i = 0; i < A_pool.size(); ++i) {
          A_pool[i] = SmallMatrix<double>(max_stencil_size, MeshType::Dimension);
          B_pool[i] = SmallMatrix<double>(max_stencil_size, DataType::Dimension);
        }

        parallel_for(
          mesh.numberOfCells(), PUGS_LAMBDA(const CellId cell_j_id) {
            if (cell_is_owned[cell_j_id]) {
              const int32_t t = tokens.acquire();

              auto stencil_cell_list = stencil[cell_j_id];
              ShrinkMatrixView B(B_pool[t], stencil_cell_list.size());

              const DataType& pj = p0_function[cell_j_id];
              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const DataType& pi_pj  = p0_function[cell_i_id] - pj;
                for (size_t k = 0; k < DataType::Dimension; ++k) {
                  B(i, k) = pi_pj[k];
                }
              }

              ShrinkMatrixView A(A_pool[t], stencil_cell_list.size());

              const Rd& Xj = xj[cell_j_id];
              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const Rd Xi_Xj         = xj[cell_i_id] - Xj;
                for (size_t l = 0; l < MeshType::Dimension; ++l) {
                  A(i, l) = Xi_Xj[l];
                }
              }

              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const double weight    = 1. / l2Norm(xj[cell_i_id] - Xj);
                for (size_t l = 0; l < MeshType::Dimension - 1; ++l) {
                  A(i, l) *= weight;
                }
                for (size_t l = 0; l < DataType::Dimension; ++l) {
                  B(i, l) *= weight;
                }
              }

              TinyMatrix<MeshType::Dimension, DataType::Dimension> X;
              Givens::solveCollectionInPlace(A, X, B);

              auto dpk_j = dpk.coefficients(cell_j_id);

              dpk_j[0] = pj;
              for (size_t i = 0; i < MeshType::Dimension; ++i) {
                auto& dpk_j_ip1 = dpk_j[i + 1];
                for (size_t k = 0; k < DataType::Dimension; ++k) {
                  dpk_j_ip1[k] = X(i, k);
                }
              }

              tokens.release(t);
            }
          });
      } else if constexpr (is_tiny_matrix_v<DataType>) {
        SmallArray<SmallMatrix<double>> A_pool(Kokkos::DefaultExecutionSpace::concurrency());
        SmallArray<SmallMatrix<double>> B_pool(Kokkos::DefaultExecutionSpace::concurrency());
        for (size_t i = 0; i < A_pool.size(); ++i) {
          A_pool[i] = SmallMatrix<double>(max_stencil_size, MeshType::Dimension);
          B_pool[i] = SmallMatrix<double>(max_stencil_size, DataType::Dimension);
        }

        parallel_for(
          mesh.numberOfCells(), PUGS_LAMBDA(const CellId cell_j_id) {
            if (cell_is_owned[cell_j_id]) {
              const int32_t t = tokens.acquire();

              auto stencil_cell_list = stencil[cell_j_id];
              ShrinkMatrixView B(B_pool[t], stencil_cell_list.size());

              const DataType& pj = p0_function[cell_j_id];
              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const DataType& pi_pj  = p0_function[cell_i_id] - pj;
                for (size_t k = 0; k < DataType::NumberOfRows; ++k) {
                  for (size_t l = 0; l < DataType::NumberOfColumns; ++l) {
                    B(i, k * DataType::NumberOfColumns + l) = pi_pj(k, l);
                  }
                }
              }

              ShrinkMatrixView A(A_pool[t], stencil_cell_list.size());

              const Rd& Xj = xj[cell_j_id];

              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const Rd Xi_Xj         = xj[cell_i_id] - Xj;
                for (size_t l = 0; l < MeshType::Dimension; ++l) {
                  A(i, l) = Xi_Xj[l];
                }
              }

              for (size_t i = 0; i < stencil_cell_list.size(); ++i) {
                const CellId cell_i_id = stencil_cell_list[i];
                const double weight    = 1. / l2Norm(xj[cell_i_id] - Xj);
                for (size_t l = 0; l < MeshType::Dimension - 1; ++l) {
                  A(i, l) *= weight;
                }
                for (size_t l = 0; l < DataType::Dimension; ++l) {
                  B(i, l) *= weight;
                }
              }

              TinyMatrix<MeshType::Dimension, DataType::Dimension> X;
              Givens::solveCollectionInPlace(A, X, B);

              auto dpk_j = dpk.coefficients(cell_j_id);
              dpk_j[0]   = pj;

              for (size_t i = 0; i < MeshType::Dimension; ++i) {
                auto& dpk_j_ip1 = dpk_j[i + 1];
                for (size_t k = 0; k < DataType::NumberOfRows; ++k) {
                  for (size_t l = 0; l < DataType::NumberOfColumns; ++l) {
                    dpk_j_ip1(k, l) = X(i, k * DataType::NumberOfColumns + l);
                  }
                }
              }

              tokens.release(t);
            }
          });
      } else {
        throw NotImplementedError("dealing with " + demangle<DataType>());
      }
    }

    synchronize(dpk.cellArrays());
    return dpk;
  }

  [[nodiscard]] std::vector<std::shared_ptr<const DiscreteFunctionDPkVariant>> build(
    const size_t degree,
    const std::vector<std::shared_ptr<const DiscreteFunctionVariant>>& discrete_function_variant_list) const;

  template <typename... DiscreteFunctionT>
  [[nodiscard]] std::vector<std::shared_ptr<const DiscreteFunctionDPkVariant>>
  build(const size_t degree, DiscreteFunctionT... input) const
  {
    std::vector<std::shared_ptr<const DiscreteFunctionVariant>> variant_vector;
    auto convert = [&variant_vector](auto&& df) {
      using DF_T = std::decay_t<decltype(df)>;
      if constexpr (is_discrete_function_v<DF_T> or std::is_same_v<DiscreteFunctionVariant, DF_T>) {
        variant_vector.push_back(std::make_shared<DiscreteFunctionVariant>(df));
      } else if constexpr (is_shared_ptr_v<DF_T>) {
        using DF_Value_T = std::decay_t<typename DF_T::element_type>;
        if constexpr (is_discrete_function_v<DF_Value_T> or std::is_same_v<DiscreteFunctionVariant, DF_Value_T>) {
          variant_vector.push_back(std::make_shared<DiscreteFunctionVariant>(*df));
        } else {
          static_assert(is_false_v<DF_T>, "unexpected type");
        }
      } else {
        static_assert(is_false_v<DF_T>, "unexpected type");
      }
    };

    (convert(std::forward<DiscreteFunctionT>(input)), ...);
    return this->build(degree, variant_vector);
  }

  PolynomialReconstruction() {}
};

#endif   // POLYNOMIAL_RECONSTRUCTION_HPP