Merge pull request #1871 from pshriwise/find_parent_cells

Allow instance specification in`Cell::get_contained_cells`

include/openmc/cell.h

@@ -6,6 +6,7 @@
 #include <limits>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 
 #include "hdf5.h"
 #include "pugixml.hpp"
@@ -166,10 +167,18 @@ public:
   void set_name(const std::string& name) { name_ = name; };
 
   //! Get all cell instances contained by this cell
-  //! \return Map with cell indexes as keys and instances as values
-  std::unordered_map<int32_t, vector<int32_t>> get_contained_cells() const;
+  //! \param[in] instance Instance of the cell for which to get contained cells
+  //! (default instance is zero)
+  //! \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;
 
 protected:
+  //! Determine the path to this cell instance in the geometry hierarchy
+  vector<ParentCell> find_parent_cells(int32_t instance) const;
+
+  //! Inner function for retrieving contained cells
   void get_contained_cells_inner(
     std::unordered_map<int32_t, vector<int32_t>>& contained_cells,
     vector<ParentCell>& parent_cells) const;
@@ -298,19 +307,11 @@ private:
   vector<vector<int32_t>> partitions_;
 };
 
-//==============================================================================
-//! Define a containing (parent) cell
-//==============================================================================
-
-struct ParentCell {
-  gsl::index cell_index;
-  gsl::index lattice_index;
-};
-
 //==============================================================================
 //! Define an instance of a particular cell
 //==============================================================================
 
+//!  Stores information used to identify a unique cell in the model
 struct CellInstance {
   //! Check for equality
   bool operator==(const CellInstance& other) const
@@ -322,6 +323,8 @@ struct CellInstance {
   gsl::index instance;
 };
 
+//! Structure necessary for inserting CellInstance into hashed STL data
+//! structures
 struct CellInstanceHash {
   std::size_t operator()(const CellInstance& k) const
   {

src/cell.cpp

@@ -340,7 +340,7 @@ void Cell::set_temperature(double T, int32_t instance, bool set_contained)
           id_)};
     }
 
-    auto contained_cells = this->get_contained_cells();
+    auto contained_cells = this->get_contained_cells(instance);
     for (const auto& entry : contained_cells) {
       auto& cell = model::cells[entry.first];
       Expects(cell->type_ == Fill::MATERIAL);
@@ -1220,10 +1220,199 @@ extern "C" int openmc_cell_set_name(int32_t index, const char* name)
   return 0;
 }
 
