Merge pull request #1995 from paulromano/lattice-instance-fix

Fix cell instance calculation for multiple lattices filled in cells in same universe

src/geometry.cpp

@@ -86,6 +86,7 @@ int cell_instance_at_level(const Particle& p, int level)
     if (c_i.type_ == Fill::UNIVERSE) {
       instance += c_i.offset_[c.distribcell_index_];
     } else if (c_i.type_ == Fill::LATTICE) {
+      instance += c_i.offset_[c.distribcell_index_];
       auto& lat {*model::lattices[p.coord(i + 1).lattice]};
       const auto& i_xyz {p.coord(i + 1).lattice_i};
       if (lat.are_valid_indices(i_xyz)) {

src/geometry_aux.cpp

@@ -412,8 +412,9 @@ void prepare_distribcell(const std::vector<int32_t>* user_distribcells)
             search_univ, target_univ_id, univ_count_memo);
 
         } else if (c.type_ == Fill::LATTICE) {
+          c.offset_[map] = offset;
           Lattice& lat = *model::lattices[c.fill_];
-          offset =
+          offset +=
             lat.fill_offset_table(offset, target_univ_id, map, univ_count_memo);
         }
       }

src/lattice.cpp

@@ -101,6 +101,18 @@ 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)
 {
+  // 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
+  // offsets_ array doesn't actually include the offset accounting for the last
+  // universe, so we get the before-last offset for the given map and then
+  // explicitly add the count for the last universe.
+  if (offsets_[map * universes_.size()] != C_NONE) {
+    int last_offset = offsets_[(map + 1) * universes_.size() - 1];
+    int last_univ = universes_.back();
+    return last_offset +
+           count_universe_instances(last_univ, target_univ_id, univ_count_memo);
+  }
+
   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);

tests/unit_tests/test_cell_instance.py

@@ -0,0 +1,73 @@
+import numpy as np
+import pytest
+
+import openmc
+import openmc.lib
+
+from tests import cdtemp
+
+
+@pytest.fixture(scope='module', autouse=True)
+def double_lattice_model():
+    model = openmc.Model()
+
+    # Create a single material
+    m = openmc.Material()
+    m.add_nuclide('U235', 1.0)
+    m.set_density('g/cm3', 10.0)
+    model.materials.append(m)
+
+    # Create a universe with a single infinite cell
+    c = openmc.Cell(fill=m)
+    u = openmc.Universe(cells=[c])
+
+    # Create a 2x2 lattice filled with above universe
+    lattice = openmc.RectLattice()
+    lattice.lower_left = (0.0, 0.0)
+    lattice.pitch = (1.0, 1.0)
+    lattice.universes = np.full((2, 2), u)
+
+    # Create two cells each filled with the same lattice, one from x=0..2 and
+    # y=0..2 and the other from x=2..4 and y=0..2
+    x0 = openmc.XPlane(0.0, boundary_type='vacuum')
+    x2 = openmc.XPlane(2.0)
+    x4 = openmc.XPlane(4.0, boundary_type='vacuum')
+    y0 = openmc.YPlane(0.0, boundary_type='vacuum')
+    y2 = openmc.YPlane(2.0, boundary_type='vacuum')
+    cell_with_lattice1 = openmc.Cell(fill=lattice, region=+x0 & -x2 & +y0 & -y2)
+    cell_with_lattice2 = openmc.Cell(fill=lattice, region=+x2 & -x4 & +y0 & -y2)
+    cell_with_lattice2.translation = (2., 0., 0.)
+    model.geometry = openmc.Geometry([cell_with_lattice1, cell_with_lattice2])
+
+    # Add necessary settings and export
+    model.settings.batches = 10
+    model.settings.inactive = 0
+    model.settings.particles = 100
+
+    with cdtemp():
+        model.export_to_xml()
+        openmc.lib.init()
+        yield
+        openmc.lib.finalize()
+
+
+# This shows the expected cell instance numbers for each lattice position:
+#      ┌─┬─┬─┬─┐
+#      │2│3│6│7│
+#      ├─┼─┼─┼─┤
+#      │0│1│4│5│
+#      └─┴─┴─┴─┘
+expected_results = [
+    ((0.5, 0.5, 0.0), 0),
+    ((1.5, 0.5, 0.0), 1),
+    ((0.5, 1.5, 0.0), 2),
+    ((1.5, 1.5, 0.0), 3),
+    ((2.5, 0.5, 0.0), 4),
+    ((3.5, 0.5, 0.0), 5),
+    ((2.5, 1.5, 0.0), 6),
+    ((3.5, 1.5, 0.0), 7),
+]
+@pytest.mark.parametrize("r,expected_cell_instance", expected_results)
+def test_cell_instance_multilattice(r, expected_cell_instance):
+    _, cell_instance = openmc.lib.find_cell(r)
+    assert cell_instance == expected_cell_instance