Add ghost Morton IDs

include/sctl/tree.hpp

@@ -383,6 +383,92 @@ template <class Real, Integer DIM> class Tree {
           }
         }
       }
+      { // Add place-holder for ghost nodes
+        Long start_idx, end_idx;
+        { // Set start_idx, end_idx
+          start_idx = std::lower_bound(node_mid.begin(), node_mid.end(), mins[rank]) - node_mid.begin();
+          end_idx = std::lower_bound(node_mid.begin(), node_mid.end(), (rank+1==np ? Morton<DIM>().Next() : mins[rank+1])) - node_mid.begin();
+        }
+        { // Set user_node_lst
+          Vector<SortPair<Long,Morton<DIM>>> user_node_lst;
+          Vector<Morton<DIM>> nlst;
+          std::set<Long> user_procs;
+          for (Long i = start_idx; i < end_idx; i++) {
+            Morton<DIM> m0 = node_mid[i];
+            Integer d0 = m0.Depth();
+            m0.NbrList(nlst, std::max<Integer>(d0-2,0), periodic);
+            user_procs.clear();
+            for (const auto& m : nlst) {
+              Morton<DIM> m_start = m.DFD();
+              Morton<DIM> m_end = m.Next();
+              Integer p_start = std::lower_bound(mins.begin(), mins.end(), m_start) - mins.begin() - 1;
+              Integer p_end   = std::lower_bound(mins.begin(), mins.end(), m_end  ) - mins.begin();
+              SCTL_ASSERT(0 <= p_start);
+              SCTL_ASSERT(p_start < p_end);
+              SCTL_ASSERT(p_end <= np);
+              for (Long p = p_start; p < p_end; p++) {
+                if (p != rank) user_procs.insert(p);
+              }
+            }
+            for (const auto p : user_procs) {
+              SortPair<Long,Morton<DIM>> pair;
+              pair.key = p;
+              pair.data = m0;
+              user_node_lst.PushBack(pair);
+            }
+          }
+          omp_par::merge_sort(user_node_lst.begin(), user_node_lst.end());
+
+          user_cnt.ReInit(np);
+          user_mid.ReInit(user_node_lst.Dim());
+          for (Integer i = 0; i < np; i++) {
+            SortPair<Long,Morton<DIM>> pair_start, pair_end;
+            pair_start.key = i;
+            pair_end.key = i+1;
+            Long cnt_start = std::lower_bound(user_node_lst.begin(), user_node_lst.end(), pair_start) - user_node_lst.begin();
+            Long cnt_end   = std::lower_bound(user_node_lst.begin(), user_node_lst.end(), pair_end  ) - user_node_lst.begin();
+            user_cnt[i] = cnt_end - cnt_start;
+            for (Long j = cnt_start; j < cnt_end; j++) {
+              user_mid[j] = user_node_lst[j].data;
+            }
+            std::sort(user_mid.begin() + cnt_start, user_mid.begin() + cnt_end);
+          }
+        }
+
+        Vector<Morton<DIM>> ghost_mid;
+        { // SendRecv user_node_lst
+          const Vector<Long> &send_cnt = user_cnt;
+          Vector<Long> send_dsp(np);
+          scan(send_dsp, send_cnt);
+
+          Vector<Long> recv_cnt(np), recv_dsp(np);
+          comm.Alltoall(send_cnt.begin(), 1, recv_cnt.begin(), 1);
+          scan(recv_dsp, recv_cnt);
+
+          const Vector<Morton<DIM>>& send_mid = user_mid;
+          Long Nsend = send_dsp[np-1] + send_cnt[np-1];
+          Long Nrecv = recv_dsp[np-1] + recv_cnt[np-1];
+          SCTL_ASSERT(send_mid.Dim() == Nsend);
+
+          ghost_mid.ReInit(Nrecv);
+          comm.Alltoallv(send_mid.begin(), send_cnt.begin(), send_dsp.begin(), ghost_mid.begin(), recv_cnt.begin(), recv_dsp.begin());
+        }
+
+        { // Update node_mid <-- ghost_mid + node_mid
+          Vector<Morton<DIM>> new_mid(end_idx-start_idx + ghost_mid.Dim());
+          Long Nsplit = std::lower_bound(ghost_mid.begin(), ghost_mid.end(), mins[rank]) - ghost_mid.begin();
+          for (Long i = 0; i < Nsplit; i++) {
+            new_mid[i] = ghost_mid[i];
+          }
+          for (Long i = 0; i < end_idx - start_idx; i++) {
+            new_mid[Nsplit + i] = node_mid[start_idx + i];
+          }
+          for (Long i = Nsplit; i < ghost_mid.Dim(); i++) {
+            new_mid[end_idx - start_idx + i] = ghost_mid[i];
+          }
+          node_mid.Swap(new_mid);
+        }
+      }
       { // Set node_mid, node_attr
         Morton<DIM> m0 = (rank      ? mins[rank]   : Morton<DIM>()       );
         Morton<DIM> m1 = (rank+1<np ? mins[rank+1] : Morton<DIM>().Next());
@@ -393,10 +479,18 @@ template <class Real, Integer DIM> class Tree {
           node_attr[i].Ghost = (node_mid[i] < m0 || node_mid[i] >= m1);
         }
       }
