lagrangedatacollector.hh

#pragma once
 
#include <cassert>
#include <map>
#include <vector>
 
#include <dune/geometry/referenceelement.hh>
#include <dune/grid/common/partitionset.hh>
#include <dune/grid/common/mcmgmapper.hh>
#include <dune/vtk/types.hh>
#include <dune/vtk/utility/lagrangepoints.hh>
 
#include "unstructureddatacollector.hh"
 
namespace Dune::Vtk
{
  template <class GridView, int ORDER = -1>
  class LagrangeDataCollector
      : public UnstructuredDataCollectorInterface<GridView, LagrangeDataCollector<GridView,ORDER>, Partitions::All>
  {
    using Self = LagrangeDataCollector;
    using Super = UnstructuredDataCollectorInterface<GridView, Self, Partitions::All>;
 
    auto pointSetMapper() const
    {
      auto pointSetLayout = [&](Dune::GeometryType gt, int dimgrid) -> unsigned int {
        if (gt.dim()==dim)
        {
          auto refElem = Dune::referenceElement<double,dim>(gt);
          return pointSets_.at(gt).size() - refElem.size(dim);
        }
        else if (gt.dim()==0)
          return 1;
        else
          return 0;
      };
      return Dune::MultipleCodimMultipleGeomTypeMapper(gridView(), pointSetLayout);
    }
 
  public:
    static_assert(ORDER != 0, "Order 0 not supported");
    using Super::dim;
    using Super::partition; // NOTE: Lagrange data-collector currently implemented for the All partition only
    using Super::gridView;
 
  public:
    LagrangeDataCollector (GridView const& gridView, int order = ORDER)
      : Super(gridView)
      , order_(order)
    {
      assert(order > 0 && "Order 0 not supported");
      assert(ORDER < 0 || order == ORDER);
    }
 
    void updateImpl ()
    {
      auto const& indexSet = gridView().indexSet();
 
      pointSets_.clear();
      for (auto gt : indexSet.types(0))
        pointSets_.emplace(gt, order_);
 
      for (auto& pointSet : pointSets_)
        pointSet.second.build(pointSet.first);
 
      numPoints_ = pointSetMapper().size();
    }
 
    std::uint64_t numPointsImpl () const
    {
      return numPoints_;
    }
 
 
    template <class T>
    std::vector<T> pointsImpl () const
    {
      std::vector<T> data(this->numPoints() * 3);
      auto const& mapper = pointSetMapper();
 
      for (auto const& element : elements(gridView(), partition)) {
        auto geometry = element.geometry();
        auto refElem = referenceElement<T,dim>(element.type());
 
        auto const& pointSet = pointSets_.at(element.type());
        unsigned int vertexDOFs = refElem.size(dim);
 
        for (std::size_t i = 0; i < pointSet.size(); ++i) {
          auto const& p = pointSet[i];
          if (i < vertexDOFs)
            assert(p.localKey().codim() == dim);
 
          auto const& localKey = p.localKey();
          std::size_t idx = 3 * (localKey.codim() == dim
              ? mapper.subIndex(element, localKey.subEntity(), dim)
              : mapper.index(element) + (i - vertexDOFs));
 
          auto v = geometry.global(p.point());
          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 gridView().size(0);
    }
 
 
    Cells cellsImpl () const
    {
      Cells cells;
      cells.connectivity.reserve(this->numPoints());
      cells.offsets.reserve(this->numCells());
      cells.types.reserve(this->numCells());
 
      auto const& mapper = pointSetMapper();
 
      std::int64_t old_o = 0;
      for (auto const& element : elements(gridView(), partition)) {
        auto refElem = referenceElement<double,dim>(element.type());
        Vtk::CellType cellType(element.type(), Vtk::CellType::LAGRANGE);
 
        auto const& pointSet = pointSets_.at(element.type());
        unsigned int vertexDOFs = refElem.size(dim);
 
        for (std::size_t i = 0; i < pointSet.size(); ++i) {
          auto const& p = pointSet[i];
          auto const& localKey = p.localKey();
          std::size_t idx = (localKey.codim() == dim
              ? mapper.subIndex(element, localKey.subEntity(), dim)
              : mapper.index(element) + (i - vertexDOFs));
          cells.connectivity.push_back(std::int64_t(idx));
        }
 
        cells.offsets.push_back(old_o += pointSet.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 const& mapper = pointSetMapper();
 
      auto localFct = localFunction(fct);
      for (auto const& element : elements(gridView(), partition)) {
        localFct.bind(element);
        auto refElem = referenceElement<T,dim>(element.type());
 
        auto const& pointSet = pointSets_.at(element.type());
        unsigned int vertexDOFs = refElem.size(dim);
 
        for (std::size_t i = 0; i < pointSet.size(); ++i) {
          auto const& p = pointSet[i];
          auto const& localKey = p.localKey();
          std::size_t idx = nComps * (localKey.codim() == dim
              ? mapper.subIndex(element, localKey.subEntity(), dim)
              : mapper.index(element) + (i - vertexDOFs));
 
          for (int comp = 0; comp < nComps; ++comp)
            data[idx + comp] = T(localFct.evaluate(comp, p.point()));
        }
        localFct.unbind();
      }
      return data;
    }
 
  private:
    unsigned int order_;
    std::uint64_t numPoints_ = 0;
 
    using PointSet = LagrangePointSet<typename GridView::ctype, GridView::dimension>;
    std::map<GeometryType, PointSet> pointSets_;
  };
 
} // end namespace Dune::Vtk