Optimize find_cell on universes with many z-planes

include/openmc/cell.h

@@ -43,6 +43,7 @@ constexpr int32_t OP_UNION        {std::numeric_limits<int32_t>::max() - 4};
 
 class Cell;
 class Universe;
+class UniversePartitioner;
 
 namespace model {
   extern std::vector<std::unique_ptr<Cell>> cells;
@@ -65,6 +66,8 @@ public:
   //! \brief Write universe information to an HDF5 group.
   //! \param group_id An HDF5 group id.
   void to_hdf5(hid_t group_id) const;
+
+  std::unique_ptr<UniversePartitioner> partitioner_;
 };
 
 //==============================================================================
@@ -193,6 +196,42 @@ public:
 };
 #endif
 
+//==============================================================================
+
+class UniversePartitioner
+{
+public:
+  UniversePartitioner(const Universe& univ);
+
+  const std::vector<int32_t>& get_cells(Position r, Direction u) const
+  {return get_cells(r, u, 0, surfs_.size()-1, (surfs_.size()-1)/2);}
+
+  const std::vector<int32_t>& get_cells(Position r, Direction u,
+    int left, int right, int middle) const
+  {
+    auto i_surf = surfs_[middle];
+    const auto& surf = *model::surfaces[i_surf];
+    if (surf.sense(r, u)) {
+      int right_leaf = right - (right - middle) / 2;
+      if (right_leaf != middle) {
+        return get_cells(r, u, middle+1, right, right_leaf);
+      } else {
+        return cells_[middle+1];
+      }
+    } else {
+      int left_leaf = left + (middle - left) / 2;
+      if (left_leaf != middle) {
+        return get_cells(r, u, left, middle-1, left_leaf);
+      } else {
+        return cells_[middle];
+      }
+    }
+  }
+
+  std::vector<std::vector<int32_t>> cells_;
+  std::vector<int32_t> surfs_;
+};
+
 //==============================================================================
 // Non-member functions
 //==============================================================================

include/openmc/geometry_aux.h

