VTKWriter.cpp

#include <output/VTKWriter.hpp>

#include <mesh/Connectivity.hpp>
#include <mesh/Mesh.hpp>
#include <utils/Messenger.hpp>
#include <utils/RevisionInfo.hpp>

#include <ctime>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <unordered_map>

std::string
VTKWriter::_getFilenamePVTU() const
{
  std::ostringstream sout;
  sout << m_base_filename;
  sout << '.' << std::setfill('0') << std::setw(4) << m_saved_times.size() << ".pvtu";
  return sout.str();
}

std::string
VTKWriter::_getDateAndVersionComment() const
{
  std::ostringstream os;

  std::time_t now = std::time(nullptr);
  os << "<!--\n";
  os << "  Generated by pugs: " << std::ctime(&now);
  os << "  version: " << RevisionInfo::version() << '\n';
  os << "  tag:  " << RevisionInfo::gitTag() << '\n';
  os << "  HEAD: " << RevisionInfo::gitHead() << '\n';
  os << "  hash: " << RevisionInfo::gitHash() << " (" << ((RevisionInfo::gitIsClean()) ? "clean" : "dirty") << ")\n";
  os << "-->\n";
  return os.str();
}

std::string
VTKWriter::_getFilenameVTU(int rank_number) const
{
  std::ostringstream sout;
  sout << m_base_filename;
  if (parallel::size() > 1) {
    sout << '-' << std::setfill('0') << std::setw(4) << rank_number;
  }
  sout << '.' << std::setfill('0') << std::setw(4) << m_saved_times.size() << ".vtu";
  return sout.str();
}

template <typename DataType>
void
VTKWriter::_write_node_pvtu(std::ofstream& os, const std::string& name, const NodeValue<const DataType>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\"/>\n";
}

template <size_t N, typename DataType>
void
VTKWriter::_write_node_pvtu(std::ofstream& os,
                            const std::string& name,
                            const NodeValue<const TinyVector<N, DataType>>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N
     << "\"/>\n";
}

template <size_t N, typename DataType>
void
VTKWriter::_write_node_pvtu(std::ofstream& os,
                            const std::string& name,
                            const NodeValue<const TinyMatrix<N, DataType>>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N * N
     << "\"/>\n";
}

template <typename DataType>
void
VTKWriter::_write_node_pvtu(std::ofstream&, const std::string&, const CellValue<const DataType>&) const
{}

template <typename DataType>
void
VTKWriter::_write_cell_pvtu(std::ofstream& os, const std::string& name, const CellValue<const DataType>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\"/>\n";
}

template <size_t N, typename DataType>
void
VTKWriter::_write_cell_pvtu(std::ofstream& os,
                            const std::string& name,
                            const CellValue<const TinyVector<N, DataType>>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N
     << "\"/>\n";
}

template <size_t N, typename DataType>
void
VTKWriter::_write_cell_pvtu(std::ofstream& os,
                            const std::string& name,
                            const CellValue<const TinyMatrix<N, DataType>>&) const
{
  os << "<PDataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N * N
     << "\"/>\n";
}

template <typename DataType>
void
VTKWriter::_write_cell_pvtu(std::ofstream&, const std::string&, const NodeValue<const DataType>&) const
{}

template <typename DataType>
struct VTKWriter::VTKType
{
  inline const static std::string name = [] {
    static_assert(std::is_arithmetic_v<DataType>, "invalid data type");

    if constexpr (std::is_integral_v<DataType>) {
      if constexpr (std::is_unsigned_v<DataType>) {
        return "UInt" + std::to_string(sizeof(DataType) * 8);
      } else {
        return "UInt" + std::to_string(sizeof(DataType) * 8);
      }
    } else if constexpr (std::is_floating_point_v<DataType>) {
      return "Float" + std::to_string(sizeof(DataType) * 8);
    }
  }();
};

