Make distribcell setup scale with the number of universes

include/openmc/cell.h

@@ -13,6 +13,7 @@
 
 #include "openmc/bounding_box.h"
 #include "openmc/constants.h"
+#include "openmc/distribcell_offsets.h"
 #include "openmc/memory.h" // for unique_ptr
 #include "openmc/neighbor_list.h"
 #include "openmc/position.h"
@@ -410,7 +411,7 @@ class Cell {
   //! also present at the end of the vector, making it of length 12.
   vector<double> rotation_;
 
-  vector<int32_t> offset_; //!< Distribcell offset table
+  DistribcellOffsets offset_; //!< Distribcell offsets
 
   // Right now, either CSG or DAGMC cells are used.
   virtual GeometryType geom_type() const = 0;

include/openmc/distribcell_offsets.h

@@ -0,0 +1,54 @@
+#ifndef OPENMC_DISTRIBCELL_OFFSETS_H
+#define OPENMC_DISTRIBCELL_OFFSETS_H
+
+#include <cstdint>
+
+#include "openmc/vector.h"
+
+namespace openmc {
+
+//==============================================================================
+//! Distributed cell offsets of a fill cell or lattice tile.
+//
+//! A map numbers the instances of a universe containing distributed cells. The
+//! offset for a map is the number of instances of the map's universe before
+//! the fill cell in its universe, or before the tile in its lattice. Offsets
+//! are only read for maps whose universe is in the cell's fill or the tile, so
+//! only those are stored, in the order of the sorted maps of the fill or tile.
+//==============================================================================
+
+class DistribcellOffsets {
+public:
+  DistribcellOffsets() = default;
+
+  //! \param maps Sorted maps in the fill or tile
+  //! \param values Offset for each map
+  //! Both must outlive this object.
+  DistribcellOffsets(const vector<int32_t>& maps, const int32_t* values);
+
+  //! \return Offset for a map, or 0 if the map's universe isn't contained
+  int32_t operator[](int32_t map) const
+  {
+    // Maps are numbered so that the maps in a universe are mostly consecutive,
+    // so first look in the consecutive maps at the start of the list
+    auto d = static_cast<uint32_t>(map - first_map_);
+    if (d < static_cast<uint32_t>(n_first_))
+      return values_[d];
+    return find(map);
+  }
+
+  //! \return Whether the map's universe is contained
+  bool contains(int32_t map) const;
+
+private:
+  int32_t find(int32_t map) const;
+
+  const int32_t* values_ {nullptr};
+  const vector<int32_t>* maps_ {nullptr};
+  int32_t first_map_ {0}; //!< First map
+  int32_t n_first_ {0};   //!< Number of consecutive maps from the first
+};
+
+} // namespace openmc
+
+#endif // OPENMC_DISTRIBCELL_OFFSETS_H

include/openmc/geometry_aux.h

@@ -84,28 +84,15 @@ void prepare_distribcell(
   const std::vector<int32_t>* user_distribcells = nullptr);
 
 //==============================================================================
-//! Recursively search through the geometry and count universe instances.
+//! Count the number of instances of every universe in the geometry.
 //!
-//! This function will update Universe.n_instances_ for each
-//! universe in the geometry.
+//! This function will update Universe.n_instances_ for each universe in the
+//! geometry. The universes are visited once each, in an order where every
+//! universe comes after the universes containing it.
 //==============================================================================
 
 void count_universe_instances();
 
-//==============================================================================
-//! Recursively search through universes and count universe instances.
-//! \param search_univ The index of the universe to begin searching from.
-//! \param target_univ_id The ID of the universe to be counted.
-//! \param univ_count_memo Memoized counts that make this function faster for
-//!   large systems.  The first call to this function for each target_univ_id
-//!   should start with an empty memo.
-//! \return The number of instances of target_univ_id in the geometry tree under
-//!   search_univ.
-//==============================================================================
-
-int count_universe_instances(int32_t search_univ, int32_t target_univ_id,
-  std::unordered_map<int32_t, int32_t>& univ_count_memo);
-
 //==============================================================================
 //! Build a character array representing the path to a distribcell instance.
 //! \param target_cell The index of the Cell in the global Cell array.

include/openmc/lattice.h

@@ -10,6 +10,7 @@
 
 #include "openmc/array.h"
 #include "openmc/constants.h"
+#include "openmc/distribcell_offsets.h"
 #include "openmc/memory.h"
 #include "openmc/position.h"
 #include "openmc/vector.h"
@@ -50,7 +51,13 @@ class Lattice {
   LatticeType type_;
   vector<int32_t> universes_;         //!< Universes filling each lattice tile
   int32_t outer_ {NO_OUTER_UNIVERSE}; //!< Universe tiled outside the lattice
-  vector<int32_t> offsets_;           //!< Distribcell offset table
+
+  //! Sorted distributed cell maps of the universes in the tiles and the outer
+  //! universe
+  vector<int32_t> distribcell_maps_;
+
+  vector<int32_t> offset_values_;      //!< Distribcell offsets of all tiles
+  vector<DistribcellOffsets> offsets_; //!< Distribcell offsets of each tile
 
   explicit Lattice(pugi::xml_node lat_node);
 
@@ -68,17 +75,6 @@ class Lattice {
   //! Convert internal universe values from IDs to indices using universe_map.
   void adjust_indices();
 
-  //! Allocate offset table for distribcell.
-  void allocate_offset_table(int n_maps)
-  {
-    offsets_.resize(n_maps * universes_.size());
-    std::fill(offsets_.begin(), offsets_.end(), C_NONE);
-  }
-
-  //! Populate the distribcell offset tables.
-  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.
   //! \param i_xyz[3] The indices for a lattice tile.
   //! \return true if the given indices fit within the lattice bounds.  False
@@ -135,14 +131,17 @@ class Lattice {
   //! \param i_xyz[3] The indices for a lattice tile.
   //! \return Distribcell offset i.e. the largest instance number for the target
   //!  cell found in the geometry tree under this lattice tile.
-  virtual int32_t& offset(int map, const array<int, 3>& i_xyz) = 0;
+  int32_t offset(int map, const array<int, 3>& i_xyz) const
+  {
+    return offset(map, get_flat_index(i_xyz));
+  }
 
   //! \brief Get the distribcell offset for a lattice tile.
   //! \param The map index for the target cell.
   //! \param indx The index for a lattice tile.
   //! \return Distribcell offset i.e. the largest instance number for the target
   //!  cell found in the geometry tree for this lattice index.
-  virtual int32_t offset(int map, int indx) const = 0;
+  int32_t offset(int map, int indx) const { return offsets_[indx][map]; }
 
   //! \brief Convert an array index to a useful human-readable string.
   //! \param indx The index for a lattice tile.
@@ -234,10 +233,6 @@ class RectLattice : public Lattice {
   Direction get_normal(
     const array<int, 3>& i_xyz, bool& is_valid) const override;
 
-  int32_t& offset(int map, const array<int, 3>& i_xyz) override;
-
-  int32_t offset(int map, int indx) const override;
-
   std::string index_to_string(int indx) const override;
 
   void to_hdf5_inner(hid_t group_id) const override;
@@ -284,10 +279,6 @@ class HexLattice : public Lattice {
 
   bool is_valid_index(int indx) const override;
 
-  int32_t& offset(int map, const array<int, 3>& i_xyz) override;
-
-  int32_t offset(int map, int indx) const override;
-
   std::string index_to_string(int indx) const override;
 
   void to_hdf5_inner(hid_t group_id) const override;

include/openmc/universe.h

@@ -31,6 +31,13 @@ class Universe {
   vector<int32_t> cells_; //!< Cells within this universe
   int32_t n_instances_;   //!< Number of instances of this universe
 
+  //! Sorted distributed cell maps of the universes in this universe,
+  //! including itself
+  vector<int32_t> distribcell_maps_;
+
+  //! Distributed cell offsets of the fill cells in this universe
+  vector<int32_t> offset_values_;
+
   //! \brief Write universe information to an HDF5 group.
   //! \param group_id An HDF5 group id.
   virtual void to_hdf5(hid_t group_id) const;