curvedgeometrydatacollector.hh

#ifndef DUNE_VTK_DATACOLLECTORS_CURVEDGEOMETRYDATACOLLECTOR_HH
#define DUNE_VTK_DATACOLLECTORS_CURVEDGEOMETRYDATACOLLECTOR_HH

#if !HAVE_DUNE_VTK
#error "CurvedGeometryDataCollector needs dune-vtk"
#endif

#include <functional>
#include <vector>

#include <dune/geometry/referenceelements.hh>
#include <dune/grid/common/partitionset.hh>
#include <dune/vtk/forward.hh>
#include <dune/vtk/types.hh>
#include <dune/vtk/datacollectors/quadraticdatacollector.hh>

namespace Dune {

/// \brief DataCollector for dune-vtk writers extracting the parametrized geometry from
/// a local-to-global coordinate mapping and writing quadratic VTK elements
/**
 * To write curved elements, a \ref CurvedGeometry is constructed on-the-fly from the given
 * mappings in the constructor.
 *
 * \tparam GridView  The grid-view to write
 * \tparam Geometry  The CurvedGeometry to use for element parametrization
 **/
template <class GridView, class Geometry>
class CurvedGeometryDataCollector
    : public Vtk::UnstructuredDataCollectorInterface<GridView,
        CurvedGeometryDataCollector<GridView,Geometry>, Partitions::All>
{
  using Self = CurvedGeometryDataCollector;
  using Super = Vtk::UnstructuredDataCollectorInterface<GridView, Self, Partitions::All>;

  using LocalCoord = typename GridView::template Codim<0>::Entity::Geometry::LocalCoordinate;
  using GlobalCoord = typename GridView::template Codim<0>::Entity::Geometry::GlobalCoordinate;

  using Mapping = std::function<typename Geometry::GlobalCoordinate(GlobalCoord)>;
  using Parametrization = std::vector<std::vector<typename Geometry::GlobalCoordinate>>;

public:
  using Super::dim;
  using Super::partition;

public:
  /// \brief Construct the curved geometry from a coordinate mapping
  /**
   * \param gridView  An instance of the GridView type
   * \param mapping   A coordinate mapping from global flat to global curved coordinates, i.e.
   *                  with signature `typename Geometry::GlobalCoordinate(GlobalCoord)`,
   *                  see \ref Mapping
   **/
  CurvedGeometryDataCollector (const GridView& gridView, Mapping mapping)
    : Super(gridView)
    , dataCollector_(gridView)
    , mapping_(mapping)
  {}

  /// \brief construct the curved geometry from a stored parametrization, i.e. a vector of
  /// coefficient vectors
  /**
   * \param gridView         An instance of the GridView type
   * \param parametrization  Vector (accessed by element index) of element parametrization
   *                         coefficients, see \ref Parametrization
   **/
  CurvedGeometryDataCollector (const GridView& gridView, const Parametrization& parametrization)
    : Super(gridView)
    , dataCollector_(gridView)
    , parametrization_(&parametrization)
  {}

  /// collect the points of the geometry. This uses the curved geometry instead of the element geometry
  template <class T>
  std::vector<T> pointsImpl () const
  {
    std::vector<T> data(Super::numPoints() * 3);
    const auto& indexSet = gridView_.indexSet();
    for (const auto& element : elements(gridView_, partition)) {
      auto refElem = referenceElement<T,Geometry::mydimension>(element.type());
      auto geometry = makeGeometry(indexSet, element);

      // vertices
      for (unsigned int i = 0; i < element.subEntities(dim); ++i) {
        std::size_t idx = 3 * indexSet.subIndex(element, i, dim);
        auto v = geometry.global(refElem.position(i,dim));
        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);
      }
      // edge centers
      for (unsigned int i = 0; i < element.subEntities(dim-1); ++i) {
        std::size_t idx = 3 * (indexSet.subIndex(element, i, dim-1) + gridView_.size(dim));
        auto v = geometry.global(refElem.position(i,dim-1));
        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;
  }

private:

  /// construct the geometry object, either from the mapping or the coefficient vector
  template <class IndexSet, class Element>
  Geometry makeGeometry(const IndexSet& indexSet, const Element& element) const
  {
    if (parametrization_ != nullptr) {
      return Geometry{element.type(), (*parametrization_)[indexSet.index(element)]};
    } else {
      // mapping from local to global coordinates
      auto X = [mapping=mapping_,geo=element.geometry()](const LocalCoord& local)
        -> typename Geometry::GlobalCoordinate
      {
        return mapping(geo.global(local));
      };

      // create a curved geometry
      return Geometry{element.type(), X};
    }
  }

public: // methods forwarded to quadratic data collector

  void updateImpl ()
  {
    dataCollector_.update();
  }

  std::uint64_t numPointsImpl () const
  {
    return dataCollector_.numPoints();
  }

  std::vector<std::uint64_t> pointIdsImpl () const
  {
    return dataCollector_.pointIds();
  }

  std::uint64_t numCellsImpl () const
  {
    return dataCollector_.numCells();
  }

  Vtk::Cells cellsImpl () const
  {
    return dataCollector_.cells();
  }

  template <class T, class GlobalFunction>
  std::vector<T> pointDataImpl (const GlobalFunction& fct) const
  {
    return dataCollector_.template pointData<T>(fct);
  }

private:
  using Super::gridView_;
  Vtk::QuadraticDataCollector<GridView> dataCollector_;
  Mapping mapping_;
  const Parametrization* parametrization_ = nullptr;
};

} // end namespace Dune

#endif // DUNE_VTK_DATACOLLECTORS_CURVEDGEOMETRYDATACOLLECTOR_HH