DiscreteFunctionP0.hpp

#ifndef DISCRETE_FUNCTION_P0_HPP
#define DISCRETE_FUNCTION_P0_HPP

#include <scheme/IDiscreteFunction.hpp>

#include <language/utils/InterpolateItemValue.hpp>
#include <mesh/Connectivity.hpp>
#include <mesh/Mesh.hpp>
#include <mesh/MeshData.hpp>
#include <mesh/MeshDataManager.hpp>
#include <scheme/DiscreteFunctionDescriptorP0.hpp>

template <size_t Dimension, typename DataType>
class DiscreteFunctionP0 : public IDiscreteFunction
{
 public:
  using data_type = DataType;
  using MeshType  = Mesh<Connectivity<Dimension>>;

  static constexpr HandledItemDataType handled_data_type = HandledItemDataType::value;

 private:
  std::shared_ptr<const MeshType> m_mesh;
  CellValue<DataType> m_cell_values;

  DiscreteFunctionDescriptorP0 m_discrete_function_descriptor;

 public:
  PUGS_INLINE
  ASTNodeDataType
  dataType() const final
  {
    return ast_node_data_type_from<DataType>;
  }

  PUGS_INLINE
  const CellValue<DataType>&
  cellValues() const
  {
    return m_cell_values;
  }

  PUGS_INLINE
  std::shared_ptr<const IMesh>
  mesh() const
  {
    return m_mesh;
  }

  PUGS_INLINE
  const IDiscreteFunctionDescriptor&
  descriptor() const final
  {
    return m_discrete_function_descriptor;
  }

  PUGS_FORCEINLINE
  operator DiscreteFunctionP0<Dimension, const DataType>() const
  {
    return DiscreteFunctionP0<Dimension, const DataType>(m_mesh, m_cell_values);
  }

  PUGS_INLINE
  void
  fill(const DataType& data) const noexcept
  {
    static_assert(not std::is_const_v<DataType>, "Cannot modify ItemValue of const");
    m_cell_values.fill(data);
  }

  friend PUGS_INLINE DiscreteFunctionP0<Dimension, std::remove_const_t<DataType>>
  copy(const DiscreteFunctionP0& source)
  {
    return DiscreteFunctionP0<Dimension, std::remove_const_t<DataType>>{source.m_mesh, copy(source.cellValues())};
  }

  friend PUGS_INLINE void
  copy_to(const DiscreteFunctionP0<Dimension, DataType>& source,
          DiscreteFunctionP0<Dimension, std::remove_const_t<DataType>>& destination)
  {
    Assert(source.m_mesh == destination.m_mesh);
    copy_to(source.m_cell_values, destination.m_cell_values);
  }

  PUGS_FORCEINLINE DataType&
  operator[](const CellId cell_id) const noexcept(NO_ASSERT)
  {
    return m_cell_values[cell_id];
  }

  PUGS_INLINE DiscreteFunctionP0<Dimension, std::remove_const_t<DataType>>
  operator-() const
  {
    Assert(m_cell_values.isBuilt());
    DiscreteFunctionP0<Dimension, std::remove_const_t<DataType>> opposite = copy(*this);
    parallel_for(
      m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { opposite[cell_id] = -opposite[cell_id]; });

    return opposite;
  }

  template <typename DataType2T>
  PUGS_INLINE DiscreteFunctionP0<Dimension, decltype(DataType{} + DataType2T{})>
  operator+(const DiscreteFunctionP0<Dimension, DataType2T>& g) const
  {
    const DiscreteFunctionP0& f = *this;
    Assert(f.mesh() == g.mesh(), "functions are not defined on the same mesh");
    DiscreteFunctionP0<Dimension, decltype(DataType{} + DataType2T{})> sum(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { sum[cell_id] = f[cell_id] + g[cell_id]; });
    return sum;
  }

