VTKWriter.hpp

#ifndef VTK_WRITER_HPP
#define VTK_WRITER_HPP

#include <string>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <TinyVector.hpp>
#include <IConnectivity.hpp>

#include <ItemValue.hpp>
#include <OutputNamedItemValueSet.hpp>

class VTKWriter
{
 private:
  template <typename DataType> struct VTKType {};

  template <> struct VTKType<int8_t>
  {
    inline const static std::string name{"Int8"};
  };

  template <> struct VTKType<uint8_t>
  {
    inline const static std::string name{"UInt8"};
  };

  template <> struct VTKType<int16_t>
  {
    inline const static std::string name{"Int16"};
  };

  template <> struct VTKType<uint16_t>
  {
    inline const static std::string name{"UInt16"};
  };

  template <> struct VTKType<int32_t>
  {
    inline const static std::string name{"Int32"};
  };

  template <> struct VTKType<uint32_t>
  {
    inline const static std::string name{"UInt32"};
  };

  template <> struct VTKType<int64_t>
  {
    inline const static std::string name{"Int64"};
  };

  template <> struct VTKType<uint64_t>
  {
    inline const static std::string name{"UInt64"};
  };

  template <> struct VTKType<float>
  {
    inline const static std::string name{"Float32"};
  };

  template <> struct VTKType<double>
  {
    inline const static std::string name{"Float64"};
  };

  const std::string m_base_filename;
  unsigned int m_file_number;
  double m_last_time;
  const double m_time_period;


  template <typename DataType>
  void _write_node_value(std::ofstream& os,
                   const std::string& name,
                   const NodeValue<const DataType>& item_value)
  {
    os << "<DataArray type=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\">\n";
    for (NodeId i=0; i<item_value.size(); ++i) {
      os << item_value[i] << ' ';
    }
    os << "\n</DataArray>\n";
  }

  template <size_t N,
            typename DataType>
  void _write_node_value(std::ofstream& os,
                         const std::string& name,
                         const NodeValue<const TinyVector<N, DataType>>& item_value)
  {
    os << "<DataArray type=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N <<  "\">\n";
    for (NodeId i=0; i<item_value.size(); ++i) {
      for (size_t j=0; j<N; ++j) {
        os << item_value[i][j] << ' ';
      }
    }
    os << "\n</DataArray>\n";
  }

  template <typename DataType>
  void _write_node_value(std::ofstream& os,
                         const std::string& name,
                         const CellValue<const DataType>& item_value) {}


  template <typename DataType>
  void _write_cell_value(std::ofstream& os,
                   const std::string& name,
                   const NodeValue<const DataType>& item_value) {}

  template <typename DataType>
  void _write_cell_value(std::ofstream& os,
                         const std::string& name,
                         const CellValue<const DataType>& item_value)
  {
    os << "<DataArray type=\"Float64\" Name=\"" << name << "\">\n";
    for (CellId i=0; i<item_value.size(); ++i) {
      os << item_value[i] << ' ';
    }
    os << "\n</DataArray>\n";
  }

  template <size_t N,
            typename DataType>
  void _write_cell_value(std::ofstream& os,
                         const std::string& name,
                         const CellValue<const TinyVector<N, DataType>>& item_value)
  {
    os << "<DataArray type=\"Float64\" Name=\"" << name << "\" NumberOfComponents=\"" << N <<  "\">\n";
    for (CellId i=0; i<item_value.size(); ++i) {
      for (size_t j=0; j<N; ++j) {
        os << item_value[i][j] << ' ';
      }
    }
    os << "\n</DataArray>\n";
  }

