basicwriter.hh

// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
 
#ifndef DUNE_GRID_IO_FILE_VTK_BASICWRITER_HH
#define DUNE_GRID_IO_FILE_VTK_BASICWRITER_HH
 
#include <fstream>
#include <iomanip>
#include <iterator>
#include <list>
#include <memory>
#include <sstream>
#include <string>
 
#include <dune/common/parallel/mpiguard.hh>
#include <dune/common/path.hh>
 
#include <dune/geometry/referenceelements.hh>
 
#include <dune/grid/io/file/vtk/common.hh>
#include <dune/grid/io/file/vtk/functionwriter.hh>
#include <dune/grid/io/file/vtk/pvtuwriter.hh>
#include <dune/grid/io/file/vtk/vtuwriter.hh>
 
namespace Dune
{
 
  namespace VTK {
 
    template<typename IteratorFactory>
    class BasicWriter {
      typedef typename IteratorFactory::CellIterator CellIterator;
      typedef typename IteratorFactory::CornerIterator CornerIterator;
      typedef typename IteratorFactory::PointIterator PointIterator;
 
      typedef typename IteratorFactory::Cell Cell;
 
    public:
      typedef FunctionWriterBase<Cell> FunctionWriter;
 
    private:
      typedef std::list<std::shared_ptr<FunctionWriter> > WriterList;
      typedef typename WriterList::const_iterator WIterator;
 
      typedef typename Cell::Geometry::ctype ctype;
      static const unsigned celldim = Cell::mydimension;
      typedef ReferenceElements<ctype, celldim> Refelems;
 
      static const FileType fileType = celldim == 1
                                       ? polyData : unstructuredGrid;
 
      const IteratorFactory& factory;
 
      WriterList cellData;
      WriterList pointData;
 
      CoordinatesWriter<Cell> coords;
      typename IteratorFactory::ConnectivityWriter connectivity;
      OffsetsWriter<Cell> offsets;
      TypesWriter<Cell> types;
 
    public:
      BasicWriter(const IteratorFactory& factory_)
        : factory(factory_), connectivity(factory.makeConnectivity())
      { }
 
      //
      //  Methods for adding data
      //
 
      void addCellData(const std::shared_ptr<FunctionWriter>& writer) {
        cellData.push_back(writer);
      }
 
      void addPointData(const std::shared_ptr<FunctionWriter>& writer) {
        pointData.push_back(writer);
      }
 
      void clear() {
        cellData.clear();
        pointData.clear();
      }
 
    protected:
      //
      //  Methods for writing single functions
      //
 
      void writeCellFunction(VTUWriter& vtuWriter,
                             FunctionWriter& functionWriter,
                             unsigned ncells) const
      {
        if(functionWriter.beginWrite(vtuWriter, ncells)) {
          const CellIterator& cellend = factory.endCells();
          for(CellIterator cellit = factory.beginCells(); cellit != cellend;
              ++cellit)
            functionWriter.write(*cellit, Refelems::general(cellit->type()).
                                 position(0,0));
        }
        functionWriter.endWrite();
      }
 
      void writePointFunction(VTUWriter& vtuWriter,
                              FunctionWriter& functionWriter,
                              unsigned npoints) const
      {
        if(functionWriter.beginWrite(vtuWriter, npoints)) {
          const PointIterator& pend = factory.endPoints();
          for(PointIterator pit = factory.beginPoints(); pit != pend; ++pit)
            functionWriter.write(pit->cell(), pit->duneIndex());
        }
        functionWriter.endWrite();
      }
 
      void writeCornerFunction(VTUWriter& vtuWriter,
                               FunctionWriter& functionWriter,
                               unsigned ncorners) const
      {
        if(functionWriter.beginWrite(vtuWriter, ncorners)) {
          const CornerIterator& cend = factory.endCorners();
          for(CornerIterator cit = factory.beginCorners(); cit != cend; ++cit)
            functionWriter.write(cit->cell(), cit->duneIndex());
        }
        functionWriter.endWrite();
      }
 
