Add max_tracks setting (maximum number of tracks per process)

include/openmc/settings.h

@@ -88,6 +88,7 @@ extern int max_order;         //!< Maximum Legendre order for multigroup data
 extern int n_log_bins;        //!< number of bins for logarithmic energy grid
 extern int n_batches;         //!< number of (inactive+active) batches
 extern int n_max_batches;     //!< Maximum number of batches
+extern int max_tracks; //!< Maximum number of particle tracks written to file
 extern ResScatMethod res_scat_method; //!< resonance upscattering method
 extern double res_scat_energy_min; //!< Min energy in [eV] for res. upscattering
 extern double res_scat_energy_max; //!< Max energy in [eV] for res. upscattering

include/openmc/track_output.h

@@ -15,6 +15,12 @@ void open_track_file();
 //! Close HDF5 resources for track file
 void close_track_file();
 
+//! Determine whether a given particle should collect/write track information
+//
+//! \param[in] p  Current particle
+//! \return Whether to collect/write track information
+bool check_track_criteria(const Particle& p);
+
 //! Create a new track state history for a primary/secondary particle
 //
 //! \param[in] p  Current particle

openmc/settings.py

@@ -105,6 +105,10 @@ class Settings:
         Maximum number of times a particle can split during a history
 
         .. versionadded:: 0.13
+    max_tracks : int
+        Maximum number of tracks written to a track file (per MPI process).
+
+        .. versionadded:: 0.13.1
     no_reduce : bool
         Indicate that all user-defined and global tallies should not be reduced
         across processes in a parallel calculation.
@@ -291,6 +295,7 @@ class Settings:
         self._weight_windows = cv.CheckedList(WeightWindows, 'weight windows')
         self._weight_windows_on = None
         self._max_splits = None
+        self._max_tracks = None
 
     @property
     def run_mode(self):
@@ -476,6 +481,10 @@ class Settings:
     def max_splits(self):
         return self._max_splits
 
+    @property
+    def max_tracks(self):
+        return self._max_tracks
+
     @run_mode.setter
     def run_mode(self, run_mode: str):
         cv.check_value('run mode', run_mode, {x.value for x in RunMode})
@@ -881,9 +890,15 @@ class Settings:
     @max_splits.setter
     def max_splits(self, value: int):
         cv.check_type('maximum particle splits', value, Integral)
-        cv.check_greater_than('max particles in flight', value, 0)
+        cv.check_greater_than('max particle splits', value, 0)
         self._max_splits = value
 
+    @max_tracks.setter
+    def max_tracks(self, value: int):
+        cv.check_type('maximum particle tracks', value, Integral)
+        cv.check_greater_than('maximum particle tracks', value, 0, True)
+        self._max_tracks = value
+
     def _create_run_mode_subelement(self, root):
         elem = ET.SubElement(root, "run_mode")
         elem.text = self._run_mode.value
@@ -1196,6 +1211,11 @@ class Settings:
             elem = ET.SubElement(root, "max_splits")
             elem.text = str(self._max_splits)
 
+    def _create_max_tracks_subelement(self, root):
+        if self._max_tracks is not None:
+            elem = ET.SubElement(root, "max_tracks")
+            elem.text = str(self._max_tracks)
+
     def _eigenvalue_from_xml_element(self, root):
         elem = root.find('eigenvalue')
         if elem is not None:
@@ -1499,6 +1519,11 @@ class Settings:
         if text is not None:
             self.max_splits = int(text)
 
