vtk_writer.h

 
#ifndef __3e48c7b32a034cb3be3dbca884ff4f6c
#define __3e48c7b32a034cb3be3dbca884ff4f6c
 
#include <lf/mesh/mesh.h>
#include <lf/mesh/utils/utils.h>
 
#include <boost/variant/variant.hpp>
#include <string>
#include <utility>
#include <vector>
 
namespace lf::io {
 
class VtkFile {
 public:
  using size_type = unsigned int;
 
  enum class Format { ASCII, BINARY };
 
  enum class CellType {
    VTK_VERTEX = 1,
    VTK_POLY_VERTEX = 2,
    VTK_LINE = 3,
    VTK_POLY_LINE = 4,
    VTK_TRIANGLE = 5,
    VTK_TRIANGLE_STRIP = 6,
    VTK_POLYGON = 7,
    VTK_PIXEL = 8,
    VTK_QUAD = 9,
    VTK_TETRA = 10,
    VTK_VOXEL = 11,
    VTK_HEXAHEDRON = 12,
    VTK_WEDGE = 13,
    VTK_PYRAMID = 14,
    VTK_QUADRATIC_EDGE = 21,
    VTK_QUADRATIC_TRIANGLE = 22,
    VTK_QUADRATIC_QUAD = 23,
    VTK_QUADRATIC_TETRA = 24,
    VTK_QUADRATIC_HEXAHEDRON = 25,
    VTK_LAGRANGE_CURVE = 68,
    VTK_LAGRANGE_TRIANGLE = 69,
    VTK_LAGRANGE_QUADRILATERAL = 70,
    VTK_LAGRANGE_TETRAHEDRON = 71,
    VTK_LAGRANGE_HEXAHEDRON = 72,
    VTK_LAGRANGE_WEDGE = 73,
    VTK_LAGRANGE_PYRAMID = 74
  };
 
  class UnstructuredGrid {
   public:
    std::vector<Eigen::Vector3f> points;
    std::vector<std::vector<size_type>> cells;
    std::vector<CellType> cell_types;
  };
 
  template <class T>
  class FieldDataArray {
   public:
    std::string name;
    std::vector<T> data;
 
    FieldDataArray() = default;
 
    explicit FieldDataArray(std::string name, std::vector<T> data)
        : name(std::move(name)), data(std::move(data)) {}
  };
 
  template <class T>
  class ScalarData {
   public:
    std::string name;
    std::vector<T> data;
    std::string lookup_table = "default";
 
    ScalarData() = default;
 
    explicit ScalarData(std::string data_name, std::vector<T> data,
                        std::string lookup_table = "default")
        : name(std::move(data_name)),
          data(std::move(data)),
          lookup_table(std::move(lookup_table)) {}
  };
 
  template <class T>
  class VectorData {
   public:
    std::string name;
    std::vector<Eigen::Matrix<T, 3, 1>> data;
 
    VectorData() = default;
 
    explicit VectorData(std::string name,
                        std::vector<Eigen::Matrix<T, 3, 1>> data)
        : name(std::move(name)), data(std::move(data)) {}
  };
 
  // WARNING: the type long is interepreted differently depending on the
  // platorm. I suggest to not use it at all.
  using Attributes = std::vector<boost::variant<
      ScalarData<char>, ScalarData<unsigned char>, ScalarData<short>,
      ScalarData<unsigned short>, ScalarData<int>, ScalarData<unsigned int>,
      ScalarData<long>, ScalarData<unsigned long>, ScalarData<float>,
      ScalarData<double>, VectorData<float>, VectorData<double>>>;
 
  using FieldData =
      std::vector<boost::variant<FieldDataArray<int>, FieldDataArray<float>,
                                 FieldDataArray<double>>>;
 
  // Actual members
 
  std::string header;
  Format format = Format::ASCII;
 
  UnstructuredGrid unstructured_grid;
 
  FieldData field_data;
 
  // Data that is attached to points
  Attributes point_data;
 
  // Data that is attached to the cells of the mesh
  Attributes cell_data;
};
 
void WriteToFile(const VtkFile& vtk_file, const std::string& filename);
 
// clang-format off
 
// clang-format on
class VtkWriter {
 public:
  using dim_t = base::dim_t;
 
  VtkWriter(const VtkWriter&) = delete;
  VtkWriter(VtkWriter&&) = delete;
  VtkWriter& operator=(const VtkWriter&) = delete;
  VtkWriter& operator=(VtkWriter&&) = delete;
 