-
-      { // Add place-holder for ghost nodes
-        // TODO
+      { // Check tree
+        Morton<DIM> m0;
+        SCTL_ASSERT(node_mid.Dim() && m0 == node_mid[0]);
+        for (Long i = 1; i < node_mid.Dim(); i++) {
+          const auto& m = node_mid[i];
+          if (m0.isAncestor(m)) m0 = m0.Ancestor(m0.Depth()+1);
+          else m0 = m0.Next();
+          SCTL_ASSERT(m0 == m);
+        }
+        SCTL_ASSERT(m0.Next() == Morton<DIM>().Next());
       }
+
       { // Update node_data, node_cnt
         Long start_idx, end_idx;
         { // Set start_idx, end_idx
@@ -570,11 +664,11 @@ template <class Real, Integer DIM> class Tree {
       omp_par::scan(cnt.begin(), dsp.begin(), cnt.Dim());
     }
 
-    //template <typename A, typename B> struct SortPair {
-    //  int operator<(const SortPair<A, B> &p1) const { return key < p1.key; }
-    //  A key;
-    //  B data;
-    //};
+    template <typename A, typename B> struct SortPair {
+      int operator<(const SortPair<A, B> &p1) const { return key < p1.key; }
+      A key;
+      B data;
+    };
 
   private:
 
@@ -585,6 +679,9 @@ template <class Real, Integer DIM> class Tree {
     std::map<std::string, Vector<Real>> node_data;
     std::map<std::string, Vector<Long>> node_cnt;
 
+    Vector<Morton<DIM>> user_mid;
+    Vector<Long> user_cnt;
+
     Comm comm;
 };
 

include/tree.f90

@@ -52,7 +52,7 @@ interface
     type(c_ptr), intent(in) :: tree_ctx
   end subroutine
 
-  !> Add some data to the tree nodes
+  !> Add some data to the tree nodes (for example multipole expansions)
   !! @param data_name   name for the data
   !! @param node_data   data array
   !! @param Ndata       length of the data array

src/test-cpp.cpp

@@ -7,82 +7,6 @@ int main(int argc, char** argv) {
   constexpr int DIM = 3;
   using Real = double;
 
-  //{
-  //  sctl::Comm comm = sctl::Comm::World();
-  //  srand48(comm.Rank());
-
-  //  sctl::Vector<Real> coord;
-  //  { // Set coord
-  //    long N_total = 2e5;
-  //    int np = comm.Size();
-  //    int myrank = comm.Rank();
-  //    long start= myrank    * N_total / np;
-  //    long end  =(myrank+1) * N_total / np;
-  //    long N_local = end - start;
-  //    coord.ReInit(N_local * DIM);
-  //    for(long i=0;i<N_local;i++){
-  //      Real phi = 2 * M_PI * drand48();
-  //      Real theta = M_PI * drand48();
-  //      coord[i*DIM+0] = 0.5 + 0.1125 * sin(theta) * cos(phi);
-  //      coord[i*DIM+1] = 0.5 + 0.1125 * sin(theta) * sin(phi);
-  //      coord[i*DIM+2] = 0.5 + 0.45 * cos(theta);
-  //    }
-  //  }
-
-  //  sctl::ParticleTreeWrapper<Real,DIM> tree;
-
-  //  { // UpdateRefinement
-  //    tree.UpdateRefinement(&coord[0], coord.Dim()/DIM, 100);
-  //  }
-
-  //  int32_t Nnodes;
-  //  { // GetTree
-  //    const Real* node_coord;
-  //    const int8_t *node_depth, *node_ghost, *node_leaf;
-  //    tree.GetTree(node_coord, node_depth, node_ghost, node_leaf, Nnodes);
-  //  }
-  //  { // AddData
-  //    sctl::Vector<Real> data(Nnodes*2);
-  //    sctl::Vector<int32_t> cnt(Nnodes);
-  //    cnt = 1;
-  //    tree.AddData("test", &data[0], data.Dim(), &cnt[0], cnt.Dim());
-  //  }
-  //  { // GetData
-  //    Real* data;
-  //    const int32_t* cnt;
-  //    int32_t Ndata, Ncnt;
-  //    tree.GetData(data, Ndata, cnt, Ncnt, "test");
-  //    SCTL_ASSERT(Ncnt == Nnodes);
-  //    SCTL_ASSERT(Ndata == Nnodes*2);
-  //  }
-  //  { // DeleteData
-  //  }
-  //  { // WriteTreeVTK
-  //    tree.WriteTreeVTK("tree", false);
-  //  }
-
-  //  { // AddParticles
-  //    tree.AddParticles("src_pts", &coord[0], coord.Dim()/DIM);
-  //  }
-  //  { // AddParticleData
-  //    tree.AddParticleData("src_coord", "src_pts", &coord[0], coord.Dim());
-  //  }
-  //  { // GetParticleData
-  //    Real* data;
-  //    int32_t Ndata;
-  //    tree.GetParticleData(data, Ndata, "src_coord");
-  //    SCTL_ASSERT(coord.Dim() == Ndata);
-  //    for (long i = 0; i < coord.Dim(); i++) {
-  //      SCTL_ASSERT(data[i] == coord[i]);
-  //    }
-  //  }
-  //  { // WriteParticleVTK
-  //    tree.WriteParticleVTK("pts", "src_coord", false);
-  //  }
-  //  { // DeleteParticleData
-  //    tree.DeleteParticleData("src_pts");
-  //  }
-  //}
   if (1) {
     sctl::Comm comm = sctl::Comm::World();
     srand48(comm.Rank());