Limit number of lost particles written to disk. (#2688)

Co-authored-by: John Tramm <john.tramm@gmail.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

docs/source/io_formats/settings.rst

@@ -276,6 +276,16 @@ then, OpenMC will only use up to the :math:`P_1` data.
   .. note:: This element is not used in the continuous-energy
     :ref:`energy_mode`.
 
+
+--------------------------------------
+``<max_write_lost_particles>`` Element
+--------------------------------------
+
+This ``<max_write_lost_particles>`` element indicates the maximum number of
+particle restart files (per MPI process) to write for lost particles.
+
+  *Default*: None
+
 .. _mesh_element:
 
 ------------------

include/openmc/settings.h

@@ -82,6 +82,9 @@ extern "C" int32_t max_lost_particles; //!< maximum number of lost particles
 extern double
   rel_max_lost_particles; //!< maximum number of lost particles, relative to the
                           //!< total number of particles
+extern "C" int32_t
+  max_write_lost_particles;       //!< maximum number of lost particles
+                                  //!< to be written to files
 extern "C" int32_t gen_per_batch; //!< number of generations per batch
 extern "C" int64_t n_particles;   //!< number of particles per generation
 

openmc/settings.py

@@ -121,6 +121,11 @@ class Settings:
         Maximum number of tracks written to a track file (per MPI process).
 
         .. versionadded:: 0.13.1
+    max_write_lost_particles : int
+        Maximum number of particle restart files (per MPI process) to write for
+        lost particles.
+
+        .. versionadded:: 0.13.4
     no_reduce : bool
         Indicate that all user-defined and global tallies should not be reduced
         across processes in a parallel calculation.
@@ -261,6 +266,7 @@ class Settings:
         self._inactive = None
         self._max_lost_particles = None
         self._rel_max_lost_particles = None
+        self._max_write_lost_particles = None
         self._particles = None
         self._keff_trigger = None
 
@@ -398,6 +404,16 @@ class Settings:
         cv.check_less_than('rel_max_lost_particles', rel_max_lost_particles, 1)
         self._rel_max_lost_particles = rel_max_lost_particles
 
+    @property
+    def max_write_lost_particles(self) -> int:
+        return self._max_write_lost_particles
+
+    @max_write_lost_particles.setter
+    def max_write_lost_particles(self, max_write_lost_particles: int):
+        cv.check_type('max_write_lost_particles', max_write_lost_particles, Integral)
+        cv.check_greater_than('max_write_lost_particles', max_write_lost_particles, 0)
+        self._max_write_lost_particles = max_write_lost_particles
+
     @property
     def particles(self) -> int:
         return self._particles
@@ -1014,6 +1030,11 @@ class Settings:
             element = ET.SubElement(root, "rel_max_lost_particles")
             element.text = str(self._rel_max_lost_particles)
 
+    def _create_max_write_lost_particles_subelement(self, root):
+        if self._max_write_lost_particles is not None:
+            element = ET.SubElement(root, "max_write_lost_particles")
+            element.text = str(self._max_write_lost_particles)
+
     def _create_particles_subelement(self, root):
         if self._particles is not None:
             element = ET.SubElement(root, "particles")
@@ -1376,6 +1397,7 @@ class Settings:
             self._inactive_from_xml_element(elem)
             self._max_lost_particles_from_xml_element(elem)
             self._rel_max_lost_particles_from_xml_element(elem)
+            self._max_write_lost_particles_from_xml_element(elem)
             self._generations_per_batch_from_xml_element(elem)
 
     def _run_mode_from_xml_element(self, root):
@@ -1408,6 +1430,11 @@ class Settings:
         if text is not None:
             self.rel_max_lost_particles = float(text)
 
+    def _max_write_lost_particles_from_xml_element(self, root):
+        text = get_text(root, 'max_write_lost_particles')
+        if text is not None:
+            self.max_write_lost_particles = int(text)
+
     def _generations_per_batch_from_xml_element(self, root):
         text = get_text(root, 'generations_per_batch')
         if text is not None:
@@ -1719,6 +1746,7 @@ class Settings:
         self._create_inactive_subelement(element)
         self._create_max_lost_particles_subelement(element)
         self._create_rel_max_lost_particles_subelement(element)
+        self._create_max_write_lost_particles_subelement(element)
         self._create_generations_per_batch_subelement(element)
         self._create_keff_trigger_subelement(element)
         self._create_source_subelement(element)
@@ -1815,6 +1843,7 @@ class Settings:
         settings._inactive_from_xml_element(elem)
         settings._max_lost_particles_from_xml_element(elem)
         settings._rel_max_lost_particles_from_xml_element(elem)
+        settings._max_write_lost_particles_from_xml_element(elem)
         settings._generations_per_batch_from_xml_element(elem)
         settings._keff_trigger_from_xml_element(elem)
         settings._source_from_xml_element(elem, meshes)

src/finalize.cpp

@@ -88,6 +88,7 @@ int openmc_finalize()
   settings::max_particles_in_flight = 100000;
   settings::max_splits = 1000;
   settings::max_tracks = 1000;
+  settings::max_write_lost_particles = -1;
   settings::n_inactive = 0;
   settings::n_particles = -1;
   settings::output_summary = true;

src/particle.cpp

@@ -709,8 +709,10 @@ void Particle::mark_as_lost(const char* message)
 {
   // Print warning and write lost particle file
   warning(message);
-  write_restart();
-
+  if (settings::max_write_lost_particles < 0 ||
+      simulation::n_lost_particles < settings::max_write_lost_particles) {
+    write_restart();
+  }
   // Increment number of lost particles
   wgt() = 0.0;
 #pragma omp atomic

src/settings.cpp

@@ -89,6 +89,7 @@ std::string weight_windows_file;
 int32_t n_inactive {0};
 int32_t max_lost_particles {10};
 double rel_max_lost_particles {1.0e-6};
+int32_t max_write_lost_particles {-1};
 int32_t gen_per_batch {1};
 int64_t n_particles {-1};
 
@@ -172,6 +173,12 @@ void get_run_parameters(pugi::xml_node node_base)
       std::stod(get_node_value(node_base, "rel_max_lost_particles"));
   }
 
+  // Get relative number of lost particles
+  if (check_for_node(node_base, "max_write_lost_particles")) {
+    max_write_lost_particles =
+      std::stoi(get_node_value(node_base, "max_write_lost_particles"));
+  }
+
   // Get number of inactive batches
   if (run_mode == RunMode::EIGENVALUE) {
     if (check_for_node(node_base, "inactive")) {

tests/unit_tests/test_lost_particles.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+
+import openmc
+import pytest
+
+from tests.testing_harness import config
+
+
+@pytest.fixture
+def model():
+    mat = openmc.Material()
+    mat.add_nuclide('N14', 1.0)
+    mat.set_density('g/cm3', 1e-5)
+
+    s1 = openmc.Sphere(r=80.0)
+    s2 = openmc.Sphere(r=90.0)
+    s3 = openmc.Sphere(r=100.0, boundary_type='vacuum')
+    cell1 = openmc.Cell(fill=mat, region=-s1)
+    cell2 = openmc.Cell(fill=mat, region=+s2 & -s3)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell1, cell2])
+
+    model.settings.run_mode = 'fixed source'
+    model.settings.batches = 10
+    model.settings.inactive = 5
+    model.settings.particles = 50
+    model.settings.max_lost_particles = 1000
+    model.settings.source = openmc.IndependentSource(space=openmc.stats.Point())
+
+    return model
+
+
+def test_max_write_lost_particles(model: openmc.Model, run_in_tmpdir):
+    # Set maximum number of lost particle restart files
+    model.settings.max_write_lost_particles = 5
+
+    # Run OpenMC to generate lost particle files. Use one thread so that we know
+    # exactly how much will be produced. If running in MPI mode, setup proper
+    # keyword arguments for run()
+    kwargs = {'openmc_exec': config['exe']}
+    if config['mpi']:
+        kwargs['mpi_args'] = [config['mpiexec'], '-n', config['mpi_np']]
+    model.run(threads=1, **kwargs)
+
+    # Make sure number of lost particle files is as expected
+    lost_particle_files = list(Path.cwd().glob('particle*.h5'))
+    n_procs = int(config['mpi_np']) if config['mpi'] else 1
+    assert len(lost_particle_files) == model.settings.max_write_lost_particles * n_procs
+