Write surface source files per batch

docs/source/io_formats/settings.rst

@@ -923,6 +923,15 @@ attributes/sub-elements:
 
     *Default*: None
 
+  :batches:
+    An Iterable of integers indicating at what batches a surface source file 
+    should be written. The surface source bank will be cleared in simulation 
+    memory each time a surface source file is written. By default, in the 
+    absence of this attribute, a ``surface_source.h5`` file will be created 
+    after the last batch.
+
+    *Default*: None
+
   :mcpl:
     An optional boolean which indicates if the banked particles should be
     written to a file in the MCPL_-format instead of the native HDF5-based

docs/source/usersguide/settings.rst

@@ -336,6 +336,17 @@ or particles going to a cell::
 .. note:: The ``cell``, ``cellfrom`` and ``cellto`` attributes cannot be
           used simultaneously.
 
+To generate surface source files in different batches, ``batches`` is available. The 
+surface source bank will be cleared in simulation memory each time a surface source 
+file is written. As an example, to write a surface source file every five batches::
+
+  settings.batches = n
+  settings.surf_source_write = {
+      'surfaces_ids': [1, 2, 3],
+      'max_particles': 10000,
+      'batches': range(5, n + 5, 5)
+  }
+
 .. _compiled_source:
 
 Compiled Sources

include/openmc/settings.h

@@ -127,8 +127,11 @@ extern std::unordered_set<int>
   sourcepoint_batch; //!< Batches when source should be written
 extern std::unordered_set<int>
   statepoint_batch; //!< Batches when state should be written
+extern std::unordered_set<int>
+  surface_source_batch; //!< Batches when surface source should be written
 extern std::unordered_set<int>
   source_write_surf_id; //!< Surface ids where sources will be written
+
 extern int
   max_history_splits; //!< maximum number of particle splits for weight windows
 extern int64_t max_surface_particles; //!< maximum number of particles to be
@@ -171,4 +174,4 @@ void free_memory_settings();
 
 } // namespace openmc
 
-#endif // OPENMC_SETTINGS_H
+#endif // OPENMC_SETTINGS_H

openmc/settings.py

@@ -217,6 +217,7 @@ class Settings:
         :cellto: Cell ID used to determine if particles crossing identified
                  surfaces are to be banked. Particles going to this declared
                  cell will be banked (int)
+        :batches: List of batches at which to write surface source files
     survival_biasing : bool
         Indicate whether survival biasing is to be used
     tabular_legendre : dict
@@ -714,14 +715,18 @@ def surf_source_write(self, surf_source_write: dict):
             cv.check_value(
                 "surface source writing key",
                 key,
-                ("surface_ids", "max_particles", "mcpl", "cell", "cellfrom", "cellto"),
+                ("surface_ids", "max_particles", "batches", "mcpl", "cell", "cellfrom", "cellto"),
             )
-            if key == "surface_ids":
-                cv.check_type(
-                    "surface ids for source banking", value, Iterable, Integral
+            if key in ("surface_ids", "batches"):
+                name = {
+                    "surface_ids": "surface ids for source banking",
+                    "batches": "surface source batches",
+                }[key]
+                cv.check_type(name, value, Iterable, Integral
                 )
-                for surf_id in value:
-                    cv.check_greater_than("surface id for source banking", surf_id, 0)
+                for i in value:
+                    cv.check_greater_than(name, i, 0)
+
             elif key == "mcpl":
                 cv.check_type("write to an MCPL-format file", value, bool)
             elif key in ("max_particles", "cell", "cellfrom", "cellto"):
@@ -1235,11 +1240,12 @@ def _create_surf_source_read_subelement(self, root):
     def _create_surf_source_write_subelement(self, root):
         if self._surf_source_write:
             element = ET.SubElement(root, "surf_source_write")
-            if "surface_ids" in self._surf_source_write:
-                subelement = ET.SubElement(element, "surface_ids")
-                subelement.text = " ".join(
-                    str(x) for x in self._surf_source_write["surface_ids"]
-                )
+            for key in ("surface_ids", "batches"):
+                if key in self._surf_source_write:
+                    subelement = ET.SubElement(element, key)
+                    subelement.text = " ".join(
+                        str(x) for x in self._surf_source_write[key]
+                    )
             if "mcpl" in self._surf_source_write:
                 subelement = ET.SubElement(element, "mcpl")
                 subelement.text = str(self._surf_source_write["mcpl"]).lower()
@@ -1642,10 +1648,10 @@ def _surf_source_write_from_xml_element(self, root):
         elem = root.find('surf_source_write')
         if elem is None:
             return
-        for key in ('surface_ids', 'max_particles', 'mcpl', 'cell', 'cellto', 'cellfrom'):
+        for key in ('surface_ids', 'max_particles', 'batches', 'mcpl', 'cell', 'cellto', 'cellfrom'):
             value = get_text(elem, key)
             if value is not None:
-                if key == 'surface_ids':
+                if key in ('surface_ids', 'batches'):
                     value = [int(x) for x in value.split()]
                 elif key == 'mcpl':
                     value = value in ('true', '1')
@@ -2088,4 +2094,4 @@ def from_xml(cls, path: PathLike = 'settings.xml'):
         tree = ET.parse(path, parser=parser)
         root = tree.getroot()
         meshes = _read_meshes(root)
-        return cls.from_xml_element(root, meshes)
+        return cls.from_xml_element(root, meshes)

src/settings.cpp

@@ -118,6 +118,7 @@ SolverType solver_type {SolverType::MONTE_CARLO};
 std::unordered_set<int> sourcepoint_batch;
 std::unordered_set<int> statepoint_batch;
 std::unordered_set<int> source_write_surf_id;
+std::unordered_set<int> surface_source_batch;
 int64_t max_surface_particles;
 int64_t ssw_cell_id {C_NONE};
 SSWCellType ssw_cell_type {SSWCellType::None};
@@ -783,6 +784,13 @@ void read_settings_xml(pugi::xml_node root)
       }
     }
 
