codimdatacollector.hh

#pragma once
 
#include <numeric>
#include <vector>
 
#include <dune/geometry/referenceelements.hh>
#include <dune/grid/common/partitionset.hh>
#include <dune/grid/utility/globalindexset.hh>
#include <dune/vtk/types.hh>
 
#include <dune/vtk/datacollectors/unstructureddatacollector.hh>
 
namespace Dune::Vtk
{
  template <class GridView, class Partition = Partitions::InteriorBorder>
  class CodimDataCollector
      : public UnstructuredDataCollectorInterface<GridView, CodimDataCollector<GridView,Partition>, Partition>
  {
    using Self = CodimDataCollector;
    using Super = UnstructuredDataCollectorInterface<GridView, Self, Partition>;
 
  public:
    using Super::dim;
    using Super::partition;
    using Super::gridView;
 
  public:
    explicit CodimDataCollector (GridView const& gridView, int codim = 1)
      : Super(gridView)
      , codim_(codim)
    {}
 
    void updateImpl ()
    {
      numPoints_ = gridView().size(dim);
      numCells_ = gridView().size(codim_);
    }
 
    std::uint64_t numPointsImpl () const
    {
      return numPoints_;
    }
 
    template <class T>
    std::vector<T> pointsImpl () const
    {
      std::vector<T> data(numPoints_ * 3);
      auto const& indexSet = gridView().indexSet();
      for (auto const& vertex : vertices(gridView(), partition)) {
        auto v = vertex.geometry().center();
        std::size_t idx = 3*indexSet.index(vertex);
        for (std::size_t j = 0; j < v.size(); ++j)
          data[idx+j] = v[j];
        for (std::size_t j = v.size(); j < 3u; ++j)
          data[idx+j] = T(0);
      }
      return data;
    }
 
    std::vector<std::uint64_t> pointIdsImpl () const
    {
      std::vector<std::uint64_t> data(numPoints_);
      auto const& indexSet = gridView().indexSet();
      GlobalIndexSet<GridView> globalIndexSet(gridView(), dim);
      for (auto const& vertex : vertices(gridView(), partition)) {
        data[indexSet.index(vertex)] = globalIndexSet.index(vertex);
      }
      return data;
    }
 
    std::uint64_t numCellsImpl () const
    {
      return numCells_;
    }
 
    Cells cellsImpl () const
    {
      auto const& indexSet = gridView().indexSet();
      auto types = indexSet.types(codim_);
      int maxVertices = std::accumulate(types.begin(), types.end(), 1, [codim=codim_](int m, GeometryType t) {
        switch (codim) {
          case 0:
            return std::max(m, referenceElement<double,dim>(t).size(dim));
          case 1: if constexpr (dim >= 1) {
            return std::max(m, referenceElement<double,dim-1>(t).size(dim-1)); } break;
          case 2: if constexpr (dim >= 2) {
            return std::max(m, referenceElement<double,dim-2>(t).size(dim-2)); } break;
          case 3: if constexpr (dim >= 3) {
            return std::max(m, referenceElement<double,dim-3>(t).size(dim-3)); } break;
          default:
            return std::max(m, referenceElement<double,0>(t).size(0));
        }
        return m;
      });
 
      Cells cells;
      cells.connectivity.reserve(numCells_ * maxVertices);
      cells.offsets.reserve(numCells_);
      cells.types.reserve(numCells_);
 
      std::vector<bool> visited(gridView().size(codim_), false);
      std::int64_t old_o = 0;
      for (auto const& c : elements(gridView(), partition)) {
        auto refElem = referenceElement(c);
        for (int i = 0; i < refElem.size(codim_); ++i) {
          Vtk::CellType cellType(refElem.type(i,codim_));
          std::size_t idx = indexSet.subIndex(c,i,codim_);
          if (!visited[idx]) {
            visited[idx] = true;
            for (int j = 0; j < refElem.size(i,codim_,dim); ++j)
              cells.connectivity.push_back(indexSet.subIndex(c,
                refElem.subEntity(i,codim_,cellType.permutation(j),dim), dim));
            cells.offsets.push_back(old_o += refElem.size(i,codim_,dim));
            cells.types.push_back(cellType.type());
          }
        }
      }
      return cells;
    }
 
    template <class T, class GlobalFunction>
    std::vector<T> pointDataImpl (GlobalFunction const& fct) const
    {
      std::vector<T> data(numPoints_ * fct.numComponents());
      auto const& indexSet = gridView().indexSet();
      auto localFct = localFunction(fct);
      std::vector<bool> visited(gridView().size(codim_), false);
      for (auto const& c : elements(gridView(), partition)) {
        localFct.bind(c);
        auto refElem = referenceElement(c);
        for (int i = 0; i < refElem.size(codim_); ++i) {
          Vtk::CellType cellType(refElem.type(i,codim_));
          std::size_t edge_idx = indexSet.subIndex(c,i,codim_);
          if (!visited[edge_idx]) {
            visited[edge_idx] = true;
            for (int j = 0; j < refElem.size(i,codim_,dim); ++j) {
              std::size_t cj = refElem.subEntity(i,codim_,cellType.permutation(j),dim);
              std::size_t idx = fct.numComponents() * indexSet.subIndex(c,cj,dim);
              auto local = refElem.position(cj,dim);
              for (int comp = 0; comp < fct.numComponents(); ++comp)
                data[idx + comp] = T(localFct.evaluate(comp, local));
            }
          }
        }
        localFct.unbind();
      }
      return data;
    }
 
    template <class T, class GlobalFunction>
    std::vector<T> cellDataImpl (GlobalFunction const& fct) const
    {
      std::vector<T> data;
      data.reserve(numCells_ * fct.numComponents());
      auto const& indexSet = gridView().indexSet();
      auto localFct = localFunction(fct);
      std::vector<bool> visited(gridView().size(codim_), false);
      for (auto const& c : elements(gridView(), partition)) {
        localFct.bind(c);
        auto refElem = referenceElement(c);
        for (int i = 0; i < refElem.size(codim_); ++i) {
          std::size_t edge_idx = indexSet.subIndex(c,i,codim_);
          if (!visited[edge_idx]) {
            visited[edge_idx] = true;
            auto local = refElem.position(i,codim_);
            for (int comp = 0; comp < fct.numComponents(); ++comp)
              data.push_back(T(localFct.evaluate(comp, local)));
          }
        }
        localFct.unbind();
      }
      return data;
    }
 
  private:
    int codim_ = 1;
    std::uint64_t numPoints_ = 0;
    std::uint64_t numCells_ = 0;
  };
 
} // end namespace Dune::Vtk