Merge pull request #1966 from pshriwise/contained_cells_hint

Acceleration for Cell::find_parent_cells.

Author: paulromano

include/openmc/cell.h

@@ -169,14 +169,32 @@ class Cell {
   //! Get all cell instances contained by this cell
   //! \param[in] instance Instance of the cell for which to get contained cells
   //! (default instance is zero)
+  //! \param[in] hint positional hint for determining the parent cells
   //! \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;
+  //! \param[in] instance of the cell to find parent cells for
+  //! \param[in] r position used to do a fast search for parent cells
+  //! \return parent cells
+  vector<ParentCell> find_parent_cells(
+    int32_t instance, const Position& r) const;
+
+  //! Determine the path to this cell instance in the geometry hierarchy
+  //! \param[in] instance of the cell to find parent cells for
+  //! \param[in] p particle used to do a fast search for parent cells
+  //! \return parent cells
+  vector<ParentCell> find_parent_cells(
+    int32_t instance, Particle& p) const;
+
+  //! Determine the path to this cell instance in the geometry hierarchy
+  //! \param[in] instance of the cell to find parent cells for
+  //! \return parent cells
+  vector<ParentCell> exhaustive_find_parent_cells(
+    int32_t instance) 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;

include/openmc/particle_data.h

@@ -341,6 +341,8 @@ class ParticleData {
   const int& cell_instance() const { return cell_instance_; }
   LocalCoord& coord(int i) { return coord_[i]; }
   const LocalCoord& coord(int i) const { return coord_[i]; }
+  const vector<LocalCoord>& coord() const { return coord_; }
+
 
   int& n_coord_last() { return n_coord_last_; }
   const int& n_coord_last() const { return n_coord_last_; }

src/cell.cpp

@@ -1316,7 +1316,59 @@ struct ParentCellStack {
     visited_cells_;
 };
 
-vector<ParentCell> Cell::find_parent_cells(int32_t instance) const
+vector<ParentCell> Cell::find_parent_cells(
+  int32_t instance, const Position& r) const {
+
+  // create a temporary particle
+  Particle dummy_particle {};
+  dummy_particle.r() = r;
+  dummy_particle.u() = {0., 0., 1.};
+
+  return find_parent_cells(instance, dummy_particle);
+}
+
+vector<ParentCell> Cell::find_parent_cells(
+  int32_t instance, Particle& p) const {
+  // look up the particle's location
+  exhaustive_find_cell(p);
+  const auto& coords = p.coord();
+
+  // build a parent cell stack from the particle coordinates
+  ParentCellStack stack;
+  bool cell_found = false;
+  for (auto it = coords.begin(); it != coords.end(); it++) {
+    const auto& coord = *it;
+    const auto& cell = model::cells[coord.cell];
+    // if the cell at this level matches the current cell, stop adding to the stack
+    if (coord.cell == model::cell_map[this->id_]) {
+      cell_found = true;
+      break;
+    }
+
+    // if filled with a lattice, get the lattice index from the next
+    // level in the coordinates to push to the stack
+    int lattice_idx = C_NONE;
+    if (cell->type_ == Fill::LATTICE) {
+      const auto& next_coord = *(it + 1);
+      lattice_idx = model::lattices[next_coord.lattice]->get_flat_index(next_coord.lattice_i);
+    }
+    stack.push(coord.universe, {coord.cell, lattice_idx});
+  }
+
+  // if this loop finished because the cell was found and
+  // the instance matches the one requested in the call
+  // we have the correct path and can return the stack
+  if (cell_found && stack.compute_instance(this->distribcell_index_) == instance) {
+    return stack.parent_cells();
+  }
+
+  // fall back on an exhaustive search for the cell's parents
+  return exhaustive_find_parent_cells(instance);
+}
+
+
+vector<ParentCell> Cell::exhaustive_find_parent_cells(
+  int32_t instance) const
 {
   ParentCellStack stack;
   // start with this cell's universe
@@ -1403,7 +1455,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 +1464,12 @@ 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;
+
+  // if a positional hint is provided, attempt to do a fast lookup
+  // of the parent cells
+  parent_cells = hint ? find_parent_cells(instance, *hint)
+                      : exhaustive_find_parent_cells(instance);
 
   // 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(