+    // User gave specific batches to write surface source files
+    if (check_for_node(node_ssw, "batches")) {
+      auto temp = get_node_array<int>(node_ssw, "batches");
+      for (const auto& b : temp) {
+        surface_source_batch.insert(b);
+      }
+    }
     // Get maximum number of particles to be banked per surface
     if (check_for_node(node_ssw, "max_particles")) {
       max_surface_particles =
@@ -1081,6 +1089,7 @@ void free_memory_settings()
   settings::statepoint_batch.clear();
   settings::sourcepoint_batch.clear();
   settings::source_write_surf_id.clear();
+  settings::surface_source_batch.clear();
   settings::res_scat_nuclides.clear();
 }
 
@@ -1134,4 +1143,4 @@ extern "C" int openmc_get_n_batches(int* n_batches, bool get_max_batches)
   return 0;
 }
 
-} // namespace openmc
+} // namespace openmc

src/simulation.cpp

@@ -345,7 +345,6 @@ void initialize_batch()
 {
   // Increment current batch
   ++simulation::current_batch;
-
   if (settings::run_mode == RunMode::FIXED_SOURCE) {
     if (settings::solver_type == SolverType::RANDOM_RAY &&
         simulation::current_batch < settings::n_inactive + 1) {
@@ -464,17 +463,34 @@ void finalize_batch()
   }
 
   // Write out surface source if requested.
-  if (settings::surf_source_write &&
-      simulation::current_batch == settings::n_batches) {
-    auto filename = settings::path_output + "surface_source";
-    auto surf_work_index =
-      mpi::calculate_parallel_index_vector(simulation::surf_source_bank.size());
-    gsl::span<SourceSite> surfbankspan(simulation::surf_source_bank.begin(),
-      simulation::surf_source_bank.size());
-    if (settings::surf_mcpl_write) {
-      write_mcpl_source_point(filename.c_str(), surfbankspan, surf_work_index);
-    } else {
-      write_source_point(filename.c_str(), surfbankspan, surf_work_index);
+  if (settings::surf_source_write) {
+    if (contains(settings::surface_source_batch, simulation::current_batch) ||
+        (simulation::current_batch == settings::n_batches &&
+          settings::surface_source_batch.size() == 0)) {
+      auto filename = settings::path_output + "surface_source." +
+                      std::to_string(simulation::current_batch);
+      // no batches specified for writing, write surface source bank
+      if (settings::surface_source_batch.size() == 0)
+        filename = settings::path_output + "surface_source";
+      auto surf_work_index = mpi::calculate_parallel_index_vector(
+        simulation::surf_source_bank.size());
+      gsl::span<SourceSite> surfbankspan(simulation::surf_source_bank.begin(),
+        simulation::surf_source_bank.size());
+      if (settings::surf_mcpl_write) {
+        write_message(
+          "Creating surface source " + filename + ".mcpl with {} particles ...",
+          simulation::surf_source_bank.size(), 5);
+        write_mcpl_source_point(
+          filename.c_str(), surfbankspan, surf_work_index);
+      } else {
+        write_message(
+          "Creating surface source " + filename + ".h5 with {} particles ...",
+          simulation::surf_source_bank.size(), 5);
+        write_source_point(filename.c_str(), surfbankspan, surf_work_index);
+      }
+      simulation::surf_source_bank.clear();
+      if (simulation::current_batch < settings::n_batches)
+        simulation::surf_source_bank.reserve(settings::max_surface_particles);
     }
   }
 }
@@ -840,4 +856,4 @@ void transport_event_based()
   }
 }
 
-} // namespace openmc
+} // namespace openmc

