Adding vtkhdf

src/openmc2dolfinx/__init__.py

@@ -6,6 +6,6 @@
     __version__ = "unknown"
 
 
-from .core import StructuredGridReader, UnstructuredMeshReader
+from .core import StructuredGridReader, UnstructuredMeshReader, VTKHDFUnstructuredMeshReader
 
-__all__ = ["StructuredGridReader, UnstructuredMeshReader"]
+__all__ = ["StructuredGridReader", "UnstructuredMeshReader", "VTKHDFUnstructuredMeshReader"]

src/openmc2dolfinx/core.py

@@ -8,7 +8,7 @@
 import ufl
 from dolfinx.mesh import create_mesh
 
-__all__ = ["StructuredGridReader", "UnstructuredMeshReader"]
+__all__ = ["StructuredGridReader", "UnstructuredMeshReader", "VTKHDFUnstructuredMeshReader"]
 
 
 class OpenMC2dolfinx(pyvista.VTKDataSetReader):
@@ -147,3 +147,83 @@ def cell_connectivity(self):
         if self._cell_connectivity is None:
             self.get_connectivity()
         return self._cell_connectivity
+
+
+class VTKHDFUnstructuredMeshReader:
+    """
+    VTKHDF Unstructured Mesh Reader
+
+    Reads an OpenMC .vtkhdf results file with unstructured meshes and converts
+    the data into a dolfinx.fem.Function
+
+    Args:
+        path: the path to the OpenMC .vtkhdf file
+
+    Attributes:
+        data: the mesh and results from the OpenMC .vtkhdf file
+        dolfinx_mesh: the dolfinx mesh
+
+    Example:
+    .. code-block:: python
+        reader = VTKHDFUnstructuredMeshReader("path/to/file.vtkhdf")
+        dolfinx_function = reader.create_dolfinx_function()
+    """
+
+    cell_type = "tetrahedron"
+
+    def __init__(self, path: str):
+        self.path = path
+        self.data: pyvista.core.pointset.UnstructuredGrid | None = None
+        self.dolfinx_mesh: dolfinx.mesh.Mesh | None = None
+
+    def read(self) -> pyvista.core.pointset.UnstructuredGrid:
+        """Reads the VTKHDF file using pyvista."""
+        self.data = pyvista.read(self.path)
+        return self.data
+
+    @property
+    def cell_connectivity(self) -> np.ndarray:
+        """Returns the cell connectivity for tetrahedra (VTK cell type 10)."""
+        if self.data is None:
+            self.read()
+        return self.data.cells_dict[10]
+
+    def create_dolfinx_mesh(self):
+        """Creates the dolfinx mesh from the VTKHDF data."""
+        if self.data is None:
+            self.read()
+
+        degree = 1
+        cell = ufl.Cell(self.cell_type)
+        mesh_element = basix.ufl.element(
+            "Lagrange", cell.cellname(), degree, shape=(3,)
+        )
+
+        mesh_ufl = ufl.Mesh(mesh_element)
+        self.dolfinx_mesh = create_mesh(
+            comm=MPI.COMM_WORLD,
+            cells=self.cell_connectivity,
+            x=self.data.points,
+            e=mesh_ufl,
+        )
+
+    def create_dolfinx_function(self, data: str = "mean") -> dolfinx.fem.Function:
+        """Converts VTKHDF cell data to a dolfinx Function.
+
+        Arguments:
+            data: the name of the cell data array to extract (default: "mean")
+
+        Returns:
+            dolfinx function with OpenMC results mapped
+        """
+        if self.dolfinx_mesh is None:
+            self.create_dolfinx_mesh()
+
+        function_space = dolfinx.fem.functionspace(self.dolfinx_mesh, ("DG", 0))
+        u = dolfinx.fem.Function(function_space)
+
+        u.x.array[:] = self.data.cell_data[data][
+            self.dolfinx_mesh.topology.original_cell_index
+        ]
+
+        return u

test/test_unstructured_mesh.py

@@ -5,7 +5,7 @@
 from mpi4py import MPI
 import pytest
 
-from openmc2dolfinx import UnstructuredMeshReader
+from openmc2dolfinx import UnstructuredMeshReader, VTKHDFUnstructuredMeshReader
 
 
 @pytest.fixture
@@ -96,3 +96,16 @@ def test_download_from_pyvista_examples(tmpdir):
     # export to vtk for visualisation
     writer = dolfinx.io.VTXWriter(MPI.COMM_WORLD, tmpdir + "/out.bp", u, "BP5")
     writer.write(t=0)
+
+
+def test_vtkhdf_unstructured_mesh_reader():
+    """Test VTKHDFUnstructuredMeshReader with OpenMC VTKHDF tally file."""
+    import pathlib
+
+    vtkhdf_file = pathlib.Path(__file__).parent / "tally.vtkhdf"
+
+    reader = VTKHDFUnstructuredMeshReader(str(vtkhdf_file))
+    dolfinx_function = reader.create_dolfinx_function()
+
+    assert isinstance(dolfinx_function, fem.Function)
+    assert dolfinx_function.function_space.mesh.topology.dim == 3