diff --git a/.gitignore b/.gitignore index 15201ac..ee7dee7 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,9 @@ cython_debug/ # PyPI configuration file .pypirc + +# versioning file +*_version.py + +# results files +*.bp diff --git a/README.md b/README.md index 6104516..e1e3297 100644 --- a/README.md +++ b/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) +``` diff --git a/src/openmc2dolfinx/__init__.py b/src/openmc2dolfinx/__init__.py new file mode 100644 index 0000000..afaa677 --- /dev/null +++ b/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"] diff --git a/src/openmc2dolfinx/core.py b/src/openmc2dolfinx/core.py new file mode 100644 index 0000000..74aa24f --- /dev/null +++ b/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 diff --git a/src/openmc2dolfinx/vtk_reader.py b/src/openmc2dolfinx/vtk_reader.py deleted file mode 100644 index d0285e6..0000000 --- a/src/openmc2dolfinx/vtk_reader.py +++ /dev/null @@ -1,23 +0,0 @@ -import dolfinx -import pyvista - - -class VTKReader: - """""" - - def __init__(self): - pass - - -class StructuredGridReader(VTKReader): - """""" - - def __init__(self): - pass - - -class UnstructuredMeshReader(VTKReader): - """""" - - def __init__(self): - pass diff --git a/test/test_example.py b/test/test_example.py deleted file mode 100644 index 592585e..0000000 --- a/test/test_example.py +++ /dev/null @@ -1,5 +0,0 @@ -import pytest - - -def test_example(): - pass diff --git a/test/test_structured_grid.py b/test/test_structured_grid.py new file mode 100644 index 0000000..3ffdb3a --- /dev/null +++ b/test/test_structured_grid.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) diff --git a/test/test_unstructured_mesh.py b/test/test_unstructured_mesh.py new file mode 100644 index 0000000..26ae187 --- /dev/null +++ b/test/test_unstructured_mesh.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)