Write surface source files per batch

include/openmc/capi.h

@@ -134,8 +134,8 @@ int openmc_reset_timers();
 int openmc_run();
 int openmc_sample_external_source(size_t n, uint64_t* seed, void* sites);
 void openmc_set_seed(int64_t new_seed);
-int openmc_set_n_batches(
-  int32_t n_batches, bool set_max_batches, bool add_statepoint_batch);
+int openmc_set_n_batches(int32_t n_batches, bool set_max_batches,
+  bool add_statepoint_batch, bool add_surface_source_batch);
 int openmc_simulation_finalize();
 int openmc_simulation_init();
 int openmc_source_bank(void** ptr, int64_t* n);

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/lib/settings.py

@@ -15,7 +15,7 @@
 _dll.openmc_get_n_batches.argtypes = [POINTER(c_int), c_bool]
 _dll.openmc_get_n_batches.restype = c_int
 _dll.openmc_get_n_batches.errcheck = _error_handler
-_dll.openmc_set_n_batches.argtypes = [c_int32, c_bool, c_bool]
+_dll.openmc_set_n_batches.argtypes = [c_int32, c_bool, c_bool, c_bool]
 _dll.openmc_set_n_batches.restype = c_int
 _dll.openmc_set_n_batches.errcheck = _error_handler
 
@@ -68,7 +68,7 @@ def seed(self):
     def seed(self, seed):
         _dll.openmc_set_seed(seed)
 
-    def set_batches(self, n_batches, set_max_batches=True, add_sp_batch=True):
+    def set_batches(self, n_batches, set_max_batches=True, add_sp_batch=True, add_surf_source_batch=True):
         """Set number of batches or maximum number of batches
 
         Parameters
@@ -83,7 +83,7 @@ def set_batches(self, n_batches, set_max_batches=True, add_sp_batch=True):
             Whether to add `n_batches` as a statepoint batch
 
         """
-        _dll.openmc_set_n_batches(n_batches, set_max_batches, add_sp_batch)
+        _dll.openmc_set_n_batches(n_batches, set_max_batches, add_sp_batch, add_surf_source_batch)
 
     def get_batches(self, get_max_batches=True):
         """Get number of batches or maximum number of batches

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,17 @@ 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);
+      }
+    } else {
+      // If no <batches> tag was present, by default write surface source at
+      // last batch only
+      surface_source_batch.insert(n_batches);
+    }
     // Get maximum number of particles to be banked per surface
     if (check_for_node(node_ssw, "max_particles")) {
       max_surface_particles =
@@ -1081,15 +1093,16 @@ 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();
 }
 
 //==============================================================================
 // C API functions
 //==============================================================================
 
-extern "C" int openmc_set_n_batches(
-  int32_t n_batches, bool set_max_batches, bool add_statepoint_batch)
+extern "C" int openmc_set_n_batches(int32_t n_batches, bool set_max_batches,
+  bool add_statepoint_batch, bool add_surface_source_batch)
 {
   if (settings::n_inactive >= n_batches) {
     set_errmsg("Number of active batches must be greater than zero.");
@@ -1124,6 +1137,12 @@ extern "C" int openmc_set_n_batches(
       !(contains(settings::statepoint_batch, n_batches)))
     settings::statepoint_batch.insert(n_batches);
 
+  // Add value of n_batches to surface_source_batch
+  // when no specific batches are defined
+  if (add_surface_source_batch &&
+      !(contains(settings::surface_source_batch, n_batches)))
+    settings::surface_source_batch.insert(n_batches);
+
   return 0;
 }
 
@@ -1134,4 +1153,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,35 @@ 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 (std::find(settings::surface_source_batch.begin(),
+          settings::surface_source_batch.end(),
+          simulation::current_batch) != settings::surface_source_batch.end()) {
+      auto filename = settings::path_output + "surface_source." +
+                      std::to_string(simulation::current_batch);
+      auto it = settings::surface_source_batch.begin();
+      if (*it == settings::n_batches &&
+          settings::surface_source_batch.size() == 1) {
+        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);
+        write_message(
+          "Creating surface source " + filename + ".mcpl with {} particles ...",
+          simulation::surf_source_bank.size(), 5);
+      } else {
+        write_source_point(filename.c_str(), surfbankspan, surf_work_index);
+        write_message(
+          "Creating surface source " + filename + ".h5 with {} particles ...",
+          simulation::surf_source_bank.size(), 5);
+      }
+      simulation::surf_source_bank.clear();
+      simulation::surf_source_bank.reserve(settings::max_surface_particles);
     }
   }
 }
@@ -840,4 +857,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

@@ -245,4 +245,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