export_vtk.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
 
// includes, FEAT
#include "kernel/shape.hpp"
#include "kernel/util/assertion.hpp"
#include <kernel/geometry/conformal_mesh.hpp>
#include <kernel/geometry/mesh_part.hpp>
#include <kernel/geometry/structured_mesh.hpp>
#include <kernel/util/dist.hpp>
#include <kernel/util/dist_file_io.hpp>
#include <kernel/util/exception.hpp>
 
// includes, STL
#include <deque>
#include <fstream>
#include <functional>
#include <vector>
 
namespace FEAT
{
  namespace Geometry
  {
    namespace Intern
    {
      template<typename Shape_>
      struct VTKShape;
 
      template<>
      struct VTKShape<Shape::Simplex<1>>
      {
        static constexpr int type = 3; // VTK_LINE
        static inline int map(int i)
        {
          return i;
        }
      };
 
      template<>
      struct VTKShape<Shape::Simplex<2>>
      {
        static constexpr int type = 5; // VTK_TRIANGLE
        static inline int map(int i)
        {
          return i;
        }
      };
 
      template<>
      struct VTKShape<Shape::Simplex<3>>
      {
        static constexpr int type = 10; // VTK_TETRA
        static inline int map(int i)
        {
          return i;
        }
      };
 
      template<>
      struct VTKShape<Shape::Hypercube<1>>
      {
        static constexpr int type = 3; // VTK_LINE
        static inline int map(int i)
        {
          return i;
        }
      };
 
      template<>
      struct VTKShape<Shape::Hypercube<2>>
      {
        static constexpr int type = 9; // VTK_QUAD
        static inline int map(int i)
        {
          // bit-stunt: {0, 1, 2, 3} -> {0, 1, 3, 2}
          return (i ^ ((i >> 1) & 1));
        }
      };
 
      template<>
      struct VTKShape<Shape::Hypercube<3>>
      {
        static constexpr int type = 12; // VTK_HEXAHEDRON
        static inline int map(int i)
        {
          // bit-stunt: {0, 1, ..., 7} -> {0, 1, 3, 2, 4, 5, 7, 6}
          return (i ^ ((i >> 1) & 1));
        }
      };
    } // namespace Intern
 
    template<typename Mesh_, int cell_dim_ = Mesh_::ShapeType::dimension>
    class ExportVTK
    {
    public:
      using MeshType = Mesh_;
      using ShapeType = typename MeshType::ShapeType;
      using CellShapeType = typename Shape::FaceTraits<ShapeType, cell_dim_>::ShapeType;
      using VTKShapeType = Intern::VTKShape<CellShapeType>;
 
    protected:
      using VarDeque = std::deque<std::pair<String, std::vector<double>>>;
 
      // NOTE(mmuegge): Depending on wether we are writing a mesh or a meshpart,
      // we may need to do some additional handling of indices.
      // The accessor functions allow us to do so,
      // while presenting a consistent interface to the actual file writing logic.
 
      std::function<double(Index, int)> _vertices;
      std::function<Index(Index, int)> _cells;
      Index _num_verts;
      Index _num_cells;
      VarDeque _vertex_scalars;
      VarDeque _vertex_vectors;
      VarDeque _cell_scalars;
      VarDeque _cell_vectors;
      int _var_prec;
 
    public:
      explicit ExportVTK(const MeshType& mesh, int var_prec = 0) :
        _num_verts(mesh.get_num_entities(0)),
        _num_cells(mesh.get_num_entities(cell_dim_)),
        _var_prec(Math::max(0, var_prec))
      {
        const auto& vertex_set = mesh.get_vertex_set();
        const auto& index_set = mesh.template get_index_set<cell_dim_, 0>();
 
        _vertices = [&](Index i, int j) { return vertex_set[i][j]; };
        _cells = [&](Index i, int j) { return index_set(i, j); };
      }
 
