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:
  const std::string m_base_filename;
  unsigned int m_file_number;
  double m_last_time;
  const double m_time_period;

  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"};
  };

  template <typename DataType>
  void _write_array(std::ofstream& os,
                    const std::string& name,
                    const Array<DataType>& item_value)
  {
    os << "<DataArray type=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\">\n";
    for (typename Array<DataType>::index_type i=0;
         i<item_value.size(); ++i) {
      // The following '+' enforces integer output for char types
      os << +item_value[i] << ' ';
    }
    os << "\n</DataArray>\n";
  }

  template <size_t N,
            typename DataType>
  void _write_array(std::ofstream& os,
                    const std::string& name,
                    const Array<TinyVector<N, DataType>>& item_value)
  {
    os << "<DataArray type=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N << "\">\n";
    for (typename Array<DataType>::index_type i=0;
         i<item_value.size(); ++i) {
      for (size_t j=0; j<N; ++j) {
        // The following '+' enforces integer output for char types
        os << +item_value[i][j] << ' ';
      }
    }
    os << "\n</DataArray>\n";
  }

  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) {
      // The following '+' enforces integer output for char types
      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) {
      // The following '+' enforces integer output for char types
        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 CellValue<const DataType>& item_value)
  {
    os << "<DataArray type=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\">\n";
    for (CellId i=0; i<item_value.size(); ++i) {
      // The following '+' enforces integer output for char types
      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=\"" << VTKType<DataType>::name
       << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N <<  "\">\n";
    for (CellId i=0; i<item_value.size(); ++i) {
      for (size_t j=0; j<N; ++j) {
      // The following '+' enforces integer output for char types
        os << +item_value[i][j] << ' ';
      }
    }
    os << "\n</DataArray>\n";
  }

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

 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";
    {
      using Rd = TinyVector<MeshType::dimension>;
      const NodeValue<const Rd>& xr = mesh.xr();
      Array<TinyVector<3>> positions(mesh.numberOfNodes());
      for (NodeId r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<MeshType::dimension; ++i) {
          positions[r][i] = xr[r][i];
        }
        for (unsigned short i=MeshType::dimension; i<3; ++i) {
          positions[r][i] = 0;
        }
      }
      _write_array(fout, "Positions", positions);
    }
    fout << "</Points>\n";

    fout << "<Cells>\n";

    {
      const auto& cell_to_node_matrix
          = mesh.connectivity().cellToNodeMatrix();
      Array<int> connectivity(cell_to_node_matrix.numberOfEntries());

      for (size_t i=0; i<cell_to_node_matrix.numberOfEntries(); ++i) {
        connectivity[i] = cell_to_node_matrix.entries()[i];
      }
      _write_array(fout, "connectivity", connectivity);
    }

    {
      const auto& cell_to_node_matrix
          = mesh.connectivity().cellToNodeMatrix();
      Array<unsigned int> offsets(mesh.numberOfCells());
      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();
        offsets[j]=offset;
      }
      _write_array(fout, "offsets", offsets);
    }

    {
      const auto& cell_to_node_matrix
          = mesh.connectivity().cellToNodeMatrix();
      Array<int8_t> types(mesh.numberOfCells());
      for (CellId j=0; j<mesh.numberOfCells(); ++j) {
        const auto& cell_nodes = cell_to_node_matrix[j];
        switch (cell_nodes.size()) {
          case 2: {
            types[j]=3;
            break;
          }
          case 3: {
            types[j]=5;
            break;
          }
          case 4: {
            if (mesh.meshDimension() == 3) {
              types[j]=10;
            } else {
              types[j]=9;
            }
            break;
          }
          case 8: {
            types[j]=12;
            break;
          }
          default: {
            types[j]=7;
            break;
          }
        }
      }
      _write_array(fout, "types", types);
    }

    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