WIP: Update ownership model to work in general forests and show tri-valent singular node fixup

example/boundary_exchange/boundary_exchange.cpp

@@ -12,9 +12,11 @@
 //========================================================================================
 
 // Standard Includes
+#include <cstdio>
 #include <iostream>
 #include <limits>
 #include <memory>
+#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -101,6 +103,185 @@ TaskStatus SetBlockValues(MeshData<Real> *md) {
   return TaskStatus::complete;
 }
 
+struct ParameterizedLine {
+  enum class ltype { straight, arc };
+
+  const Real x1, y1;
+  const Real x2, y2;
+  Real r, delta, phi;
+  ltype type;
+
+  using node = parthenon::forest::Node;
+  ParameterizedLine(std::shared_ptr<node> start, std::shared_ptr<node> end)
+      : x1{start->x[0]}, y1{start->x[1]}, x2{end->x[0]}, y2{end->x[1]},
+        type{ltype::straight} {
+    Real d1 = std::sqrt(x1 * x1 + y1 * y1);
+    Real d2 = std::sqrt(x2 * x2 + y2 * y2);
+    if (std::abs(d1 - d2) < 1.e-8 && d1 > 1.1) {
+      type = ltype::arc;
+      r = d1;
+      delta = M_PI / 4.0;
+      phi = 0.0;
+      if (y1 > 1.e-5) phi = delta;
+    }
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  Real GetX(Real u) const {
+    if (type == ltype::straight)
+      return x1 * (1.0 - u) + x2 * u;
+    else
+      return r * cos(delta * u + phi);
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  Real GetY(Real u) const {
+    if (type == ltype::straight)
+      return y1 * (1.0 - u) + y2 * u;
+    else
+      return r * sin(delta * u + phi);
+  }
+};
+
+TaskStatus SetCoordinates(MeshData<Real> *md) {
+  using TE = parthenon::TopologicalElement;
+  auto pmesh = md->GetMeshPointer();
+  auto desc = parthenon::MakePackDescriptor<position>(md);
+
+  for (auto &ptree : pmesh->forest.GetTrees()) {
+    auto tree_id = ptree->GetId();
+    auto pack = desc.GetPack(md, parthenon::GetBlockSelector::OnTree(ptree->GetId()));
+    int i{0};
+    Real posx[4], posy[4];
+    for (auto &pnode : ptree->forest_nodes) {
+      posx[i] = pnode->x[0];
+      posy[i] = pnode->x[1];
+    }
+
+    auto &pnodes = ptree->forest_nodes;
+    ParameterizedLine c1(pnodes[0], pnodes[1]);
+    ParameterizedLine c3(pnodes[2], pnodes[3]);
+    ParameterizedLine c2(pnodes[0], pnodes[2]);
+    ParameterizedLine c4(pnodes[1], pnodes[3]);
+
+    IndexRange ib = md->GetBoundsI(IndexDomain::interior, TE::NN);
+    IndexRange jb = md->GetBoundsJ(IndexDomain::interior, TE::NN);
+    IndexRange kb = md->GetBoundsK(IndexDomain::interior, TE::NN);
+
+    parthenon::par_for(
+        parthenon::loop_pattern_mdrange_tag, "SetPosition", DevExecSpace(), 0,
+        pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+        KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
+          auto &coords = pack.GetCoordinates(b);
+          const Real u = coords.X<X1DIR, TE::NN>(k, j, i);
+          const Real v = coords.X<X2DIR, TE::NN>(k, j, i);
+          const Real mu = 1.0 - u;
+          const Real mv = 1.0 - v;
+
+          Real x = mv * c1.GetX(u) + v * c3.GetX(u) + mu * c2.GetX(v) + u * c4.GetX(v) -
+                   (mu * mv * c1.GetX(0.0) + u * mv * c1.GetX(1.0) +
+                    mu * v * c2.GetX(1.0) + u * v * c3.GetX(1.0));
+          Real y = mv * c1.GetY(u) + v * c3.GetY(u) + mu * c2.GetY(v) + u * c4.GetY(v) -
+                   (mu * mv * c1.GetY(0.0) + u * mv * c1.GetY(1.0) +
+                    mu * v * c2.GetY(1.0) + u * v * c3.GetY(1.0));
+          pack(b, TE::NN, position(0), k, j, i) = x; // x-position
+          pack(b, TE::NN, position(1), k, j, i) = y; // y-position
+        });
+  }
+  return TaskStatus::complete;
+}
+
+TaskStatus FixTrivalentNodes2D(MeshData<Real> *md, bool coarse) {
+  using TE = parthenon::TopologicalElement;
+  auto pmesh = md->GetMeshPointer();
+  std::set<parthenon::PDOpt> opts{};
+  if (coarse) opts.insert(parthenon::PDOpt::Coarse);
+
+  auto desc = parthenon::MakePackDescriptor<position>(md, {}, opts);
+
+  const parthenon::CellLevel clevel =
+      coarse ? parthenon::CellLevel::coarse : parthenon::CellLevel::same;
+  IndexRange ib_in = md->GetBoundsI(clevel, IndexDomain::interior, TE::NN);
+  IndexRange jb_in = md->GetBoundsJ(clevel, IndexDomain::interior, TE::NN);
+  IndexRange kb = md->GetBoundsK(clevel, IndexDomain::interior, TE::NN);
+
+  for (auto &ptree : pmesh->forest.GetTrees()) {
+    auto tree_id = ptree->GetId();
+    int pos{0};
+    for (auto &pnode : ptree->forest_nodes) {
+      bool trivalent =
+          (pnode->associated_faces.size() == 3) && !pnode->on_physical_boundary;
+      if (trivalent) {
+        parthenon::CellCentOffsets offset(2 * (pos % 2) - 1, 2 * (pos / 2) - 1, 0);
+        auto pack =
+            desc.GetPack(md, parthenon::GetBlockSelector::OnTree(ptree->GetId(), offset));
+        if (pack.GetNBlocks() == 0) return TaskStatus::complete;
+        // Copy shared boundary data
+
+        // This pack has to have only one block
+        auto gid = pack.GetGIDHost(0);
+        parthenon::MeshBlock *pmb;
+        for (auto &pmbd : md->GetAllBlockData())
+          if (pmbd->GetParentPointer()->gid == gid) pmb = pmbd->GetParentPointer();
+
+        // Now check which row of corner shared elements has been set by an owning block
+        parthenon::CellCentOffsets offsetX1(2 * (pos % 2) - 1, 0, 0);
+        parthenon::CellCentOffsets offsetX2(0, 2 * (pos / 2) - 1, 0);
+        int dir = X2DIR;
+        for (auto &neighbor : pmb->neighbors) {
+          if (neighbor.offsets == offsetX1 && neighbor.origin_ownership(offsetX2))
+            dir = X1DIR;
+          if (neighbor.offsets == offsetX2 && neighbor.origin_ownership(offsetX1))
+            dir = X2DIR;
+        }
+
+        // Select the reference location on the node shared by all three blocks
+        const int icorner = (offset(X1DIR) == parthenon::Offset::Low) ? ib_in.s : ib_in.e;
+        const int jcorner = (offset(X2DIR) == parthenon::Offset::Low) ? jb_in.s : jb_in.e;
+        // Select the index space of elements that need to get overwritten
+        IndexRange ib = offset(X1DIR) == parthenon::Offset::Low
+                            ? IndexRange{0, ib_in.s - (dir == X2DIR)}
+                            : IndexRange{ib_in.e + (dir == X2DIR),
+                                         ib_in.e + parthenon::Globals::nghost};
+        IndexRange jb = offset(X2DIR) == parthenon::Offset::Low
+                            ? IndexRange{0, jb_in.s - (dir == X1DIR)}
+                            : IndexRange{jb_in.e + (dir == X1DIR),
+                                         ib_in.e + parthenon::Globals::nghost};
+        parthenon::par_for(
+            parthenon::loop_pattern_mdrange_tag, "SetPosition", DevExecSpace(), 0,
+            pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+            KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
+              const int enclosed_area = std::abs((i - icorner) * (j - jcorner));
+              int ioff{0};
+              int joff{0};
+              // This is an element shared by both of the neighbors, but only one
+              // communicated it, so copy it to the other side of the corner
+              if (enclosed_area == 0) {
+                joff = std::abs(icorner - i) * offset(X2DIR);
+                ioff = std::abs(jcorner - j) * offset(X1DIR);
+              } else {
+                // These are non-existent elements that we fill with data copied from
+                // their nearest neighbors
+                int off = 1;
+                while (off * off < enclosed_area)
+                  off++;
+                ioff = offset(X1DIR) * off * (dir == X1DIR);
+                joff = offset(X2DIR) * off * (dir == X2DIR);
+              }
+              pack(b, TE::NN, position(0), k, j, i) =
+                  pack(b, TE::NN, position(0), k, jcorner + joff,
+                       icorner + ioff); // x-position
+              pack(b, TE::NN, position(1), k, j, i) =
+                  pack(b, TE::NN, position(1), k, jcorner + joff,
+                       icorner + ioff); // y-position
+            });
+      }
+      pos++;
+    }
+  }
+  return TaskStatus::complete;
+}
+
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   auto package = std::make_shared<StateDescriptor>("boundary_exchange");
   Params &params = package->AllParams();
