Multi subdomain support

src/foam2dolfinx/open_foam_reader.py

@@ -6,7 +6,7 @@
 import dolfinx
 import numpy as np
 import pyvista
-import ufl  # type: ignore
+import ufl
 from dolfinx.mesh import create_mesh
 
 __all__ = ["OpenFOAMReader", "find_closest_value"]
@@ -25,36 +25,41 @@ class OpenFOAMReader:
             filename: the filename
             cell_type: cell type id (12 corresponds to HEXAHEDRON)
             reader: pyvista OpenFOAM reader for .foam files
-            OF_mesh: the mesh from the openfoam file
-            OF_cells: an array of the cells with associated vertices
-            connectivity: The OpenFOAM mesh cell connectivity with vertices reordered
-                in a sorted order for mapping with the dolfinx mesh.
-            mesh_element: the basix element used in the dolfinx mesh
-            dolfinx_mesh: the dolfinx mesh
-            function_space: the function space of the dolfinx function returned in
-                create_dolfinx_function()
+            multidomain: boolean indicating if the mesh is multi-domain
+            OF_meshes_dict: dictionary of meshes from the openfoam file
+            OF_cells_dict: dictionary of arrays of the cells with associated vertices
+            connectivities_dict: dictionary of the OpenFOAM mesh cell connectivity with
+                vertices reordered in a sorted order for mapping with the dolfinx mesh.
+            dolfinx_meshes_dict: dictionary of dolfinx meshes
 
         Notes:
             The cell type refers to the VTK cell type, a full list of cells and their
             respective integers can be found at: https://vtk.org/doc/nightly/html/vtkCellType_8h_source.html
+
+            If only one mesh is present in the OpenFOAM file, all data will be under the
+            key: "default"
     """
 
     filename: str
     cell_type: int
 
     reader: pyvista.POpenFOAMReader
-    OF_mesh: pyvista.pyvista_ndarray | pyvista.DataSet
-    OF_cells: np.ndarray
-    connectivity: np.ndarray
-    mesh_element: basix.ufl._BlockedElement
-    dolfinx_mesh: dolfinx.mesh.Mesh
-    function_space: dolfinx.fem.FunctionSpace
+    multidomain: bool
+    OF_meshes_dict: dict[str, pyvista.pyvista_ndarray | pyvista.DataSet]
+    OF_cells_dict: dict[str, np.ndarray]
+    connectivities_dict: dict[str, np.ndarray]
+    dolfinx_meshes_dict: dict[str, dolfinx.mesh.Mesh]
 
     def __init__(self, filename, cell_type: int = 12):
         self.filename = filename
         self.cell_type = cell_type
 
         self.reader = pyvista.POpenFOAMReader(self.filename)
+        self.multidomain = False
+        self.OF_meshes_dict = {}
+        self.OF_cells_dict = {}
+        self.connectivities_dict = {}
+        self.dolfinx_meshes_dict = {}
 
     @property
     def cell_type(self):
@@ -66,44 +71,89 @@ def cell_type(self, value):
             raise TypeError("cell_type value should be an int")
         self._cell_type = value
 
-    def _read_with_pyvista(self, t: float):
-        """reads the filename dolfinx.fem.Function from the OpenFOAM file.
+    def _read_with_pyvista(self, t: float, subdomain: Optional[str] = "default"):
+        """
+        Reads the openfoam data in the filename provided, passes details of the
+        openfoam mesh to OF_mesh and details of the cells to OF_cells.
 
         Args:
             t: timestamp of the data to read
-            name: Name of the field in the openfoam file, defaults to "U" for velocity
+            subdomain: Name of the subdmain in the openfoam file, from which a field is
+                extracted
 
-        Returns:
-            the dolfinx function
         """
         self.reader.set_active_time_value(t)  # Set the time value to read data from
         OF_multiblock = self.reader.read()  # Read the data from the OpenFOAM file
-        self.OF_mesh = OF_multiblock["internalMesh"]  # Extract the internal mesh
 
-        # Dictionary mapping cell type to connectivity
-        assert hasattr(self.OF_mesh, "cells_dict")  # Ensure the mesh has cell data
-        OF_cells_dict = self.OF_mesh.cells_dict  # Get the cell dictionary
+        # Check if the reader has a multiblock dataset block named "internalMesh"
+        if "internalMesh" not in OF_multiblock.keys():
+            self.multidomain = True
+            if subdomain not in OF_multiblock.keys():
+                raise ValueError(
+                    f"Subdomain {subdomain} not found in the OpenFOAM file. "
+                    f"Available subdomains: {OF_multiblock.keys()}"
+                )
+
+        # Extract the internal mesh
+        if self.multidomain:
+            for cell_array_name in OF_multiblock.keys():
+                self.OF_meshes_dict[cell_array_name] = OF_multiblock[cell_array_name][
+                    "internalMesh"
+                ]
+        else:
+            self.OF_meshes_dict[subdomain] = OF_multiblock["internalMesh"]
+
+        # Ensure the mesh has cell data
+        assert hasattr(self.OF_meshes_dict[subdomain], "cells_dict")
 
