From 9686851e7a1bf1dc49ad313673829bbbabd3945e Mon Sep 17 00:00:00 2001
From: Zoe Prieto <101403129+zoeprieto@users.noreply.github.com>
Date: Thu, 3 Oct 2024 19:32:03 -0300
Subject: [PATCH] Write surface source files per batch (#3124)

Co-authored-by: Patrick Shriwise <pshriwise@gmail.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
---
 docs/source/io_formats/settings.rst           |  9 +++
 docs/source/usersguide/settings.rst           | 11 ++++
 include/openmc/settings.h                     | 15 +++--
 include/openmc/simulation.h                   |  1 +
 include/openmc/state_point.h                  |  8 ++-
 openmc/settings.py                            | 13 ++--
 src/finalize.cpp                              |  5 ++
 src/settings.cpp                              | 13 +++-
 src/simulation.cpp                            | 60 +++++++++++--------
 src/state_point.cpp                           | 20 ++++++-
 tests/regression_tests/dagmc/legacy/test.py   |  3 +-
 tests/regression_tests/surface_source/test.py |  2 +-
 .../surface_source_write/test.py              |  2 +-
 tests/unit_tests/test_surface_source_write.py | 53 ++++++++++++++++
 14 files changed, 167 insertions(+), 48 deletions(-)

diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst
index e395ada6826..2c6c3265649 100644
--- a/docs/source/io_formats/settings.rst
+++ b/docs/source/io_formats/settings.rst
@@ -923,6 +923,15 @@ attributes/sub-elements:
 
     *Default*: None
 
+  :max_source_files:
+    An integer value indicating the number of surface source files to be written
+    containing the maximum number of particles each. The surface source bank
+    will be cleared in simulation memory each time a surface source file is
+    written. By default a ``surface_source.h5`` file will be created when the
+    maximum number of saved particles is reached.
+
+    *Default*: 1
+
   :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
diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst
index d73d90cbb3a..4111481a9ea 100644
--- a/docs/source/usersguide/settings.rst
+++ b/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 more than one surface source files when the maximum number of stored
+particles is reached, ``max_source_files`` 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 maximum of three surface source files:::
+
+  settings.surf_source_write = {
+      'surfaces_ids': [1, 2, 3],
+      'max_particles': 10000,
+      'max_source_files': 3
+  }
+
 .. _compiled_source:
 
 Compiled Sources
diff --git a/include/openmc/settings.h b/include/openmc/settings.h
index a95f1ced9f1..1e44c08801d 100644
--- a/include/openmc/settings.h
+++ b/include/openmc/settings.h
@@ -129,14 +129,17 @@ extern std::unordered_set<int>
   statepoint_batch; //!< Batches when state 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
-                                      //!< banked on surfaces per process
-extern int64_t ssw_cell_id;           //!< Cell id for the surface source
-                                      //!< write setting
-extern SSWCellType ssw_cell_type;     //!< Type of option for the cell
-                                      //!< argument of surface source write
+extern int64_t ssw_max_particles; //!< maximum number of particles to be
+                                  //!< banked on surfaces per process
+extern int64_t ssw_max_files;     //!< maximum number of surface source files
+                                  //!<  to be created
+extern int64_t ssw_cell_id;       //!< Cell id for the surface source
+                                  //!< write setting
+extern SSWCellType ssw_cell_type; //!< Type of option for the cell
+                                  //!< argument of surface source write
 extern TemperatureMethod
   temperature_method; //!< method for choosing temperatures
 extern double
diff --git a/include/openmc/simulation.h b/include/openmc/simulation.h
index dd8bb7cdfd7..3e4e24e1d06 100644
--- a/include/openmc/simulation.h
+++ b/include/openmc/simulation.h
@@ -37,6 +37,7 @@ extern "C" int n_lost_particles;   //!< cumulative number of lost particles
 extern "C" bool need_depletion_rx; //!< need to calculate depletion rx?
 extern "C" int restart_batch;      //!< batch at which a restart job resumed
 extern "C" bool satisfy_triggers;  //!< have tally triggers been satisfied?
+extern int ssw_current_file;       //!< current surface source file
 extern "C" int total_gen;          //!< total number of generations simulated
 extern double total_weight;        //!< Total source weight in a batch
 extern int64_t work_per_rank;      //!< number of particles per MPI rank
diff --git a/include/openmc/state_point.h b/include/openmc/state_point.h
index 95524f6b622..f0d41e1697f 100644
--- a/include/openmc/state_point.h
+++ b/include/openmc/state_point.h
@@ -2,6 +2,7 @@
 #define OPENMC_STATE_POINT_H
 
 #include <cstdint>
+#include <string>
 
 #include <gsl/gsl-lite.hpp>
 
@@ -33,8 +34,11 @@ void load_state_point();
 // values on each rank, used to create global indexing. This vector
 // can be created by calling calculate_parallel_index_vector on
 // source_bank.size() if such a vector is not already available.
-void write_source_point(const char* filename, gsl::span<SourceSite> source_bank,
-  const vector<int64_t>& bank_index);
+void write_h5_source_point(const char* filename,
+  gsl::span<SourceSite> source_bank, const vector<int64_t>& bank_index);
+
+void write_source_point(std::string, gsl::span<SourceSite> source_bank,
+  const vector<int64_t>& bank_index, bool use_mcpl);
 
 // This appends a source bank specification to an HDF5 file
 // that's already open. It is used internally by write_source_point.
diff --git a/openmc/settings.py b/openmc/settings.py
index ddaac040191..96e6368e4f3 100644
--- a/openmc/settings.py
+++ b/openmc/settings.py
@@ -211,6 +211,7 @@ class Settings:
                    banked (int)
         :max_particles: Maximum number of particles to be banked on surfaces per
                    process (int)
+        :max_source_files: Maximum number of surface source files to be created (int)
         :mcpl: Output in the form of an MCPL-file (bool)
         :cell: Cell ID used to determine if particles crossing identified
                surfaces are to be banked. Particles coming from or going to this
@@ -718,7 +719,7 @@ 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", "max_source_files", "mcpl", "cell", "cellfrom", "cellto"),
             )
             if key == "surface_ids":
                 cv.check_type(
@@ -726,11 +727,13 @@ def surf_source_write(self, surf_source_write: dict):
                 )
                 for surf_id in value:
                     cv.check_greater_than("surface id for source banking", surf_id, 0)