@@ -109,8 +290,11 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
              std::vector<int>{8});
   m.RegisterRefinementOps<parthenon::refinement_ops::ProlongatePiecewiseConstant,
                           parthenon::refinement_ops::RestrictAverage>();
-  package->AddField(neighbor_info::name(), m);
+  package->AddField<neighbor_info>(m);
 
+  Metadata m_node({Metadata::Node, Metadata::Independent, Metadata::FillGhost},
+                  std::vector<int>{2});
+  package->AddField<position>(m_node);
   return package;
 }
 

example/boundary_exchange/boundary_exchange.hpp

@@ -34,7 +34,16 @@ struct neighbor_info : public parthenon::variable_names::base_t<false, 4> {
   static std::string name() { return "neighbor_info"; }
 };
 
+struct position : public parthenon::variable_names::base_t<false, 4> {
+  template <class... Ts>
+  KOKKOS_INLINE_FUNCTION position(Ts &&...args)
+      : parthenon::variable_names::base_t<false, 4>(std::forward<Ts>(args)...) {}
+  static std::string name() { return "position"; }
+};
+
 TaskStatus SetBlockValues(MeshData<Real> *rc);
+TaskStatus SetCoordinates(MeshData<Real> *rc);
+TaskStatus FixTrivalentNodes2D(MeshData<Real> *rc, bool coarse);
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);
 
 } // namespace boundary_exchange