-        self.OF_cells = OF_cells_dict.get(self.cell_type)
+        # obtain dictionary of cell types in OF_mesh
+        OF_cell_type_dict = self.OF_meshes_dict[subdomain].cells_dict
+
+        cell_types_in_mesh = [int(k) for k in OF_cell_type_dict.keys()]
 
         # Raise error if OF_mesh is mixed topology
-        if len(OF_cells_dict.keys()) > 1:
+        if len(cell_types_in_mesh) > 1:
             raise NotImplementedError("Cannot support mixed-topology meshes")
 
+        self.OF_cells_dict[subdomain] = OF_cell_type_dict.get(self.cell_type)
+
         # Raise error if no cells of the specified type are found in the OF_mesh
-        if self.OF_cells is None:
+        if self.OF_cells_dict[subdomain] is None:
             raise ValueError(
-                f"No {self.cell_type} cells found in the mesh. Found "
-                f"{OF_cells_dict.keys()}"
+                f"No cell type {self.cell_type} found in the mesh. Found "
+                f"{cell_types_in_mesh}"
             )
 
-    def _create_dolfinx_mesh(self):
+    def _create_dolfinx_mesh(self, subdomain: Optional[str] = "default"):
         """Creates a dolfinx.mesh.Mesh based on the elements within the OpenFOAM mesh"""
 
+        # Define mesh element and define args conn based on the OF cell type
+        if self.cell_type == 12:
+            shape = "hexahedron"
+            args_conn = np.tile(
+                np.array([0, 1, 3, 2, 4, 5, 7, 6]),
+                (len(self.OF_cells_dict[subdomain]), 1),
+            )
+
+        elif self.cell_type == 10:
+            shape = "tetrahedron"
+            args_conn = np.argsort(
+                self.OF_cells_dict[subdomain], axis=1
+            )  # Sort the cell connectivity
+
+        else:
+            raise ValueError(
+                f"Cell type: {self.cell_type}, not supported, please use"
+                " either 12 (hexahedron) or 10 (tetrahedron) cells in OF mesh"
+            )
+
         # create the connectivity between the OpenFOAM and dolfinx meshes
-        args_conn = np.argsort(self.OF_cells, axis=1)  # Sort the cell connectivity
-        rows = np.arange(self.OF_cells.shape[0])[:, None]  # Create row indices
-        self.connectivity = self.OF_cells[rows, args_conn]  # Reorder connectivity
+        # Create row indices
+        rows = np.arange(self.OF_cells_dict[subdomain].shape[0])[:, None]
+        # Reorder connectivity
+        self.connectivities_dict[subdomain] = self.OF_cells_dict[subdomain][
+            rows, args_conn
+        ]
 
         # Define mesh element
         if self.cell_type == 12:
@@ -126,68 +176,72 @@ def _create_dolfinx_mesh(self):
 
         # Create dolfinx Mesh
         mesh_ufl = ufl.Mesh(self.mesh_vector_element)