  template <typename LHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(LHSDataType{} + DataType{})>
  operator+(const LHSDataType& a, const DiscreteFunctionP0& g)
  {
    using ProductDataType = decltype(LHSDataType{} + DataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> sum(g.m_mesh);
    parallel_for(
      g.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { sum[cell_id] = a + g[cell_id]; });
    return sum;
  }

  template <typename RHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(DataType{} + RHSDataType{})>
  operator+(const DiscreteFunctionP0& f, const RHSDataType& b)
  {
    using ProductDataType = decltype(DataType{} + RHSDataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> sum(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { sum[cell_id] = f[cell_id] + b; });
    return sum;
  }

  template <typename DataType2T>
  PUGS_INLINE DiscreteFunctionP0<Dimension, decltype(DataType{} - DataType2T{})>
  operator-(const DiscreteFunctionP0<Dimension, DataType2T>& g) const
  {
    const DiscreteFunctionP0& f = *this;
    Assert(f.mesh() == g.mesh(), "functions are not defined on the same mesh");
    DiscreteFunctionP0<Dimension, decltype(DataType{} - DataType2T{})> difference(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { difference[cell_id] = f[cell_id] - g[cell_id]; });
    return difference;
  }

  template <typename LHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(LHSDataType{} - DataType{})>
  operator-(const LHSDataType& a, const DiscreteFunctionP0& g)
  {
    using ProductDataType = decltype(LHSDataType{} - DataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> difference(g.m_mesh);
    parallel_for(
      g.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { difference[cell_id] = a - g[cell_id]; });
    return difference;
  }

  template <typename RHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(DataType{} - RHSDataType{})>
  operator-(const DiscreteFunctionP0& f, const RHSDataType& b)
  {
    using ProductDataType = decltype(DataType{} - RHSDataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> difference(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { difference[cell_id] = f[cell_id] - b; });
    return difference;
  }

  template <typename DataType2T>
  PUGS_INLINE DiscreteFunctionP0<Dimension, decltype(DataType{} * DataType2T{})>
  operator*(const DiscreteFunctionP0<Dimension, DataType2T>& g) const
  {
    const DiscreteFunctionP0& f = *this;
    Assert(f.mesh() == g.mesh(), "functions are not defined on the same mesh");
    DiscreteFunctionP0<Dimension, decltype(DataType{} * DataType2T{})> product(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { product[cell_id] = f[cell_id] * g[cell_id]; });
    return product;
  }

  template <typename LHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(LHSDataType{} * DataType{})>
  operator*(const LHSDataType& a, const DiscreteFunctionP0& f)
  {
    using ProductDataType = decltype(LHSDataType{} * DataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> product(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { product[cell_id] = a * f[cell_id]; });
    return product;
  }

  template <typename RHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(DataType{} * RHSDataType{})>
  operator*(const DiscreteFunctionP0& f, const RHSDataType& b)
  {
    using ProductDataType = decltype(DataType{} * RHSDataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> product(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { product[cell_id] = f[cell_id] * b; });
    return product;
  }

  template <typename DataType2T>
  PUGS_INLINE DiscreteFunctionP0<Dimension, decltype(DataType{} / DataType2T{})>
  operator/(const DiscreteFunctionP0<Dimension, DataType2T>& g) const
  {
    const DiscreteFunctionP0& f = *this;
    Assert(f.mesh() == g.mesh(), "functions are not defined on the same mesh");
    DiscreteFunctionP0<Dimension, decltype(DataType{} / DataType2T{})> ratio(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { ratio[cell_id] = f[cell_id] / g[cell_id]; });
    return ratio;
  }

  template <typename LHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(LHSDataType{} / DataType{})>
  operator/(const LHSDataType& a, const DiscreteFunctionP0& f)
  {
    using ProductDataType = decltype(LHSDataType{} / DataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> ratio(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { ratio[cell_id] = a / f[cell_id]; });
    return ratio;
  }

  template <typename RHSDataType>
  PUGS_INLINE friend DiscreteFunctionP0<Dimension, decltype(DataType{} / RHSDataType{})>
  operator/(const DiscreteFunctionP0& f, const RHSDataType& b)
  {
    using ProductDataType = decltype(DataType{} / RHSDataType{});
    DiscreteFunctionP0<Dimension, ProductDataType> ratio(f.m_mesh);
    parallel_for(
      f.m_mesh->numberOfCells(), PUGS_LAMBDA(CellId cell_id) { ratio[cell_id] = f[cell_id] / b; });
    return ratio;
  }

  DiscreteFunctionP0(const std::shared_ptr<const MeshType>& mesh, const FunctionSymbolId& function_id) : m_mesh(mesh)
  {
    using MeshDataType      = MeshData<Dimension>;
    MeshDataType& mesh_data = MeshDataManager::instance().getMeshData(*mesh);

    m_cell_values =
      InterpolateItemValue<DataType(TinyVector<Dimension>)>::template interpolate<ItemType::cell>(function_id,
                                                                                                  mesh_data.xj());
  }

  DiscreteFunctionP0(const std::shared_ptr<const MeshType>& mesh) : m_mesh{mesh}, m_cell_values{mesh->connectivity()} {}

  DiscreteFunctionP0(const std::shared_ptr<const MeshType>& mesh, const CellValue<DataType>& cell_value)
    : m_mesh{mesh}, m_cell_values{cell_value}
  {
    Assert(mesh->connectivity().shared_ptr() == cell_value.connectivity_ptr());
  }

  DiscreteFunctionP0(const DiscreteFunctionP0&) noexcept = default;
  DiscreteFunctionP0(DiscreteFunctionP0&&) noexcept      = default;

  ~DiscreteFunctionP0() = default;
};

#endif   // DISCRETE_FUNCTION_P0_HPP