+
             elif key == "mcpl":
                 cv.check_type("write to an MCPL-format file", value, bool)
-            elif key in ("max_particles", "cell", "cellfrom", "cellto"):
+            elif key in ("max_particles", "max_source_files", "cell", "cellfrom", "cellto"):
                 name = {
                     "max_particles": "maximum particle banks on surfaces per process",
+                    "max_source_files": "maximun surface source files to be written",
                     "cell": "Cell ID for source banking (from or to)",
                     "cellfrom": "Cell ID for source banking (from only)",
                     "cellto": "Cell ID for source banking (to only)",
@@ -1251,7 +1254,7 @@ def _create_surf_source_write_subelement(self, root):
             if "mcpl" in self._surf_source_write:
                 subelement = ET.SubElement(element, "mcpl")
                 subelement.text = str(self._surf_source_write["mcpl"]).lower()
-            for key in ("max_particles", "cell", "cellfrom", "cellto"):
+            for key in ("max_particles", "max_source_files", "cell", "cellfrom", "cellto"):
                 if key in self._surf_source_write:
                     subelement = ET.SubElement(element, key)
                     subelement.text = str(self._surf_source_write[key])
@@ -1650,14 +1653,14 @@ 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', 'max_source_files', 'mcpl', 'cell', 'cellto', 'cellfrom'):
             value = get_text(elem, key)
             if value is not None:
                 if key == 'surface_ids':
                     value = [int(x) for x in value.split()]
                 elif key == 'mcpl':
                     value = value in ('true', '1')
-                elif key in ('max_particles', 'cell', 'cellfrom', 'cellto'):
+                elif key in ('max_particles', 'max_source_files', 'cell', 'cellfrom', 'cellto'):
                     value = int(value)
                 self.surf_source_write[key] = value
 
diff --git a/src/finalize.cpp b/src/finalize.cpp
index 2bde0e3387c..08c2fced308 100644
--- a/src/finalize.cpp
+++ b/src/finalize.cpp
@@ -120,6 +120,10 @@ int openmc_finalize()
   settings::source_latest = false;
   settings::source_separate = false;
   settings::source_write = true;
+  settings::ssw_cell_id = C_NONE;
+  settings::ssw_cell_type = SSWCellType::None;
+  settings::ssw_max_particles = 0;
+  settings::ssw_max_files = 1;
   settings::survival_biasing = false;
   settings::temperature_default = 293.6;
   settings::temperature_method = TemperatureMethod::NEAREST;
@@ -141,6 +145,7 @@ int openmc_finalize()
 
   simulation::keff = 1.0;
   simulation::need_depletion_rx = false;
+  simulation::ssw_current_file = 1;
   simulation::total_gen = 0;
 
   simulation::entropy_mesh = nullptr;
diff --git a/src/settings.cpp b/src/settings.cpp
index ba451e219c5..3a905bdf961 100644
--- a/src/settings.cpp
+++ b/src/settings.cpp
@@ -118,7 +118,8 @@ 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;
-int64_t max_surface_particles;
+int64_t ssw_max_particles;
+int64_t ssw_max_files;
 int64_t ssw_cell_id {C_NONE};
 SSWCellType ssw_cell_type {SSWCellType::None};
 TemperatureMethod temperature_method {TemperatureMethod::NEAREST};
@@ -799,14 +800,20 @@ void read_settings_xml(pugi::xml_node root)
 
     // Get maximum number of particles to be banked per surface
     if (check_for_node(node_ssw, "max_particles")) {
-      max_surface_particles =
-        std::stoll(get_node_value(node_ssw, "max_particles"));
+      ssw_max_particles = std::stoll(get_node_value(node_ssw, "max_particles"));
     } else {
       fatal_error("A maximum number of particles needs to be specified "
                   "using the 'max_particles' parameter to store surface "
                   "source points.");
     }
 
+    // Get maximum number of surface source files to be created
+    if (check_for_node(node_ssw, "max_source_files")) {
+      ssw_max_files = std::stoll(get_node_value(node_ssw, "max_source_files"));
+    } else {
+      ssw_max_files = 1;
+    }
+
     if (check_for_node(node_ssw, "mcpl")) {
       surf_mcpl_write = get_node_value_bool(node_ssw, "mcpl");
 
diff --git a/src/simulation.cpp b/src/simulation.cpp
index be03e48553c..09b3764732b 100644
--- a/src/simulation.cpp
+++ b/src/simulation.cpp
@@ -117,6 +117,7 @@ int openmc_simulation_init()
   // Reset global variables -- this is done before loading state point (as that
   // will potentially populate k_generation and entropy)
   simulation::current_batch = 0;
+  simulation::ssw_current_file = 1;
   simulation::k_generation.clear();
   simulation::entropy.clear();
   openmc_reset();
@@ -308,6 +309,7 @@ int n_lost_particles {0};
 bool need_depletion_rx {false};
 int restart_batch;
 bool satisfy_triggers {false};
+int ssw_current_file;
 int total_gen {0};
 double total_weight;
 int64_t work_per_rank;
@@ -337,7 +339,7 @@ void allocate_banks()
 
   if (settings::surf_source_write) {
     // Allocate surface source bank
-    simulation::surf_source_bank.reserve(settings::max_surface_particles);
+    simulation::surf_source_bank.reserve(settings::ssw_max_particles);
   }
 }
 
@@ -345,7 +347,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) {
@@ -439,42 +440,49 @@ void finalize_batch()
       std::string source_point_filename = fmt::format("{0}source.{1:0{2}}",
         settings::path_output, simulation::current_batch, w);
       gsl::span<SourceSite> bankspan(simulation::source_bank);
-      if (settings::source_mcpl_write) {
-        write_mcpl_source_point(
-          source_point_filename.c_str(), bankspan, simulation::work_index);
-      } else {
-        write_source_point(
-          source_point_filename.c_str(), bankspan, simulation::work_index);
-      }
+      write_source_point(source_point_filename, bankspan,
+        simulation::work_index, settings::source_mcpl_write);
     }
 
     // Write a continously-overwritten source point if requested.
     if (settings::source_latest) {
-
       // note: correct file extension appended automatically
       auto filename = settings::path_output + "source";
       gsl::span<SourceSite> bankspan(simulation::source_bank);
-      if (settings::source_mcpl_write) {
-        write_mcpl_source_point(
-          filename.c_str(), bankspan, simulation::work_index);
-      } else {
-        write_source_point(filename.c_str(), bankspan, simulation::work_index);
-      }
+      write_source_point(filename.c_str(), bankspan, simulation::work_index,
+        settings::source_mcpl_write);
     }
   }
 
   // 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);