@@ -10,6 +10,11 @@
 
 namespace openmc {
 
+extern int search_numerator1;
+extern int search_denominator1;
+extern int search_numerator2;
+extern int search_denominator2;
+
 void read_geometry_xml();
 
 //==============================================================================

include/openmc/surface.h

@@ -220,8 +220,8 @@ public:
 
 class SurfaceZPlane : public PeriodicSurface
 {
-  double z0_;
 public:
+  double z0_;
   explicit SurfaceZPlane(pugi::xml_node surf_node);
   double evaluate(Position r) const;
   double distance(Position r, Direction u, bool coincident) const;

src/cell.cpp

@@ -635,6 +635,84 @@ void DAGCell::to_hdf5(hid_t group_id) const { return; }
 
 #endif
 
+//==============================================================================
+//==============================================================================
+
+UniversePartitioner::UniversePartitioner(const Universe& univ)
+{
+  std::unordered_set<int32_t> surf_set;
+  for (auto i_cell : univ.cells_) {
+    for (auto token : model::cells[i_cell]->rpn_) {
+      if (token < OP_UNION) {
+        auto i_surf = std::abs(token) - 1;
+        const auto* surf = model::surfaces[i_surf].get();
+        if (const auto* zplane = dynamic_cast<const SurfaceZPlane*>(surf)) {
+          if (surf_set.find(i_surf) == surf_set.end()) {
+            surf_set.insert(i_surf);
+            surfs_.push_back(i_surf);
+          }
+        }
+      }
+    }
+  }
+
+  struct {
+    bool operator()(int32_t i_surf, int32_t j_surf) const
+    {
+      const auto* surf = model::surfaces[i_surf].get();
+      const auto* zplane = dynamic_cast<const SurfaceZPlane*>(surf);
+      double zi = zplane->z0_;
+      surf = model::surfaces[j_surf].get();
+      zplane = dynamic_cast<const SurfaceZPlane*>(surf);
+      double zj = zplane->z0_;
+      return zi < zj;
+    }
+  } compare_surfs;
+  std::sort(surfs_.begin(), surfs_.end(), compare_surfs);
+
+  cells_.resize(surfs_.size() + 1);
+
+  for (auto i_cell : univ.cells_) {
+    int32_t min_token = 0, max_token = 0;
+    double min_z, max_z;
+    for (auto token : model::cells[i_cell]->rpn_) {
+      if (token < OP_UNION) {
+        const auto* surf = model::surfaces[std::abs(token) - 1].get();
+        if (const auto* zplane = dynamic_cast<const SurfaceZPlane*>(surf)) {
+          if (min_token == 0 || zplane->z0_ < min_z) {
+            min_token = token;
+            min_z = zplane->z0_;
+          }
+          if (max_token == 0 || zplane->z0_ > max_z) {
+            max_token = token;
+            max_z = zplane->z0_;
+          }
+        }
+      }
+    }
+
+    if (min_token == 0) {
+      min_token = -(surfs_.front() + 1);
+      max_token = surfs_.back() + 1;
+    }
+
+    bool in_partition = min_token < 0;
+
+    for (auto i = 0; i < surfs_.size(); ++i) {
+      if (in_partition) {
+        cells_[i].push_back(i_cell);
+
+        if (max_token == -(surfs_[i] + 1))
+          in_partition = false;
+      } else {
+        in_partition = min_token == surfs_[i] + 1;
+      }
+    }
+    if (in_partition)
+      cells_.back().push_back(i_cell);
+  }
+}
+
 //==============================================================================
 // Non-method functions
 //==============================================================================

src/geometry.cpp

@@ -11,6 +11,8 @@
 #include "openmc/simulation.h"
 #include "openmc/surface.h"
 
+#include "openmc/geometry_aux.h"
+
 
 namespace openmc {
 
@@ -69,6 +71,7 @@ find_cell_inner(Particle* p, const NeighborList* neighbor_list)
   bool found = false;
   int32_t i_cell;
   if (neighbor_list) {
+    ++search_numerator1;
     for (auto it = neighbor_list->cbegin(); it != neighbor_list->cend(); ++it) {
       i_cell = *it;
 
@@ -80,6 +83,7 @@ find_cell_inner(Particle* p, const NeighborList* neighbor_list)
       Position r {p->r_local()};
       Direction u {p->u_local()};
       auto surf = p->surface_;
+      ++search_denominator1;
       if (model::cells[i_cell]->contains(r, u, surf)) {
         p->coord_[p->n_coord_-1].cell = i_cell;
         found = true;
@@ -88,8 +92,15 @@ find_cell_inner(Particle* p, const NeighborList* neighbor_list)
     }
 
   } else {
+    ++search_numerator2;
     int i_universe = p->coord_[p->n_coord_-1].universe;
-    const auto& cells {model::universes[i_universe]->cells_};
+    //const auto& cells {model::universes[i_universe]->cells_};
+    const auto& univ {*model::universes[i_universe]};
+    const auto& cells {
+      !univ.partitioner_ ?
+      model::universes[i_universe]->cells_
+      : univ.partitioner_->get_cells(p->r_local(), p->u_local())
+    };
     for (auto it = cells.cbegin(); it != cells.cend(); it++) {
       i_cell = *it;
 
@@ -101,6 +112,7 @@ find_cell_inner(Particle* p, const NeighborList* neighbor_list)
       Position r {p->r_local()};
       Direction u {p->u_local()};
       auto surf = p->surface_;
+      ++search_denominator2;
       if (model::cells[i_cell]->contains(r, u, surf)) {
         p->coord_[p->n_coord_-1].cell = i_cell;
         found = true;

src/geometry_aux.cpp

@@ -23,6 +23,11 @@
 
 namespace openmc {
 
+int search_denominator1 {0};
+int search_numerator1 {0};
+int search_denominator2 {0};
+int search_numerator2 {0};
+
 void read_geometry_xml()
 {
 #ifdef DAGMC
@@ -219,6 +224,28 @@ void finalize_geometry(std::vector<std::vector<double>>& nuc_temps,
 
   // Determine number of nested coordinate levels in the geometry
   model::n_coord_levels = maximum_levels(model::root_universe);
+
+  for (const auto& univ : model::universes) {
+    if (univ->cells_.size() > 10) {
+      std::unordered_set<int32_t> surf_inds;
+      for (auto i_cell : univ->cells_) {
+        for (auto token : model::cells[i_cell]->rpn_) {
+          if (token < OP_UNION) surf_inds.insert(std::abs(token) - 1);
+        }
+      }
+      int n_zplanes = 0;
+      for (auto i_surf : surf_inds) {
+        if (dynamic_cast<const SurfaceZPlane*>(model::surfaces[i_surf].get())) {
+          ++n_zplanes;
+          if (n_zplanes > 5) {
+            univ->partitioner_ = std::make_unique<UniversePartitioner>(
+              UniversePartitioner(*univ));
+            break;
+          }
+        }
+      }
+    }
+  }
 }
 
 //==============================================================================
@@ -519,6 +546,8 @@ maximum_levels(int32_t univ)
 void
 free_memory_geometry()
 {
+  std::cout << "SEARCH EFFICIENCY = " << static_cast<double>(search_numerator1) / search_denominator1 << "\n";
+  std::cout << "SEARCH EFFICIENCY = " << static_cast<double>(search_numerator2) / search_denominator2 << "\n";
   model::cells.clear();
   model::cell_map.clear();