diff --git a/src/openmc2dolfinx/__init__.py b/src/openmc2dolfinx/__init__.py index afaa677..263b458 100644 --- a/src/openmc2dolfinx/__init__.py +++ b/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"] diff --git a/src/openmc2dolfinx/core.py b/src/openmc2dolfinx/core.py index dbc8373..c6e336a 100644 --- a/src/openmc2dolfinx/core.py +++ b/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 diff --git a/test/tally.vtkhdf b/test/tally.vtkhdf new file mode 100644 index 0000000..ba1931b Binary files /dev/null and b/test/tally.vtkhdf differ diff --git a/test/test_unstructured_mesh.py b/test/test_unstructured_mesh.py index 26ae187..9ae62e0 100644 --- a/test/test_unstructured_mesh.py +++ b/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