[Trivial] Reorganize LogicalLocation code

CHANGELOG.md

@@ -15,6 +15,7 @@
 - [[PR 1031]](https://github.com/parthenon-hpc-lab/parthenon/pull/1031) Fix bug in non-cell centered AMR
 
 ### Infrastructure (changes irrelevant to downstream codes)
+- [[PR 1028]](https://github.com/parthenon-hpc-lab/parthenon/pull/1028) Internal reorganization of LogicalLocation files
 - [[PR 1009]](https://github.com/parthenon-hpc-lab/parthenon/pull/1009) Move from a single octree to a forest of octrees
 
 

src/CMakeLists.txt

@@ -156,18 +156,20 @@ add_library(parthenon
   interface/variable.cpp
   interface/variable.hpp
 
-  mesh/amr_loadbalance.cpp
   mesh/domain.hpp
+  mesh/forest/block_ownership.cpp
+  mesh/forest/block_ownership.hpp
   mesh/forest/forest.cpp
   mesh/forest/forest.hpp
   mesh/forest/relative_orientation.hpp
   mesh/forest/relative_orientation.cpp
   mesh/forest/tree.hpp
   mesh/forest/tree.cpp
-  mesh/logical_location.cpp
-  mesh/logical_location.hpp
+  mesh/forest/logical_location.cpp
+  mesh/forest/logical_location.hpp
   mesh/mesh_refinement.cpp
   mesh/mesh_refinement.hpp
+  mesh/mesh-amr_loadbalance.cpp
   mesh/mesh-gmg.cpp
   mesh/mesh.cpp
   mesh/mesh.hpp

src/bvals/neighbor_block.cpp

@@ -30,63 +30,13 @@
 #include <unordered_set>
 
 #include "globals.hpp"
-#include "mesh/logical_location.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "mesh/mesh.hpp"
 #include "utils/buffer_utils.hpp"
 #include "utils/error_checking.hpp"
 
 namespace parthenon {
 
-//----------------------------------------------------------------------------------------
-// \!fn void NeighborBlock::SetNeighbor(int irank, int ilevel, int igid, int ilid,
-//                          int iox1, int iox2, int iox3, NeighborConnect itype,
-//                          int ibid, int itargetid, int ifi1=0, int ifi2=0)
-// \brief Set neighbor information
-
-void NeighborBlock::SetNeighbor(LogicalLocation inloc, int irank, int ilevel, int igid,
-                                int ilid, int iox1, int iox2, int iox3,
-                                NeighborConnect itype, int ibid, int itargetid,
-                                int ifi1, // =0
-                                int ifi2  // =0
-) {
-  snb.rank = irank;
-  snb.level = ilevel;
-  snb.gid = igid;
-  snb.lid = ilid;
-  ni.ox1 = iox1;
-  ni.ox2 = iox2;
-  ni.ox3 = iox3;
-  ni.type = itype;
-  ni.fi1 = ifi1;
-  ni.fi2 = ifi2;
-  bufid = ibid;
-  targetid = itargetid;
-  loc = inloc;
-  if (ni.type == NeighborConnect::face) {
-    if (ni.ox1 == -1)
-      fid = BoundaryFace::inner_x1;
-    else if (ni.ox1 == 1)
-      fid = BoundaryFace::outer_x1;
-    else if (ni.ox2 == -1)
-      fid = BoundaryFace::inner_x2;
-    else if (ni.ox2 == 1)
-      fid = BoundaryFace::outer_x2;
-    else if (ni.ox3 == -1)
-      fid = BoundaryFace::inner_x3;
-    else if (ni.ox3 == 1)
-      fid = BoundaryFace::outer_x3;
-  }
-  if (ni.type == NeighborConnect::edge) {
-    if (ni.ox3 == 0)
-      eid = ((((ni.ox1 + 1) >> 1) | ((ni.ox2 + 1) & 2)));
-    else if (ni.ox2 == 0)
-      eid = (4 + (((ni.ox1 + 1) >> 1) | ((ni.ox3 + 1) & 2)));
-    else if (ni.ox1 == 0)
-      eid = (8 + (((ni.ox2 + 1) >> 1) | ((ni.ox3 + 1) & 2)));
-  }
-  return;
-}
-
 NeighborConnect NCFromOffsets(const std::array<int, 3> offsets) {
   int connect_indicator =
       std::abs(offsets[0]) + std::abs(offsets[1]) + std::abs(offsets[2]);

src/bvals/neighbor_block.hpp

@@ -29,6 +29,8 @@
 #include "parthenon_mpi.hpp"
 
 #include "defs.hpp"
+#include "mesh/forest/block_ownership.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "parthenon_arrays.hpp"
 #include "utils/error_checking.hpp"
 
@@ -144,9 +146,6 @@ struct NeighborBlock {
   block_ownership_t ownership;
   RegionSize block_size;
 
-  void SetNeighbor(LogicalLocation inloc, int irank, int ilevel, int igid, int ilid,
-                   int iox1, int iox2, int iox3, NeighborConnect itype, int ibid,
-                   int itargetid, int ifi1 = 0, int ifi2 = 0);
   NeighborBlock() = default;
   NeighborBlock(Mesh *mesh, LogicalLocation loc, int rank, int gid, int lid,
                 std::array<int, 3> offsets, NeighborConnect type, int bid, int target_id,

src/defs.hpp

@@ -25,7 +25,7 @@
 
 #include "basic_types.hpp"
 #include "config.hpp"
-#include "mesh/logical_location.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "parthenon_arrays.hpp"
 
 namespace parthenon {

src/mesh/forest/block_ownership.cpp

@@ -0,0 +1,183 @@
+//========================================================================================
+// Athena++ astrophysical MHD code
+// Copyright(C) 2014 James M. Stone <jmstone@princeton.edu> and other code contributors
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// Parthenon performance portable AMR framework
+// Copyright(C) 2020-2023 The Parthenon collaboration
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved.
+//
+// This program was produced under U.S. Government contract 89233218CNA000001 for Los
+// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
+// for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+// in the program are reserved by Triad National Security, LLC, and the U.S. Department
+// of Energy/National Nuclear Security Administration. The Government is granted for
+// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// license in this material to reproduce, prepare derivative works, distribute copies to
+// the public, perform publicly and display publicly, and to permit others to do so.
+//========================================================================================
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "mesh/forest/block_ownership.hpp"
+#include "mesh/forest/logical_location.hpp"
+#include "utils/error_checking.hpp"
+#include "utils/morton_number.hpp"
+
+namespace parthenon {
+
+// TODO(LFR): Remove this
+block_ownership_t
+DetermineOwnership(const LogicalLocation &main_block,
+                   const std::unordered_set<LogicalLocation> &allowed_neighbors,
+                   const RootGridInfo &rg_info,
+                   const std::unordered_set<LogicalLocation> &newly_refined) {
+  block_ownership_t main_owns;
+
+  auto ownership_level = [&](const LogicalLocation &a) {
+    // Newly-refined blocks are treated as higher-level than blocks at their
+    // parent level, but lower-level than previously-refined blocks at their
+    // current level.
+    if (newly_refined.count(a)) return 2 * a.level() - 1;
+    return 2 * a.level();
+  };
+
+  auto ownership_less_than = [ownership_level](const LogicalLocation &a,
+                                               const LogicalLocation &b) {
+    // Ownership is first determined by block with the highest level, then by maximum
+    // Morton number this is reversed in precedence from the normal comparators where
+    // Morton number takes precedence
+    if (ownership_level(a) == ownership_level(b)) return a.morton() < b.morton();
+    return ownership_level(a) < ownership_level(b);
+  };
+
+  for (int ox1 : {-1, 0, 1}) {
+    for (int ox2 : {-1, 0, 1}) {
+      for (int ox3 : {-1, 0, 1}) {
+        main_owns(ox1, ox2, ox3) = true;
+        for (auto &n : allowed_neighbors) {
+          if (ownership_less_than(main_block, n) &&
+              main_block.IsNeighborOfTE(n, ox1, ox2, ox3, rg_info)) {
+            main_owns(ox1, ox2, ox3) = false;
+            break;
+          }
+        }
+      }
+    }
+  }
+  return main_owns;
+}
+
+block_ownership_t
+DetermineOwnershipForest(const LogicalLocation &main_block,
+                         const std::vector<NeighborLocation> &allowed_neighbors,
+                         const std::unordered_set<LogicalLocation> &newly_refined) {
+  block_ownership_t main_owns;
+
+  auto ownership_level = [&](const LogicalLocation &a) {
+    // Newly-refined blocks are treated as higher-level than blocks at their
+    // parent level, but lower-level than previously-refined blocks at their
+    // current level.
+    if (newly_refined.count(a)) return 2 * a.level() - 1;
+    return 2 * a.level();
+  };
+
+  auto ownership_less_than = [ownership_level](const LogicalLocation &a,
+                                               const LogicalLocation &b) {
+    // Ownership is first determined by block with the highest level, then by maximum
+    // (tree, Morton) number this is reversed in precedence from the normal comparators
+    // where (tree, Morton) number takes precedence
+    if (ownership_level(a) != ownership_level(b))
+      return ownership_level(a) < ownership_level(b);
+    if (a.tree() != b.tree()) return a.tree() < b.tree();
+    return a.morton() < b.morton();
+  };
+
+  for (int ox1 : {-1, 0, 1}) {
+    for (int ox2 : {-1, 0, 1}) {
+      for (int ox3 : {-1, 0, 1}) {
+        main_owns(ox1, ox2, ox3) = true;
+        for (const auto &n : allowed_neighbors) {
+          if (ownership_less_than(main_block, n.global_loc) &&
+              main_block.IsNeighborOfTEForest(n.origin_loc, {ox1, ox2, ox3})) {
+            main_owns(ox1, ox2, ox3) = false;
+            break;
+          }
+        }
+      }
+    }
+  }
+  return main_owns;
+}
+
+// Given a topological element, ownership array of the sending block, and offset indices
+// defining the location of an index region within the block (i.e. the ghost zones passed
+// across the x-face or the ghost zones passed across the z-edge), return the index range
+// masking array required for masking out unowned regions of the index space. ox? defines
+// buffer location on the owner block
+block_ownership_t
+GetIndexRangeMaskFromOwnership(TopologicalElement el,
+                               const block_ownership_t &sender_ownership, int ox1,
+                               int ox2, int ox3) {
+  using vp_t = std::vector<std::pair<int, int>>;
+
+  // Transform general block ownership to element ownership over entire block. For
+  // instance, x-faces only care about block ownership in the x-direction First index of
+  // the pair is the element index and the second index is the block index that is copied
+  // to that element index
+  block_ownership_t element_ownership = sender_ownership;
+  auto x1_idxs = TopologicalOffsetI(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  auto x2_idxs = TopologicalOffsetJ(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  auto x3_idxs = TopologicalOffsetK(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  for (auto [iel, ibl] : x1_idxs) {
+    for (auto [jel, jbl] : x2_idxs) {
+      for (auto [kel, kbl] : x3_idxs) {
+        element_ownership(iel, jel, kel) = sender_ownership(ibl, jbl, kbl);
+      }
+    }
+  }
+
+  // Now, the ownership status is correct for the entire interior index range of the
+  // block, but the offsets ox? define a subset of these indices (e.g. one edge of the
+  // interior). Therefore, we need to set the index ownership to true for edges of the
+  // index range that are contained in the interior of the sending block
+  if (ox1 != 0) {
+    for (auto j : {-1, 0, 1}) {
+      for (auto k : {-1, 0, 1}) {
+        element_ownership(-ox1, j, k) = element_ownership(0, j, k);
+      }
+    }
+  }
+  if (ox2 != 0) {
+    for (auto i : {-1, 0, 1}) {
+      for (auto k : {-1, 0, 1}) {
+        element_ownership(i, -ox2, k) = element_ownership(i, 0, k);
+      }
+    }
+  }
+  if (ox3 != 0) {
+    for (auto i : {-1, 0, 1}) {
+      for (auto j : {-1, 0, 1}) {
+        element_ownership(i, j, -ox3) = element_ownership(i, j, 0);
+      }
+    }
+  }
+
+  return element_ownership;
+}
+
+} // namespace parthenon