Adding roundtrip test for mesh sampling

openmc/mesh.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
+from collections import OrderedDict
 from math import pi
 from numbers import Real, Integral
 from pathlib import Path
@@ -1937,3 +1938,15 @@ class UnstructuredMesh(MeshBase):
         length_multiplier = float(get_text(elem, 'length_multiplier', 1.0))
 
         return cls(filename, library, mesh_id, '', length_multiplier)
+
+
+def read_meshes(tree):
+    """Reads all mesh nodes under an XML tree
+    """
+    out = OrderedDict()
+    root = tree.getroot()
+    for mesh_elem in root.iter('mesh'):
+        mesh = MeshBase.from_xml_element(mesh_elem)
+        out[mesh.id] = mesh
+
+    return out

openmc/settings.py

@@ -16,7 +16,7 @@ from openmc.stats.multivariate import MeshSpatial
 from . import RegularMesh, Source, VolumeCalculation, WeightWindows
 from ._xml import clean_indentation, get_text, reorder_attributes
 from openmc.checkvalue import PathLike
-from .mesh import MeshBase
+from .mesh import MeshBase, read_meshes
 
 
 class RunMode(Enum):
@@ -1328,19 +1328,7 @@ class Settings:
 
     def _source_from_xml_element(self, root, meshes=None):
         for elem in root.findall('source'):
-            src = Source.from_xml_element(elem)
-            if isinstance(src.space, MeshSpatial):
-                mesh_id = int(get_text(elem, 'mesh'))
-                if mesh_id not in meshes:
-                    path = f"./mesh[@id='{mesh_id}']"
-                    mesh_elem = root.find(path)
-                    if mesh_elem is not None:
-                        mesh = MeshBase.from_xml_element(mesh_elem)
-                    meshes[mesh.id] = mesh
-                try:
-                    src.space.mesh = meshes[mesh_id]
-                except KeyError as e:
-                    raise e(f'Mesh with ID {mesh_id} was not found.')
+            src = Source.from_xml_element(elem, meshes)
             # add newly constructed source object to the list
             self.source.append(src)
 
@@ -1774,4 +1762,5 @@ class Settings:
         """
         tree = ET.parse(path)
         root = tree.getroot()
-        return cls.from_xml_element(root)
+        meshes = read_meshes(tree)
+        return cls.from_xml_element(root, meshes)

openmc/source.py

@@ -258,13 +258,16 @@ class Source:
         return element
 
     @classmethod
-    def from_xml_element(cls, elem: ET.Element) -> 'openmc.Source':
+    def from_xml_element(cls, elem: ET.Element, meshes=None) -> 'openmc.Source':
         """Generate source from an XML element
 
         Parameters
         ----------
         elem : xml.etree.ElementTree.Element
             XML element
+        meshes : dict
+            Dictionary with mesh IDs as keys and openmc.MeshBase instaces as
+            values
 
         Returns
         -------
@@ -313,7 +316,7 @@ class Source:
 
         space = elem.find('space')
         if space is not None:
-            source.space = Spatial.from_xml_element(space)
+            source.space = Spatial.from_xml_element(space, meshes)
 
         angle = elem.find('angle')
         if angle is not None:

openmc/stats/multivariate.py

@@ -262,7 +262,7 @@ class Spatial(ABC):
 
     @classmethod
     @abstractmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem, meshes=None):
         distribution = get_text(elem, 'type')
         if distribution == 'cartesian':
             return CartesianIndependent.from_xml_element(elem)
@@ -275,7 +275,7 @@ class Spatial(ABC):
         elif distribution == 'point':
             return Point.from_xml_element(elem)
         elif distribution == 'mesh':
-            return MeshSpatial.from_xml_element(elem)
+            return MeshSpatial.from_xml_element(elem, meshes)
 
 
 class CartesianIndependent(Spatial):
@@ -714,13 +714,16 @@ class MeshSpatial(Spatial):
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem, meshes):
         """Generate spatial distribution from an XML element
 
         Parameters
         ----------
         elem : xml.etree.ElementTree.Element
             XML element
+        meshes : dict
+            A dictionary with mesh IDs as keys and openmc.MeshBase instances as
+            values
 
         Returns
         -------
@@ -732,16 +735,16 @@ class MeshSpatial(Spatial):
         mesh_id = int(elem.get('mesh_id'))
 
         # check if this mesh has been read in from another location already
-        if mesh_id not in MESHES:
+        if mesh_id not in meshes:
             raise RuntimeError(f'Could not locate mesh with ID "{mesh_id}"')
 
         volume_normalized = elem.get("volume_normalized")
         volume_normalized = get_text(elem, 'volume_normalized').lower() == 'true'
-        if elem.get('strengths') is not None:
+        strengths = get_text(elem, 'strengths')
+        if strengths is not None:
             strengths = [float(b) for b in get_text(elem, 'strengths').split()]
-        else:
-            strengths = None
-        return cls(MESHES[mesh_id], strengths, volume_normalized)
+
+        return cls(meshes[mesh_id], strengths, volume_normalized)
 
 
 class Box(Spatial):

tests/unit_tests/test_source_mesh.py

@@ -106,10 +106,8 @@ def test_unstructured_mesh_sampling(model, request, test_cases):
         cell_counts = np.zeros((n_cells, n_measurements))
 
         # This model contains 1000 geometry cells. Each cell is a hex
-        # corresponding to 12 of the tets. This test runs 10000 particles. This
+        # corresponding to 12 of the tets. This test runs 1000 samples. This
         #  results in the following average for each cell
-        average_in_hex = n_samples / n_cells
-
         openmc.lib.init([])
 
         # perform many sets of samples and track counts for each cell
@@ -174,4 +172,32 @@ def test_strengths_size_failure(request, model):
 
     with pytest.raises(RuntimeError, match=r'strengths array'), cdtemp([mesh_filename]):
         model.export_to_xml()
-        openmc.run()
+        openmc.run()
+
+def test_roundtrip(run_in_tmpdir, model, request):
+    if not openmc.lib._libmesh_enabled() and not openmc.lib._dagmc_enabled():
+        pytest.skip("Unstructured mesh is not enabled in this build.")
+
+    mesh_filename = Path(request.fspath).parent / 'test_mesh_tets.e'
+    ucd_mesh = openmc.UnstructuredMesh(mesh_filename, library='libmesh')
+
+    if not openmc.lib._libmesh_enabled():
+        ucd_mesh.library = 'moab'
+
+    n_cells = len(model.geometry.get_all_cells())
+
+    space_out = openmc.MeshSpatial(ucd_mesh)
+    space_out.strengths = np.random.rand(n_cells*TETS_PER_VOXEL)
+    model.settings.source = openmc.Source(space=space_out)
+
+    # write out the model
+    model.export_to_xml()
+
+    model_in = openmc.Model.from_xml()
+
+    space_in = model_in.settings.source[0].space
+
+    np.testing.assert_equal(space_out.strengths, space_in.strengths)
+
+    assert space_in.mesh.id == space_out.mesh.id
+    assert space_in.volume_normalized == space_out.volume_normalized