      explicit ExportVTK(const MeshType& mesh, const MeshPart<MeshType>& part, int var_prec = 0) :
        _num_verts(part.get_num_entities(0)),
        _num_cells(part.get_num_entities(cell_dim_)),
        _var_prec(Math::max(0, var_prec))
      {
        const auto& vertex_set = mesh.get_vertex_set();
        const auto& vertex_target_set = part.template get_target_set<0>();
        const auto& cell_target_set = part.template get_target_set<cell_dim_>();
 
        _vertices = [&](Index i, int j) { return vertex_set[vertex_target_set[i]][j]; };
 
        if(part.has_topology())
        {
          // The meshpart has its own topology. Use that.
          const auto& index_set = part.template get_index_set<cell_dim_, 0>();
          _cells = [&](Index i, int j) { return index_set(i, j); };
        }
        else
        {
          // The meshpart has no own topology. Use the mesh topology.
          const auto& index_set = mesh.template get_index_set<cell_dim_, 0>();
 
          // The mesh index-set returns indices of vertices the mesh's index-space,
          // but the .vtu will be written in the meshparts index-space.
          // We thus need to search for the corresponding index on the meshpart for any vertex index.
          _cells = [&](Index i, int j)
          {
            const Index vertex = index_set(cell_target_set[i], j);
            for(Index k(0); k < vertex_target_set.get_num_entities(); k++)
            {
              if(vertex_target_set[k] == vertex)
              {
                return k;
              }
            }
 
            XABORTM("Lookup of vertex index failed!");
          };
        }
      }
 
      virtual ~ExportVTK() = default;
 
      ExportVTK(const ExportVTK& other) = default;
 
      ExportVTK(ExportVTK&& other) noexcept = default;
 
      ExportVTK& operator=(const ExportVTK& other) = default;
 
      ExportVTK& operator=(ExportVTK&& other) = default;
 
      void clear()
      {
        _cell_scalars.clear();
        _vertex_scalars.clear();
        _vertex_vectors.clear();
        _cell_vectors.clear();
      }
 
      template<typename T_>
      void add_vertex_scalar(const String& name, const T_* data, double scaling_factor = 1.0)
      {
        XASSERTM(data != nullptr, "data array is nullptr");
        std::vector<double> d(_num_verts);
        for(Index i(0); i < _num_verts; ++i)
        {
          d[i] = scaling_factor * double(data[i]);
        }
        _vertex_scalars.emplace_back(name, std::move(d));
      }
 
      template<typename T_>
      void add_vertex_vector(
        const String& name,
        const T_* x,
        const T_* y = nullptr,
        const T_* z = nullptr,
        double scaling_factor = 1.0)
      {
        XASSERTM(x != nullptr, "x-data array is nullptr");
        std::vector<double> d(3 * _num_verts);
 
        if(y != nullptr && z != nullptr)
        {
          for(Index i(0); i < _num_verts; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = scaling_factor * double(y[i]);
            d[(3 * i) + 2] = scaling_factor * double(z[i]);
          }
        }
        else if(y != nullptr)
        {
          for(Index i(0); i < _num_verts; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = scaling_factor * double(y[i]);
            d[(3 * i) + 2] = 0.0;
          }
        }
        else
        {
          for(Index i(0); i < _num_verts; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = 0.0;
            d[(3 * i) + 2] = 0.0;
          }
        }
        _vertex_vectors.emplace_back(name, std::move(d));
      }
 
      template<typename VectorType_>
      void add_vertex_vector(const String& name, const VectorType_& v, double scaling_factor = 1.0)
      {
        std::vector<double> d(3 * _num_verts);
 
        for(Index i(0); i < _num_verts; ++i)
        {
          for(Index j(0); j < 3; ++j)
          {
            d[(Index(3) * i) + j] = double(0);
          }
 
          for(int j(0); j < v.BlockSize; ++j)
          {
            d[(Index(3) * i) + Index(j)] = scaling_factor * double(v(i)[j]);
          }
        }
        _vertex_vectors.emplace_back(name, std::move(d));
      }
 
      template<typename T_>
      void add_cell_scalar(const String& name, const T_* data, double scaling_factor = 1.0)
      {
        XASSERTM(data != nullptr, "data array is nullptr");
        std::vector<double> d(_num_cells);
        for(Index i(0); i < _num_cells; ++i)
        {
          d[i] = scaling_factor * double(data[i]);
        }
        _cell_scalars.emplace_back(name, std::move(d));
      }
 
