Add VTU export for Unstructured meshes

openmc/mesh.py

@@ -9,6 +9,7 @@
 import h5py
 import lxml.etree as ET
 import numpy as np
+from pathlib import Path
 
 import openmc
 import openmc.checkvalue as cv
@@ -2358,17 +2359,21 @@ def write_vtk_mesh(self, **kwargs):
         self.write_data_to_vtk(**kwargs)
 
     def write_data_to_vtk(
-            self,
-            filename: PathLike | None  = None,
-            datasets: dict | None = None,
-            volume_normalization: bool = True
+        self,
+        filename: PathLike | None = None,
+        datasets: dict | None = None,
+        volume_normalization: bool = True,
     ):
         """Map data to unstructured VTK mesh elements.
 
+        If filename is None, then a filename will be generated based on the mesh ID,
+        and exported to VTK format.
+
         Parameters
         ----------
         filename : str or pathlib.Path
-            Name of the VTK file to write
+            Name of the VTK file to write. If the filename ends in '.vtu' then a VTU
+            format file will be written, if the filename ends in '.vtk' then a legacy VTK file will be written.
         datasets : dict
             Dictionary whose keys are the data labels
             and values are numpy appropriately sized arrays
@@ -2377,36 +2382,48 @@ def write_data_to_vtk(
             Whether or not to normalize the data by the
             volume of the mesh elements
         """
-        import vtk
-        from vtk.util import numpy_support as nps
+        from vtkmodules.util import numpy_support
+        from vtkmodules import vtkCommonCore
+        from vtkmodules import vtkCommonDataModel
+        from vtkmodules import vtkIOLegacy
+        from vtkmodules import vtkIOXML
 
         if self.connectivity is None or self.vertices is None:
-            raise RuntimeError('This mesh has not been '
-                               'loaded from a statepoint file.')
+            raise RuntimeError(
+                "This mesh has not been " "loaded from a statepoint file."
+            )
 
         if filename is None:
-            filename = f'mesh_{self.id}.vtk'
+            filename = f"mesh_{self.id}.vtk"
+
+        if Path(filename).suffix == ".vtk":
+            writer = vtkIOLegacy.vtkUnstructuredGridWriter()
 
-        writer = vtk.vtkUnstructuredGridWriter()
+        elif Path(filename).suffix == ".vtu":
+            writer = vtkIOXML.vtkXMLUnstructuredGridWriter()
+            writer.SetCompressorTypeToZLib()
+            writer.SetDataModeToBinary()
 
         writer.SetFileName(str(filename))
 
-        grid = vtk.vtkUnstructuredGrid()
+        grid = vtkCommonDataModel.vtkUnstructuredGrid()
 
-        vtk_pnts = vtk.vtkPoints()
-        vtk_pnts.SetData(nps.numpy_to_vtk(self.vertices))
-        grid.SetPoints(vtk_pnts)
+        points = vtkCommonCore.vtkPoints()
+        points.SetData(numpy_support.numpy_to_vtk(self.vertices))
+        grid.SetPoints(points)
 
         n_skipped = 0
         for elem_type, conn in zip(self.element_types, self.connectivity):
             if elem_type == self._LINEAR_TET:
-                elem = vtk.vtkTetra()
+                elem = vtkCommonDataModel.vtkTetra()
             elif elem_type == self._LINEAR_HEX:
-                elem = vtk.vtkHexahedron()
+                elem = vtkCommonDataModel.vtkHexahedron()
             elif elem_type == self._UNSUPPORTED_ELEM:
                 n_skipped += 1
+                continue
             else:
-                raise RuntimeError(f'Invalid element type {elem_type} found')
+                raise RuntimeError(f"Invalid element type {elem_type} found")
+
             for i, c in enumerate(conn):
                 if c == -1:
                     break
