.

include/sctl/tree.hpp

@@ -4,6 +4,7 @@
 #include SCTL_INCLUDE(common.hpp)
 #include SCTL_INCLUDE(morton.hpp)
 #include SCTL_INCLUDE(comm.hpp)
+#include SCTL_INCLUDE(matrix.hpp) // TODO: fix issues when this is before #Include <comm.hpp>
 
 #include <fstream>
 #include <algorithm>
@@ -22,18 +23,14 @@ struct VTUData {
   Vector<int32_t> offset;
   Vector<uint8_t> types;
 
-  void WriteVTK(const std::string& fname, Comm comm = Comm::Self()) const {
+  void WriteVTK(const std::string& fname, const Comm& comm = Comm::Self()) const {
     typedef typename VTUData::VTKReal VTKReal;
-
-    Integer rank = comm.Rank();
-    Integer np = comm.Size();
-
     Long value_dof = 0;
     {  // Write vtu file.
       std::ofstream vtufile;
       {  // Open file for writing.
         std::stringstream vtufname;
-        vtufname << fname << std::setfill('0') << std::setw(6) << rank << ".vtu";
+        vtufname << fname << std::setfill('0') << std::setw(6) << comm.Rank() << ".vtu";
         vtufile.open(vtufname.str().c_str());
         if (vtufile.fail()) return;
       }
@@ -44,7 +41,7 @@ struct VTUData {
 
         Vector<int32_t> mpi_rank;
         {  // Set  mpi_rank
-          Integer new_myrank = rank;
+          Integer new_myrank = comm.Rank();
           mpi_rank.ReInit(pt_cnt);
           for (Long i = 0; i < mpi_rank.Dim(); i++) mpi_rank[i] = new_myrank;
         }
@@ -86,7 +83,7 @@ struct VTUData {
         vtufile << "        <DataArray type=\"Int32\" Name=\"offsets\" format=\"appended\" offset=\"" << data_size << "\" />\n";
         data_size += sizeof(uint32_t) + offset.Dim() * sizeof(int32_t);
         vtufile << "        <DataArray type=\"UInt8\" Name=\"types\" format=\"appended\" offset=\"" << data_size << "\" />\n";
-        data_size += sizeof(uint32_t) + types.Dim() * sizeof(uint8_t);
+        //data_size += sizeof(uint32_t) + types.Dim() * sizeof(uint8_t);
         vtufile << "      </Cells>\n";
         //---------------------------------------------------------------------------
         vtufile << "    </Piece>\n";
@@ -134,7 +131,7 @@ struct VTUData {
       }
       vtufile.close();  // close file
     }
-    if (!rank) {  // Write pvtu file
+    if (!comm.Rank()) {  // Write pvtu file
       std::ofstream pvtufile;
       {  // Open file for writing
         std::stringstream pvtufname;
@@ -166,7 +163,7 @@ struct VTUData {
           // char *fname_ = (char*)strrchr(vtupath.str().c_str(), '/') + 1;
           // std::string fname_ =
           // boost::filesystem::path(fname).filename().string().
-          for (Integer i = 0; i < np; i++) pvtufile << "      <Piece Source=\"" << fname_ << std::setfill('0') << std::setw(6) << i << ".vtu\"/>\n";
+          for (Integer i = 0; i < comm.Size(); i++) pvtufile << "      <Piece Source=\"" << fname_ << std::setfill('0') << std::setw(6) << i << ".vtu\"/>\n";
         }
         pvtufile << "  </PUnstructuredGrid>\n";
         pvtufile << "</VTKFile>\n";
@@ -174,6 +171,103 @@ struct VTUData {
       pvtufile.close();  // close file
     }
   };
+
+  template <class ElemLst> void AddElems(const ElemLst elem_lst, Integer order, const Comm& comm = Comm::Self()) {
+    constexpr Integer COORD_DIM = ElemLst::CoordDim();
+    constexpr Integer ElemDim = ElemLst::ElemDim();
+    using CoordBasis = typename ElemLst::CoordBasis;
+    using CoordType = typename ElemLst::CoordType;
+    Long N0 = coord.Dim() / COORD_DIM;
+    Long NElem = elem_lst.NElem();
+
+    Matrix<CoordType> nodes = VTK_Nodes<CoordType, ElemDim>(order);
+    Integer Nnodes = sctl::pow<ElemDim,Integer>(order);
+    SCTL_ASSERT(nodes.Dim(0) == ElemDim);
+    SCTL_ASSERT(nodes.Dim(1) == Nnodes);
+    { // Set coord
+      Matrix<CoordType> vtk_coord;
+      auto M = CoordBasis::SetupEval(nodes);
+      CoordBasis::Eval(vtk_coord, elem_lst.ElemVector(), M);
+      for (Long k = 0; k < NElem; k++) {
+        for (Integer i = 0; i < Nnodes; i++) {
+          constexpr Integer dim = (COORD_DIM < 3 ? COORD_DIM : 3);
+          for (Integer j = 0; j < dim; j++) {
+            coord.PushBack((VTUData::VTKReal)vtk_coord[k*COORD_DIM+j][i]);
+          }
+          for (Integer j = dim; j < 3; j++) {
+            coord.PushBack((VTUData::VTKReal)0);
+          }
+        }
+      }
+    }
+
+    if (ElemLst::ElemDim() == 2) {
+      for (Long k = 0; k < NElem; k++) {
+        for (Integer i = 0; i < order-1; i++) {
+          for (Integer j = 0; j < order-1; j++) {
+            Long idx = k*Nnodes + i*order + j;
+            connect.PushBack(N0+idx);
+            connect.PushBack(N0+idx+1);
+            connect.PushBack(N0+idx+order+1);
+            connect.PushBack(N0+idx+order);
+            offset.PushBack(connect.Dim());
+            types.PushBack(9);
+          }
+        }
+      }
+    } else {
+      // TODO
+      SCTL_ASSERT(false);
+    }
+  }
+  template <class ElemLst, class ValueBasis> void AddElems(const ElemLst elem_lst, const Vector<ValueBasis>& elem_value, Integer order, const Comm& comm = Comm::Self()) {
+    constexpr Integer ElemDim = ElemLst::ElemDim();
+    using ValueType = typename ValueBasis::ValueType;
+    Long NElem = elem_lst.NElem();
+
+    Integer dof = (NElem==0 ? 0 : elem_value.Dim() / NElem);
+    SCTL_ASSERT(elem_value.Dim() == NElem * dof);
+    AddElems(elem_lst, order, comm);
+
+    Matrix<ValueType> nodes = VTK_Nodes<ValueType, ElemDim>(order);
+    Integer Nnodes = sctl::pow<ElemDim,Integer>(order);
+    SCTL_ASSERT(nodes.Dim(0) == ElemDim);
+    SCTL_ASSERT(nodes.Dim(1) == Nnodes);
+
+    { // Set value
+      Matrix<ValueType> vtk_value;
+      auto M = ValueBasis::SetupEval(nodes);
+      ValueBasis::Eval(vtk_value, elem_value, M);
+      for (Long k = 0; k < NElem; k++) {
+        for (Integer i = 0; i < Nnodes; i++) {
+          for (Integer j = 0; j < dof; j++) {
+            value.PushBack((VTUData::VTKReal)vtk_value[k*dof+j][i]);
+          }
+        }
+      }
+    }
+  }
+
+  private:
+    template <class CoordType, Integer ELEM_DIM> static Matrix<CoordType> VTK_Nodes(Integer order) {
+      Matrix<CoordType> nodes;
+      if (ELEM_DIM == 2) {
+        Integer Nnodes = order*order;
+        nodes.ReInit(ELEM_DIM, Nnodes);
+        for (Integer i = 0; i < order; i++) {
+          for (Integer j = 0; j < order; j++) {
+            //nodes[0][i*order+j] = i / (CoordType)(order-1);
+            //nodes[1][i*order+j] = j / (CoordType)(order-1);
+            nodes[0][i*order+j] = 0.5 - 0.5 * sctl::cos<CoordType>((2*i+1) * const_pi<CoordType>() / (2*order));
+            nodes[1][i*order+j] = 0.5 - 0.5 * sctl::cos<CoordType>((2*j+1) * const_pi<CoordType>() / (2*order));
+          }
+        }
+      } else {
+        // TODO
+        SCTL_ASSERT(false);
+      }
+      return nodes;
+    }
 };
 
 template <Integer DIM> class Tree {