Enabling MCPL source files to be read when using surf_source_read (#3472)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

include/openmc/source.h

@@ -173,7 +173,7 @@ protected:
   SourceSite sample(uint64_t* seed) const override;
 
 private:
-  vector<SourceSite> sites_; //!< Source sites from a file
+  vector<SourceSite> sites_; //!< Source sites
 };
 
 //==============================================================================

src/source.cpp

@@ -417,11 +417,7 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
 FileSource::FileSource(pugi::xml_node node) : Source(node)
 {
   auto path = get_node_value(node, "file", false, true);
-  if (ends_with(path, ".mcpl") || ends_with(path, ".mcpl.gz")) {
-    sites_ = mcpl_source_sites(path);
-  } else {
-    this->load_sites_from_file(path);
-  }
+  load_sites_from_file(path);
 }
 
 FileSource::FileSource(const std::string& path)
@@ -431,30 +427,33 @@ FileSource::FileSource(const std::string& path)
 
 void FileSource::load_sites_from_file(const std::string& path)
 {
-  // Check if source file exists
-  if (!file_exists(path)) {
-    fatal_error(fmt::format("Source file '{}' does not exist.", path));
+  // If MCPL file, use the dedicated file reader
+  if (ends_with(path, ".mcpl") || ends_with(path, ".mcpl.gz")) {
+    sites_ = mcpl_source_sites(path);
+  } else {
+    // Check if source file exists
+    if (!file_exists(path)) {
+      fatal_error(fmt::format("Source file '{}' does not exist.", path));
+    }
+
+    write_message(6, "Reading source file from {}...", path);
+
+    // Open the binary file
+    hid_t file_id = file_open(path, 'r', true);
+
+    // Check to make sure this is a source file
+    std::string filetype;
+    read_attribute(file_id, "filetype", filetype);
+    if (filetype != "source" && filetype != "statepoint") {
+      fatal_error("Specified starting source file not a source file type.");
+    }
+
+    // Read in the source particles
+    read_source_bank(file_id, sites_, false);
+
+    // Close file
+    file_close(file_id);
   }
-
-  // Read the source from a binary file instead of sampling from some
-  // assumed source distribution
-  write_message(6, "Reading source file from {}...", path);
-
-  // Open the binary file
-  hid_t file_id = file_open(path, 'r', true);
-
-  // Check to make sure this is a source file
-  std::string filetype;
-  read_attribute(file_id, "filetype", filetype);
-  if (filetype != "source" && filetype != "statepoint") {
-    fatal_error("Specified starting source file not a source file type.");
-  }
-
-  // Read in the source particles
-  read_source_bank(file_id, sites_, false);
-
-  // Close file
-  file_close(file_id);
 }
 
 SourceSite FileSource::sample(uint64_t* seed) const

tests/regression_tests/surface_source/inputs_true_read_mcpl.dat

@@ -0,0 +1,33 @@
+<?xml version='1.0' encoding='utf-8'?>
+<model>
+  <materials>
+  </materials>
+  <geometry>
+    <cell id="1" material="void" region="-1" universe="1"/>
+    <cell id="2" material="void" region="1 -2" universe="1"/>
+    <cell id="3" material="void" region="2 -3" universe="1"/>
+    <cell id="4" material="void" region="3 -4" universe="1"/>
+    <surface coeffs="0.0 0.0 0.0 1.0" id="1" type="sphere"/>
+    <surface coeffs="0.0 0.0 0.0 2.0" id="2" type="sphere"/>
+    <surface coeffs="0.0 0.0 0.0 2.5" id="3" type="sphere"/>
+    <surface boundary="vacuum" coeffs="0.0 0.0 0.0 4.0" id="4" type="sphere"/>
+  </geometry>
+  <settings>
+    <run_mode>fixed source</run_mode>
+    <particles>1000</particles>
+    <batches>10</batches>
+    <surf_source_read>
+      <path>surface_source_true.mcpl</path>
+    </surf_source_read>
+    <seed>1</seed>
+  </settings>
+  <tallies>
+    <filter id="1" type="cell">
+      <bins>3</bins>
+    </filter>
+    <tally id="1">
+      <filters>1</filters>
+      <scores>flux</scores>
+    </tally>
+  </tallies>
+</model>

tests/regression_tests/surface_source/inputs_true_write_mcpl.dat

@@ -0,0 +1,40 @@
+<?xml version='1.0' encoding='utf-8'?>
+<model>
+  <materials>
+  </materials>
+  <geometry>
+    <cell id="1" material="void" region="-1" universe="1"/>
+    <cell id="2" material="void" region="1 -2" universe="1"/>
+    <cell id="3" material="void" region="2 -3" universe="1"/>
+    <cell id="4" material="void" region="3 -4" universe="1"/>
+    <surface coeffs="0.0 0.0 0.0 1.0" id="1" type="sphere"/>
+    <surface coeffs="0.0 0.0 0.0 2.0" id="2" type="sphere"/>
+    <surface coeffs="0.0 0.0 0.0 2.5" id="3" type="sphere"/>
+    <surface boundary="vacuum" coeffs="0.0 0.0 0.0 4.0" id="4" type="sphere"/>
+  </geometry>
+  <settings>
+    <run_mode>fixed source</run_mode>
+    <particles>1000</particles>
+    <batches>10</batches>
+    <source particle="neutron" strength="1.0" type="independent">
+      <space type="point">
+        <parameters>0 0 0</parameters>
+      </space>
+    </source>
+    <surf_source_write>
+      <surface_ids>1</surface_ids>
+      <mcpl>true</mcpl>
+      <max_particles>1000</max_particles>
+    </surf_source_write>
+    <seed>1</seed>
+  </settings>
+  <tallies>
+    <filter id="1" type="cell">
+      <bins>3</bins>
+    </filter>
+    <tally id="1">
+      <filters>1</filters>
+      <scores>flux</scores>
+    </tally>
+  </tallies>
+</model>

tests/regression_tests/surface_source/test.py

@@ -10,6 +10,25 @@ from tests.testing_harness import PyAPITestHarness
 from tests.regression_tests import config
 
 
+def mcpl_to_array(filepath):
+    import mcpl
+
+    source = []
+    with mcpl.MCPLFile(filepath) as f:
+        for p in f.particles:
+            source.append(
+                [
+                    *tuple(p.position),
+                    *tuple(p.direction),
+                    1.0e6 * p.ekin,
+                    1.0e-3 * p.time,
+                    p.weight,
+                    p.pdgcode,
+                ]
+            )
+    return np.sort(np.array(source), axis=0)
+
+
 def assert_structured_arrays_close(arr1, arr2, rtol=1e-5, atol=1e-8):
     assert arr1.dtype == arr2.dtype
 
@@ -24,8 +43,7 @@ def assert_structured_arrays_close(arr1, arr2, rtol=1e-5, atol=1e-8):
 @pytest.fixture
 def model(request):
     openmc.reset_auto_ids()
-    marker = request.node.get_closest_marker("surf_source_op")
-    surf_source_op = marker.args[0]
+    operation, file_format = request.node.get_closest_marker("params").args
 
     openmc_model = openmc.model.Model()
 
@@ -54,15 +72,19 @@ def model(request):
     openmc_model.settings.batches = 10
     openmc_model.settings.seed = 1
 
-    if surf_source_op == 'write':
+    if operation == 'write':
         point = openmc.stats.Point((0, 0, 0))
         pt_src = openmc.IndependentSource(space=point)
         openmc_model.settings.source = pt_src
 
-        openmc_model.settings.surf_source_write = {'surface_ids': [1],
-                                                   'max_particles': 1000}
-    elif surf_source_op == 'read':
-        openmc_model.settings.surf_source_read = {'path': 'surface_source_true.h5'}
+        surf_source_write_settings = {'surface_ids': [1],
+                                      'max_particles': 1000}
+        if file_format == "mcpl":
+            surf_source_write_settings["mcpl"] = True
+
+        openmc_model.settings.surf_source_write = surf_source_write_settings
+    elif operation == 'read':
+        openmc_model.settings.surf_source_read = {'path': f"surface_source_true.{file_format}"}
 
     # Tallies
     tal = openmc.Tally()
@@ -75,22 +97,30 @@ def model(request):
 
 
 class SurfaceSourceTestHarness(PyAPITestHarness):
+    def __init__(self, statepoint_name, model=None, inputs_true=None, file_format="h5"):
+        super().__init__(statepoint_name, model, inputs_true)
+        self.file_format = file_format
+
     def _test_output_created(self):
-        """Make sure surface_source.h5 has also been created."""
+        """Make sure the surface_source file has also been created."""
         super()._test_output_created()
-        # Check if 'surface_source.h5' has been created.
         if self._model.settings.surf_source_write:
-            assert os.path.exists('surface_source.h5'), \
+            assert os.path.exists(f"surface_source.{self.file_format}"), \
                 'Surface source file does not exist.'
 
     def _compare_output(self):
         """Make sure the current surface_source.h5 agree with the reference."""
         if self._model.settings.surf_source_write:
-            with h5py.File("surface_source_true.h5", 'r') as f:
-                source_true = np.sort(f['source_bank'][()])
-            with h5py.File("surface_source.h5", 'r') as f:
-                source_test = np.sort(f['source_bank'][()])
-            assert_structured_arrays_close(source_true, source_test, atol=1e-07)
+            if self.file_format == "h5":
+                with h5py.File("surface_source_true.h5", 'r') as f:
+                    source_true = np.sort(f['source_bank'][()])
+                with h5py.File("surface_source.h5", 'r') as f:
+                    source_test = np.sort(f['source_bank'][()])
+                assert_structured_arrays_close(source_true, source_test, atol=1e-07)
+            elif self.file_format == "mcpl":
+                source_true = mcpl_to_array("surface_source_true.mcpl")
+                source_test = mcpl_to_array("surface_source.mcpl")
+                np.testing.assert_allclose(source_true, source_test, rtol=1e-5, atol=1e-7)
 
     def execute_test(self):
         """Build input XMLs, run OpenMC, check output and results."""
@@ -111,30 +141,56 @@ class SurfaceSourceTestHarness(PyAPITestHarness):
     def _overwrite_results(self):
         """Overwrite the results_true with the results_test."""
         shutil.copyfile('results_test.dat', 'results_true.dat')
-        if os.path.exists('surface_source.h5'):
-            shutil.copyfile('surface_source.h5', 'surface_source_true.h5')
+        if os.path.exists(f"surface_source.{self.file_format}"):
+            shutil.copyfile(f"surface_source.{self.file_format}", f"surface_source_true.{self.file_format}")
 
     def _cleanup(self):
         """Delete statepoints, tally, and test files."""
         super()._cleanup()
-        fs = 'surface_source.h5'
+        fs = f"surface_source.{self.file_format}"
         if os.path.exists(fs):
             os.remove(fs)
 
 
-@pytest.mark.surf_source_op('write')
-def test_surface_source_write(model, monkeypatch):
-    # Test result is based on 1 MPI process
-    monkeypatch.setitem(config, "mpi_np", "1")
-    harness = SurfaceSourceTestHarness('statepoint.10.h5',
-                                       model,
-                                       'inputs_true_write.dat')
+@pytest.mark.params('write', 'h5')
+def test_surface_source_write(model, monkeypatch, request):
+    monkeypatch.setitem(config, "mpi_np", "1")  # Results generated with 1 MPI process
+    operation, file_format = request.node.get_closest_marker("params").args
+    harness = SurfaceSourceTestHarness(
+        "statepoint.10.h5", model, f"inputs_true_{operation}_{file_format}.dat",
+        file_format=file_format
+    )
     harness.main()
 
 
-@pytest.mark.surf_source_op('read')
-def test_surface_source_read(model):
-    harness = SurfaceSourceTestHarness('statepoint.10.h5',
-                                       model,
-                                       'inputs_true_read.dat')
+@pytest.mark.params('read', 'h5')
+def test_surface_source_read(model, request):
+    operation, file_format = request.node.get_closest_marker("params").args
+    harness = SurfaceSourceTestHarness(
+        "statepoint.10.h5", model, f"inputs_true_{operation}_{file_format}.dat",
+        file_format=file_format
+    )
+    harness.main()
+
+
+@pytest.mark.skipif(shutil.which("mcpl-config") is None, reason="MCPL is not available.")
+@pytest.mark.params('write', 'mcpl')
+def test_surface_source_write_mcpl(model, monkeypatch, request):
+    monkeypatch.setitem(config, "mpi_np", "1")  # Results generated with 1 MPI process
+    operation, file_format = request.node.get_closest_marker("params").args
+    harness = SurfaceSourceTestHarness(
+        "statepoint.10.h5", model, f"inputs_true_{operation}_{file_format}.dat",
+        file_format=file_format
+    )
+    harness.main()
+
+
+@pytest.mark.skipif(shutil.which("mcpl-config") is None, reason="MCPL is not available.")
+@pytest.mark.params('read', 'mcpl')
+def test_surface_source_read_mcpl(model, request):
+    operation, file_format = request.node.get_closest_marker("params").args
+    harness = SurfaceSourceTestHarness(
+        "statepoint.10.h5", model, f"inputs_true_{operation}_{file_format}.dat",
+        file_format=file_format
+    )
     harness.main()