+      simulation::ssw_current_file <= settings::ssw_max_files) {
+    bool last_batch = (simulation::current_batch == settings::n_batches);
+    if (simulation::surf_source_bank.full() || last_batch) {
+      // Determine appropriate filename
+      auto filename = fmt::format("{}surface_source.{}", settings::path_output,
+        simulation::current_batch);
+      if (settings::ssw_max_files == 1 ||
+          (simulation::ssw_current_file == 1 && last_batch)) {
+        filename = settings::path_output + "surface_source";
+      }
+
+      // Get span of source bank and calculate parallel index vector
+      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());
+
+      // Write surface source file
+      write_source_point(
+        filename, surfbankspan, surf_work_index, settings::surf_mcpl_write);
+
+      // Reset surface source bank and increment counter
+      simulation::surf_source_bank.clear();
+      if (!last_batch && settings::ssw_max_files >= 1) {
+        simulation::surf_source_bank.reserve(settings::ssw_max_particles);
+      }
+      ++simulation::ssw_current_file;
     }
   }
 }
diff --git a/src/state_point.cpp b/src/state_point.cpp
index c7b7d6ad85c..adc026fa3c2 100644
--- a/src/state_point.cpp
+++ b/src/state_point.cpp
@@ -15,6 +15,7 @@
 #include "openmc/error.h"
 #include "openmc/file_utils.h"
 #include "openmc/hdf5_interface.h"
