Merge pull request #5 from festim-dev/readers

Readers for structured and unstructured meshes

Author: jhdark

.gitignore

@@ -169,3 +169,9 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# versioning file
+*_version.py
+
+# results files
+*.bp

README.md

@@ -1,2 +1,31 @@
 # openmc2dolfinx
 A repository to handle the conversion of results in OpenMC vtk files to dolfinx functions
+
+
+## Usage
+```python
+from openmc2dolfinx import StructuredGridReader, UnstructuredMeshReader
+import pyvista as pv
+import numpy as np
+import dolfinx
+from mpi4py import MPI
+
+# download an example tetmesh
+filename = pv.examples.download_tetrahedron(load=False)
+
+grid = pv.read(filename)
+
+# assign random cell data
+grid.cell_data["mean"] = np.arange(grid.n_cells)
+grid.save("out.vtk")
+
+# read the vtk file
+reader = UnstructuredMeshReader("out.vtk")
+
+# make a dolfinx function
+u = reader.create_dolfinx_function("mean")
+
+# export to vtk for visualisation
+writer = dolfinx.io.VTXWriter(MPI.COMM_WORLD, "out.bp", u, "BP5")
+writer.write(t=0)
+```

src/openmc2dolfinx/__init__.py

@@ -0,0 +1,11 @@
+from importlib import metadata
+
+try:
+    __version__ = metadata.version("openmc2dolfinx")
+except Exception:
+    __version__ = "unknown"
+
+
+from .core import StructuredGridReader, UnstructuredMeshReader
+
+__all__ = ["StructuredGridReader, UnstructuredMeshReader"]

src/openmc2dolfinx/core.py

@@ -0,0 +1,146 @@
+from mpi4py import MPI
+
+import basix
+import dolfinx
+import numpy as np
+import pyvista
+import pyvista.core.pointset
+import ufl
+from dolfinx.mesh import create_mesh
+
+__all__ = ["StructuredGridReader", "UnstructuredMeshReader"]
+
+
+class OpenMC2dolfinx(pyvista.VTKDataSetReader):
+    """
+    Base OpenMC2Dolfinx Mesh Reader
+
+    Converts OpenMC results data into a dolfinx.fem.Function
+
+    Args:
+        path: the path to the OpenMC .vtk file
+
+    Attributes:
+        data: the mesh and results from the OpenMC .vtk file
+        connectivity: The OpenMC mesh cell connectivity
+        dolfinx_mesh: the dolfinx mesh
+    """
+
+    data: pyvista.core.pointset.UnstructuredGrid | pyvista.core.pointset.StructuredGrid
+    connectivity: np.ndarray
+    dolfinx_mesh: dolfinx.mesh.Mesh = None
+
+    def create_dolfinx_mesh(self):
+        """Creates the dolfinx mesh depending on the type of cell provided
+
+        args:
+            cell_type: the cell type for the dolfinx mesh, defaults to "tetrahedron"
+        """
+
+        # TODO find a way to fix this with abstractmethod and property
+        if not hasattr(self, "cell_type"):
+            raise AttributeError("cell_type must be defined in the child class")
+
+        self.data = self.read()
+        if not hasattr(self, "cell_connectivity"):
+            raise AttributeError("cell_connectivity must be defined in the child class")
+
+        degree = 1  # Set polynomial degree
+
+        cell = ufl.Cell(f"{self.cell_type}")
+        mesh_element = basix.ufl.element(
+            "Lagrange", cell.cellname(), degree, shape=(3,)
+        )
+
+        # Create dolfinx Mesh
+        mesh_ufl = ufl.Mesh(mesh_element)
+        self.dolfinx_mesh = create_mesh(
+            MPI.COMM_WORLD, self.cell_connectivity, self.data.points, mesh_ufl
+        )
+
+    def create_dolfinx_function(self, data: str = "mean") -> dolfinx.fem.Function:
+        """reads the filename of the OpenMC file
+
+        Arguments:
+            data: the name of the data to extract from the vtk file
+
+        Returns:
+            dolfinx function with openmc results mapped
+        """
+
+        if not self.dolfinx_mesh:
+            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[f"{data}"][
+            self.dolfinx_mesh.topology.original_cell_index
+        ]
+
+        return u
+
+
+class UnstructuredMeshReader(OpenMC2dolfinx):
+    """
+    Unstructured Mesh Reader
+
+    Reads an OpenMC .vtk results file with unstructured meshes and converts the data
+    into a dolfinx.fem.Function
+
+    Args:
+        path: the path to the OpenMC .vtk file
+
+    Example:
+    .. code-block:: python
+        reader = UnstructuredMeshReader("path/to/file.vtk")
+        dolfinx_function = reader.create_dolfinx_function()
+    """
+
+    cell_type = "tetrahedron"
+
+    @property
+    def cell_connectivity(self):
+        return self.data.cells_dict[10]
+
+
+class StructuredGridReader(OpenMC2dolfinx):
+    """
+    Structured Mesh Reader
+
+    Reads an OpenMC .vtk results file with Structured meshes and converts the data
+    into a dolfinx.fem.Function
+
+    Args:
+        path: the path to the OpenMC .vtk file
+
+    Example:
+    .. code-block:: python
+        reader = StructuredGridReader("path/to/file.vtk")
+        dolfinx_function = reader.create_dolfinx_function()
+    """
+
+    cell_type = "hexahedron"
+    _cell_connectivity = None
+
+    def get_connectivity(self):
+        num_cells = self.data.GetNumberOfCells()
+        assert self.data.GetCellType(0) == 12, "Only hexahedron cells are supported"
+
+        # Extract connectivity information
+        ordering = [0, 1, 3, 2, 4, 5, 7, 6]
+
+        self._cell_connectivity = []
+
+        # TODO numpify this
+        # Extract all cell connectivity data at once
+        for i in range(num_cells):
+            cell = self.data.GetCell(i)  # Get the i-th cell
+            point_ids = [cell.GetPointId(j) for j in ordering]  # Extract connectivity
+            self._cell_connectivity.append(point_ids)
+
+    @property
+    def cell_connectivity(self):
+        if self._cell_connectivity is None:
+            self.get_connectivity()
+        return self._cell_connectivity