-        self.dolfinx_mesh = create_mesh(
-            MPI.COMM_WORLD, self.connectivity, self.OF_mesh.points, mesh_ufl
+        self.dolfinx_meshes_dict[subdomain] = create_mesh(
+            MPI.COMM_WORLD,
+            self.connectivities_dict[subdomain],
+            self.OF_meshes_dict[subdomain].points,
+            mesh_ufl,
         )
 
     def create_dolfinx_function(
-        self, t: Optional[float] = None, name: Optional[str] = "U"
+        self, t: float, name: str = "U", subdomain: Optional[str] = "default"
     ) -> dolfinx.fem.Function:
         """Creates a dolfinx.fem.Function from the OpenFOAM file.
 
         Args:
             t: timestamp of the data to read
             name: Name of the field in the openfoam file, defaults to "U" for velocity
+            subdomain: Name of the subdmain in the openfoam file, from which a field is
+                extracted
 
         Returns:
             the dolfinx function
         """
-        self._read_with_pyvista(t=t)
 
-        # Create dolfinx mesh if it doesn't exist
-        if not hasattr(self, "dolfinx_mesh"):
-            self._create_dolfinx_mesh()
+        # read the openfoam data in the filename provided
+        self._read_with_pyvista(t=t, subdomain=subdomain)
+
+        # create the dolfinx mesh
+        if subdomain not in self.dolfinx_meshes_dict:
+            self._create_dolfinx_mesh(subdomain=subdomain)
+
+        mesh = self.dolfinx_meshes_dict[subdomain]
 
         if name == "U":
             element = self.mesh_vector_element
         else:
             element = self.mesh_scalar_element
 
-        self.function_space = dolfinx.fem.functionspace(self.dolfinx_mesh, element)
-        u = dolfinx.fem.Function(self.function_space)
+        function_space = dolfinx.fem.functionspace(mesh, element)
+        u = dolfinx.fem.Function(function_space)
 
         num_vertices = (
-            self.dolfinx_mesh.topology.index_map(0).size_local
-            + self.dolfinx_mesh.topology.index_map(0).num_ghosts
+            mesh.topology.index_map(0).size_local
+            + mesh.topology.index_map(0).num_ghosts
         )
         vertex_map = np.empty(num_vertices, dtype=np.int32)
 
         # Get cell-to-vertex connectivity
-        c_to_v = self.dolfinx_mesh.topology.connectivity(
-            self.dolfinx_mesh.topology.dim, 0
-        )
-
+        c_to_v = mesh.topology.connectivity(mesh.topology.dim, 0)
         # Map the OF_mesh vertices to dolfinx_mesh vertices
         num_cells = (
-            self.dolfinx_mesh.topology.index_map(
-                self.dolfinx_mesh.topology.dim
-            ).size_local
-            + self.dolfinx_mesh.topology.index_map(
-                self.dolfinx_mesh.topology.dim
-            ).num_ghosts
+            mesh.topology.index_map(mesh.topology.dim).size_local
+            + mesh.topology.index_map(mesh.topology.dim).num_ghosts
         )
         vertices = np.array([c_to_v.links(cell) for cell in range(num_cells)])
         flat_vertices = np.concatenate(vertices)
         cell_indices = np.repeat(np.arange(num_cells), [len(v) for v in vertices])
         vertex_positions = np.concatenate([np.arange(len(v)) for v in vertices])
 
-        vertex_map[flat_vertices] = self.connectivity[
-            self.dolfinx_mesh.topology.original_cell_index
+        vertex_map[flat_vertices] = self.connectivities_dict[subdomain][
+            mesh.topology.original_cell_index
         ][cell_indices, vertex_positions]
 
         # Assign values in OF_mesh to dolfinx_mesh
-        assert hasattr(self.OF_mesh, "point_data")
-        u.x.array[:] = self.OF_mesh.point_data[name][vertex_map].flatten()
+        assert hasattr(self.OF_meshes_dict[subdomain], "point_data")
+        u.x.array[:] = (
+            self.OF_meshes_dict[subdomain].point_data[name][vertex_map].flatten()
+        )
 
         return u
 

test/test_create_dolfinx_mesh.py

@@ -17,7 +17,7 @@ def test_error_rasied_when_using_mixed_topology_mesh():
     rng = np.random.default_rng()
 
     # Create a 400x8 array filled with random values
-    my_reader.OF_cells = rng.random((400, 8))
+    my_reader.OF_cells_dict["default"] = rng.random((400, 8))
 
     error_message = "Cell type: 1, not supported, please use either 12 (hexahedron) or 10 (tetrahedron) cells in OF mesh"
     pattern = re.escape(error_message)

test/test_read_with_pyvista.py

@@ -1,3 +1,6 @@
+import zipfile
+from pathlib import Path
+
 import pytest
 from pyvista import examples
 
@@ -20,6 +23,47 @@ def test_error_rasied_when_cells_wanted_are_not_in_file_provided():
 
     with pytest.raises(
         ValueError,
-        match="No 1 cells found in the mesh. Found dict_keys([np.uint8(12)])",
+        match=r"No cell type 1 found in the mesh\. Found \[.*12]",
     ):
         my_reader._read_with_pyvista(t=0)
+
+
+def test_error_rasied_when_subdomain_is_not_given_in_multidomain_case(tmpdir):
+    zip_path = Path("test/data/test_2Regions.zip")
+    extract_path = Path(tmpdir) / "test_2Regions"
+
+    # Unzip the file
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(extract_path)
+
+    # Construct the path to the .foam file
+    foam_file = extract_path / "test_2Regions/pv.foam"
+
+    # read the .foam file
+    my_of_reader = OpenFOAMReader(filename=str(foam_file), cell_type=12)
+
+    with pytest.raises(
+        ValueError,
+        match="Subdomain None not found in the OpenFOAM file\. Available subdomains: \['defaultRegion', 'fluid', 'solid']",
+    ):
+        my_of_reader._read_with_pyvista(t=20.0, subdomain=None)
+
+
+@pytest.mark.parametrize("subdomain", ["fluid", "solid"])
+def test_read_with_pyvista_finds_all_mesh_data(tmpdir, subdomain):
+    zip_path = Path("test/data/test_2Regions.zip")
+    extract_path = Path(tmpdir) / "test_2Regions"
+
+    # Unzip the file
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(extract_path)
+
+    # Construct the path to the .foam file
+    foam_file = extract_path / "test_2Regions/pv.foam"
+
+    # read the .foam file
+    my_of_reader = OpenFOAMReader(filename=str(foam_file), cell_type=12)
+
+    my_of_reader._read_with_pyvista(t=20.0, subdomain=subdomain)
+
+    assert subdomain in my_of_reader.OF_meshes_dict