+#include "openmc/mcpl_interface.h"
 #include "openmc/mesh.h"
 #include "openmc/message_passing.h"
 #include "openmc/mgxs_interface.h"
@@ -568,8 +569,23 @@ hid_t h5banktype()
   return banktype;
 }
 
-void write_source_point(const char* filename, gsl::span<SourceSite> source_bank,
-  const vector<int64_t>& bank_index)
+void write_source_point(std::string filename, gsl::span<SourceSite> source_bank,
+  const vector<int64_t>& bank_index, bool use_mcpl)
+{
+  std::string ext = use_mcpl ? "mcpl" : "h5";
+  write_message("Creating source file {}.{} with {} particles ...", filename,
+    ext, source_bank.size(), 5);
+
+  // Dispatch to appropriate function based on file type
+  if (use_mcpl) {
+    write_mcpl_source_point(filename.c_str(), source_bank, bank_index);
+  } else {
+    write_h5_source_point(filename.c_str(), source_bank, bank_index);
+  }
+}
+
+void write_h5_source_point(const char* filename,
+  gsl::span<SourceSite> source_bank, const vector<int64_t>& bank_index)
 {
   // When using parallel HDF5, the file is written to collectively by all
   // processes. With MPI-only, the file is opened and written by the master
diff --git a/tests/regression_tests/dagmc/legacy/test.py b/tests/regression_tests/dagmc/legacy/test.py
index 884264f30ef..b6b1e376b00 100644
--- a/tests/regression_tests/dagmc/legacy/test.py
+++ b/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()
\ No newline at end of file
diff --git a/tests/regression_tests/surface_source/test.py b/tests/regression_tests/surface_source/test.py
index 81bba0c9463..251e9e9ad4a 100644
--- a/tests/regression_tests/surface_source/test.py
+++ b/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()
\ No newline at end of file
diff --git a/tests/regression_tests/surface_source_write/test.py b/tests/regression_tests/surface_source_write/test.py
index 38581736ec4..fe11d68a6e4 100644
--- a/tests/regression_tests/surface_source_write/test.py
+++ b/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()
\ No newline at end of file
diff --git a/tests/unit_tests/test_surface_source_write.py b/tests/unit_tests/test_surface_source_write.py
index 47236307330..6f18d32b718 100644
--- a/tests/unit_tests/test_surface_source_write.py
+++ b/tests/unit_tests/test_surface_source_write.py
@@ -32,6 +32,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, "surface_ids": [2], "max_source_files": 1},
     ],
 )
 def test_xml_serialization(parameter, run_in_tmpdir):
@@ -44,6 +45,58 @@ 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
+    h1 = openmc.Material(name="H1")
+    h1.add_nuclide("H1", 1.0)
+    h1.set_density('g/cm3', 1e-7)
+
+    # Geometry
+    radius = 1.0
+    sphere = openmc.Sphere(r=radius, boundary_type="vacuum")
+    cell = openmc.Cell(region=-sphere, fill=h1)
+    model.geometry = openmc.Geometry([cell])
+
+    # Settings
+    model.settings = openmc.Settings()
+    model.settings.run_mode = "fixed source"
+    model.settings.particles = 100
+    model.settings.batches = 3
+    model.settings.seed = 1
+
+    distribution = openmc.stats.Point()
+    model.settings.source = openmc.IndependentSource(space=distribution)
+    return model
+
+
+@pytest.mark.parametrize(
+    "max_particles, max_source_files",
+    [
+        (100, 2),
+        (100, 3),
+        (100, 1),
+    ],
+)
+def test_number_surface_source_file_created(max_particles, max_source_files,
+                                            run_in_tmpdir, model):
+    """Check the number of surface source files written."""
+    model.settings.surf_source_write = {
+        "max_particles": max_particles,
+        "max_source_files": max_source_files
+    }
+    model.run()
+    should_be_numbered = max_source_files > 1
+    for i in range(1, max_source_files + 1):
+        if should_be_numbered:
+            assert Path(f"surface_source.{i}.h5").exists()
+    if not should_be_numbered:
+        assert Path("surface_source.h5").exists()
+
 ERROR_MSG_1 = (
     "A maximum number of particles needs to be specified "
     "using the 'max_particles' parameter to store surface "