  VtkWriter(std::shared_ptr<const mesh::Mesh> mesh, std::string filename,
            dim_t codim = 0, unsigned char order = 1);
 
  void setBinary(bool binary) {
    if (binary) {
      vtk_file_.format = VtkFile::Format::BINARY;
    } else {
      vtk_file_.format = VtkFile::Format::ASCII;
    }
  }
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<unsigned char>& mds,
                      unsigned char undefined_value = 0);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<char>& mds,
                      char undefined_value = 0);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<unsigned int>& mds,
                      unsigned int undefined_value = 0);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<int>& mds,
                      int undefined_value = 0);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<float>& mds,
                      float undefined_value = 0.F);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<double>& mds,
                      double undefined_value = 0.);
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<Eigen::Vector2d>& mds,
                      const Eigen::Vector2d& undefined_value = {0, 0});
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<Eigen::Vector2f>& mds,
                      const Eigen::Vector2f& undefined_value = {0, 0});
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<Eigen::Vector3d>& mds,
                      const Eigen::Vector3d& undefined_value = {0, 0, 0});
 
  void WritePointData(const std::string& name,
                      const mesh::utils::MeshDataSet<Eigen::Vector3f>& mds,
                      const Eigen::Vector3f& undefined_value = {0, 0, 0});
 
  void WritePointData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::VectorXd>& mds,
      const Eigen::VectorXd& undefined_value = Eigen::Vector3d(0, 0, 0));
 
  void WritePointData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::VectorXf>& mds,
      const Eigen::VectorXf& undefined_value = Eigen::Vector3f(0, 0, 0));
 
  template <
      class MESH_FUNCTION,
      class = std::enable_if_t<lf::mesh::utils::isMeshFunction<MESH_FUNCTION>>>
  void WritePointData(const std::string& name,
                      const MESH_FUNCTION& mesh_function);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<unsigned char>& mds,
                     unsigned char undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<char>& mds,
                     char undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<unsigned int>& mds,
                     unsigned int undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<int>& mds,
                     int undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<float>& mds,
                     float undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<double>& mds,
                     double undefined_value = 0);
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<Eigen::Vector2d>& mds,
                     const Eigen::Vector2d& undefined_value = {0, 0});
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<Eigen::Vector2f>& mds,
                     const Eigen::Vector2f& undefined_value = {0, 0});
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<Eigen::Vector3d>& mds,
                     const Eigen::Vector3d& undefined_value = {0, 0, 0});
 
  void WriteCellData(const std::string& name,
                     const mesh::utils::MeshDataSet<Eigen::Vector3f>& mds,
                     const Eigen::Vector3f& undefined_value = {0, 0, 0});
 
  void WriteCellData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::VectorXd>& mds,
      const Eigen::VectorXd& undefined_value = Eigen::Vector3d(0, 0, 0));
 
  void WriteCellData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::VectorXf>& mds,
      const Eigen::VectorXf& undefined_value = Eigen::Vector3f(0, 0, 0));
 
  template <
      class MESH_FUNCTION,
      class = std::enable_if_t<lf::mesh::utils::isMeshFunction<MESH_FUNCTION>>>
  void WriteCellData(const std::string& name,
                     const MESH_FUNCTION& mesh_function);
 
  void WriteGlobalData(const std::string& name, std::vector<int> data);
 
  void WriteGlobalData(const std::string& name, std::vector<float> data);
 
  void WriteGlobalData(const std::string& name, std::vector<double> data);
 
  ~VtkWriter() { WriteToFile(vtk_file_, filename_); }
 
 private:
  std::shared_ptr<const mesh::Mesh> mesh_;
  VtkFile vtk_file_;
  std::string filename_;
  dim_t codim_;
  unsigned char order_;
 
  // entry [i] stores the reference coordinates for the auxilliary nodes for
  // RefEl.Id() == i
  std::array<Eigen::MatrixXd, 5> aux_nodes_;
 
  // entry[cd] contains the index of the first auxiliary node for the codim=cd
  // entity
  std::array<mesh::utils::CodimMeshDataSet<unsigned int>, 2> aux_node_offset_;
 
  template <class T>
  void WriteScalarPointData(const std::string& name,
                            const mesh::utils::MeshDataSet<T>& mds,
                            T undefined_value);
 
  template <int ROWS, class T>
  void WriteVectorPointData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::Matrix<T, ROWS, 1>>& mds,
      const Eigen::Matrix<T, ROWS, 1>& undefined_value);
 
  template <class T>
  void WriteScalarCellData(const std::string& name,
                           const mesh::utils::MeshDataSet<T>& mds,
                           T undefined_value);
 
  template <int ROWS, class T>
  void WriteVectorCellData(
      const std::string& name,
      const mesh::utils::MeshDataSet<Eigen::Matrix<T, ROWS, 1>>& mds,
      const Eigen::Matrix<T, ROWS, 1>& undefined_value);
 
  template <class T>
  void WriteFieldData(const std::string& name, std::vector<T> data);
 
  template <class DATA>
  void CheckAttributeSetName(const DATA& data, const std::string& name);
 
  template <int ROWS, class T>
  static void PadWithZeros(Eigen::Matrix<T, 3, 1>& out,
                           const Eigen::Matrix<T, ROWS, 1>& in);
};
 