      template<typename T_>
      void add_cell_vector(
        const String& name,
        const T_* x,
        const T_* y = nullptr,
        const T_* z = nullptr,
        double scaling_factor = 1.0)
      {
        XASSERTM(x != nullptr, "x-data array is nullptr");
        std::vector<double> d(3 * _num_cells);
 
        if(y != nullptr && z != nullptr)
        {
          for(Index i(0); i < _num_cells; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = scaling_factor * double(y[i]);
            d[(3 * i) + 2] = scaling_factor * double(z[i]);
          }
        }
        else if(y != nullptr)
        {
          for(Index i(0); i < _num_cells; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = scaling_factor * double(y[i]);
            d[(3 * i) + 2] = 0.0;
          }
        }
        else
        {
          for(Index i(0); i < _num_cells; ++i)
          {
            d[(3 * i) + 0] = scaling_factor * double(x[i]);
            d[(3 * i) + 1] = 0.0;
            d[(3 * i) + 2] = 0.0;
          }
        }
        _cell_vectors.emplace_back(name, std::move(d));
      }
 
      template<typename VectorType_>
      void add_cell_vector(const String& name, const VectorType_& v, double scaling_factor = 1.0)
      {
        std::vector<double> d(3 * _num_cells);
 
        for(Index i(0); i < _num_cells; ++i)
        {
          for(Index j(0); j < 3; ++j)
          {
            d[(Index(3) * i) + j] = double(0);
          }
 
          for(int j(0); j < v.BlockSize; ++j)
          {
            d[(Index(3) * i) + Index(j)] = scaling_factor * double(v(i)[j]);
          }
        }
        _cell_vectors.emplace_back(name, std::move(d));
      }
 
      void write(const String& filename) const
      {
        // try to open the output file
        String vtu_name(filename + ".vtu");
        std::ofstream ofs(vtu_name.c_str());
        if(!(ofs.is_open() && ofs.good()))
        {
          throw FileError("Failed to create '" + vtu_name + "'");
        }
 
        // write
        write_vtu(ofs);
 
        // and close
        ofs.close();
      }
 
      void write(const String& filename, const int rank, const int nparts)
      {
        // call the standard serial write version if nparts < 1
        if(nparts <= 1)
        {
          write(filename);
          return;
        }
 
        // verify rank
        XASSERTM((rank >= 0) && (rank < nparts), "Invalid rank");
 
        // Add rank cell array since we're parallel if we come to here
        std::vector<double> rank_array((std::size_t(_num_cells)), double(rank));
        add_cell_scalar("rank", rank_array.data());
 
        // compute number of non-zero digits in (nparts-1) for padding
        const std::size_t ndigits = Math::ilog10(std::size_t(nparts - 1));
 
        // write serial VTU file: "filename.#rank.vtu"
        write(filename + "." + stringify(rank).pad_front(ndigits, '0'));
 
        // we're done unless we have rank = 0
        if(rank != 0)
        {
          return;
        }
 
        // try to open our output file
        String pvtu_name(filename + ".pvtu");
        std::ofstream ofs(pvtu_name.c_str());
        if(!(ofs.is_open() && ofs.good()))
        {
          throw FileError("Failed to create '" + pvtu_name + "'");
        }
 
        // extract the file title from our filename
        std::size_t p = filename.find_last_of("\\/");
        String file_title = filename.substr(p == FEAT::String::npos ? 0 : ++p);
 
        // write PVTU file
        write_pvtu(ofs, file_title, nparts);
 
        // and close
        ofs.close();
      }
 