      //
      //  Methods for writing whole sections
      //
 
      static std::string getFirstScalar(const WriterList& data) {
        const WIterator& wend = data.end();
        for(WIterator wit = data.begin(); wit != wend; ++wit)
          if((*wit)->ncomps() == 1)
            return (*wit)->name();
        return "";
      }
 
      static std::string getFirstVector(const WriterList& data) {
        const WIterator& wend = data.end();
        for(WIterator wit = data.begin(); wit != wend; ++wit)
          if((*wit)->ncomps() == 3)
            return (*wit)->name();
        return "";
      }
 
      void writeCellData(VTUWriter& vtuWriter, unsigned ncells) const {
        if(cellData.empty()) return;
 
        vtuWriter.beginCellData(getFirstScalar(cellData),
                                getFirstVector(cellData));
        const WIterator& wend = cellData.end();
        for(WIterator wit = cellData.begin(); wit != wend; ++wit)
          writeCellFunction(vtuWriter, **wit, ncells);
        vtuWriter.endCellData();
      }
 
      void writePointData(VTUWriter& vtuWriter, unsigned npoints) const {
        if(pointData.empty()) return;
 
        vtuWriter.beginPointData(getFirstScalar(pointData),
                                 getFirstVector(pointData));
        const WIterator& wend = pointData.end();
        for(WIterator wit = pointData.begin(); wit != wend; ++wit)
          writePointFunction(vtuWriter, **wit, npoints);
        vtuWriter.endPointData();
      }
 
      void writeGrid(VTUWriter& vtuWriter, unsigned ncells, unsigned npoints,
                     unsigned ncorners) {
        vtuWriter.beginPoints();
        writePointFunction(vtuWriter, coords, npoints);
        vtuWriter.endPoints();
 
        vtuWriter.beginCells();
        writeCornerFunction(vtuWriter, connectivity, ncorners);
        writeCellFunction(vtuWriter, offsets, ncells);
        if(fileType != polyData)
          writeCellFunction(vtuWriter, types, ncells);
        vtuWriter.endCells();
      }
 
      void writeAll(VTUWriter& vtuWriter, unsigned ncells, unsigned npoints,
                    unsigned ncorners) {
        writeCellData(vtuWriter, ncells);
        writePointData(vtuWriter, npoints);
        writeGrid(vtuWriter, ncells, npoints, ncorners);
      }
 
    public:
      void writePiece(const std::string& filename, OutputType outputType) {
        std::ofstream stream;
        stream.exceptions(std::ios_base::badbit | std::ios_base::failbit |
                          std::ios_base::eofbit);
        stream.open(filename.c_str(), std::ios::binary);
 
        VTUWriter vtuWriter(stream, outputType, fileType);
 
        unsigned ncells = std::distance(factory.beginCells(),
                                        factory.endCells());
        unsigned npoints = std::distance(factory.beginPoints(),
                                         factory.endPoints());
        unsigned ncorners = std::distance(factory.beginCorners(),
                                          factory.endCorners());
 
        vtuWriter.beginMain(ncells, npoints);
        writeAll(vtuWriter, ncells, npoints, ncorners);
        vtuWriter.endMain();
 
        if(vtuWriter.beginAppended())
          writeAll(vtuWriter, ncells, npoints, ncorners);
        vtuWriter.endAppended();
 
      }
 
 
      void writeCollection(const std::string name,
                           const std::string& piecename,
                           const std::string& piecepath)
      {
        std::ofstream stream;
        stream.exceptions(std::ios_base::badbit | std::ios_base::failbit |
                          std::ios_base::eofbit);
        stream.open(name.c_str(), std::ios::binary);
 
        PVTUWriter writer(stream, fileType);
 
        writer.beginMain();
 
        // PPointData
        writer.beginPointData(getFirstScalar(pointData),
                              getFirstVector(pointData));
        for(WIterator it=pointData.begin(); it!=pointData.end(); ++it)
          (*it)->addArray(writer);
        writer.endPointData();
 
        // PCellData
        writer.beginCellData(getFirstScalar(cellData),
                             getFirstVector(cellData));
        for(WIterator it=cellData.begin(); it!=cellData.end(); ++it)
          (*it)->addArray(writer);
        writer.endCellData();
 
        // PPoints
        writer.beginPoints();
        coords.addArray(writer);
        writer.endPoints();
 
        // Pieces
        for( int i = 0; i < factory.comm().size(); ++i )
          writer.addPiece(getParallelPieceName(piecename, piecepath, i));
 
        writer.endMain();
      }
 
