Fixed a bug in distribcell offsets logic (openmc-dev#3424)

Author: GuySten

include/openmc/cell.h

@@ -226,6 +226,8 @@ class Cell {
   void set_temperature(
     double T, int32_t instance = -1, bool set_contained = false);
 
+  int32_t n_instances() const;
+
   //! Set the rotation matrix of a cell instance
   //! \param[in] rot The rotation matrix of length 3 or 9
   void set_rotation(const vector<double>& rot);
@@ -312,12 +314,11 @@ class Cell {
   //----------------------------------------------------------------------------
   // Data members
 
-  int32_t id_;              //!< Unique ID
-  std::string name_;        //!< User-defined name
-  Fill type_;               //!< Material, universe, or lattice
-  int32_t universe_;        //!< Universe # this cell is in
-  int32_t fill_;            //!< Universe # filling this cell
-  int32_t n_instances_ {0}; //!< Number of instances of this cell
+  int32_t id_;       //!< Unique ID
+  std::string name_; //!< User-defined name
+  Fill type_;        //!< Material, universe, or lattice
+  int32_t universe_; //!< Universe # this cell is in
+  int32_t fill_;     //!< Universe # filling this cell
 
   //! \brief Index corresponding to this cell in distribcell arrays
   int distribcell_index_ {C_NONE};

include/openmc/geometry_aux.h

@@ -15,8 +15,6 @@
 namespace openmc {
 
 namespace model {
-extern std::unordered_map<int32_t, std::unordered_map<int32_t, int32_t>>
-  universe_cell_counts;
 extern std::unordered_map<int32_t, int32_t> universe_level_counts;
 } // namespace model
 
@@ -80,15 +78,13 @@ void prepare_distribcell(
   const std::vector<int32_t>* user_distribcells = nullptr);
 
 //==============================================================================
-//! Recursively search through the geometry and count cell instances.
+//! Recursively search through the geometry and count universe instances.
 //!
-//! This function will update the Cell::n_instances value for each cell in the
-//! geometry.
-//! \param univ_indx The index of the universe to begin searching from (probably
-//!   the root universe).
+//! This function will update Universe.n_instances_ for each
+//! universe in the geometry.
 //==============================================================================
 
-void count_cell_instances(int32_t univ_indx);
+void count_universe_instances();
 
 //==============================================================================
 //! Recursively search through universes and count universe instances.

include/openmc/lattice.h

@@ -76,7 +76,7 @@ class Lattice {
   }
 
   //! Populate the distribcell offset tables.
-  int32_t fill_offset_table(int32_t offset, int32_t target_univ_id, int map,
+  int32_t fill_offset_table(int32_t target_univ_id, int map,
     std::unordered_map<int32_t, int32_t>& univ_count_memo);
 
   //! \brief Check lattice indices.

include/openmc/universe.h

@@ -29,6 +29,7 @@ class Universe {
 public:
   int32_t id_;            //!< Unique ID
   vector<int32_t> cells_; //!< Cells within this universe
+  int32_t n_instances_;   //!< Number of instances of this universe
 
   //! \brief Write universe information to an HDF5 group.
   //! \param group_id An HDF5 group id.

src/cell.cpp

@@ -40,6 +40,11 @@ vector<unique_ptr<Cell>> cells;
 // Cell implementation
 //==============================================================================
 
+int32_t Cell::n_instances() const
+{
+  return model::universes[universe_]->n_instances_;
+}
+
 void Cell::set_rotation(const vector<double>& rot)
 {
   if (fill_ == C_NONE) {
@@ -115,8 +120,8 @@ void Cell::set_temperature(double T, int32_t instance, bool set_contained)
   if (type_ == Fill::MATERIAL) {
     if (instance >= 0) {
       // If temperature vector is not big enough, resize it first
-      if (sqrtkT_.size() != n_instances_)
-        sqrtkT_.resize(n_instances_, sqrtkT_[0]);
+      if (sqrtkT_.size() != n_instances())
+        sqrtkT_.resize(n_instances(), sqrtkT_[0]);
 
       // Set temperature for the corresponding instance
       sqrtkT_.at(instance) = std::sqrt(K_BOLTZMANN * T);
@@ -170,7 +175,7 @@ void Cell::import_properties_hdf5(hid_t group)
 
   // Ensure number of temperatures makes sense
   auto n_temps = temps.size();
-  if (n_temps > 1 && n_temps != n_instances_) {
+  if (n_temps > 1 && n_temps != n_instances()) {
     throw std::runtime_error(fmt::format(
       "Number of temperatures for cell {} doesn't match number of instances",
       id_));
@@ -1618,7 +1623,7 @@ extern "C" int openmc_cell_get_num_instances(
     set_errmsg("Index in cells array is out of bounds.");
     return OPENMC_E_OUT_OF_BOUNDS;
   }
-  *num_instances = model::cells[index]->n_instances_;
+  *num_instances = model::cells[index]->n_instances();
   return 0;
 }
 

src/finalize.cpp

@@ -185,7 +185,6 @@ int openmc_finalize()
 int openmc_reset()
 {
 
-  model::universe_cell_counts.clear();
   model::universe_level_counts.clear();
 
   for (auto& t : model::tallies) {

src/geometry_aux.cpp

@@ -25,21 +25,9 @@
 namespace openmc {
 
 namespace model {
-std::unordered_map<int32_t, std::unordered_map<int32_t, int32_t>>
-  universe_cell_counts;
 std::unordered_map<int32_t, int32_t> universe_level_counts;
 } // namespace model
 
-// adds the cell counts of universe b to universe a
-void update_universe_cell_count(int32_t a, int32_t b)
-{
-  auto& universe_a_counts = model::universe_cell_counts[a];
-  const auto& universe_b_counts = model::universe_cell_counts[b];
-  for (const auto& it : universe_b_counts) {
-    universe_a_counts[it.first] += it.second;
-  }
-}
-
 void read_geometry_xml()
 {
   // Display output message
@@ -263,7 +251,7 @@ void finalize_geometry()
 {
   // Perform some final operations to set up the geometry
   adjust_indices();
-  count_cell_instances(model::root_universe);
+  count_universe_instances();
   partition_universes();
 
   // Assign temperatures to cells that don't have temperatures already assigned
@@ -356,35 +344,38 @@ void prepare_distribcell(const std::vector<int32_t>* user_distribcells)
     Cell& c {*model::cells[i]};
 
     if (c.material_.size() > 1) {
-      if (c.material_.size() != c.n_instances_) {
+      if (c.material_.size() != c.n_instances()) {
         fatal_error(fmt::format(
           "Cell {} was specified with {} materials but has {} distributed "
           "instances. The number of materials must equal one or the number "
           "of instances.",
-          c.id_, c.material_.size(), c.n_instances_));
+          c.id_, c.material_.size(), c.n_instances()));
       }
     }
 
     if (c.sqrtkT_.size() > 1) {
-      if (c.sqrtkT_.size() != c.n_instances_) {
+      if (c.sqrtkT_.size() != c.n_instances()) {
         fatal_error(fmt::format(
           "Cell {} was specified with {} temperatures but has {} distributed "
           "instances. The number of temperatures must equal one or the number "
           "of instances.",
-          c.id_, c.sqrtkT_.size(), c.n_instances_));
+          c.id_, c.sqrtkT_.size(), c.n_instances()));
       }
     }
   }
 
   // Search through universes for material cells and assign each one a
-  // unique distribcell array index.
-  int distribcell_index = 0;
+  // distribcell array index according to the containing universe.
   vector<int32_t> target_univ_ids;
   for (const auto& u : model::universes) {
     for (auto idx : u->cells_) {
       if (distribcells.find(idx) != distribcells.end()) {
-        model::cells[idx]->distribcell_index_ = distribcell_index++;
-        target_univ_ids.push_back(u->id_);
+        if (!contains(target_univ_ids, u->id_)) {
+          target_univ_ids.push_back(u->id_);
+        }
+        model::cells[idx]->distribcell_index_ =
+          std::find(target_univ_ids.begin(), target_univ_ids.end(), u->id_) -
+          target_univ_ids.begin();
       }
     }
   }
@@ -419,8 +410,7 @@ void prepare_distribcell(const std::vector<int32_t>* user_distribcells)
         } else if (c.type_ == Fill::LATTICE) {
           c.offset_[map] = offset;
           Lattice& lat = *model::lattices[c.fill_];
-          offset +=
-            lat.fill_offset_table(offset, target_univ_id, map, univ_count_memo);
+          offset += lat.fill_offset_table(target_univ_id, map, univ_count_memo);
         }
       }
     }
@@ -429,32 +419,12 @@ void prepare_distribcell(const std::vector<int32_t>* user_distribcells)
 
 //==============================================================================
 
-void count_cell_instances(int32_t univ_indx)
+void count_universe_instances()
 {
-  const auto univ_counts = model::universe_cell_counts.find(univ_indx);
-  if (univ_counts != model::universe_cell_counts.end()) {
-    for (const auto& it : univ_counts->second) {
-      model::cells[it.first]->n_instances_ += it.second;
-    }
-  } else {
-    for (int32_t cell_indx : model::universes[univ_indx]->cells_) {
-      Cell& c = *model::cells[cell_indx];
-      ++c.n_instances_;
-      model::universe_cell_counts[univ_indx][cell_indx] += 1;
-
-      if (c.type_ == Fill::UNIVERSE) {
-        // This cell contains another universe.  Recurse into that universe.
-        count_cell_instances(c.fill_);
-        update_universe_cell_count(univ_indx, c.fill_);
-      } else if (c.type_ == Fill::LATTICE) {
-        // This cell contains a lattice.  Recurse into the lattice universes.
-        Lattice& lat = *model::lattices[c.fill_];
-        for (auto it = lat.begin(); it != lat.end(); ++it) {
-          count_cell_instances(*it);
-          update_universe_cell_count(univ_indx, *it);
-        }
-      }
-    }
+  for (auto& univ : model::universes) {
+    std::unordered_map<int32_t, int32_t> univ_count_memo;
+    univ->n_instances_ = count_universe_instances(
+      model::root_universe, univ->id_, univ_count_memo);
   }
 }
 
@@ -528,19 +498,16 @@ std::string distribcell_path_inner(int32_t target_cell, int32_t map,
 
     // Material cells don't contain other cells so ignore them.
     if (c.type_ != Fill::MATERIAL) {
-      int32_t temp_offset;
-      if (c.type_ == Fill::UNIVERSE) {
-        temp_offset =
-          offset + c.offset_[map]; // TODO: should also apply to lattice fills?
-      } else {
+      int32_t temp_offset = offset + c.offset_[map];
+      if (c.type_ == Fill::LATTICE) {
         Lattice& lat = *model::lattices[c.fill_];
         int32_t indx = lat.universes_.size() * map + lat.begin().indx_;
-        temp_offset = offset + lat.offsets_[indx];
+        temp_offset += lat.offsets_[indx];
       }
 
       // The desired cell is the first cell that gives an offset smaller or
       // equal to the target offset.
-      if (temp_offset <= target_offset - c.offset_[map])
+      if (temp_offset <= target_offset)
         break;
     }
   }
@@ -570,12 +537,12 @@ std::string distribcell_path_inner(int32_t target_cell, int32_t map,
     path << "l" << lat.id_;
     for (ReverseLatticeIter it = lat.rbegin(); it != lat.rend(); ++it) {
       int32_t indx = lat.universes_.size() * map + it.indx_;
-      int32_t temp_offset = offset + lat.offsets_[indx];
-      if (temp_offset <= target_offset - c.offset_[map]) {
+      int32_t temp_offset = offset + lat.offsets_[indx] + c.offset_[map];
+      if (temp_offset <= target_offset) {
         offset = temp_offset;
         path << "(" << lat.index_to_string(it.indx_) << ")->";
-        path << distribcell_path_inner(target_cell, map, target_offset,
-          *model::universes[*it], offset + c.offset_[map]);
+        path << distribcell_path_inner(
+          target_cell, map, target_offset, *model::universes[*it], offset);
         return path.str();
       }
     }

src/lattice.cpp

@@ -103,8 +103,8 @@ void Lattice::adjust_indices()
 
 //==============================================================================
 
-int32_t Lattice::fill_offset_table(int32_t offset, int32_t target_univ_id,
-  int map, std::unordered_map<int32_t, int32_t>& univ_count_memo)
+int32_t Lattice::fill_offset_table(int32_t target_univ_id, int map,
+  std::unordered_map<int32_t, int32_t>& univ_count_memo)
 {
   // If the offsets have already been determined for this "map", don't bother
   // recalculating all of them and just return the total offset. Note that the
@@ -119,6 +119,7 @@ int32_t Lattice::fill_offset_table(int32_t offset, int32_t target_univ_id,
            count_universe_instances(last_univ, target_univ_id, univ_count_memo);
   }
 
+  int32_t offset = 0;
   for (LatticeIter it = begin(); it != end(); ++it) {
     offsets_[map * universes_.size() + it.indx_] = offset;
     offset += count_universe_instances(*it, target_univ_id, univ_count_memo);

src/random_ray/flat_source_domain.cpp

@@ -46,7 +46,7 @@ FlatSourceDomain::FlatSourceDomain() : negroups_(data::mg.num_energy_groups_)
       source_region_offsets_.push_back(-1);
     } else {
       source_region_offsets_.push_back(base_source_regions);
-      base_source_regions += c->n_instances_;
+      base_source_regions += c->n_instances();
     }
   }
 
@@ -61,7 +61,7 @@ FlatSourceDomain::FlatSourceDomain() : negroups_(data::mg.num_energy_groups_)
   for (int i = 0; i < model::cells.size(); i++) {
     Cell& cell = *model::cells[i];
     if (cell.type_ == Fill::MATERIAL) {
-      for (int j = 0; j < cell.n_instances_; j++) {
+      for (int j = 0; j < cell.n_instances(); j++) {
         source_regions_.material(source_region_id++) = cell.material(j);
       }
     }
@@ -1014,7 +1014,7 @@ void FlatSourceDomain::apply_external_source_to_cell_and_children(
   Cell& cell = *model::cells[i_cell];
 
   if (cell.type_ == Fill::MATERIAL) {
-    vector<int> instances(cell.n_instances_);
+    vector<int> instances(cell.n_instances());
     std::iota(instances.begin(), instances.end(), 0);
     apply_external_source_to_cell_instances(
       i_cell, discrete, strength_factor, target_material_id, instances);
@@ -1343,12 +1343,12 @@ void FlatSourceDomain::apply_mesh_to_cell_and_children(int32_t i_cell,
   Cell& cell = *model::cells[i_cell];
 
   if (cell.type_ == Fill::MATERIAL) {
-    vector<int> instances(cell.n_instances_);
+    vector<int> instances(cell.n_instances());
     std::iota(instances.begin(), instances.end(), 0);
     apply_mesh_to_cell_instances(
       i_cell, mesh_idx, target_material_id, instances, is_target_void);
   } else if (target_material_id == C_NONE && !is_target_void) {
-    for (int j = 0; j < cell.n_instances_; j++) {
+    for (int j = 0; j < cell.n_instances(); j++) {
       std::unordered_map<int32_t, vector<int32_t>> cell_instance_list =
         cell.get_contained_cells(j, nullptr);
       for (const auto& pair : cell_instance_list) {

src/tallies/filter_distribcell.cpp

@@ -34,7 +34,7 @@ void DistribcellFilter::set_cell(int32_t cell)
   assert(cell >= 0);
   assert(cell < model::cells.size());
   cell_ = cell;
-  n_bins_ = model::cells[cell]->n_instances_;
+  n_bins_ = model::cells[cell]->n_instances();
 }
 
 void DistribcellFilter::get_all_bins(