template <class MESH_FUNCTION, class>
void VtkWriter::WritePointData(const std::string& name,
                               const MESH_FUNCTION& mesh_function) {
  //  we have to evaluate the mesh function at all points
  // and write this into the vtk file:
  CheckAttributeSetName(vtk_file_.point_data, name);
  using T = mesh::utils::MeshFunctionReturnType<MESH_FUNCTION>;
  auto dim_mesh = mesh_->DimMesh();
 
  Eigen::Matrix<double, 0, 1> origin;
 
  if constexpr (std::is_same_v<T, unsigned char> || std::is_same_v<T, char> ||
                std::is_same_v<T, unsigned> || std::is_same_v<T, int> ||
                std::is_same_v<T, float> ||
                std::is_same_v<T, double>) {  // NOLINT
    // MeshFunction is scalar valued:
    VtkFile::ScalarData<T> data{};
    data.data.resize(vtk_file_.unstructured_grid.points.size());
    data.name = name;
 
    // evaluate at nodes:
    for (const auto* n : mesh_->Entities(dim_mesh)) {
      data.data[mesh_->Index(*n)] = mesh_function(*n, origin)[0];
    }
 
    for (int codim = static_cast<int>(dim_mesh - 1);
         codim >= static_cast<char>(codim_); --codim) {
      for (const auto* e : mesh_->Entities(codim)) {
        auto ref_el = e->RefEl();
        if (order_ == 1 || (order_ == 2 && ref_el == base::RefEl::kTria())) {
          continue;
        }
        auto values = mesh_function(*e, aux_nodes_[ref_el.Id()]);
        auto offset = aux_node_offset_[codim](*e);
        for (typename std::vector<T>::size_type i = 0; i < values.size(); ++i) {
          data.data[offset + i] = values[i];
        }
      }
    }
    vtk_file_.point_data.push_back(std::move(data));
  } else if constexpr (base::is_eigen_matrix<T>) {  // NOLINT
    static_assert(T::ColsAtCompileTime == 1,
                  "The MeshFunction must return row-vectors");
    static_assert(
        T::RowsAtCompileTime == Eigen::Dynamic ||
            (T::RowsAtCompileTime > 1 && T::RowsAtCompileTime < 4),
        "The Row vectors returned by the MeshFunction must be either dynamic "
        "or must have 2 or 3 rows (at compile time).");
    using Scalar = typename T::Scalar;
    static_assert(
        std::is_same_v<double, Scalar> || std::is_same_v<float, Scalar>,
        "The RowVectors returned by the MeshFunction must be either double "
        "or float valued.");
    VtkFile::VectorData<Scalar> data{};
    data.data.resize(vtk_file_.unstructured_grid.points.size());
    data.name = name;
 
    // evaluate at nodes:
    for (const auto* n : mesh_->Entities(dim_mesh)) {
      PadWithZeros<T::RowsAtCompileTime, Scalar>(data.data[mesh_->Index(*n)],
                                                 mesh_function(*n, origin)[0]);
    }
 
    for (int codim = static_cast<int>(dim_mesh - 1);
         codim >= static_cast<char>(codim_); --codim) {
      for (const auto* e : mesh_->Entities(codim)) {
        auto ref_el = e->RefEl();
        if (order_ < 2 || (order_ < 3 && ref_el == base::RefEl::kTria())) {
          continue;
        }
        auto values = mesh_function(*e, aux_nodes_[ref_el.Id()]);
        auto offset = aux_node_offset_[codim](*e);
        for (unsigned long i = 0; i < values.size(); ++i) {
          PadWithZeros<T::RowsAtCompileTime, Scalar>(data.data[offset + i],
                                                     values[i]);
        }
      }
    }
    vtk_file_.point_data.push_back(std::move(data));
  } else {
    LF_VERIFY_MSG(false,
                  "MeshFunction values must be one of: unsigned char, char, "
                  "unsigned, int, float, double, Eigen::Vector<double, ...> "
                  "or Eigen::Vector<float, ...>");
  }
}
 