      //
      //  Filename generators
      //
 
 
      std::string getParallelPieceName(const std::string& name,
                                       const std::string& path, int rank) const
      {
        std::ostringstream s;
        if(path.size() > 0) {
          s << path;
          if(path[path.size()-1] != '/')
            s << '/';
        }
        s << 's' << std::setw(4) << std::setfill('0') << factory.comm().size()
          << ':';
        s << 'p' << std::setw(4) << std::setfill('0') << rank << ':';
        s << name;
        switch(fileType) {
        case polyData :         s << ".vtp"; break;
        case unstructuredGrid : s << ".vtu"; break;
        }
        return s.str();
      }
 
 
      std::string getParallelHeaderName(const std::string& name,
                                        const std::string& path) const
      {
        std::ostringstream s;
        if(path.size() > 0) {
          s << path;
          if(path[path.size()-1] != '/')
            s << '/';
        }
        s << 's' << std::setw(4) << std::setfill('0') << factory.comm().size()
          << ':';
        s << name;
        switch(fileType) {
        case polyData :         s << ".pvtp"; break;
        case unstructuredGrid : s << ".pvtu"; break;
        }
        return s.str();
      }
 
 
      std::string getSerialPieceName(const std::string& name,
                                     const std::string& path) const
      {
        switch(fileType) {
        case polyData :         return concatPaths(path, name+".vtp");
        case unstructuredGrid : return concatPaths(path, name+".vtu");
        }
      }
 
      //
      //  User interface functions for writing
      //
 
 
      std::string pwrite(const std::string& name, const std::string& path,
                         const std::string& extendpath, OutputType outputType)
      {
        MPIGuard guard(factory.comm());
 
        // do some magic because paraview can only cope with relative paths to
        // piece files
        std::ofstream file;
        file.exceptions(std::ios_base::badbit | std::ios_base::failbit |
                        std::ios_base::eofbit);
        std::string piecepath = concatPaths(path, extendpath);
        std::string relpiecepath = relativePath(path, piecepath);
 
        // write this processes .vtu/.vtp piece file
        std::string fullname = getParallelPieceName(name, piecepath,
                                                    factory.comm().rank());
        writePiece(fullname, outputType);
 
        // if we are rank 0, write .pvtu/.pvtp parallel header
        fullname = getParallelHeaderName(name, path);
        if(factory.comm().rank() == 0)
          writeCollection(fullname, name, relpiecepath);
 
        guard.finalize();
 
        return fullname;
      }
 
      std::string write(const std::string &name, OutputType outputType)
      {
        // in the parallel case, just use pwrite, it has all the necessary
        // stuff, so we don't need to reimplement it here.
        if(factory.comm().size() > 1)
          return pwrite(name, "", "", outputType);
 
        // generate filename for process data
        std::string pieceName = getSerialPieceName(name, "");
 
        writePiece(pieceName, outputType);
 
        return pieceName;
      }
 
    };
 
  } // namespace VTK
 
 
} // namespace Dune
 
#endif // DUNE_GRID_IO_FILE_VTK_BASICWRITER_HH