VtuPlot.cpp

#include "VtuPlot.hpp"
#include "Config.hpp"

#include <fstream>


using namespace std;

int vtkCellType(CELLTYPE cType) {
  switch (cType) {
    case INTERVAL:
      return 3;
    case TRIANGLE:
      return 5;
    case QUADRILATERAL:
      return 9;
    case TETRAHEDRON:
      return 10;
    case HEXAHEDRON:
      return 12;
    default:
      return -1;
  }
}

string VtuPlot::plotIndent(int i) {
  string indent;
  for (int j = 0; j < i; j++)
    indent += standardIndent;
  return indent;
}

string VtuPlot::plotLineBreak(int i, int limit, int indent) {
  return (i % limit != 0 ? " " : ("\n" + plotIndent(indent)));
}

void VtuPlot::createVtu(const string &filename, const string& filecontent) {
  string filepath = dataPath + filename + string(".vtu");
  std::ofstream out(filepath.c_str());
  out << FileString() << std::endl;
}

string VtuPlot::FileString() {
  std::stringstream out;

  // == Header ==
  out << R"(<VTKFile type ="UnstructuredGrid" version="0.1" byte_order="LittleEndian">)" << "\n"
      << plotIndent(1) << "<UnstructuredGrid>" << "\n"
      << plotIndent(2) << "<Piece NumberOfPoints=\"" << pointN
      << "\" NumberOfCells=\"" << cellN << "\">" << "\n";

  // == Points ==
  out << plotIndent(3) << "<Points>" << "\n"
      << plotIndent(4) << R"(<DataArray type="Float32" NumberOfComponents="3" format="ascii">)";
  for (int i = 0; i < pointN; i++) {
    out << plotLineBreak(i, lineBreakP, 5);
    plotPoint(out, points.at(i) + displacement.at(i));
  }
  out << "\n"
      << plotIndent(4) << "</DataArray>" << "\n"
      << plotIndent(3) << "</Points>" << "\n";

  // == PointData ==
  out << plotIndent(3) << "<PointData";
  if (!activeArrays.first.empty()) {
    string aArray = " ";
    aArray += (pointDataSize.at(activeArrays.first) > 1 ? R"(Vectors=")" : R"(Scalars=")");
    aArray += activeArrays.first;
    aArray += R"(")";
    out << aArray;
  }
  out << ">" << "\n";
  for (auto pData : pointData) {
    if (pData.first == "ID") continue;

    int dataSize = pointDataSize.at(pData.first);
    out << plotIndent(4)
        << R"(<DataArray type="Float32" Name=")" << pData.first
        << R"(" NumberOfComponents=")" << dataSize
        << R"(" format="ascii">)";
    for (int i = 0; i < pointN; i++) {
      for (int j = 0; j < dataSize; j++)
        out << plotLineBreak(i + j, lineBreakZ, 5)
            << to_string(pData.second->at(points[i]).GetData(j));
    }
    out << "\n" << plotIndent(4) << "</DataArray>" << "\n";
  }
  out << plotIndent(3) << "</PointData>" << "\n";

  // == Cells ==
  out << plotIndent(3) << "<Cells>" << "\n"
      << plotIndent(4) << R"(<DataArray type="Int64" Name="connectivity" format="ascii">)"
      << cellConnectivity << "\n"
      << plotIndent(4) << "</DataArray>" << "\n";
  out << plotIndent(4) << R"(<DataArray type="Int64" Name="offsets" format="ascii">)"
      << cellOffset << "\n"
      << plotIndent(4) << "</DataArray>" << "\n";
  out << plotIndent(4) << R"(<DataArray type="Int32" Name="types" format="ascii">)"
      << cellTypes << "\n"
      << plotIndent(4) << "</DataArray>" << "\n"
      << plotIndent(3) << "</Cells>" << "\n";

  // == CellData ==
  out << plotIndent(3) << "<CellData";
  if (!activeArrays.second.empty()) {
    string aArray = " ";
    aArray += (cellDataSize.at(activeArrays.second) > 1 ? R"(Vectors=")" : R"(Scalars=")");
    aArray += activeArrays.second;
    aArray += R"(")";
    out << aArray;
  }
  out << ">" << "\n";
  for (const auto &cData : cellData) {
    if (cData.first == "ID") continue;

    out << plotIndent(4)
        << R"(<DataArray type="Float32" Name=")" << cData.first
        << R"(" NumberOfComponents=")" << cellDataSize.at(cData.first)
        << R"(" format="ascii">)";

    int dataSize = cellDataSize.at(cData.first);
    for (int i = 0; i < cellN; i++) {
      for (int j = 0; j < dataSize; j++)
        out << plotLineBreak(i + j, lineBreakZ, 5)
            << to_string(cData.second->at(cells[i]).GetData(j));
    }
    out << "\n" << plotIndent(4) << "</DataArray>" << "\n";
  }
  out << plotIndent(3) << "</CellData>" << "\n";
  // == End of File ==
  out << plotIndent(2) << "</Piece>" << "\n"
      << plotIndent(1) << "</UnstructuredGrid>" << "\n"
      << "</VTKFile>";

  return out.str();
}

void VtuPlot::Erase() {
  ClearData();
  ClearMesh();
}

void VtuPlot::ClearData() {
  vector<string> dataToDrop;
  for (const auto &pData : pointData) {
    if (pData.first == "ID") continue;
    dataToDrop.push_back(pData.first);
  }
  for (const auto &dataName : dataToDrop) {
    pointData.erase(dataName);
    pointDataSize.erase(dataName);
  }

  dataToDrop.clear();
  for (const auto &cData : cellData) {
    if (std::string{"IDSubdomainProcLoad"}.find(cData.first) != std::string::npos) continue;
    dataToDrop.push_back(cData.first);
  }
  for (const auto &dataName : dataToDrop) {
    cellData.erase(dataName);
    cellDataSize.erase(dataName);
  }

  activeArrays = {"", ""};
  for (auto & i : displacement) {
    i.second = Origin;
  }
}

void VtuPlot::ClearMesh(){
  displacement.clear();
  pointData.erase("ID");
  cellData.erase("ID");
  cellData.erase("Subdomain");
  cellData.erase("ProcLoad");

  plotMesh.Reset();
}

void VtuPlot::plotPoint(std::ostream &os, const Point &p) {
  os << p[0];
  for (int i = 1; i < SpaceDimension; ++i)
    os << " " << p[i];
  for (int i = SpaceDimension; i < 3; ++i)
    os << " 0.0";
}

void VtuPlot::plotData(std::ostream &os, const Point &p) {
  os << p[0];
  for (int i = 1; i < SpaceDimension; ++i)
    os << " " << p[i];
  for (int i = SpaceDimension; i < 3; ++i)
    os << " 0.0";
}

void VtuPlot::Flush(PlotStatus status) {
  if(status==PlotStatus::FLUSH){
    ClearData();
  }
  if(status==PlotStatus::CLEAR){
    Erase();
  }
}


VtuPlot &operator<<(VtuPlot &plot, const Vector &v) {
  plot.AddData("value" + std::to_string(plot.DataCount()), v);
  return plot;
}

VtuPlot &operator<<(VtuPlot &plot, PlotStatus status) {
  plot.PlotFile();
  plot.Flush(status);
  return plot;
}

VtuPlot &operator<<(VtuPlot &plot, const string &name) {
  plot.FileName(name);
  return plot;
}

VtuPlot &operator<<(VtuPlot &plot, const std::pair<std::string, PlotStatus>& saveoptions) {
  return plot << saveoptions.first << saveoptions.second;
}