src/openmc2dolfinx/vtk_reader.py

@@ -0,0 +1,30 @@
+import numpy as np
+import pyvista as pv
+from dolfinx import fem
+
+from openmc2dolfinx import StructuredGridReader
+
+
+def test_read_and_generation_of_dolfinx_function_from_structured_grid(tmpdir):
+    """Test StructuredGridReader"""
+    xrng = np.arange(0, 20, 5, dtype=np.float32)
+    yrng = np.arange(0, 20, 5, dtype=np.float32)
+    zrng = np.arange(0, 20, 5, dtype=np.float32)
+    x, y, z = np.meshgrid(xrng, yrng, zrng, indexing="ij")
+    grid = pv.StructuredGrid(x, y, z)
+
+    # add cell data
+    grid.cell_data["mean"] = np.arange(grid.n_cells)
+
+    # save to vtk file
+    filename = str(tmpdir.join("original_structured.vtk"))
+    grid.save(filename)
+
+    # save to vtk file
+    filename = str(tmpdir.join("original.vtk"))
+    grid.save(filename)
+
+    reader = StructuredGridReader(filename)
+    dolfinx_function = reader.create_dolfinx_function()
+
+    assert isinstance(dolfinx_function, fem.Function)

test/test_example.py

@@ -0,0 +1,98 @@
+import numpy as np
+import pyvista as pv
+import dolfinx
+from dolfinx import fem
+from mpi4py import MPI
+import pytest
+
+from openmc2dolfinx import UnstructuredMeshReader
+
+
+@pytest.fixture
+def unstructured_mesh():
+    points = [
+        [1.0, 1.0, 1.0],
+        [1.0, -1.0, -1.0],
+        [-1.0, 1.0, -1.0],
+        [-1.0, -1.0, 1.0],
+        [2.0, 2.0, -1.0],  # Additional point for the second tetrahedron
+    ]
+
+    cells = [4, 0, 1, 2, 3, 4, 0, 1, 2, 4]  # First tetrahedron  # Second tetrahedron
+
+    celltypes = [pv.CellType.TETRA, pv.CellType.TETRA]
+
+    grid = pv.UnstructuredGrid(cells, celltypes, points)
+    grid.cell_data["mean"] = np.arange(grid.n_cells)
+
+    return grid
+
+
+def test_read_and_generation_of_dolfinx_function_from_unstructured_mesh(
+    tmpdir, unstructured_mesh
+):
+    """Test UnstructuredMeshReader"""
+
+    # save to vtk file
+    filename = str(tmpdir.join("original_unstructured.vtk"))
+    unstructured_mesh.save(filename)
+
+    reader = UnstructuredMeshReader(filename)
+    dolfinx_function = reader.create_dolfinx_function()
+
+    assert isinstance(dolfinx_function, fem.Function)
+
+
+def test_cell_type_raises_error_if_not_defined(tmpdir, unstructured_mesh):
+    # save to vtk file
+    filename = str(tmpdir.join("original_unstructured.vtk"))
+    unstructured_mesh.save(filename)
+
+    from openmc2dolfinx.core import OpenMC2dolfinx
+
+    class Temp(OpenMC2dolfinx):
+        pass
+
+    with pytest.raises(
+        AttributeError, match="cell_type must be defined in the child class"
+    ):
+        my_temp = Temp(filename)
+        my_temp.create_dolfinx_mesh()
+
+
+def test_cell_connectivity_raises_error_if_not_defined(tmpdir, unstructured_mesh):
+    # save to vtk file
+    filename = str(tmpdir.join("original_unstructured.vtk"))
+    unstructured_mesh.save(filename)
+
+    from openmc2dolfinx.core import OpenMC2dolfinx
+
+    class Temp(OpenMC2dolfinx):
+        cell_type = "tetrahedron"
+
+    with pytest.raises(
+        AttributeError, match="cell_connectivity must be defined in the child class"
+    ):
+        my_temp = Temp(filename)
+        my_temp.create_dolfinx_mesh()
+
+
+def test_download_from_pyvista_examples(tmpdir):
+    # download an example tetmesh
+    filename = pv.examples.download_tetrahedron(load=False)
+
+    grid = pv.read(filename)
+
+    # assign random cell data
+    grid.cell_data["mean"] = np.arange(grid.n_cells)
+    grid.save("out.vtk")
+
+    # read the vtk file
+    reader = UnstructuredMeshReader("out.vtk")
+
+    # make a dolfinx function
+    u = reader.create_dolfinx_function("mean")
+
+    # export to vtk for visualisation
+    writer = dolfinx.io.VTXWriter(MPI.COMM_WORLD, tmpdir + "/out.bp", u, "BP5")
+    writer.write(t=0)