      void write(const String& filename, const Dist::Comm& comm)
      {
        // Add rank cell array since we're parallel if we come to here
        std::vector<double> rank_array(std::size_t(_num_cells), double(comm.rank()));
        add_cell_scalar("rank", rank_array.data());
 
        // compute number of non-zero digits in (nparts-1) for padding
        const auto ndigits = std::size_t(Math::max(Math::ilog10(comm.size() - 1), 1));
 
        // write serial VTU file into a stringstream
        std::stringstream stream;
        write_vtu(stream);
 
        // generate pattern for filename: "filename.#rank.vtu"
        String pattern = filename + "." + String(ndigits, '*') + ".vtu";
 
        // write distributed VTU files
        DistFileIO::write_sequence(stream, pattern, comm);
 
        // we're done unless we have rank = 0
        if(comm.rank() != 0)
        {
          return;
        }
 
        // try to open our output file
        String pvtu_name(filename + ".pvtu");
        std::ofstream ofs(pvtu_name.c_str());
        if(!(ofs.is_open() && ofs.good()))
        {
          throw FileError("Failed to create '" + pvtu_name + "'");
        }
 
        // extract the file title from our filename
        std::size_t p = filename.find_last_of("\\/");
        String file_title = filename.substr(p == FEAT::String::npos ? 0 : ++p);
 
        // write PVTU file
        write_pvtu(ofs, file_title, comm.size());
 
        // and close
        ofs.close();
      }
 
      void write_vtu(std::ostream& os) const
      {
        // fetch basic information
        const int num_coords = MeshType::world_dim;
        const int verts_per_cell = Shape::FaceTraits<CellShapeType, 0>::count;
 
        // write VTK header
        os << "<VTKFile type=\"UnstructuredGrid\" version=\"0.1\">\n)";
        os << "<!-- Generated by FEAT v" << version_major << "." << version_minor;
        os << "." << version_patch << " -->\n";
 
        // write mesh header
        os << "<UnstructuredGrid>\n";
        os << "<Piece NumberOfPoints=\"" << _num_verts << "\" NumberOfCells=\"" << _num_cells << "\">\n";
 
        // write point data
        if((!_vertex_scalars.empty()) || (!_vertex_vectors.empty()))
        {
          os << "<PointData>\n";
 
          // write vertex variables
          for(const auto& var : _vertex_scalars)
          {
            os << "<DataArray type=\"Float64\" Name=\"" << var.first << "\" Format=\"ascii\">\n";
            for(Index j(0); j < _num_verts; ++j)
            {
              os << stringify_fp_sci(var.second[j], _var_prec) << "\n";
            }
            os << "</DataArray>\n";
          }
          // write vertex fields
          for(const auto& var : _vertex_vectors)
          {
            os << "<DataArray type=\"Float64\" Name=\"" << var.first;
            os << "\" NumberOfComponents=\"3\" Format=\"ascii\">\n";
            for(Index j(0); j < _num_verts; ++j)
            {
              os << stringify_fp_sci(var.second[(3 * j) + 0], _var_prec) << " ";
              os << stringify_fp_sci(var.second[(3 * j) + 1], _var_prec) << " ";
              os << stringify_fp_sci(var.second[(3 * j) + 2], _var_prec) << "\n";
            }
            os << "</DataArray>\n";
          }
 
          os << "</PointData>\n";
        }
 
        // write cell data
        if(!_cell_scalars.empty() || !_cell_vectors.empty())
        {
          os << "<CellData>\n";
          if(!_cell_scalars.empty())
          {
            for(const auto& var : _cell_scalars)
            {
              os << "<DataArray type=\"Float64\" Name=\"" << var.first << "\" Format=\"ascii\">\n";
              for(Index j(0); j < _num_cells; ++j)
              {
                os << stringify_fp_sci(var.second[j], _var_prec) << "\n";
              }
              os << "</DataArray>\n";
            }
          }
 
          if(!_cell_vectors.empty())
          {
            // write cell fields
            for(const auto& var : _cell_vectors)
            {
              os << "<DataArray type=\"Float64\" Name=\"" << var.first;
              os << "\" NumberOfComponents=\"3\" Format=\"ascii\">\n";
              for(Index j(0); j < _num_cells; ++j)
              {
                os << stringify_fp_sci(var.second[(3 * j) + 0], _var_prec) << " ";
                os << stringify_fp_sci(var.second[(3 * j) + 1], _var_prec) << " ";
                os << stringify_fp_sci(var.second[(3 * j) + 2], _var_prec) << "\n";
              }
              os << "</DataArray>\n";
            }
          }
          os << "</CellData>\n";
        }
 
        // write vertices
        os << "<Points>\n";
        os << "<DataArray type=\"Float32\" NumberOfComponents=\"3\" Format=\"ascii\">\n";
        for(Index i(0); i < _num_verts; ++i)
        {
          os << _vertices(i, 0);
          for(int j(1); j < num_coords; ++j)
          {
            os << " " << _vertices(i, j);
          }
          for(int j(num_coords); j < 3; ++j)
          {
            os << " 0";
          }
          os << "\n";
        }
        os << "</DataArray>\n";
        os << "</Points>\n";
 
        // write cells
        os << "<Cells>\n";
        os << "<DataArray type=\"UInt32\" Name=\"connectivity\">\n";
        for(Index i(0); i < _num_cells; ++i)
        {
          os << _cells(i, VTKShapeType::map(0));
          for(int j(1); j < verts_per_cell; ++j)
          {
            os << " " << _cells(i, VTKShapeType::map(j));
          }
          os << "\n";
        }
        os << "</DataArray>\n";
        os << "<DataArray type=\"UInt32\" Name=\"offsets\">\n";
        for(Index i(0); i < _num_cells; ++i)
        {
          os << ((i + 1) * verts_per_cell) << "\n";
        }
        os << "</DataArray>\n";
        os << "<DataArray type=\"UInt32\" Name=\"types\">\n";
        for(Index i(0); i < _num_cells; ++i)
        {
          os << VTKShapeType::type << "\n";
        }
        os << "</DataArray>\n";
        os << "</Cells>\n";
 
        // finish
        os << "</Piece>\n";
        os << "</UnstructuredGrid>\n";
        os << "</VTKFile>\n";
      }
 
