Add level-restriction, parent and child lists

include/sctl/morton.hpp

@@ -25,7 +25,7 @@ template <Integer DIM = 3> class Morton {
   typedef uint64_t UINT_T;
   #endif
 
-  static const Integer MAX_DEPTH = SCTL_MAX_DEPTH;
+  static constexpr Integer MAX_DEPTH = SCTL_MAX_DEPTH;
 
   Morton() {
     depth = 0;

include/sctl/tree.hpp

@@ -239,7 +239,7 @@ template <class Real, Integer DIM> class Tree {
       return comm;
     }
 
-    void UpdateRefinement(const Vector<Real>& coord, Long M = 1) {
+    void UpdateRefinement(const Vector<Real>& coord, Long M = 1, bool balance21 = 0, bool periodic = 0) {
       Integer np = comm.Size();
       Integer rank = comm.Rank();
 
@@ -311,7 +311,10 @@ template <class Real, Integer DIM> class Tree {
         while (m0 < mend) {
           Integer d = m0.Depth();
           Morton<DIM> m1 = (idx + M < pt_mid.Dim() ? pt_mid[idx+M] : Morton<DIM>().Next());
-          while (d < Morton<DIM>::MAX_DEPTH && m0.Ancestor(d) == m1.Ancestor(d)) d++;
+          while (d < Morton<DIM>::MAX_DEPTH && m0.Ancestor(d) == m1.Ancestor(d)) {
+            node_mid.PushBack(m0.Ancestor(d));
+            d++;
+          }
           m0 = m0.Ancestor(d);
           node_mid.PushBack(m0);
           m0 = m0.Next();
@@ -325,18 +328,73 @@ template <class Real, Integer DIM> class Tree {
         Morton<DIM> m0 = coarsest_ancestor_mid(node_mid[min_idx]);
         comm.Allgather(Ptr2ConstItr<Morton<DIM>>(&m0,1), 1, mins.begin(), 1);
       }
+      if (balance21) { // 2:1 balance refinement // TODO: optimize
+        Vector<Morton<DIM>> parent_mid;
+        { // add balancing Morton IDs
+          Vector<std::set<Morton<DIM>>> parent_mid_set(Morton<DIM>::MAX_DEPTH+1);
+          Vector<Morton<DIM>> nlst;
+          for (const auto& m0 : node_mid) {
+            Integer d0 = m0.Depth();
+            parent_mid_set[m0.Depth()].insert(m0.Ancestor(d0-1));
+          }
+          for (Integer d = Morton<DIM>::MAX_DEPTH; d > 0; d--) {
+            for (const auto& m : parent_mid_set[d]) {
+              m.NbrList(nlst, d-1, periodic);
+              parent_mid_set[d-1].insert(nlst.begin(), nlst.end());
+              parent_mid.PushBack(m);
+            }
+          }
+        }
+
+        Vector<Morton<DIM>> parent_mid_sorted;
+        { // sort and repartition
+          comm.HyperQuickSort(parent_mid, parent_mid_sorted);
+          comm.PartitionS(parent_mid_sorted, mins[comm.Rank()]);
+        }
+
+        Vector<Morton<DIM>> tmp_mid;
+        { // add children
+          Vector<Morton<DIM>> clst;
+          tmp_mid.PushBack(Morton<DIM>()); // include root node
+          for (Long i = 0; i < parent_mid_sorted.Dim(); i++) {
+            if (i+1 == parent_mid_sorted.Dim() || parent_mid_sorted[i] != parent_mid_sorted[i+1]) {
+              const auto& m = parent_mid_sorted[i];
+              tmp_mid.PushBack(m);
+              m.Children(clst);
+              for (const auto& c : clst) tmp_mid.PushBack(c);
+            }
+          }
+          auto insert_ancestor_children = [](Vector<Morton<DIM>>& mvec, const Morton<DIM>& m0) {
+            Integer d0 = m0.Depth();
+            Vector<Morton<DIM>> clst;
+            for (Integer d = 0; d < d0; d++) {
+              m0.Ancestor(d).Children(clst);
+              for (const auto& m : clst) mvec.PushBack(m);
+            }
+          };
+          insert_ancestor_children(tmp_mid, mins[rank]);
+          omp_par::merge_sort(tmp_mid.begin(), tmp_mid.end());
+        }
+
+        node_mid.ReInit(0);
+        for (Long i = 0; i < tmp_mid.Dim(); i++) { // remove duplicates
+          if (i+1 == tmp_mid.Dim() || tmp_mid[i] != tmp_mid[i+1]) {
+            node_mid.PushBack(tmp_mid[i]);
+          }
+        }
+      }
       { // 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());
         Long Nnodes = node_mid.Dim();
         node_attr.ReInit(Nnodes);
         for (Long i = 0; i < Nnodes; i++) {
-          node_attr[i].Leaf = 1;
+          node_attr[i].Leaf = !(i+1<Nnodes && node_mid[i].isAncestor(node_mid[i+1]));
           node_attr[i].Ghost = (node_mid[i] < m0 || node_mid[i] >= m1);
         }
       }
 
-      { // Add non-leaf nodes and place-holder for ghost nodes
+      { // Add place-holder for ghost nodes
         // TODO
       }
       { // Update node_data, node_cnt
@@ -348,6 +406,24 @@ template <class Real, Integer DIM> class Tree {
 
         comm.PartitionS(node_mid_orig, mins[comm.Rank()]);
 
+        Vector<Long> new_cnt_range0(node_mid.Dim()), new_cnt_range1(node_mid.Dim());
+        { // Set new_cnt_range0, new_cnt_range1
+          for (Long i = 0; i < start_idx; i++) {
+            new_cnt_range0[i] = 0;
+            new_cnt_range1[i] = 0;
+          }
+          for (Long i = start_idx; i < end_idx; i++) {
+            auto m0 = (node_mid[i+0]);
+            auto m1 = (i+1==end_idx ? Morton<DIM>().Next() : (node_mid[i+1]));
+            new_cnt_range0[i] = std::lower_bound(node_mid_orig.begin(), node_mid_orig.begin() + node_mid_orig.Dim(), m0) - node_mid_orig.begin();
+            new_cnt_range1[i] = std::lower_bound(node_mid_orig.begin(), node_mid_orig.begin() + node_mid_orig.Dim(), m1) - node_mid_orig.begin();
+          }
+          for (Long i = end_idx; i < node_mid.Dim(); i++) {
+            new_cnt_range0[i] = 0;
+            new_cnt_range1[i] = 0;
+          }
+        }
+
         Vector<Long> cnt_tmp;
         Vector<Real> data_tmp;
         for (const auto& pair : node_data) {
@@ -375,21 +451,12 @@ template <class Real, Integer DIM> class Tree {
           comm.PartitionN(cnt_tmp, node_mid_orig.Dim());
 
           cnt_->ReInit(node_mid.Dim());
-          for (Long i = 0; i < start_idx; i++) {
-            cnt_[0][i] = 0;
-          }
-          for (Long i = start_idx; i < end_idx; i++) {
-            auto m0 = coarsest_ancestor_mid(node_mid[i+0]);
-            auto m1 = (i+1==end_idx ? Morton<DIM>().Next() : coarsest_ancestor_mid(node_mid[i+1]));
-            Long a = std::lower_bound(node_mid_orig.begin(), node_mid_orig.begin() + node_mid_orig.Dim(), m0) - node_mid_orig.begin();
-            Long b = std::lower_bound(node_mid_orig.begin(), node_mid_orig.begin() + node_mid_orig.Dim(), m1) - node_mid_orig.begin();
-            // TODO: precompute a and b
-
-            cnt_[0][i] = 0;
-            for (Long j = a; j < b; j++) cnt_[0][i] += cnt_tmp[j];
-          }
-          for (Long i = end_idx; i < node_mid.Dim(); i++) {
-            cnt_[0][i] = 0;
+          for (Long i = 0; i < node_mid.Dim(); i++) {
+            Long sum = 0;
+            Long j0 = new_cnt_range0[i];
+            Long j1 = new_cnt_range1[i];
+            for (Long j = j0; j < j1; j++) sum += cnt_tmp[j];
+            cnt_[0][i] = sum;
           }
           SCTL_ASSERT(omp_par::reduce(cnt_->begin(), cnt_->Dim()) == omp_par::reduce(cnt_tmp.begin(), cnt_tmp.Dim()));
 
@@ -499,7 +566,7 @@ template <class Real, Integer DIM> class Tree {
 
     static void scan(Vector<Long>& dsp, const Vector<Long>& cnt) {
       dsp.ReInit(cnt.Dim());
-      dsp[0] = 0;
+      if (cnt.Dim()) dsp[0] = 0;
       omp_par::scan(cnt.begin(), dsp.begin(), cnt.Dim());
     }
 
@@ -537,9 +604,9 @@ template <class Real, Integer DIM, class BaseTree = Tree<Real,DIM>> class PtTree
       #endif
     }
 
-    void UpdateRefinement(const Vector<Real>& coord, Long M = 1) {
+    void UpdateRefinement(const Vector<Real>& coord, Long M = 1, bool balance21 = 0, bool periodic = 0) {
       const auto& comm = this->GetComm();
-      BaseTree::UpdateRefinement(coord, M);
+      BaseTree::UpdateRefinement(coord, M, balance21, periodic);
 
       Long start_node_idx, end_node_idx;
       { // Set start_node_idx, end_node_idx

include/tree.f90

@@ -10,21 +10,25 @@ interface
     type(c_ptr), intent(in) :: tree_ctx
   end subroutine
 
-  subroutine GetTree(node_coord, node_depth, node_ghost, node_leaf, Nnodes, tree_ctx)
+  subroutine GetTree(node_coord, parent_lst, child_lst, node_depth, node_ghost, node_leaf, Nnodes, tree_ctx)
     use, intrinsic :: ISO_C_BINDING
-    type(c_ptr), intent(out) :: node_coord ! real*8, dimension(Nnode*3)
-    type(c_ptr), intent(out) :: node_depth ! integer*1, dimension(Nnode)
-    type(c_ptr), intent(out) :: node_ghost ! integer*1, dimension(Nnode)
-    type(c_ptr), intent(out) :: node_leaf  ! integer*1, dimension(Nnode)
+    type(c_ptr), intent(out) :: node_coord  ! real*8, dimension(Nnode*3)
+    type(c_ptr), intent(out) :: parent_lst  ! integer*4, dimension(Nnode)
+    type(c_ptr), intent(out) :: child_lst   ! integer*4, dimension(Nnode*8)
+    type(c_ptr), intent(out) :: node_depth  ! integer*1, dimension(Nnode)
+    type(c_ptr), intent(out) :: node_ghost  ! integer*1, dimension(Nnode)
+    type(c_ptr), intent(out) :: node_leaf   ! integer*1, dimension(Nnode)
     integer*4  , intent(out) :: Nnodes
     type(c_ptr), intent(in)  :: tree_ctx
   end subroutine
 
-  subroutine UpdateRefinement(pt_coord, Npt, max_pts, tree_ctx)
+  subroutine UpdateRefinement(pt_coord, Npt, max_pts, level_restrict, periodic, tree_ctx)
     use, intrinsic :: ISO_C_BINDING
     real*8     , intent(in) :: pt_coord(*)
     integer*4  , intent(in) :: Npt
     integer*4  , intent(in) :: max_pts
+    logical    , intent(in) :: level_restrict
+    logical    , intent(in) :: periodic
     type(c_ptr), intent(in) :: tree_ctx
   end subroutine
 

src/libtree.cpp

@@ -7,32 +7,54 @@ template <class Real, Integer DIM> class ParticleTreeWrapper {
 
     ParticleTreeWrapper(Comm comm = Comm::World()) : tree_(comm) {}
 
-    void GetTree(const Real*& node_coord, const int8_t*& node_depth, const int8_t*& node_ghost, const int8_t*& node_leaf, int32_t& Nnodes) const {
+    void GetTree(const Real*& node_coord, const int32_t*& parent_lst, const int32_t*& child_lst, const int8_t*& node_depth, const int8_t*& node_ghost, const int8_t*& node_leaf, int32_t& Nnodes) const {
+      static constexpr Integer MAX_CHILD = (1u << DIM);
       const auto& node_mid_  = tree_.GetNodeMID();
       const auto& node_attr_ = tree_.GetNodeAttr();
       SCTL_ASSERT(node_mid_.Dim() == node_attr_.Dim());
       Nnodes = node_mid_.Dim();
 
-      if (node_coord_.Dim() != Nnodes*DIM) node_coord_.ReInit(Nnodes*DIM);
-      if (node_depth_.Dim() != Nnodes    ) node_depth_.ReInit(Nnodes    );
-      if (node_ghost_.Dim() != Nnodes    ) node_ghost_.ReInit(Nnodes    );
-      if (node_leaf_ .Dim() != Nnodes    ) node_leaf_ .ReInit(Nnodes    );
+      if (node_coord_.Dim() != Nnodes*DIM      ) node_coord_.ReInit(Nnodes*DIM      );
+      if (parent_lst_.Dim() != Nnodes          ) parent_lst_.ReInit(Nnodes          );
+      if (child_lst_ .Dim() != Nnodes*MAX_CHILD) child_lst_ .ReInit(Nnodes*MAX_CHILD);
+      if (node_depth_.Dim() != Nnodes          ) node_depth_.ReInit(Nnodes          );
+      if (node_ghost_.Dim() != Nnodes          ) node_ghost_.ReInit(Nnodes          );
+      if (node_leaf_ .Dim() != Nnodes          ) node_leaf_ .ReInit(Nnodes          );
+      Vector<Long> ancestors(Morton<DIM>::MAX_DEPTH);
+      Vector<Long> child_cnt(Morton<DIM>::MAX_DEPTH);
       #pragma omp parallel for schedule(static)
       for (Long i = 0; i < Nnodes; i++) {
         node_mid_[i].template Coord<Real>(node_coord_.begin() + i*DIM);
         node_depth_[i] = node_mid_[i].Depth();
         node_ghost_[i] = node_attr_[i].Ghost;
         node_leaf_ [i] = node_attr_[i].Leaf;
+        parent_lst_[i] = -1;
+        child_lst_ [i] = -1;
       }
+      for (Long i = 0; i < Nnodes; i++) { // Set parent_lst, child_lst_
+        Integer depth = node_mid_[i].Depth();
+        ancestors[depth] = i;
+        child_cnt[depth] = 0;
+        if (depth) {
+          Long p = ancestors[depth-1];
+          Integer& c = child_cnt[depth-1];
+          parent_lst_[i] = p+1; // +1 because indexing starts from 1 in fortran
+          child_lst_[p*MAX_CHILD+c] = i+1; // +1 because indexing starts from 1 in fortran
+          c++;
+        }
+      }
+
       node_coord = &node_coord_[0];
+      parent_lst = &parent_lst_[0];
+      child_lst  = &child_lst_ [0];
       node_depth = &node_depth_[0];
       node_ghost = &node_ghost_[0];
       node_leaf  = &node_leaf_ [0];
     }
 
-    void UpdateRefinement(const Real* pt_coord_, int32_t Npt, int32_t max_pts) {
+    void UpdateRefinement(const Real* pt_coord_, int32_t Npt, int32_t max_pts, bool level_restrict, bool periodic) {
       const Vector<Real> pt_coord(Npt*DIM, Ptr2Itr<Real>((Real*)pt_coord_,Npt*DIM), false);
-      tree_.UpdateRefinement(pt_coord, max_pts);
+      tree_.UpdateRefinement(pt_coord, max_pts, level_restrict, periodic);
     }
 
     void AddData(const char* name, const Real* data, int32_t Ndata, const int32_t* cnt, int32_t Ncnt) {
@@ -95,6 +117,8 @@ template <class Real, Integer DIM> class ParticleTreeWrapper {
     PtTree<Real,DIM> tree_;
 
     mutable Vector<Real> node_coord_;
+    mutable Vector<int32_t> parent_lst_;
+    mutable Vector<int32_t> child_lst_;
     mutable Vector<int8_t> node_depth_;
     mutable Vector<int8_t> node_ghost_;
     mutable Vector<int8_t> node_leaf_;
@@ -118,14 +142,14 @@ void deletetree_(void** tree) {
   tree[0] = nullptr;
 }
 
-void gettree_(const double** node_coord, const int8_t** node_depth, const int8_t** node_ghost, const int8_t** node_leaf, int32_t* Nnodes, const void** tree_ctx) {
+void gettree_(const double** node_coord, const int32_t** parent_lst, const int32_t** child_lst, const int8_t** node_depth, const int8_t** node_ghost, const int8_t** node_leaf, int32_t* Nnodes, const void** tree_ctx) {
   const sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
-  tree.GetTree(*node_coord, *node_depth, *node_ghost, *node_leaf, *Nnodes);
+  tree.GetTree(*node_coord, *parent_lst, *child_lst, *node_depth, *node_ghost, *node_leaf, *Nnodes);
 }
 
-void updaterefinement_(const double* pt_coord_, const int32_t* Npt, const int32_t* max_pts, void** tree_ctx) {
+void updaterefinement_(const double* pt_coord_, const int32_t* Npt, const int32_t* max_pts, bool* level_restrict, bool* periodic, void** tree_ctx) {
   sctl::ParticleTreeWrapper<double,3>& tree = *(sctl::ParticleTreeWrapper<double,3>*)tree_ctx[0];
-  tree.UpdateRefinement(pt_coord_, *Npt, *max_pts);
+  tree.UpdateRefinement(pt_coord_, *Npt, *max_pts, *level_restrict, *periodic);
 }
 
 void adddata_(const char* name, const double* data, const int32_t* Ndata, const int32_t* cnt, const int32_t* Ncnt, void** tree_ctx) {

src/test-cpp.cpp

@@ -110,7 +110,7 @@ int main(int argc, char** argv) {
 
     sctl::Profile::Tic("Refine",&comm);
     sctl::PtTree<Real,DIM> t(comm);
-    t.UpdateRefinement(coord_unif, 1000);
+    t.UpdateRefinement(coord_unif, 1000, true, false);
     sctl::Profile::Toc();
 
     sctl::Profile::Tic("AddPts",&comm);
@@ -118,7 +118,7 @@ int main(int argc, char** argv) {
     sctl::Profile::Toc();
 
     sctl::Profile::Tic("UpdateRefine",&comm);
-    t.UpdateRefinement(coord, 100);
+    t.UpdateRefinement(coord, 100, true, false);
     sctl::Profile::Toc();
 
     sctl::Profile::Tic("GetPts",&comm);

src/test-fortran.f90

@@ -15,8 +15,9 @@ program main
 
   integer*4 :: Nnodes
   real*4, pointer :: node_coord(:)
+  integer*4, pointer :: parent_lst(:), child_lst(:)
   integer*1, pointer :: node_depth(:), node_ghost(:), node_leaf(:)
-  type(c_ptr) :: node_coord_, node_depth_, node_ghost_, node_leaf_
+  type(c_ptr) :: node_coord_, parent_lst_, child_lst_, node_depth_, node_ghost_, node_leaf_
 
   real*8, allocatable :: mult_coeff(:)
   integer*4, allocatable :: mult_cnt(:)
@@ -36,13 +37,15 @@ program main
 
   ! Create tree and refine using pt_coord with max 100 points per box
   call CreateTree(tree_ctx)
-  call UpdateRefinement(pt_coord, Npt, 100, tree_ctx)
+  call UpdateRefinement(pt_coord, Npt, 100, .true., .false., tree_ctx)
 
 
 
   ! Read tree node structure
-  call GetTree(node_coord_, node_depth_, node_ghost_, node_leaf_, Nnodes, tree_ctx)
+  call GetTree(node_coord_, parent_lst_, child_lst_, node_depth_, node_ghost_, node_leaf_, Nnodes, tree_ctx)
   call c_f_pointer(node_coord_, node_coord, [Nnodes*3])
+  call c_f_pointer(parent_lst_, parent_lst, [Nnodes])
+  call c_f_pointer(child_lst_ , child_lst , [Nnodes*8])
   call c_f_pointer(node_depth_, node_depth, [Nnodes])
   call c_f_pointer(node_ghost_, node_ghost, [Nnodes])
   call c_f_pointer(node_leaf_ , node_leaf , [Nnodes])
@@ -82,7 +85,7 @@ program main
   allocate(mult_coeff(Nnodes*10))
   allocate(mult_cnt(Nnodes))
   do i = 1, Nnodes
-    mult_cnt(i) = 10
+    mult_cnt(i) = 1
   enddo
   call AddData('mult_coeff', mult_coeff, Nnodes*10, mult_cnt, Nnodes, tree_ctx)
   call DeleteData('mult_coeff', tree_ctx)