template <class MESH_FUNCTION, class>
void VtkWriter::WriteCellData(const std::string& name,
                              const MESH_FUNCTION& mesh_function) {
  CheckAttributeSetName(vtk_file_.cell_data, name);
  using T = mesh::utils::MeshFunctionReturnType<MESH_FUNCTION>;
 
  // maps from RefEl::Id() -> barycenter of the reference element
  std::vector<Eigen::VectorXd> barycenters(5);
  barycenters[base::RefEl::kPoint().Id()] = Eigen::Matrix<double, 0, 1>();
  barycenters[base::RefEl::kSegment().Id()] =
      base::RefEl::kSegment().NodeCoords().rowwise().sum() / 2.;
  barycenters[base::RefEl::kTria().Id()] =
      base::RefEl::kTria().NodeCoords().rowwise().sum() / 3.;
  barycenters[base::RefEl::kQuad().Id()] =
      base::RefEl::kTria().NodeCoords().rowwise().sum() / 4.;
 
  if constexpr (base::is_eigen_array<T> || base::is_eigen_matrix<T>) {
    static_assert(T::ColsAtCompileTime == 1,
                  "The MeshFunction must return row-vectors");
    static_assert(
        T::RowsAtCompileTime == Eigen::Dynamic ||
            (T::RowsAtCompileTime > 1 && T::RowsAtCompileTime < 4),
        "The Row vectors returned by the MeshFunction must be either dynamic "
        "or must have 2 or 3 rows (at compile time).");
    using Scalar = typename T::Scalar;
    static_assert(
        std::is_same_v<double, Scalar> || std::is_same_v<float, Scalar>,
        "The RowVectors returned by the MeshFunction must be either double "
        "or float valued.");
    VtkFile::VectorData<Scalar> data{};
    data.data.resize(mesh_->NumEntities(codim_));
    data.name = name;
    constexpr int rows = T::RowsAtCompileTime;
    for (const auto* p : mesh_->Entities(codim_)) {
      this->PadWithZeros<rows, Scalar>(
          data.data[mesh_->Index(*p)],
          mesh_function(*p, barycenters[p->RefEl().Id()])[0]);
    }
    vtk_file_.cell_data.push_back(std::move(data));
  } else {
    VtkFile::ScalarData<T> data{};
    data.data.resize(mesh_->NumEntities(codim_));
    data.name = name;
    for (const auto* e : mesh_->Entities(codim_)) {
      data.data[mesh_->Index(*e)] =
          mesh_function(*e, barycenters[e->RefEl().Id()])[0];
    }
    vtk_file_.cell_data.push_back(std::move(data));
  }
}
 
template <class DATA>
void VtkWriter::CheckAttributeSetName(const DATA& data,
                                      const std::string& name) {
  if (std::find_if(data.begin(), data.end(), [&](auto& d) {
        return boost::apply_visitor([&](auto&& d2) { return d2.name; }, d) ==
               name;
      }) != data.end()) {
    throw base::LfException(
        "There is already another Point/Cell Attribute Set with the "
        "name " +
        name);
  }
  if (name.find(' ') != std::string::npos) {
    throw base::LfException(
        "The name of the attribute set cannot contain spaces!");
  }
}
 
template <int ROWS, class T>
void VtkWriter::PadWithZeros(Eigen::Matrix<T, 3, 1>& out,
                             const Eigen::Matrix<T, ROWS, 1>& in) {
  if constexpr (ROWS == 2) {  // NOLINT
    out.template block<2, 1>(0, 0) = in;
    out(2) = T(0);
  } else if constexpr (ROWS == 3) {  // NOLINT
    out = in;
  } else if constexpr (ROWS == Eigen::Dynamic) {  // NOLINT
    LF_ASSERT_MSG(in.rows() == 2 || in.rows() == 3,
                  "VtkWriter can only write out 2d or 3d vectors.");
    if (in.rows() == 2) {
      out(2) = T(0);
      out.topRows(2) = in;
    } else {
      out = in;
    }
  }
}
 
}  // namespace lf::io
 
#endif  // __3e48c7b32a034cb3be3dbca884ff4f6c