tests/regression_tests/dagmc/legacy/test.py

@@ -104,5 +104,4 @@ def test_surf_source(model):
 
 def test_dagmc(model):
     harness = PyAPITestHarness('statepoint.5.h5', model)
-    harness.main()
-
+    harness.main()

tests/regression_tests/surface_source/test.py

@@ -132,4 +132,4 @@ def test_surface_source_read(model):
     harness = SurfaceSourceTestHarness('statepoint.10.h5',
                                        model,
                                        'inputs_true_read.dat')
-    harness.main()
+    harness.main()

tests/regression_tests/surface_source_write/test.py

@@ -1129,4 +1129,4 @@ def test_surface_source_cell_dagmc(
     harness = SurfaceSourceWriteTestHarness(
         "statepoint.5.h5", model=model, workdir=folder
     )
-    harness.main()
+    harness.main()

tests/unit_tests/test_surface_source_write.py

@@ -8,6 +8,7 @@
 import pytest
 import h5py
 import numpy as np
+import os
 
 
 @pytest.fixture(scope="module")
@@ -32,6 +33,7 @@ def geometry():
         {"max_particles": 200, "surface_ids": [2], "cell": 1},
         {"max_particles": 200, "surface_ids": [2], "cellto": 1},
         {"max_particles": 200, "surface_ids": [2], "cellfrom": 1},
+        {"max_particles": 200, "batches": [1,2,3]},
     ],
 )
 def test_xml_serialization(parameter, run_in_tmpdir):
@@ -44,6 +46,59 @@ def test_xml_serialization(parameter, run_in_tmpdir):
     assert read_settings.surf_source_write == parameter
 
 
+@pytest.fixture(scope="module")
+def model():
+    """Simple hydrogen sphere geometry"""
+    openmc.reset_auto_ids()
+    model = openmc.Model()
+
+    # Material
+    material = openmc.Material(name="H1")
+    material.add_element("H", 1.0)
+
+    # Geometry
+    radius = 1.0
+    sphere = openmc.Sphere(r=radius, boundary_type="reflective")
+    cell = openmc.Cell(region=-sphere, fill=material)
+    root = openmc.Universe(cells=[cell])
+    model.geometry = openmc.Geometry(root)
+
+    # Settings
+    model.settings = openmc.Settings()
+    model.settings.run_mode = "fixed source"
+    model.settings.particles = 100
+    model.settings.batches = 3
+    model.settings.seed = 1
+
+    bounds = [-radius, -radius, -radius, radius, radius, radius]
+    distribution = openmc.stats.Box(bounds[:3], bounds[3:])
+    model.settings.source = openmc.IndependentSource(space=distribution)
+
+    return model
+
+@pytest.mark.parametrize(
+    "parameter",
+    [
+        {"max_particles": 200, "batches": [1]},
+        {"max_particles": 200, "batches": [2]},
+        {"max_particles": 200, "batches": [1,2]},
+        {"max_particles": 200, "batches": [2,3]},
+        {"max_particles": 200, "batches": [1,3]},
+        {"max_particles": 200, "batches": [1,2,3]},
+    ],
+)
+
+def test_number_surface_source_file_created(parameter, run_in_tmpdir, model):
+    """Check the number of surface source files written."""
+    model.settings.surf_source_write = parameter
+    model.run()
+    for batch in parameter["batches"]:
+        filename = "surface_source."+str(batch)+".h5"
+        if not os.path.exists(filename):
+            assert False
+    if os.path.exists("surface_source.h5"):
+        assert False
+
 ERROR_MSG_1 = (
     "A maximum number of particles needs to be specified "
     "using the 'max_particles' parameter to store surface "
@@ -245,4 +300,4 @@ def test_particle_direction_dagmc(parameter, run_in_tmpdir, model_dagmc):
                 elif "cellto" in parameter.keys():
                     assert ux * x + uy * y < 0
                 else:
-                    assert False
+                    assert False