-std::unordered_map<int32_t, vector<int32_t>> Cell::get_contained_cells() const
+//==============================================================================
+//! Define a containing (parent) cell
+//==============================================================================
+
+//! Used to locate a universe fill in the geometry
+struct ParentCell {
+  bool operator==(const ParentCell& other) const
+  {
+    return cell_index == other.cell_index &&
+           lattice_index == other.lattice_index;
+  }
+
+  bool operator<(const ParentCell& other) const
+  {
+    return cell_index < other.cell_index ||
+           (cell_index == other.cell_index &&
+             lattice_index < other.lattice_index);
+  }
+
+  gsl::index cell_index;
+  gsl::index lattice_index;
+};
+
+//! Structure used to insert ParentCell into hashed STL data structures
+struct ParentCellHash {
+  std::size_t operator()(const ParentCell& p) const
+  {
+    return 4096 * p.cell_index + p.lattice_index;
+  }
+};
+
+//! Used to manage a traversal stack when locating parent cells of a cell
+//! instance in the model
+struct ParentCellStack {
+
+  //! push method that adds to the parent_cells visited cells for this search
+  //! universe
+  void push(int32_t search_universe, const ParentCell& pc)
+  {
+    parent_cells_.push_back(pc);
+    // add parent cell to the set of cells we've visited for this search
+    // universe
+    visited_cells_[search_universe].insert(pc);
+  }
+
+  //! removes the last parent_cell and clears the visited cells for the popped
+  //! cell's universe
+  void pop()
+  {
+    visited_cells_[this->current_univ()].clear();
+    parent_cells_.pop_back();
+  }
+
+  //! checks whether or not the parent cell has been visited already for this
+  //! search universe
+  bool visited(int32_t search_universe, const ParentCell& parent_cell)
+  {
+    return visited_cells_[search_universe].count(parent_cell) != 0;
+  }
+
+  //! return the next universe to search for a parent cell
+  int32_t current_univ() const
+  {
+    return model::cells[parent_cells_.back().cell_index]->universe_;
+  }
+
+  //! indicates whether nor not parent cells are present on the stack
+  bool empty() const { return parent_cells_.empty(); }
+
+  //! compute an instance for the provided distribcell index
+  int32_t compute_instance(int32_t distribcell_index) const
+  {
+    int32_t instance = 0;
+    for (const auto& parent_cell : this->parent_cells_) {
+      auto& cell = model::cells[parent_cell.cell_index];
+      if (cell->type_ == Fill::UNIVERSE) {
+        instance += cell->offset_[distribcell_index];
+      } else if (cell->type_ == Fill::LATTICE) {
+        auto& lattice = model::lattices[cell->fill_];
+        instance +=
+          lattice->offset(distribcell_index, parent_cell.lattice_index);
+      }
+    }
+    return instance;
+  }
+
+  // Accessors
+  vector<ParentCell>& parent_cells() { return parent_cells_; }
+  const vector<ParentCell>& parent_cells() const { return parent_cells_; }
+
+  // Data Members
+  vector<ParentCell> parent_cells_;
+  std::unordered_map<int32_t, std::unordered_set<ParentCell, ParentCellHash>>
+    visited_cells_;
+};
+
+vector<ParentCell> Cell::find_parent_cells(int32_t instance) const
+{
+  ParentCellStack stack;
+  // start with this cell's universe
+  int32_t prev_univ_idx;
+  int32_t univ_idx = this->universe_;
+
+  while (true) {
+    const auto& univ = model::universes[univ_idx];
+    prev_univ_idx = univ_idx;
+
+    // search for a cell that is filled w/ this universe
+    for (const auto& cell : model::cells) {
+      // if this is a material-filled cell, move on
+      if (cell->type_ == Fill::MATERIAL)
+        continue;
+
+      if (cell->type_ == Fill::UNIVERSE) {
+        // if this is in the set of cells previously visited for this universe,
+        // move on
+        if (stack.visited(univ_idx, {model::cell_map[cell->id_], C_NONE}))
+          continue;
+
+        // if this cell contains the universe we're searching for, add it to the
+        // stack
+        if (cell->fill_ == univ_idx) {
+          stack.push(univ_idx, {model::cell_map[cell->id_], C_NONE});
+          univ_idx = cell->universe_;
+        }
+      } else if (cell->type_ == Fill::LATTICE) {
+        // retrieve the lattice and lattice universes
+        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())
+            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
+      if (prev_univ_idx != univ_idx)
+        break;
+    } // end cell loop search for universe
+
+    // if we're at the top of the geometry and the instance matches, we're done
+    if (univ_idx == model::root_universe &&
+        stack.compute_instance(this->distribcell_index_) == instance)
+      break;
+
+    // if there is no match on the original cell's universe, report an error
+    if (univ_idx == this->universe_) {
+      fatal_error(
+        fmt::format("Could not find the parent cells for cell {}, instance {}.",
+          this->id_, instance));
+    }
+
+    // if we don't find a suitable update, adjust the stack and continue
+    if (univ_idx == model::root_universe || univ_idx == prev_univ_idx) {
+      stack.pop();
+      univ_idx = stack.empty() ? this->universe_ : stack.current_univ();
+    }
+
+  } // end while
+
+  // reverse the stack so the highest cell comes first
+  std::reverse(stack.parent_cells().begin(), stack.parent_cells().end());
+  return stack.parent_cells();
+}
+
+std::unordered_map<int32_t, vector<int32_t>> Cell::get_contained_cells(
+  int32_t instance) const
 {
   std::unordered_map<int32_t, vector<int32_t>> contained_cells;
-  vector<ParentCell> parent_cells;
+
+  // if this is a material-filled cell it has no contained cells
+  if (this->type_ == Fill::MATERIAL)
+    return contained_cells;
+
+  // find the pathway through the geometry to this cell
+  vector<ParentCell> parent_cells = this->find_parent_cells(instance);
 
   // if this cell is filled w/ a material, it contains no other cells
   if (type_ != Fill::MATERIAL) {
@@ -1244,7 +1433,7 @@ void Cell::get_contained_cells_inner(
     int instance = 0;
     if (this->distribcell_index_ >= 0) {
       for (auto& parent_cell : parent_cells) {
-        auto& cell = openmc::model::cells[parent_cell.cell_index];
+        auto& cell = model::cells[parent_cell.cell_index];
         if (cell->type_ == Fill::UNIVERSE) {
           instance += cell->offset_[distribcell_index_];
         } else if (cell->type_ == Fill::LATTICE) {

tests/regression_tests/cpp_driver/driver.cpp

@@ -6,6 +6,7 @@
 #include "openmc/cell.h"
 #include "openmc/error.h"
 #include "openmc/geometry.h"
+#include "openmc/geometry_aux.h"
 #include "openmc/message_passing.h"
 #include "openmc/summary.h"
 #include "openmc/tallies/filter.h"
@@ -31,11 +32,14 @@ int main(int argc, char** argv) {
   for (auto& entry : openmc::model::cell_map) {
       cell_indices.push_back(entry.second);
   }
+  // enable distribcells offsets for all cells
+  prepare_distribcell(&cell_indices);
   // sort to make sure the cell bins appear in the same
   // order as the test relying on the openmc exe
   std::sort(cell_indices.begin(), cell_indices.end());
   cell_filter->set_cells(cell_indices);
 
+
   // create a new tally
   auto tally = Tally::create();
   std::vector<Filter*> filters = {cell_filter};
@@ -46,7 +50,7 @@ int main(int argc, char** argv) {
   // the lattice
   auto& root_univ = openmc::model::universes[openmc::model::root_universe];
   auto& lattice_cell = openmc::model::cells[root_univ->cells_[0]];
-  lattice_cell->set_temperature(300, 1, true);
+  lattice_cell->set_temperature(300.0, 0, true);
 
   // check that material-filled cells return no contained cells
   for (auto& cell : openmc::model::cells) {
@@ -56,6 +60,9 @@ int main(int argc, char** argv) {
     }
   }
 
+  // set a higher temperature for only one of the lattice cells (ID is 4 in the model)
+  model::cells[model::cell_map[4]]->set_temperature(400.0, 3, true);
+
   // the summary file will be used to check that
   // temperatures were set correctly so clear
   // error output can be provided

tests/regression_tests/cpp_driver/inputs_true.dat

@@ -3,14 +3,15 @@
   <cell id="1" material="1" region="-1" universe="1" />
   <cell id="2" material="2" region="1 -2" universe="1" />
   <cell id="3" material="3" region="2" universe="1" />
-  <cell fill="2" id="4" region="3 -4 5 -6" universe="3" />
-  <lattice id="2">
+  <cell fill="1" id="4" universe="2" />
+  <cell fill="3" id="5" region="3 -4 5 -6" universe="4" />
+  <lattice id="3">
     <pitch>4.0 4.0</pitch>
     <dimension>2 2</dimension>
     <lower_left>-4.0 -4.0</lower_left>
     <universes>
-1 1 
-1 1 </universes>
+2 2 
+2 2 </universes>
   </lattice>
   <surface coeffs="0.0 0.0 1.5" id="1" type="z-cylinder" />
   <surface coeffs="0.0 0.0 1.7" id="2" type="z-cylinder" />

tests/regression_tests/cpp_driver/results_true.dat

@@ -9,3 +9,5 @@ tally 1:
 9.919885E+02
 2.174849E+02
 5.292948E+03
+2.174849E+02
+5.292948E+03

tests/regression_tests/cpp_driver/test.py

@@ -81,11 +81,15 @@ def model():
 
     pincell_univ = openmc.Universe(cells=[fuel, cladding, moderator])
 
+    # insert an additional cell to add another level to the geometry
+    extra_cell = openmc.Cell(fill=pincell_univ)
+    extra_univ = openmc.Universe(cells=[extra_cell])
+
     # lattice
     lattice = openmc.RectLattice()
     lattice.pitch = (4.0, 4.0)
     lattice.lower_left = (-4.0, -4.0)
-    lattice.universes = [[pincell_univ, pincell_univ], [pincell_univ, pincell_univ]]
+    lattice.universes = [[extra_univ, extra_univ], [extra_univ, extra_univ]]
     lattice_region = openmc.model.rectangular_prism(8.0,
                                                     8.0,
                                                     boundary_type='reflective')
@@ -130,7 +134,8 @@ class ExternalDriverTestHarness(PyAPITestHarness):
         for cell in cells.values():
             if isinstance(cell.fill, openmc.Material):
                 assert len(cell.temperature) == 4
-                assert_allclose(cell.temperature, 300.0)
+                assert_allclose(cell.temperature[:3], 300.0)
+                assert_allclose(cell.temperature[-1:], 400.0)
 
 
 def test_cpp_driver(cpp_driver, model):