@@ -2415,30 +2432,36 @@ def write_data_to_vtk(
             grid.InsertNextCell(elem.GetCellType(), elem.GetPointIds())
 
         if n_skipped > 0:
-            warnings.warn(f'{n_skipped} elements were not written because '
-                          'they are not of type linear tet/hex')
+            warnings.warn(
+                f"{n_skipped} elements were not written because "
+                "they are not of type linear tet/hex"
+            )
 
         # check that datasets are the correct size
         datasets_out = []
         if datasets is not None:
             for name, data in datasets.items():
                 if data.shape != self.dimension:
-                    raise ValueError(f'Cannot apply dataset "{name}" with '
-                                     f'shape {data.shape} to mesh {self.id} '
-                                     f'with dimensions {self.dimension}')
+                    raise ValueError(
+                        f'Cannot apply dataset "{name}" with '
+                        f"shape {data.shape} to mesh {self.id} "
+                        f"with dimensions {self.dimension}"
+                    )
 
             if volume_normalization:
                 for name, data in datasets.items():
                     if np.issubdtype(data.dtype, np.integer):
-                        warnings.warn(f'Integer data set "{name}" will '
-                                      'not be volume-normalized.')
+                        warnings.warn(
+                            f'Integer data set "{name}" will '
+                            "not be volume-normalized."
+                        )
                         continue
                     data /= self.volumes
 
             # add data to the mesh
             for name, data in datasets.items():
                 datasets_out.append(data)
-                arr = vtk.vtkDoubleArray()
+                arr = vtkCommonCore.vtkDoubleArray()
                 arr.SetName(name)
                 arr.SetNumberOfTuples(data.size)
 

tests/regression_tests/unstructured_mesh/test_mesh_dagmc_tets.vtk

@@ -0,0 +1,159 @@
+# vtk DataFile Version 2.0
+made_with_cad_to_dagmc_package, Created by Gmsh 4.12.1 
+ASCII
+DATASET UNSTRUCTURED_GRID
+POINTS 14 double
+-0.5 -0.5 0.5
+-0.5 -0.5 -0.5
+-0.5 0.5 0.5
+-0.5 0.5 -0.5
+0.5 -0.5 0.5
+0.5 -0.5 -0.5
+0.5 0.5 0.5
+0.5 0.5 -0.5
+-0.5 0 0
+0.5 0 0
+0 -0.5 0
+0 0.5 0
+0 0 -0.5
+0 0 0.5
+
+CELLS 68 268
+1 0
+1 1
+1 2
+1 3
+1 4
+1 5
+1 6
+1 7
+2 1 0
+2 0 2
+2 3 2
+2 1 3
+2 5 4
+2 4 6
+2 7 6
+2 5 7
+2 1 5
+2 0 4
+2 3 7
+2 2 6
+3 1 0 8
+3 0 2 8
+3 3 1 8
+3 2 3 8
+3 5 9 4
+3 4 9 6
+3 7 9 5
+3 6 9 7
+3 0 1 10
+3 4 0 10
+3 1 5 10
+3 5 4 10
+3 2 11 3
+3 6 11 2
+3 3 11 7
+3 7 11 6
+3 1 3 12
+3 5 1 12
+3 3 7 12
+3 7 5 12
+3 0 13 2
+3 4 13 0
+3 2 13 6
+3 6 13 4
+4 13 8 12 10
+4 11 8 12 13
+4 12 11 13 9
+4 10 12 13 9
+4 12 3 11 7
+4 13 2 8 0
+4 8 12 1 3
+4 11 3 8 2
+4 10 8 1 0
+4 0 10 13 4
+4 1 12 10 5
+4 11 2 13 6
+4 4 9 13 6
+4 6 9 11 7
+4 10 9 4 5
+4 7 9 12 5
+4 3 8 12 11
+4 1 12 8 10
+4 8 11 2 13
+4 10 13 8 0
+4 13 10 9 4
+4 11 13 9 6
+4 12 11 9 7
+4 9 10 12 5
+
+CELL_TYPES 68
+1
+1
+1
+1
+1
+1
+1
+1
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10

tests/unit_tests/test_mesh.py

@@ -3,6 +3,9 @@
 import numpy as np
 import pytest
 import openmc
+from pathlib import Path
+import tempfile
+from vtkmodules import vtkIOLegacy, vtkIOXML
 
 
 @pytest.mark.parametrize("val_left,val_right", [(0, 0), (-1., -1.), (2.0, 2)])
@@ -398,6 +401,81 @@ def test_umesh_roundtrip(run_in_tmpdir, request):
     assert umesh.id == xml_mesh.id
 
 
+@pytest.mark.parametrize('export_type', ('.vtk','.vtu'))
+def test_umesh(request, export_type):
+    """Performs a minimal UnstructuredMesh simulation, reads in the resulting
+    statepoint file and writes the mesh data to vtk and vtkhdf files. It is
+    necessary to read in the unstructured mesh from a statepoint file to ensure
+    it has all the required attributes
+    """
+    surf1 = openmc.Sphere(r=1.0, boundary_type="vacuum")
+    material1 = openmc.Material()
+    material1.add_components(components={"H" : 1.0})
+    material1.set_density('g/cm3', 1.0)
+
+    materials = openmc.Materials([material1])
+    cell1 = openmc.Cell(region=-surf1, fill=material1)
+    geometry = openmc.Geometry([cell1])
+
+    umesh = openmc.UnstructuredMesh(
+       filename= request.path.parent.parent
+        / "regression_tests/unstructured_mesh/test_mesh_dagmc_tets.vtk",
+       library= "moab",
+       mesh_id = 1
+    )
+    # setting ID to make it easier to get the mesh from the statepoint later
+    mesh_filter = openmc.MeshFilter(umesh)
+
+    # Create flux mesh tally to score alpha production
+    mesh_tally = openmc.Tally(name="test_tally")
+    mesh_tally.filters = [mesh_filter]
+    mesh_tally.scores = ["total"]
+
+    tallies = openmc.Tallies([mesh_tally])
+
+    source = openmc.IndependentSource()
+
+    settings = openmc.Settings()
+    settings.run_mode = "fixed source"
+    settings.batches = 2
+    settings.particles = 100
+    settings.source = source
+
+    model = openmc.Model(
+        materials=materials, geometry=geometry, settings=settings, tallies=tallies
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+
+        statepoint_file = model.run(cwd=tmpdir)
+
+        statepoint = openmc.StatePoint(statepoint_file)
+        tally: openmc.Tally = statepoint.get_tally(name="test_tally")
+        umesh_from_sp: openmc.UnstructuredMesh = statepoint.meshes[1]
+
+        datasets={
+            "mean": tally.mean.squeeze(),
+            "std_dev": tally.std_dev.squeeze()
+        }
+
+        filename = Path(tmpdir) / ("test_mesh" + export_type)
+        umesh_from_sp.write_data_to_vtk(datasets=datasets, filename=filename)
+
+        assert Path(filename).exists()
+
+        if export_type == ".vtk":
+            reader = vtkIOLegacy.vtkGenericDataObjectReader()
+            reader.SetFileName(str(filename))
+            reader.Update()
+            assert reader.GetOutput().GetCellData().GetArray("mean")
+
+        elif export_type == ".vtk":
+            reader = vtkIOXML.vtkGenericDataObjectReader()
+            reader.SetFileName(str(filename))
+            reader.Update()
+            assert reader.GetOutput().GetCellData().GetArray("mean")
+
+
 def test_mesh_get_homogenized_materials():
     """Test the get_homogenized_materials method"""
     # Simple model with 1 cm of Fe56 next to 1 cm of H1