      void write_pvtu(std::ostream& os, const String& file_title, const int nparts) const
      {
        // write VTK header
        os << "<VTKFile type=\"PUnstructuredGrid\" version=\"0.1\">\n";
        os << "<!-- Generated by FEAT v" << version_major << "." << version_minor;
        os << "." << version_patch << " -->\n";
        os << "<PUnstructuredGrid GhostLevel=\"0\">\n";
 
        // write vertex data
        if((!_vertex_scalars.empty()) || (!_vertex_vectors.empty()))
        {
          os << "<PPointData>\n";
 
          // write vertex variables
          for(const auto& vertex_scalar : _vertex_scalars)
          {
            os << "<PDataArray type=\"Float64\" Name=\"" << vertex_scalar.first << "\" />\n";
          }
          // write vertex fields
          for(const auto& vertex_vector : _vertex_vectors)
          {
            os << "<PDataArray type=\"Float64\" Name=\"" << vertex_vector.first;
            os << "\" NumberOfComponents=\"3\" />\n";
          }
 
          os << "</PPointData>\n";
        }
 
        // write cell variables
        if(!_cell_scalars.empty() || !_cell_vectors.empty())
        {
          os << "<PCellData>\n";
          // write cell scalars
          for(const auto& cell_scalar : _cell_scalars)
          {
            os << "<PDataArray type=\"Float64\" Name=\"" << cell_scalar.first << "\" />\n";
          }
          // write cell fields
          for(const auto& cell_vector : _cell_vectors)
          {
            os << "<PDataArray type=\"Float64\" Name=\"" << cell_vector.first;
            os << "\" NumberOfComponents=\"3\" />\n";
          }
          os << "</PCellData>\n";
        }
 
        // write vertices
        os << "<PPoints>\n";
        os << "<PDataArray type=\"Float32\" NumberOfComponents=\"3\" />\n";
        os << "</PPoints>\n";
 
        // compute number of non-zero digits in (nparts-1) for padding
        const std::size_t ndigits = Math::ilog10(std::size_t(nparts - 1));
 
        // now let's write our piece data
        for(int i(0); i < nparts; ++i)
        {
          os << "<Piece Source=\"" << file_title << "." << stringify(i).pad_front(ndigits, '0');
          os << ".vtu\" />\n";
        }
 
        // finish
        os << "</PUnstructuredGrid>\n";
        os << "</VTKFile>\n";
      }
    }; // class ExportVTK
  } // namespace Geometry
} // namespace FEAT