Adding support for a position hint to get_contained_cells for acceleration

Author: pshriwise

include/openmc/cell.h

@@ -172,11 +172,12 @@ class Cell {
   //! \return Map with cell indexes as keys and
   //! instances as values
   std::unordered_map<int32_t, vector<int32_t>> get_contained_cells(
-    int32_t instance = 0) const;
+    int32_t instance = 0, Position* hint = nullptr) const;
 
 protected:
   //! Determine the path to this cell instance in the geometry hierarchy
-  vector<ParentCell> find_parent_cells(int32_t instance) const;
+  vector<ParentCell> find_parent_cells(
+    int32_t instance, Position* hint = nullptr) const;
 
   //! Inner function for retrieving contained cells
   void get_contained_cells_inner(

include/openmc/lattice.h

@@ -99,6 +99,11 @@ class Lattice {
   virtual void get_indices(
     Position r, Direction u, array<int, 3>& result) const = 0;
 
+  //! \brief Compute the the flat index for a set of lattice cell indices
+  //! \param i_xyz The indices for a lattice cell
+  //! \return Flat index into the universes vector
+  virtual int get_flat_index(const array<int, 3>& i_xyz) const = 0;
+
   //! \brief Get coordinates local to a lattice tile.
   //! \param r A 3D Cartesian coordinate.
   //! \param i_xyz The indices for a lattice tile.
@@ -210,6 +215,8 @@ class RectLattice : public Lattice {
 
   void get_indices(Position r, Direction u, array<int, 3>& result) const;
 
+  int get_flat_index(const array<int, 3>& i_xyz) const;
+
   Position get_local_position(Position r, const array<int, 3>& i_xyz) const;
 
   int32_t& offset(int map, array<int, 3> const& i_xyz);
@@ -245,6 +252,8 @@ class HexLattice : public Lattice {
 
   void get_indices(Position r, Direction u, array<int, 3>& result) const;
 
+  int get_flat_index(const array<int, 3>& i_xyz) const;
+
   Position get_local_position(Position r, const array<int, 3>& i_xyz) const;
 
   bool is_valid_index(int indx) const;

src/cell.cpp

@@ -1316,12 +1316,14 @@ struct ParentCellStack {
     visited_cells_;
 };
 
-vector<ParentCell> Cell::find_parent_cells(int32_t instance) const
+vector<ParentCell> Cell::find_parent_cells(
+  int32_t instance, Position* hint) const
 {
   ParentCellStack stack;
   // start with this cell's universe
   int32_t prev_univ_idx;
   int32_t univ_idx = this->universe_;
+  bool hint_used = false;
 
   while (true) {
     const auto& univ = model::universes[univ_idx];
@@ -1350,26 +1352,38 @@ vector<ParentCell> Cell::find_parent_cells(int32_t instance) const
         const auto& lattice = model::lattices[cell->fill_];
         const auto& lattice_univs = lattice->universes_;
 
-        // start search for universe
-        auto lat_it = lattice_univs.begin();
-        while (true) {
-          // find the next lattice cell with this universe
-          lat_it = std::find(lat_it, lattice_univs.end(), univ_idx);
-          if (lat_it == lattice_univs.end())
+        // if a hint position is provided, look up the lattice cell
+        // for that position
+        if (hint && !hint_used) {
+          std::array<int, 3> hint_idx;
+          lattice->get_indices(*hint, Direction(0.0, 0.0, 0.0), hint_idx);
+          int hint_offset = lattice->get_flat_index(hint_idx);
+          if (univ_idx == lattice_univs[hint_offset]) {
+            hint_used = true;
+            stack.push(univ_idx, {model::cell_map[cell->id_], hint_offset});
+          }
+        } else {
+          // start search for universe
+          auto lat_it = lattice_univs.begin();
+          while (true) {
+            // find the next lattice cell with this universe
+            lat_it = std::find(lat_it, lattice_univs.end(), univ_idx);
+            if (lat_it == lattice_univs.end())
+              break;
+
+            int lattice_idx = lat_it - lattice_univs.begin();
+
+            // move iterator forward one to avoid finding the same entry
+            lat_it++;
+            if (stack.visited(
+                  univ_idx, {model::cell_map[cell->id_], lattice_idx}))
+              continue;
+
+            // add this cell and lattice index to the stack and exit loop
+            stack.push(univ_idx, {model::cell_map[cell->id_], lattice_idx});
+            univ_idx = cell->universe_;
             break;
-
-          int lattice_idx = lat_it - lattice_univs.begin();
-
-          // move iterator forward one to avoid finding the same entry
-          lat_it++;
-          if (stack.visited(
-                univ_idx, {model::cell_map[cell->id_], lattice_idx}))
-            continue;
-
-          // add this cell and lattice index to the stack and exit loop
-          stack.push(univ_idx, {model::cell_map[cell->id_], lattice_idx});
-          univ_idx = cell->universe_;
-          break;
+          }
         }
       }
       // if we've updated the universe, break
@@ -1403,7 +1417,7 @@ vector<ParentCell> Cell::find_parent_cells(int32_t instance) const
 }
 
 std::unordered_map<int32_t, vector<int32_t>> Cell::get_contained_cells(
-  int32_t instance) const
+  int32_t instance, Position* hint) const
 {
   std::unordered_map<int32_t, vector<int32_t>> contained_cells;
 
@@ -1412,7 +1426,7 @@ std::unordered_map<int32_t, vector<int32_t>> Cell::get_contained_cells(
     return contained_cells;
 
   // find the pathway through the geometry to this cell
-  vector<ParentCell> parent_cells = this->find_parent_cells(instance);
+  vector<ParentCell> parent_cells = this->find_parent_cells(instance, hint);
 
   // if this cell is filled w/ a material, it contains no other cells
   if (type_ != Fill::MATERIAL) {

src/lattice.cpp

@@ -215,9 +215,7 @@ RectLattice::RectLattice(pugi::xml_node lat_node) : Lattice {lat_node}
 
 int32_t const& RectLattice::operator[](array<int, 3> const& i_xyz)
 {
-  int indx =
-    n_cells_[0] * n_cells_[1] * i_xyz[2] + n_cells_[0] * i_xyz[1] + i_xyz[0];
-  return universes_[indx];
+  return universes_[get_flat_index(i_xyz)];
 }
 
 //==============================================================================
@@ -323,6 +321,12 @@ void RectLattice::get_indices(
   }
 }
 
+int RectLattice::get_flat_index(const array<int, 3>& i_xyz) const
+{
+  return n_cells_[0] * n_cells_[1] * i_xyz[2] + n_cells_[0] * i_xyz[1] +
+         i_xyz[0];
+}
+
 //==============================================================================
 
 Position RectLattice::get_local_position(
@@ -662,9 +666,7 @@ void HexLattice::fill_lattice_y(const vector<std::string>& univ_words)
 
 int32_t const& HexLattice::operator[](array<int, 3> const& i_xyz)
 {
-  int indx = (2 * n_rings_ - 1) * (2 * n_rings_ - 1) * i_xyz[2] +
-             (2 * n_rings_ - 1) * i_xyz[1] + i_xyz[0];
-  return universes_[indx];
+  return universes_[get_flat_index(i_xyz)];
 }
 
 //==============================================================================
@@ -929,6 +931,12 @@ void HexLattice::get_indices(
   result[1] += i2_chg;
 }
 
+int HexLattice::get_flat_index(const array<int, 3>& i_xyz) const
+{
+  return (2 * n_rings_ - 1) * (2 * n_rings_ - 1) * i_xyz[2] +
+         (2 * n_rings_ - 1) * i_xyz[1] + i_xyz[0];
+}
+
 //==============================================================================
 
 Position HexLattice::get_local_position(