unit test added and suggested changes

docs/source/io_formats/settings.rst

@@ -924,10 +924,12 @@ attributes/sub-elements:
     *Default*: None
 
   :batches:
-    A list of integers indicating at what batches a surface source file 
-    should be written.
+    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 a 
+    ``surface_source.h5`` file will be created in the last batch.
 
-    *Default*: Last batch only
+    *Default*: None
 
   :mcpl:
     An optional boolean which indicates if the banked particles should be

docs/source/usersguide/settings.rst

@@ -336,8 +336,9 @@ 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. As an 
-example, to write a surface source file every five batches::
+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 = {

src/simulation.cpp

@@ -465,7 +465,8 @@ void finalize_batch()
   // Write out surface source if requested.
   if (settings::surf_source_write) {
     if (contains(settings::surface_source_batch, simulation::current_batch) ||
-        simulation::current_batch == settings::n_batches) {
+        (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
@@ -488,7 +489,8 @@ void finalize_batch()
         write_source_point(filename.c_str(), surfbankspan, surf_work_index);
       }
       simulation::surf_source_bank.clear();
-      simulation::surf_source_bank.reserve(settings::max_surface_particles);
+      if (simulation::current_batch < settings::n_batches)
+        simulation::surf_source_bank.reserve(settings::max_surface_particles);
     }
   }
 }

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 "