subdivisiondatacollector.hh

#pragma once
 
#include <cassert>
#include <map>
#include <optional>
#include <vector>
 
#include <dune/geometry/referenceelement.hh>
#include <dune/geometry/refinement.hh>
#include <dune/geometry/virtualrefinement.hh>
#include <dune/grid/common/partitionset.hh>
#include <dune/grid/common/mcmgmapper.hh>
#include <dune/vtk/types.hh>
 
#include "unstructureddatacollector.hh"
 
namespace Dune::Vtk
{
  enum class Subdivision {
    REFINEMENT,             // refine cubes into sub-cubes
    TRIANGULATION           // refine all elements into simplices
  };
 
  template <class GridView>
  class SubdivisionDataCollector
      : public UnstructuredDataCollectorInterface<GridView, SubdivisionDataCollector<GridView>, Partitions::All>
  {
    using Self = SubdivisionDataCollector;
    using Super = UnstructuredDataCollectorInterface<GridView, Self, Partitions::All>;
 
    static constexpr int dim = GridView::dimension;
 
    // GeometryType of the divided cell, depends on the Subdivision property
    static Dune::GeometryType refinementGeometryType (Dune::GeometryType const& gt, Vtk::Subdivision subdivision)
    {
      if (subdivision == Vtk::Subdivision::TRIANGULATION)
        return Dune::GeometryTypes::simplex(gt.dim());
      else if (gt.isPrism() || gt.isPyramid())
        return Dune::GeometryTypes::simplex(3);
      else
        return gt;
    }
 
    static auto& virtualRefinement (Dune::GeometryType const& gt, Vtk::Subdivision subdivision)
    {
      Dune::GeometryType coerceTo = refinementGeometryType(gt, subdivision);
      return Dune::buildRefinement<dim, typename GridView::ctype>(gt, coerceTo);
    }
 
    // The mapper contains n nodes per element interior but no nodes on other codimensions, where n
    // refers to the number of vertices per element in a virtual refinement.
    auto makeMapper () const
    {
      auto layout = [subdivision=subdivision_, ivals=ivals_](Dune::GeometryType gt, int dimgrid) -> unsigned int {
        if (int(gt.dim()) == dimgrid)
          return virtualRefinement(gt, subdivision).nVertices(ivals);
        else
          return 0;
      };
      return Dune::MultipleCodimMultipleGeomTypeMapper(gridView(), layout);
    }
 
  public:
    using Super::partition; // NOTE: subdivision data-collector currently implemented for the All partition only
    using Super::gridView;
 
  public:
    SubdivisionDataCollector (GridView const& gridView,
                              Dune::RefinementIntervals ivals = Dune::RefinementIntervals(1),
                              Vtk::Subdivision subdivision = Vtk::Subdivision::REFINEMENT)
      : Super(gridView)
      , ivals_(ivals)
      , subdivision_(subdivision)
    {
      assert(ivals_.intervals() > 0);
    }
 
    void updateImpl ()
    {
      auto const& indexSet = gridView().indexSet();
 
      mapper_.emplace(makeMapper());
      numPoints_ = mapper_->size();
      numCells_ = 0;
      for (auto gt : indexSet.types(0)) {
        const auto& ref = virtualRefinement(gt, subdivision_);
        numCells_ += indexSet.size(gt)*ref.nElements(ivals_);
      }
    }
 
    std::uint64_t numPointsImpl () const
    {
      return numPoints_;
    }
 
 
    template <class T>
    std::vector<T> pointsImpl () const
    {
      std::vector<T> data(this->numPoints() * 3);
      for (auto const& element : elements(gridView(), partition)) {
        auto geometry = element.geometry();
        std::size_t element_idx = mapper_->index(element);
 
        const auto& ref = virtualRefinement(element.type(), subdivision_);
        auto vEnd = ref.vEnd(ivals_);
        for (auto vIt = ref.vBegin(ivals_); vIt != vEnd; ++vIt) {
          std::size_t idx = 3 * (element_idx + vIt.index());
 
          auto v = geometry.global(vIt.coords());
          for (std::size_t j = 0; j < v.size(); ++j)
            data[idx + j] = T(v[j]);
          for (std::size_t j = v.size(); j < 3u; ++j)
            data[idx + j] = T(0);
        }
      }
      return data;
    }
 
    std::uint64_t numCellsImpl () const
    {
      return numCells_;
    }
 
 
    Cells cellsImpl () const
    {
      Cells cells;
      cells.connectivity.reserve(this->numPoints());
      cells.offsets.reserve(this->numCells());
      cells.types.reserve(this->numCells());
 
      std::int64_t old_o = 0;
      for (auto const& element : elements(gridView(), partition)) {
        std::size_t element_idx = mapper_->index(element);
        Vtk::CellType cellType(refinementGeometryType(element.type(),subdivision_), Vtk::CellType::LINEAR);
 
        const auto& ref = virtualRefinement(element.type(), subdivision_);
        auto eEnd = ref.eEnd(ivals_);
        for (auto eIt = ref.eBegin(ivals_); eIt != eEnd; ++eIt) {
          auto vertices = eIt.vertexIndices();
          for (std::size_t i = 0; i < vertices.size(); ++i) {
            std::size_t idx = (element_idx + vertices[cellType.permutation(i)]);
            cells.connectivity.push_back(std::int64_t(idx));
          }
          cells.offsets.push_back(old_o += vertices.size());
          cells.types.push_back(cellType.type());
        }
      }
      return cells;
    }
 
    template <class T, class GlobalFunction>
    std::vector<T> pointDataImpl (GlobalFunction const& fct) const
    {
      int nComps = fct.numComponents();
      std::vector<T> data(this->numPoints() * nComps);
 
      auto localFct = localFunction(fct);
      for (auto const& element : elements(gridView(), partition)) {
        std::size_t element_idx = mapper_->index(element);
        localFct.bind(element);
 
        const auto& ref = virtualRefinement(element.type(), subdivision_);
        auto vEnd = ref.vEnd(ivals_);
        for (auto vIt = ref.vBegin(ivals_); vIt != vEnd; ++vIt) {
          std::size_t idx = nComps * (element_idx + vIt.index());
 
          for (int comp = 0; comp < nComps; ++comp)
            data.at(idx + comp) = T(localFct.evaluate(comp, vIt.coords()));
        }
        localFct.unbind();
      }
      return data;
    }
 
    template <class T, class GlobalFunction>
    std::vector<T> cellDataImpl (GlobalFunction const& fct) const
    {
      int nComps = fct.numComponents();
      std::vector<T> data;
      data.reserve(this->numCells() * nComps);
 
      auto localFct = localFunction(fct);
      for (auto const& element : elements(gridView(), partition)) {
        localFct.bind(element);
 
        const auto& ref = virtualRefinement(element.type(), subdivision_);
        auto eEnd = ref.eEnd(ivals_);
        for (auto eIt = ref.eBegin(ivals_); eIt != eEnd; ++eIt) {
          for (int comp = 0; comp < fct.numComponents(); ++comp)
            data.emplace_back(localFct.evaluate(comp, eIt.coords()));
        }
        localFct.unbind();
      }
      return data;
    }
 
  private:
    Dune::RefinementIntervals ivals_;
    Vtk::Subdivision subdivision_;
    std::optional<Dune::MultipleCodimMultipleGeomTypeMapper<GridView>> mapper_;
 
    std::uint64_t numPoints_ = 0;
    std::uint64_t numCells_ = 0;
  };
 
} // end namespace Dune::Vtk