+    def _max_tracks_from_xml_element(self, root):
+        text = get_text(root, 'max_tracks')
+        if text is not None:
+            self.max_tracks = int(text)
+
     def export_to_xml(self, path: Union[str, os.PathLike] = 'settings.xml'):
         """Export simulation settings to an XML file.
 
@@ -1555,6 +1580,7 @@ class Settings:
         self._create_write_initial_source_subelement(root_element)
         self._create_weight_windows_subelement(root_element)
         self._create_max_splits_subelement(root_element)
+        self._create_max_tracks_subelement(root_element)
 
         # Clean the indentation in the file to be user-readable
         clean_indentation(root_element)
@@ -1634,6 +1660,7 @@ class Settings:
         settings._write_initial_source_from_xml_element(root)
         settings._weight_windows_from_xml_element(root)
         settings._max_splits_from_xml_element(root)
+        settings._max_tracks_from_xml_element(root)
 
         # TODO: Get volume calculations
 

src/finalize.cpp

@@ -85,6 +85,7 @@ int openmc_finalize()
   settings::material_cell_offsets = true;
   settings::max_particles_in_flight = 100000;
   settings::max_splits = 1000;
+  settings::max_tracks = std::numeric_limits<int>::max();
   settings::n_inactive = 0;
   settings::n_particles = -1;
   settings::output_summary = true;

src/settings.cpp

@@ -95,6 +95,7 @@ int n_log_bins {8000};
 int n_batches;
 int n_max_batches;
 int max_splits {1000};
+int max_tracks {std::numeric_limits<int>::max()};
 ResScatMethod res_scat_method {ResScatMethod::rvs};
 double res_scat_energy_min {0.01};
 double res_scat_energy_max {1000.0};
@@ -877,6 +878,10 @@ void read_settings_xml()
   if (check_for_node(root, "max_splits")) {
     settings::max_splits = std::stoi(get_node_value(root, "max_splits"));
   }
+
+  if (check_for_node(root, "max_tracks")) {
+    settings::max_tracks = std::stoi(get_node_value(root, "max_tracks"));
+  }
 }
 
 void free_memory_settings()

src/simulation.cpp

@@ -517,18 +517,7 @@ void initialize_history(Particle& p, int64_t index_source)
     p.trace() = true;
 
   // Set particle track.
-  p.write_track() = false;
-  if (settings::write_all_tracks) {
-    p.write_track() = true;
-  } else if (settings::track_identifiers.size() > 0) {
-    for (const auto& t : settings::track_identifiers) {
-      if (simulation::current_batch == t[0] &&
-          simulation::current_gen == t[1] && p.id() == t[2]) {
-        p.write_track() = true;
-        break;
-      }
-    }
-  }
+  p.write_track() = check_track_criteria(p);
 
   // Display message if high verbosity or trace is on
   if (settings::verbosity >= 9 || p.trace()) {

src/track_output.cpp

@@ -21,8 +21,9 @@ namespace openmc {
 // Global variables
 //==============================================================================
 
-hid_t track_file;  //! HDF5 identifier for track file
-hid_t track_dtype; //! HDF5 identifier for track datatype
+hid_t track_file;     //! HDF5 identifier for track file
+hid_t track_dtype;    //! HDF5 identifier for track datatype
+int n_tracks_written; //! Number of tracks written
 
 //==============================================================================
 // Non-member functions
@@ -81,6 +82,33 @@ void close_track_file()
 {
   H5Tclose(track_dtype);
   file_close(track_file);
+
+  // Reset number of tracks written
+  n_tracks_written = 0;
+}
+
+bool check_track_criteria(const Particle& p)
+{
+  if (settings::write_all_tracks) {
+    // Increment number of tracks written and get previous value
+    int n;
+#pragma omp atomic capture
+    n = n_tracks_written++;
+
+    // Indicate that track should be written for this particle
+    return n < settings::max_tracks;
+  }
+
+  // Check for match from explicit track identifiers
+  if (settings::track_identifiers.size() > 0) {
+    for (const auto& t : settings::track_identifiers) {
+      if (simulation::current_batch == t[0] &&
+          simulation::current_gen == t[1] && p.id() == t[2]) {
+        return true;
+      }
+    }
+  }
+  return false;
 }
 
 void finalize_particle_track(Particle& p)

tests/unit_tests/test_settings.py

@@ -14,6 +14,7 @@ def test_export_to_xml(run_in_tmpdir):
     s.keff_trigger = {'type': 'std_dev', 'threshold': 0.001}
     s.energy_mode = 'continuous-energy'
     s.max_order = 5
+    s.max_tracks = 1234
     s.source = openmc.Source(space=openmc.stats.Point())
     s.output = {'summary': True, 'tallies': False, 'path': 'here'}
     s.verbosity = 7
@@ -71,6 +72,7 @@ def test_export_to_xml(run_in_tmpdir):
     assert s.keff_trigger == {'type': 'std_dev', 'threshold': 0.001}
     assert s.energy_mode == 'continuous-energy'
     assert s.max_order == 5
+    assert s.max_tracks == 1234
     assert isinstance(s.source[0], openmc.Source)
     assert isinstance(s.source[0].space, openmc.stats.Point)
     assert s.output == {'summary': True, 'tallies': False, 'path': 'here'}

tests/unit_tests/test_tracks.py

@@ -21,17 +21,16 @@ def sphere_model():
     model.settings.run_mode = 'fixed source'
     model.settings.batches = 1
     model.settings.particles = 100
-    model.settings.track = [(1, 1, 1), (1, 1, 10), (1, 1, 75)]
 
     return model
 
 
-def test_tracks(sphere_model, run_in_tmpdir):
+def generate_track_file(model, **kwargs):
     # If running in MPI mode, setup proper keyword arguments for run()
-    kwargs = {'openmc_exec': config['exe']}
+    kwargs.setdefault('openmc_exec', config['exe'])
     if config['mpi']:
         kwargs['mpi_args'] = [config['mpiexec'], '-n', config['mpi_np']]
-    sphere_model.run(**kwargs)
+    model.run(**kwargs)
 
     if config['mpi'] and int(config['mpi_np']) > 1:
         # With MPI, we need to combine track files
@@ -41,6 +40,14 @@ def test_tracks(sphere_model, run_in_tmpdir):
         track_file = Path('tracks.h5')
         assert track_file.is_file()
 
+
+def test_tracks(sphere_model, run_in_tmpdir):
+    # Set track identifiers
+    sphere_model.settings.track = [(1, 1, 1), (1, 1, 10), (1, 1, 75)]
+
+    # Run OpenMC to generate tracks.h5 file
+    generate_track_file(sphere_model)
+
     # Open track file and make sure we have correct number of tracks
     tracks = openmc.TrackFile('tracks.h5')
     assert len(tracks) == len(sphere_model.settings.track)
@@ -79,3 +86,17 @@ def test_tracks(sphere_model, run_in_tmpdir):
         assert x.time == state['time']
         assert x.wgt == state['wgt']
         assert x.particle == particle_track.particle
+
+
+def test_max_tracks(sphere_model, run_in_tmpdir):
+    # Set maximum number of tracks per process to write
+    sphere_model.settings.max_tracks = expected_num_tracks = 10
+    if config['mpi']:
+        expected_num_tracks *= int(config['mpi_np'])
+
+    # Run OpenMC to generate tracks.h5 file
+    generate_track_file(sphere_model, tracks=True)
+
+    # Open track file and make sure we have correct number of tracks
+    tracks = openmc.TrackFile('tracks.h5')
+    assert len(tracks) == expected_num_tracks