template <typename DataType>
void
VTKWriter::_write_array(std::ofstream& os, const std::string& name, const Array<DataType>& item_value) const
{
  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
VTKWriter::_write_array(std::ofstream& os,
                        const std::string& name,
                        const Array<TinyVector<N, DataType>>& item_value) const
{
  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
VTKWriter::_write_node_value(std::ofstream& os,
                             const std::string& name,
                             const NodeValue<const DataType>& item_value) const
{
  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
VTKWriter::_write_node_value(std::ofstream& os,
                             const std::string& name,
                             const NodeValue<const TinyVector<N, DataType>>& item_value) const
{
  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 <size_t N, typename DataType>
void
VTKWriter::_write_node_value(std::ofstream& os,
                             const std::string& name,
                             const NodeValue<const TinyMatrix<N, DataType>>& item_value) const
{
  os << "<DataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N * N
     << "\">\n";
  for (NodeId i = 0; i < item_value.size(); ++i) {
    for (size_t j = 0; j < N; ++j) {
      for (size_t k = 0; k < N; ++k) {
        // The following '+' enforces integer output for char types
        os << +item_value[i](j, k) << ' ';
      }
    }
  }
  os << "\n</DataArray>\n";
}

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

template <typename DataType>
void
VTKWriter::_write_cell_value(std::ofstream& os,
                             const std::string& name,
                             const CellValue<const DataType>& item_value) const
{
  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
VTKWriter::_write_cell_value(std::ofstream& os,
                             const std::string& name,
                             const CellValue<const TinyVector<N, DataType>>& item_value) const
{
  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 <size_t N, typename DataType>
void
VTKWriter::_write_cell_value(std::ofstream& os,
                             const std::string& name,
                             const CellValue<const TinyMatrix<N, DataType>>& item_value) const
{
  os << "<DataArray type=\"" << VTKType<DataType>::name << "\" Name=\"" << name << "\" NumberOfComponents=\"" << N * N
     << "\">\n";
  for (CellId i = 0; i < item_value.size(); ++i) {
    for (size_t j = 0; j < N; ++j) {
      for (size_t k = 0; k < N; ++k) {
        // The following '+' enforces integer output for char types
        os << +item_value[i](j, k) << ' ';
      }
    }
  }
  os << "\n</DataArray>\n";
}

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

template <typename MeshType>
void
VTKWriter::_write(const std::shared_ptr<const MeshType>& mesh,
                  const OutputNamedItemValueSet& given_output_named_item_value_set,
                  double time) const
{
  OutputNamedItemValueSet output_named_item_value_set{given_output_named_item_value_set};
  // Adding basic mesh information
  output_named_item_value_set.add(NamedItemValue{"cell_number", mesh->connectivity().cellNumber()});
  output_named_item_value_set.add(NamedItemValue{"node_number", mesh->connectivity().nodeNumber()});

  if (parallel::rank() == 0) {   // write PVTK file
    std::ofstream fout(_getFilenamePVTU());
    fout << "<?xml version=\"1.0\"?>\n";
    fout << _getDateAndVersionComment();
    fout << "<VTKFile type=\"PUnstructuredGrid\">\n";
    fout << "<PUnstructuredGrid GhostLevel=\"0\">\n";

    fout << "<PPoints>\n";
    fout << "<PDataArray Name=\"Positions\" NumberOfComponents=\"3\" "
            "type=\"Float64\"/>\n";
    fout << "</PPoints>\n";

    fout << "<PCells>\n";
    fout << "<PDataArray type=\"Int32\" Name=\"connectivity\" "
            "NumberOfComponents=\"1\"/>\n";
    fout << "<PDataArray type=\"UInt32\" Name=\"offsets\" "
            "NumberOfComponents=\"1\"/>\n";
    fout << "<PDataArray type=\"Int8\" Name=\"types\" "
            "NumberOfComponents=\"1\"/>\n";
    for (const auto& [name, item_value_variant] : output_named_item_value_set) {
      std::visit([&, name = name](auto&& item_value) { return this->_write_cell_pvtu(fout, name, item_value); },
                 item_value_variant);
    }
    if constexpr (MeshType::Dimension == 3) {
      fout << "<PDataArray type=\"Int64\" IdType=\"1\" Name=\"faces\"/>\n";
      fout << "<PDataArray type=\"Int64\" IdType=\"1\" Name=\"faceoffsets\"/>\n";
    }
    fout << "</PCells>\n";

    fout << "<PPointData>\n";
    for (const auto& [name, item_value_variant] : output_named_item_value_set) {
      std::visit([&, name = name](auto&& item_value) { return this->_write_node_pvtu(fout, name, item_value); },
                 item_value_variant);
    }
    fout << "</PPointData>\n";

    fout << "<PCellData>\n";
    for (const auto& [name, item_value_variant] : output_named_item_value_set) {
      std::visit([&, name = name](auto&& item_value) { return this->_write_cell_pvtu(fout, name, item_value); },
                 item_value_variant);
    }
    fout << "</PCellData>\n";

    for (size_t i_rank = 0; i_rank < parallel::size(); ++i_rank) {
      fout << "<Piece Source=\"" << _getFilenameVTU(i_rank) << "\"/>\n";
    }
    fout << "</PUnstructuredGrid>\n";
    fout << "</VTKFile>\n";
  }

  {   // write VTK files
    std::ofstream fout(_getFilenameVTU(parallel::rank()));
    fout << "<?xml version=\"1.0\"?>\n";
    fout << _getDateAndVersionComment();
    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());
      parallel_for(
        mesh->numberOfNodes(), PUGS_LAMBDA(NodeId 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();

      _write_array(fout, "connectivity", cell_to_node_matrix.entries());
    }

    {
      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);
    }

    {
      Array<int8_t> types(mesh->numberOfCells());
      const auto& cell_type           = mesh->connectivity().cellType();
      const auto& cell_to_node_matrix = mesh->connectivity().cellToNodeMatrix();

      parallel_for(
        mesh->numberOfCells(), PUGS_LAMBDA(CellId j) {
          switch (cell_type[j]) {
          case CellType::Line: {
            types[j] = 3;
            break;
          }
          case CellType::Triangle: {
            types[j] = 5;
            break;
          }
          case CellType::Quadrangle: {
            types[j] = 9;
            break;
          }
          case CellType::Tetrahedron: {
            types[j] = 10;
            break;
          }
          case CellType::Pyramid: {
            if (cell_to_node_matrix[j].size() == 5) {
              types[j] = 14;   // quadrangle basis
            } else {
              types[j] = 41;   // polygonal basis
            }
            break;
          }
          case CellType::Prism: {
            types[j] = 13;
            break;
          }
          case CellType::Hexahedron: {
            types[j] = 12;
            break;
          }
          case CellType::Diamond: {
            types[j] = 41;
            break;
          }
          default: {
            std::ostringstream os;
            os << __FILE__ << ':' << __LINE__ << ": unknown cell type";
            throw UnexpectedError(os.str());
          }
          }
        });
      _write_array(fout, "types", types);
      if constexpr (MeshType::Dimension == 3) {
        const bool has_general_polyhedron = [&] {
          for (size_t i = 0; i < types.size(); ++i) {
            if (types[i] == 41)
              return true;
          }
          return false;
        }();
        if (has_general_polyhedron) {
          const auto& cell_to_face_matrix   = mesh->connectivity().cellToFaceMatrix();
          const auto& face_to_node_matrix   = mesh->connectivity().faceToNodeMatrix();
          const auto& cell_face_is_reversed = mesh->connectivity().cellFaceIsReversed();

          Array<size_t> faces_offsets(mesh->numberOfCells());
          size_t next_offset = 0;
          fout << "<DataArray type=\"Int64\" IdType=\"1\" Name=\"faces\">\n";
          for (CellId cell_id = 0; cell_id < mesh->numberOfCells(); ++cell_id) {
            const auto& cell_nodes = cell_to_node_matrix[cell_id];
            std::unordered_map<size_t, size_t> node_id_to_local_number_in_cell;
            for (size_t i_cell_node = 0; i_cell_node < cell_nodes.size(); ++i_cell_node) {
              node_id_to_local_number_in_cell[cell_nodes[i_cell_node]] = i_cell_node;
            }

            const auto& cell_faces = cell_to_face_matrix[cell_id];
            fout << cell_faces.size() << '\n';
            next_offset++;
            for (size_t i_cell_face = 0; i_cell_face < cell_faces.size(); ++i_cell_face) {
              const FaceId& face_id  = cell_faces[i_cell_face];
              const auto& face_nodes = face_to_node_matrix[face_id];
              fout << face_nodes.size();
              next_offset++;
              Array<size_t> face_node_in_cell(face_nodes.size());

              for (size_t i_face_node = 0; i_face_node < face_nodes.size(); ++i_face_node) {
                const NodeId& node_id                  = face_nodes[i_face_node];
                auto i_node_id_to_local_number_in_cell = node_id_to_local_number_in_cell.find(node_id);
                Assert(i_node_id_to_local_number_in_cell != node_id_to_local_number_in_cell.end());
                face_node_in_cell[i_face_node] = i_node_id_to_local_number_in_cell->second;
              }

              if (cell_face_is_reversed(cell_id, i_cell_face)) {
                for (size_t i = 0; i < face_node_in_cell.size(); ++i) {
                  fout << ' ' << face_node_in_cell[face_node_in_cell.size() - 1 - i];
                }
              } else {
                for (size_t i = 0; i < face_node_in_cell.size(); ++i) {
                  fout << ' ' << face_node_in_cell[i];
                }
              }

              next_offset += face_nodes.size();

              fout << '\n';
            }
            faces_offsets[cell_id] = next_offset;
          }
          fout << "</DataArray>\n";
          fout << "<DataArray type=\"Int64\" IdType=\"1\" Name=\"faceoffsets\">\n";
          for (size_t i_face_offsets = 0; i_face_offsets < faces_offsets.size(); ++i_face_offsets) {
            fout << faces_offsets[i_face_offsets] << '\n';
          }
          fout << "</DataArray>\n";
        }
      }
    }

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

  if (parallel::rank() == 0) {   // write PVD file
    std::ofstream fout(m_base_filename + ".pvd");

    fout << "<?xml version=\"1.0\"?>\n";
    fout << "<VTKFile type=\"Collection\" version=\"0.1\">\n";

    fout << "<Collection>\n";
    for (size_t i_time = 0; i_time < m_saved_times.size(); ++i_time) {
      std::ostringstream sout;
      sout << m_base_filename;
      sout << '.' << std::setfill('0') << std::setw(4) << i_time << ".pvtu";

      fout << "<DataSet timestep=\"" << m_saved_times[i_time] << "\" "
           << "file=\"" << sout.str() << "\"/>\n";
    }
    fout << "<DataSet timestep=\"" << time << "\" group=\"\" part=\"0\" file=\"" << _getFilenamePVTU() << "\"/>\n";

    fout << "</Collection>\n";
    fout << "</VTKFile>\n";
  }
}

void
VTKWriter::writeMesh(const std::shared_ptr<const IMesh>& mesh) const
{
  OutputNamedItemValueSet output_named_item_value_set;

  switch (mesh->dimension()) {
  case 1: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<1>>>(mesh), output_named_item_value_set, 0);
    break;
  }
  case 2: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<2>>>(mesh), output_named_item_value_set, 0);
    break;
  }
  case 3: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<3>>>(mesh), output_named_item_value_set, 0);
    break;
  }
  default: {
    throw UnexpectedError("invalid mesh dimension");
  }
  }
}

void
VTKWriter::write(const std::vector<std::shared_ptr<const NamedDiscreteFunction>>& named_discrete_function_list,
                 double time) const
{
  std::shared_ptr mesh = this->_getMesh(named_discrete_function_list);

  OutputNamedItemValueSet output_named_item_value_set = this->_getOutputNamedItemValueSet(named_discrete_function_list);

  switch (mesh->dimension()) {
  case 1: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<1>>>(mesh), output_named_item_value_set, time);
    break;
  }
  case 2: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<2>>>(mesh), output_named_item_value_set, time);
    break;
  }
  case 3: {
    this->_write(std::dynamic_pointer_cast<const Mesh<Connectivity<3>>>(mesh), output_named_item_value_set, time);
    break;
  }
  default: {
    throw UnexpectedError("invalid mesh dimension");
  }
  }
}