 public:
  template <typename MeshType>
  void write(const MeshType& mesh,
             const OutputNamedItemValueSet& output_named_item_value_set,
             const double& time,
             const bool& forced_output = false)
  {
    if (time == m_last_time) return; // output already performed
    if ((time - m_last_time >= m_time_period) or forced_output) {
      m_last_time = time;
    } else {
      return;
    }
    std::ostringstream sout;
    sout << m_base_filename << '.' << std::setfill('0') << std::setw(4) << m_file_number << ".vtu" << std::ends;
    std::ofstream fout(sout.str());
    fout << "<?xml version=\"1.0\"?>\n";
    fout << "<VTKFile type=\"UnstructuredGrid\">\n";
    fout << "<UnstructuredGrid>\n";
    fout << "<Piece NumberOfPoints=\""<< mesh.numberOfNodes()
         << "\" NumberOfCells=\"" << mesh.numberOfCells() << "\">\n";
    fout << "<CellData>\n";
    for(const auto& [name, item_value_variant] : output_named_item_value_set) {
      std::visit([&, name=name](auto&& item_value) {
                   return this->_write_cell_value(fout, name, item_value);
                 }, item_value_variant);
    }
    fout << "</CellData>\n";
    fout << "<PointData>\n";
    for(const auto& [name, item_value_variant] : output_named_item_value_set) {
      std::visit([&, name=name](auto&& item_value) {
                   return this->_write_node_value(fout, name, item_value);
                 }, item_value_variant);
    }
    fout << "</PointData>\n";
    fout << "<Points>\n";
    fout << "<DataArray Name=\"Positions\" NumberOfComponents=\"3\" type=\"Float64\" format=\"ascii\">\n";
    using Rd = TinyVector<MeshType::dimension>;
    const NodeValue<const Rd>& xr = mesh.xr();
    if constexpr(MeshType::dimension == 1) {
      for (NodeId r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<1; ++i) {
          fout << xr[r][i] << ' ';
        }
        fout << "0 0 "; // VTK requires 3 components
      }
    } else if constexpr (MeshType::dimension == 2) {
      for (NodeId r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<2; ++i) {
          fout << xr[r][i] << ' ';
        }
        fout << "0 "; // VTK requires 3 components
      }
    } else {
      for (NodeId r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<3; ++i) {
          fout << xr[r][i] << ' ';
        }
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";
    fout << "</Points>\n";

    fout << "<Cells>\n";

    fout << "<DataArray type=\"Int32\" Name=\"connectivity\" NumberOfComponents=\"1\" format=\"ascii\">\n";

    const auto& cell_to_node_matrix
        = mesh.connectivity().cellToNodeMatrix();

    for (CellId j=0; j<mesh.numberOfCells(); ++j) {
      const auto& cell_nodes = cell_to_node_matrix[j];
      for (unsigned short r=0; r<cell_nodes.size(); ++r) {
        fout << cell_nodes[r] << ' ';
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "<DataArray type=\"UInt32\" Name=\"offsets\" NumberOfComponents=\"1\" format=\"ascii\">\n";
    {
      unsigned int offset=0;
      for (CellId j=0; j<mesh.numberOfCells(); ++j) {
        const auto& cell_nodes = cell_to_node_matrix[j];
        offset += cell_nodes.size();
        fout << offset << ' ';
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "<DataArray type=\"Int8\" Name=\"types\" NumberOfComponents=\"1\" format=\"ascii\">\n";
    for (CellId j=0; j<mesh.numberOfCells(); ++j) {
      const auto& cell_nodes = cell_to_node_matrix[j];
      switch (cell_nodes.size()) {
        case 2: {
          fout << "3 ";
          break;
        }
        case 3: {
          fout << "5 ";
          break;
        }
        case 4: {
          if (mesh.meshDimension() == 3) {
            fout << "10 ";
          } else {
            fout << "9 ";
          }
          break;
        }
        case 8: {
          fout << "12 ";
          break;
        }
        default: {
          fout << "7 ";
          break;
        }
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "</Cells>\n";
    fout << "</Piece>\n";
    fout << "</UnstructuredGrid>\n";
    fout << "</VTKFile>\n";

    m_file_number++;
  }
  VTKWriter(const std::string& base_filename,
            const double time_period)
      : m_base_filename(base_filename),
        m_file_number  (0),
        m_last_time    (-std::numeric_limits<double>::max()),
        m_time_period  (time_period)
  {}

  ~VTKWriter() = default;
};

#endif // VTK_WRITER_HPP