Adds support for virtual lattice geometry handling

Introduces functionality to calculate distances and locate cells within virtual lattice geometries. Refactors existing geometry handling logic to reduce code duplication and improve maintainability.

Key updates include:
- Addition of `distance_in_virtual_lattice` and `find_cell_in_virtual_lattice` methods for specialized handling.
- Integration of virtual lattice checks in standard cell distance computations.
- Refactoring of repetitive code in surface crossing handlers for improved clarity.

Enhances support for TRISO particle geometries within virtual lattices, ensuring accurate material and cell identification during particle transport.

Author: skywalker_cn

include/openmc/cell.h

@@ -381,6 +381,9 @@ class CSGCell : public Cell {
   std::pair<double, int32_t> distance(Position r, Direction u,
     int32_t on_surface, GeometryState* p) const override;
 
+  std::pair<double, int32_t> distance_in_virtual_lattice(
+    Position r, Direction u, int32_t on_surface, GeometryState* p) const;
+
   bool contains(Position r, Direction u, int32_t on_surface) const override
   {
     return region_.contains(r, u, on_surface);

include/openmc/geometry.h

@@ -64,6 +64,10 @@ bool exhaustive_find_cell(GeometryState& p, bool verbose = false);
 bool neighbor_list_find_cell(
   GeometryState& p, bool verbose = false); // Only usable on surface crossings
 
+bool find_cell_in_virtual_lattice(GeometryState& p,
+  bool verbose =
+    false); // Only usable on triso surface crossings in virtual lattice
+
 //==============================================================================
 //! Move a particle into a new lattice tile.
 //==============================================================================

include/openmc/universe.h

@@ -38,6 +38,8 @@ class Universe {
 
   virtual bool find_cell(GeometryState& p) const;
 
+  virtual bool find_cell_in_virtual_lattice(GeometryState& p) const;
+
   BoundingBox bounding_box() const;
 
   /* By default, universes are CSG universes. The DAGMC

src/cell.cpp

@@ -505,96 +505,101 @@ vector<int32_t>::iterator CSGCell::find_left_parenthesis(
 }
 std::pair<double, int32_t> CSGCell::distance(
   Position r, Direction u, int32_t on_surface, GeometryState* p) const
+{
+  if (virtual_lattice_) {
+    return distance_in_virtual_lattice(r, u, on_surface, p);
+  } else {
+    return region_.distance(r, u, on_surface);
+  }
+}
+
+std::pair<double, int32_t> CSGCell::distance_in_virtual_lattice(
+  Position r, Direction u, int32_t on_surface, GeometryState* p) const
 {
   double min_dist {INFTY};
   int32_t i_surf {std::numeric_limits<int32_t>::max()};
   double min_dis_vl;
   int32_t i_surf_vl;
-  if (virtual_lattice_) {
-    double max_dis = p->collision_distance();
-    double tol_dis = 0;
-    vector<double> dis_to_bou(3), dis_to_bou_max(3);
-    double u_value = sqrt(pow(u.x, 2) + pow(u.y, 2) +
-                          pow(u.z, 2)); // don't know if u has been normalized
-    vector<double> norm_u = {u.x / u_value, u.y / u_value, u.z / u_value};
-    vector<int> lat_ind(3);
-    vector<double> temp_pos = {r.x, r.y, r.z};
-    int loop_time;
-    for (int i = 0; i < 3; i++) {
-      lat_ind[i] = floor((temp_pos[i] - vl_lower_left_[i]) / vl_pitch_[i]);
-      if (lat_ind[i] == vl_shape_[i] && norm_u[i] < 0) {
-        lat_ind[i] = vl_shape_[i] - 1;
-      }
-      if (lat_ind[i] == -1 && norm_u[i] > 0) {
-        lat_ind[i] = 0;
-      }
+
+  double max_dis = p->collision_distance();
+  double tol_dis = 0;
+  vector<double> dis_to_bou(3), dis_to_bou_max(3);
+  double u_value = sqrt(pow(u.x, 2) + pow(u.y, 2) +
+                        pow(u.z, 2)); // don't know if u has been normalized
+  vector<double> norm_u = {u.x / u_value, u.y / u_value, u.z / u_value};
+  vector<int> lat_ind(3);
+  vector<double> temp_pos = {r.x, r.y, r.z};
+  int loop_time;
+  for (int i = 0; i < 3; i++) {
+    lat_ind[i] = floor((temp_pos[i] - vl_lower_left_[i]) / vl_pitch_[i]);
+    if (lat_ind[i] == vl_shape_[i] && norm_u[i] < 0) {
+      lat_ind[i] = vl_shape_[i] - 1;
+    }
+    if (lat_ind[i] == -1 && norm_u[i] > 0) {
+      lat_ind[i] = 0;
     }
+  }
 
-    dis_to_bou = {INFTY, INFTY, INFTY};
-    for (int i = 0; i < 3; i++) {
-      if (norm_u[i] > 0) {
-        dis_to_bou[i] = std::abs(
-          ((lat_ind[i] + 1) * vl_pitch_[i] + vl_lower_left_[i] - temp_pos[i]) /
-          norm_u[i]);
-        dis_to_bou_max[i] = vl_pitch_[i] / norm_u[i];
-      } else if (norm_u[i] < 0) {
-        dis_to_bou[i] = std::abs(
-          (lat_ind[i] * vl_pitch_[i] + vl_lower_left_[i] - temp_pos[i]) /
-          norm_u[i]);
-        dis_to_bou_max[i] = -vl_pitch_[i] / norm_u[i];
-      }
+  dis_to_bou = {INFTY, INFTY, INFTY};
+  for (int i = 0; i < 3; i++) {
+    if (norm_u[i] > 0) {
+      dis_to_bou[i] = std::abs(
+        ((lat_ind[i] + 1) * vl_pitch_[i] + vl_lower_left_[i] - temp_pos[i]) /
+        norm_u[i]);
+      dis_to_bou_max[i] = vl_pitch_[i] / norm_u[i];
+    } else if (norm_u[i] < 0) {
+      dis_to_bou[i] =
+        std::abs((lat_ind[i] * vl_pitch_[i] + vl_lower_left_[i] - temp_pos[i]) /
+                 norm_u[i]);
+      dis_to_bou_max[i] = -vl_pitch_[i] / norm_u[i];
     }
+  }
 
-    while (true) {
-      if (lat_ind[0] < 0 || lat_ind[0] >= vl_shape_[0] || lat_ind[1] < 0 ||
-          lat_ind[1] >= vl_shape_[1] || lat_ind[2] < 0 ||
-          lat_ind[2] >= vl_shape_[2])
-        break;
+  while (true) {
+    if (lat_ind[0] < 0 || lat_ind[0] >= vl_shape_[0] || lat_ind[1] < 0 ||
+        lat_ind[1] >= vl_shape_[1] || lat_ind[2] < 0 ||
+        lat_ind[2] >= vl_shape_[2])
+      break;
 
-      for (int token :
-        vl_triso_distribution_[lat_ind[0] + lat_ind[1] * vl_shape_[0] +
-                               lat_ind[2] * vl_shape_[0] * vl_shape_[1]]) {
-        bool coincident {std::abs(token) == std::abs(on_surface)};
-        double d {model::surfaces[abs(token) - 1]->distance(r, u, coincident)};
-        if (d < min_dist) {
-          if (min_dist - d >= FP_PRECISION * min_dist) {
-            min_dist = d;
-            i_surf = -token;
-          }
+    for (int token :
+      vl_triso_distribution_[lat_ind[0] + lat_ind[1] * vl_shape_[0] +
+                             lat_ind[2] * vl_shape_[0] * vl_shape_[1]]) {
+      bool coincident {std::abs(token) == std::abs(on_surface)};
+      double d {model::surfaces[abs(token) - 1]->distance(r, u, coincident)};
+      if (d < min_dist) {
+        if (min_dist - d >= FP_PRECISION * min_dist) {
+          min_dist = d;
+          i_surf = -token;
         }
       }
+    }
 
-      int mes_bou_crossed = 0;
-      if (dis_to_bou[1] < dis_to_bou[0]) {
-        mes_bou_crossed = 1;
-      }
-      if (dis_to_bou[2] < dis_to_bou[mes_bou_crossed]) {
-        mes_bou_crossed = 2;
-      }
+    int mes_bou_crossed = 0;
+    if (dis_to_bou[1] < dis_to_bou[0]) {
+      mes_bou_crossed = 1;
+    }
+    if (dis_to_bou[2] < dis_to_bou[mes_bou_crossed]) {
+      mes_bou_crossed = 2;
+    }
 
-      tol_dis = dis_to_bou[mes_bou_crossed];
+    tol_dis = dis_to_bou[mes_bou_crossed];
 
-      if (min_dist < tol_dis) {
-        break;
-      }
+    if (min_dist < tol_dis) {
+      break;
+    }
 
-      if (norm_u[mes_bou_crossed] > 0) {
-        lat_ind[mes_bou_crossed] += 1;
-      } else {
-        lat_ind[mes_bou_crossed] += -1;
-      }
+    if (norm_u[mes_bou_crossed] > 0) {
+      lat_ind[mes_bou_crossed] += 1;
+    } else {
+      lat_ind[mes_bou_crossed] += -1;
+    }
 
-      dis_to_bou[mes_bou_crossed] += dis_to_bou_max[mes_bou_crossed];
+    dis_to_bou[mes_bou_crossed] += dis_to_bou_max[mes_bou_crossed];
 
-      if (tol_dis > max_dis) {
-        break;
-      }
+    if (tol_dis > max_dis) {
+      break;
     }
-
-  } else {
-    return region_.distance(r, u, on_surface);
   }
-
   return {min_dist, i_surf};
 }
 

src/geometry.cpp

@@ -294,6 +294,98 @@ bool exhaustive_find_cell(GeometryState& p, bool verbose)
   return find_cell_inner(p, nullptr, verbose);
 }
 
+bool find_cell_in_virtual_lattice(GeometryState& p, bool verbose)
+{
+  int i_surface = std::abs(p.surface());
+  if (p.surface() > 0) {
+    for (int i = p.n_coord(); i < model::n_coord_levels; i++) {
+      p.coord(i).reset();
+    }
+    p.coord(p.n_coord() - 1).cell() =
+      model::cell_map[model::surfaces[i_surface - 1]->triso_base_index_];
+  } else if (p.surface() < 0) {
+    for (int i = p.n_coord(); i < model::n_coord_levels; i++) {
+      p.coord(i).reset();
+    }
+    if (model::surfaces[i_surface - 1]->triso_particle_index_ == -1) {
+      fatal_error(fmt::format("Particle cell of surface {} is not defined",
+        model::surfaces[i_surface - 1]->id_));
+    }
+    p.lowest_coord().cell() =
+      model::cell_map[model::surfaces[i_surface - 1]->triso_particle_index_];
+  }
+
+  // find material
+  bool found = true;
+  int i_cell = p.lowest_coord().cell();
+  for (;; ++p.n_coord()) {
+    if (i_cell == C_NONE) {
+      int i_universe = p.lowest_coord().universe();
+      const auto& univ {model::universes[i_universe]};
+
+      if (univ->filled_with_triso_base_ != -1) {
+        p.lowest_coord().cell() =
+          model::cell_map[univ->filled_with_triso_base_];
+        found = true;
+      } else {
+        found = univ->find_cell(p);
+      }
+      if (!found) {
+        return found;
+      }
+    }
+
+    i_cell = p.lowest_coord().cell();
+
+    Cell& c {*model::cells[i_cell]};
+    if (c.type_ == Fill::MATERIAL) {
+      // Found a material cell which means this is the lowest coord level.
+
+      p.cell_instance() = 0;
+      // Find the distribcell instance number.
+      if (c.distribcell_index_ >= 0) {
+        p.cell_instance() = cell_instance_at_level(p, p.n_coord() - 1);
+      }
+
+      // Set the material and temperature.
+      p.material_last() = p.material();
+      if (c.material_.size() > 1) {
+        p.material() = c.material_[p.cell_instance()];
+      } else {
+        p.material() = c.material_[0];
+      }
+      p.sqrtkT_last() = p.sqrtkT();
+      if (c.sqrtkT_.size() > 1) {
+        p.sqrtkT() = c.sqrtkT_[p.cell_instance()];
+      } else {
+        p.sqrtkT() = c.sqrtkT_[0];
+      }
+      return found;
+
+    } else if (c.type_ == Fill::UNIVERSE) {
+      //========================================================================
+      //! Found a lower universe, update this coord level then search the
+      //! next.
+
+      // Set the lower coordinate level universe.
+      auto& coor {p.coord(p.n_coord())};
+      coor.universe() = c.fill_;
+
+      // Set the position and direction.
+      coor.r() = p.r_local();
+      coor.u() = p.u_local();
+
+      // Apply translation.
+      coor.r() -= c.translation_;
+
+      // Apply rotation.
+      if (!c.rotation_.empty()) {
+        coor.rotate(c.rotation_);
+      }
+      i_cell = C_NONE;
+    }
+  }
+}
 //==============================================================================
 
 void cross_lattice(GeometryState& p, const BoundaryInfo& boundary, bool verbose)