example/boundary_exchange/boundary_exchange_driver.cpp

@@ -34,24 +34,7 @@ using boundary_exchange::BoundaryExchangeDriver;
 
 Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin);
 
-int main(int argc, char *argv[]) {
-  ParthenonManager pman;
-
-  pman.app_input->ProcessPackages = ProcessPackages;
-
-  // This is called on each mesh block whenever the mesh changes.
-  // pman.app_input->InitMeshBlockUserData = &calculate_pi::SetInOrOutBlock;
-
-  auto manager_status = pman.ParthenonInitEnv(argc, argv);
-  if (manager_status == ParthenonStatus::complete) {
-    pman.ParthenonFinalize();
-    return 0;
-  }
-  if (manager_status == ParthenonStatus::error) {
-    pman.ParthenonFinalize();
-    return 1;
-  }
-
+parthenon::forest::ForestDefinition MakeFourTreeForestWithOneTreeRotated() {
   // Create the nodes for the forest, the x-y positions are only used for
   // visualizing the forest configuration and *do not* determine the global
   // coordinates of the trees.
@@ -95,7 +78,100 @@ int main(int argc, char *argv[]) {
   forest_def.AddBC(edge_t({n[7], n[8]}));
 
   forest_def.AddInitialRefinement(parthenon::LogicalLocation(0, 1, 0, 0, 0));
-  pman.ParthenonInitPackagesAndMesh(forest_def);
+
+  return forest_def;
+}
+
+parthenon::forest::ForestDefinition n_blocks(int nblocks) {
+  using namespace parthenon::forest;
+  std::unordered_map<uint64_t, std::shared_ptr<Node>> nodes;
+  ForestDefinition forest_def;
+  parthenon::Real xoffset = 0.1;
+  parthenon::Real yoffset = 0.1;
+
+  for (int point = 0; point < 2 * nblocks; ++point) {
+    nodes[point] = Node::create(point, {-std::sin(point * M_PI / nblocks) + xoffset,
+                                        -std::cos(point * M_PI / nblocks) + yoffset});
+  }
+
+  using edge_t = parthenon::forest::Edge;
+  for (int point = 0; point < 2 * nblocks; ++point) {
+    forest_def.AddBC(
+        edge_t({nodes[point % (2 * nblocks)], nodes[(point + 1) % (2 * nblocks)]}));
+  }
+
+  nodes[2 * nblocks] = Node::create(2 * nblocks, {xoffset, yoffset});
+  auto &n = nodes;
+  for (int t = 0; t < nblocks; ++t)
+    forest_def.AddFace(
+        t, {n[2 * t + 1], n[2 * t], n[(2 * t + 2) % (2 * nblocks)], n[2 * nblocks]});
+
+  int level = 1;
+  forest_def.AddInitialRefinement(
+      parthenon::LogicalLocation(0, level, (1 << level) - 1, (1 << level) - 1, 0));
+  forest_def.AddInitialRefinement(
+      parthenon::LogicalLocation(1, level, (1 << level) - 1, (1 << level) - 1, 0));
+  forest_def.AddInitialRefinement(
+      parthenon::LogicalLocation(2, level, (1 << level) - 1, (1 << level) - 1, 0));
+
+  return forest_def;
+}
+
+parthenon::forest::ForestDefinition cubed_circle(std::vector<double> radii) {
+  using namespace parthenon::forest;
+  std::vector<std::shared_ptr<Node>> nodes_x, nodes_y, nodes_c;
+  std::size_t node_idx{0};
+  std::shared_ptr<Node> origin = Node::create(node_idx++, {0.0, 0.0});
+  double fac = 1.0 / sqrt(2.0);
+  for (auto rad : radii) {
+    nodes_x.push_back(Node::create(node_idx++, {rad, 0.0}));
+    nodes_c.push_back(Node::create(node_idx++, {rad * fac, rad * fac}));
+    nodes_y.push_back(Node::create(node_idx++, {0.0, rad}));
+  }
+
+  ForestDefinition forest_def;
+  std::size_t face_idx{0};
+  forest_def.AddFace(face_idx++, {origin, nodes_x[0], nodes_y[0], nodes_c[0]});
+  forest_def.AddBC(Edge({origin, nodes_x[0]}));
+  forest_def.AddBC(Edge({origin, nodes_y[0]}));
+  for (int i = 0; i < nodes_x.size() - 1; ++i) {
+    forest_def.AddFace(face_idx++,
+                       {nodes_x[i], nodes_x[i + 1], nodes_c[i], nodes_c[i + 1]});
+    forest_def.AddFace(face_idx++,
+                       {nodes_c[i], nodes_c[i + 1], nodes_y[i], nodes_y[i + 1]});
+    forest_def.AddBC(Edge({nodes_x[i], nodes_x[i + 1]}));
+    forest_def.AddBC(Edge({nodes_y[i], nodes_y[i + 1]}));
+  }
+  forest_def.AddBC(Edge({nodes_x.back(), nodes_c.back()}));
+  forest_def.AddBC(Edge({nodes_c.back(), nodes_y.back()}));
+
+  int level = 2;
+  forest_def.AddInitialRefinement(
+      parthenon::LogicalLocation(1, level, 0, (1 << level) - 1, 0));
+
+  return forest_def;
+}
+
+int main(int argc, char *argv[]) {
+  ParthenonManager pman;
+
+  pman.app_input->ProcessPackages = ProcessPackages;
+
+  // This is called on each mesh block whenever the mesh changes.
+  // pman.app_input->InitMeshBlockUserData = &calculate_pi::SetInOrOutBlock;
+
+  auto manager_status = pman.ParthenonInitEnv(argc, argv);
+  if (manager_status == ParthenonStatus::complete) {
+    pman.ParthenonFinalize();
+    return 0;
+  }
+  if (manager_status == ParthenonStatus::error) {
+    pman.ParthenonFinalize();
+    return 1;
+  }
+
+  pman.ParthenonInitPackagesAndMesh(cubed_circle({1.0, 2.0, 3.0}));
+  // pman.ParthenonInitPackagesAndMesh(n_blocks(3));
 
   // This needs to be scoped so that the driver object is destructed before Finalize
   {
@@ -150,7 +226,14 @@ TaskCollection BoundaryExchangeDriver::MakeTaskCollection(T &blocks) {
       auto &md = pmesh->mesh_data.GetOrAdd("base", i);
       TaskID none(0);
       auto fill = tl.AddTask(none, SetBlockValues, md.get());
-      auto bound = AddBoundaryExchangeTasks(fill, tl, md, true);
+      auto set_coords = tl.AddTask(none, SetCoordinates, md.get());
+      auto bound = AddBoundaryExchangeTasks(
+          fill | set_coords, tl, md, true,
+          [](TaskID depend, TaskList *ptl, std::shared_ptr<MeshData<Real>> md,
+             bool coarse) {
+            return ptl->AddTask(depend, FixTrivalentNodes2D, md.get(), coarse);
+          });
+      // auto fix = tl.AddTask(bound, FixTrivalentNodes2D, md.get(), false);
     }
   }
 

src/bvals/comms/bnd_info.cpp

@@ -244,7 +244,12 @@ CalcIndices(const NeighborBlock &nb, MeshBlock *pmb,
       if (sox2 == 0) sox2 = loc.l(1) % 2 == 1 ? 1 : -1;
       if (sox3 == 0) sox3 = loc.l(2) % 2 == 1 ? 1 : -1;
     }
-    owns = GetIndexRangeMaskFromOwnership(el, nb.ownership, sox1, sox2, sox3);
+    // We are working in the frame of the origin block (i.e. the receiving block),
+    // so we use the ownership array in that frame to get the masking initially.
+    // This mask must then be transformed to the frame of the neighbor block (i.e.
+    // sending block) since the index range we are masking is defined in that frame.
+    owns = lcoord_trans.Transform(
+        GetIndexRangeMaskFromOwnership(el, nb.origin_ownership, sox1, sox2, sox3));
   }
   return SpatiallyMaskedIndexer6D(owns, {0, tensor_shape[0] - 1},
                                   {0, tensor_shape[1] - 1}, {0, tensor_shape[2] - 1},