Wrap C++ Mesh from Python (fixes support for Python 3.11) (#2500)

* Wrap mesh in place of too clever casting

* Wrap Mesh on Python side.

* flake8 fixes

* Update for gmsh interface

* Code improvement

* Fix function space clone

* Various fixes

* Flake8 fix

* Re-enable mypy checks

* Small edit

* mypy work-around

* Merge fix

* flake8 fixes

* Syntax fixes

* _mesh -> _cpp_object

* Small fixes

* Updates

* Add a hint

* Simplifications

* Add missing import

* Tidy up

* Simplifications

* Tidy up

* Flake8 fixes

* Eliminate test warning

* Wrap cell size function on Python side

Author: garth-wells

python/demo/demo_lagrange_variants.py

@@ -157,15 +157,15 @@ def saw_tooth(x):
 # elements, and plot the finite element interpolation.
 
 # +
-mesh = mesh.create_unit_interval(MPI.COMM_WORLD, 10)
+msh = mesh.create_unit_interval(MPI.COMM_WORLD, 10)
 
-x = ufl.SpatialCoordinate(mesh)
+x = ufl.SpatialCoordinate(msh)
 u_exact = saw_tooth(x[0])
 
 for variant in [basix.LagrangeVariant.equispaced, basix.LagrangeVariant.gll_warped]:
     element = basix.create_element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
     ufl_element = basix.ufl_wrapper.BasixElement(element)
-    V = fem.FunctionSpace(mesh, ufl_element)
+    V = fem.FunctionSpace(msh, ufl_element)
     uh = fem.Function(V)
     uh.interpolate(lambda x: saw_tooth(x[0]))
     if MPI.COMM_WORLD.size == 1:  # Skip this plotting in parallel
@@ -205,11 +205,11 @@ def saw_tooth(x):
 for variant in [basix.LagrangeVariant.equispaced, basix.LagrangeVariant.gll_warped]:
     element = basix.create_element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
     ufl_element = basix.ufl_wrapper.BasixElement(element)
-    V = fem.FunctionSpace(mesh, ufl_element)
+    V = fem.FunctionSpace(msh, ufl_element)
     uh = fem.Function(V)
     uh.interpolate(lambda x: saw_tooth(x[0]))
     M = fem.form((u_exact - uh)**2 * dx)
-    error = mesh.comm.allreduce(fem.assemble_scalar(M), op=MPI.SUM)
+    error = msh.comm.allreduce(fem.assemble_scalar(M), op=MPI.SUM)
     print(f"Computed L2 interpolation error ({variant.name}):", error ** 0.5)
 # -
 

python/demo/demo_types.py

@@ -26,7 +26,6 @@
 from dolfinx import fem, la, mesh, plot
 
 from mpi4py import MPI
-
 # -
 
 # SciPy solvers do not support MPI, so all computations will be

python/dolfinx/fem/function.py

@@ -1,9 +1,9 @@
-# Copyright (C) 2009-2022 Chris N. Richardson, Garth N. Wells and Michal Habera
+# Copyright (C) 2009-2023 Chris N. Richardson, Garth N. Wells and Michal Habera
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
-"""Collection of functions and function spaces"""
+"""Finite element function spaces and functions"""
 
 from __future__ import annotations
 
@@ -14,20 +14,19 @@
 
 from functools import singledispatch
 
-import numpy as np
-import numpy.typing as npt
-
 import basix
 import basix.ufl_wrapper
+import numpy as np
+import numpy.typing as npt
 import ufl
 import ufl.algorithms
 import ufl.algorithms.analysis
-from dolfinx import cpp as _cpp
-from dolfinx import jit, la
 from dolfinx.fem import dofmap
+from petsc4py import PETSc
 from ufl.domain import extract_unique_domain
 
-from petsc4py import PETSc
+from dolfinx import cpp as _cpp
+from dolfinx import jit, la
 
 
 class Constant(ufl.Constant):
@@ -41,7 +40,6 @@ def __init__(self, domain, c: typing.Union[np.ndarray, typing.Sequence, float]):
         """
         c = np.asarray(c)
         super().__init__(domain, c.shape)
-
         try:
             if c.dtype == np.complex64:
                 self._cpp_object = _cpp.fem.Constant_complex64(c)
@@ -373,13 +371,11 @@ def _(expr: Expression, cells: typing.Optional[np.ndarray] = None):
         except TypeError:
             # u is callable
             assert callable(u)
-            x = _cpp.fem.interpolation_coords(
-                self._V.element, self._V.mesh, cells)
-            self._cpp_object.interpolate(
-                np.asarray(u(x), dtype=self.dtype), cells)
+            x = _cpp.fem.interpolation_coords(self._V.element, self._V.mesh._cpp_object, cells)
+            self._cpp_object.interpolate(np.asarray(u(x), dtype=self.dtype), cells)
 
     def copy(self) -> Function:
-        """Return a copy of the Function. The FunctionSpace is shared and the
+        """Create a copy of the Function. The FunctionSpace is shared and the
         degree-of-freedom vector is copied.
 
         """
@@ -445,8 +441,7 @@ def split(self) -> tuple[Function, ...]:
 
     def collapse(self) -> Function:
         u_collapsed = self._cpp_object.collapse()
-        V_collapsed = FunctionSpace(None, self.ufl_element(),
-                                    u_collapsed.function_space)
+        V_collapsed = FunctionSpace(self.function_space._mesh, self.ufl_element(), u_collapsed.function_space)
         return Function(V_collapsed, u_collapsed.x)
 
 
@@ -459,23 +454,13 @@ class ElementMetaData(typing.NamedTuple):
 class FunctionSpace(ufl.FunctionSpace):
     """A space on which Functions (fields) can be defined."""
 
-    def __init__(self, mesh: typing.Union[None, Mesh],
+    def __init__(self, mesh: Mesh,
                  element: typing.Union[ufl.FiniteElementBase, ElementMetaData, typing.Tuple[str, int]],
                  cppV: typing.Optional[_cpp.fem.FunctionSpace] = None,
                  form_compiler_options: dict[str, typing.Any] = {}, jit_options: dict[str, typing.Any] = {}):
         """Create a finite element function space."""
 
-        # Create function space from a UFL element and existing cpp
-        # FunctionSpace
-        if cppV is not None:
-            assert mesh is None
-            ufl_domain = cppV.mesh.ufl_domain()
-            super().__init__(ufl_domain, element)
-            self._cpp_object = cppV
-            return
-
-        if mesh is not None:
-            assert cppV is None
+        if cppV is None:
             # Initialise the ufl.FunctionSpace
             if isinstance(element, ufl.FiniteElementBase):
                 super().__init__(mesh.ufl_domain(), element)
@@ -491,17 +476,26 @@ def __init__(self, mesh: typing.Union[None, Mesh],
                 jit_options=jit_options)
 
             ffi = module.ffi
-            cpp_element = _cpp.fem.FiniteElement(
-                ffi.cast("uintptr_t", ffi.addressof(self._ufcx_element)))
+            cpp_element = _cpp.fem.FiniteElement(ffi.cast("uintptr_t", ffi.addressof(self._ufcx_element)))
             cpp_dofmap = _cpp.fem.create_dofmap(mesh.comm, ffi.cast(
                 "uintptr_t", ffi.addressof(self._ufcx_dofmap)), mesh.topology, cpp_element)
 
-            # Initialize the cpp.FunctionSpace
-            self._cpp_object = _cpp.fem.FunctionSpace(
-                mesh, cpp_element, cpp_dofmap)
+            # Initialize the cpp.FunctionSpace and store mesh
+            self._cpp_object = _cpp.fem.FunctionSpace(mesh._cpp_object, cpp_element, cpp_dofmap)
+            self._mesh = mesh
+        else:
+            # Create function space from a UFL element and an existing
+            # C++ FunctionSpace
+            if mesh._cpp_object is not cppV.mesh:
+                raise RecursionError("Meshes do not match in FunctionSpace initialisation.")
+            ufl_domain = mesh.ufl_domain()
+            super().__init__(ufl_domain, element)
+            self._cpp_object = cppV
+            self._mesh = mesh
+            return
 
     def clone(self) -> FunctionSpace:
-        """Return a new FunctionSpace :math:`W` which shares data with this
+        """Create a new FunctionSpace :math:`W` which shares data with this
         FunctionSpace :math:`V`, but with a different unique integer ID.
 
         This function is helpful for defining mixed problems and using
@@ -513,10 +507,12 @@ def clone(self) -> FunctionSpace:
         diagonal blocks. This is relevant for the handling of boundary
         conditions.
 
+        Returns:
+            A new function space that shares data
+
         """
-        Vcpp = _cpp.fem.FunctionSpace(
-            self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap)
-        return FunctionSpace(None, self.ufl_element(), Vcpp)
+        Vcpp = _cpp.fem.FunctionSpace(self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap)
+        return FunctionSpace(self._mesh, self.ufl_element(), Vcpp)
 
     @property
     def num_sub_spaces(self) -> int:
@@ -536,7 +532,7 @@ def sub(self, i: int) -> FunctionSpace:
         assert self.ufl_element().num_sub_elements() > i
         sub_element = self.ufl_element().sub_elements()[i]
         cppV_sub = self._cpp_object.sub([i])
-        return FunctionSpace(None, sub_element, cppV_sub)
+        return FunctionSpace(self._mesh, sub_element, cppV_sub)
 
     def component(self):
         """Return the component relative to the parent space."""
@@ -562,15 +558,13 @@ def __ne__(self, other):
         """Comparison for inequality."""
         return super().__ne__(other) or self._cpp_object != other._cpp_object
 
-    def ufl_cell(self):
-        return self._cpp_object.mesh.ufl_cell()
-
     def ufl_function_space(self) -> ufl.FunctionSpace:
         """UFL function space"""
         return self
 
     @property
     def element(self):
+        """Function space finite element."""
         return self._cpp_object.element
 
     @property
@@ -579,9 +573,9 @@ def dofmap(self) -> dofmap.DofMap:
         return dofmap.DofMap(self._cpp_object.dofmap)
 
     @property
-    def mesh(self) -> _cpp.mesh.Mesh:
-        """Return the mesh on which the function space is defined."""
-        return self._cpp_object.mesh
+    def mesh(self) -> Mesh:
+        """Mesh on which the function space is defined."""
+        return self._mesh
 
     def collapse(self) -> tuple[FunctionSpace, np.ndarray]:
         """Collapse a subspace and return a new function space and a map from
@@ -592,31 +586,37 @@ def collapse(self) -> tuple[FunctionSpace, np.ndarray]:
 
         """
         cpp_space, dofs = self._cpp_object.collapse()
-        V = FunctionSpace(None, self.ufl_element(), cpp_space)
+        V = FunctionSpace(self._mesh, self.ufl_element(), cpp_space)
         return V, dofs
 
-    def tabulate_dof_coordinates(self) -> np.ndarray:
+    def tabulate_dof_coordinates(self) -> npt.NDArray[np.float64]:
+        """Tabulate the coordinates of the degrees-of-freedom in the function space.
+
+            Returns:
+                Coordinates of the degrees-of-freedom.
+
+            Notes:
+                This method should be used only for elements with point
+                evaluation degrees-of-freedom.
+
+         """
         return self._cpp_object.tabulate_dof_coordinates()
 
 
-def VectorFunctionSpace(mesh: Mesh, element: typing.Union[ElementMetaData, typing.Tuple[str, int]], dim=None,
-                        restriction=None) -> FunctionSpace:
+def VectorFunctionSpace(mesh: Mesh, element: typing.Union[ElementMetaData, typing.Tuple[str, int]],
+                        dim=None) -> FunctionSpace:
     """Create vector finite element (composition of scalar elements) function space."""
-
     e = ElementMetaData(*element)
-    ufl_element = basix.ufl_wrapper.create_vector_element(
-        e.family, mesh.ufl_cell().cellname(), e.degree, dim=dim,
-        gdim=mesh.geometry.dim)
-
+    ufl_element = basix.ufl_wrapper.create_vector_element(e.family, mesh.ufl_cell().cellname(), e.degree,
+                                                          dim=dim, gdim=mesh.geometry.dim)
     return FunctionSpace(mesh, ufl_element)
 
 
 def TensorFunctionSpace(mesh: Mesh, element: typing.Union[ElementMetaData, typing.Tuple[str, int]], shape=None,
-                        symmetry: typing.Optional[bool] = None, restriction=None) -> FunctionSpace:
+                        symmetry: typing.Optional[bool] = None) -> FunctionSpace:
     """Create tensor finite element (composition of scalar elements) function space."""
-
     e = ElementMetaData(*element)
-    ufl_element = basix.ufl_wrapper.create_tensor_element(
-        e.family, mesh.ufl_cell().cellname(), e.degree, shape=shape, symmetry=symmetry,
-        gdim=mesh.geometry.dim)
+    ufl_element = basix.ufl_wrapper.create_tensor_element(e.family, mesh.ufl_cell().cellname(),
+                                                          e.degree, shape=shape, symmetry=symmetry,
+                                                          gdim=mesh.geometry.dim)
     return FunctionSpace(mesh, ufl_element)

python/dolfinx/geometry.py

@@ -9,15 +9,17 @@
 
 import typing
 
+import numpy as np
+import numpy.typing as npt
+
 if typing.TYPE_CHECKING:
     from dolfinx.mesh import Mesh
     from dolfinx.cpp.graph import AdjacencyList_int32
 
 import numpy
+from dolfinx.cpp.geometry import compute_collisions, compute_distance_gjk
 
 from dolfinx import cpp as _cpp
-from dolfinx.cpp.geometry import (compute_closest_entity, compute_collisions,
-                                  compute_distance_gjk, create_midpoint_tree)
 
 __all__ = ["compute_colliding_cells", "squared_distance", "compute_closest_entity", "compute_collisions",
            "compute_distance_gjk", "create_midpoint_tree"]
@@ -44,7 +46,30 @@ def __init__(self, mesh: Mesh, dim: int, entities=None, padding: float = 0.0):
         if entities is None:
             entities = range(0, map.size_local + map.num_ghosts)
 
-        super().__init__(mesh, dim, entities, padding)
+        super().__init__(mesh._cpp_object, dim, entities, padding)
+
+
+def compute_closest_entity(tree: BoundingBoxTree, midpoint_tree: BoundingBoxTree, mesh: Mesh,
+                           points: numpy.ndarray) -> npt.NDArray[np.int32]:
+    """Compute closest mesh entity to a point.
+
+        Args:
+            tree: bounding box tree for the entities
+            midpoint_tree: A bounding box tree with the midpoints of all
+                the mesh entities. This is used to accelerate the search.
+            mesh: The mesh
+            points: The points to check for collision, shape=(num_points, 3)
+
+        Returns:
+            Mesh entity index for each point in `points`. Returns -1 for
+            a point if the bounding box tree is empty.
+
+    """
+    return _cpp.geometry.compute_closest_entity(tree, midpoint_tree, mesh._cpp_object, points)
+
+
+def create_midpoint_tree(mesh: Mesh, dim: int, entities: numpy.ndarray):
+    return _cpp.geometry.create_midpoint_tree(mesh._cpp_object, dim, entities)
 
 
 def compute_colliding_cells(mesh: Mesh, candidates: AdjacencyList_int32, x: numpy.ndarray):
@@ -60,7 +85,7 @@ def compute_colliding_cells(mesh: Mesh, candidates: AdjacencyList_int32, x: nump
         Adjacency list where the ith node is the list of entities that
         collide with the ith point
     """
-    return _cpp.geometry.compute_colliding_cells(mesh, candidates, x)
+    return _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates, x)
 
 
 def squared_distance(mesh: Mesh, dim: int, entities: typing.List[int], points: numpy.ndarray):
@@ -80,4 +105,4 @@ def squared_distance(mesh: Mesh, dim: int, entities: typing.List[int], points: n
         Squared shortest distance from points[i] to entities[i]
 
     """
-    return _cpp.geometry.squared_distance(mesh, dim, entities, points)
+    return _cpp.geometry.squared_distance(mesh._cpp_object, dim, entities, points)

python/dolfinx/io/gmshio.py

@@ -230,7 +230,6 @@ def model_to_mesh(model: gmsh.model, comm: _MPI.Comm, rank: int, gdim: int = 3,
 
             cells = np.asarray(topologies[cell_id]["topology"], dtype=np.int64)
             cell_values = np.asarray(topologies[cell_id]["cell_data"], dtype=np.int32)
-
         else:
             cell_id, num_nodes = comm.bcast([None, None], root=rank)
             cells, x = np.empty([0, num_nodes], dtype=np.int32), np.empty([0, gdim])
@@ -248,7 +247,8 @@ def model_to_mesh(model: gmsh.model, comm: _MPI.Comm, rank: int, gdim: int = 3,
         mesh = create_mesh(comm, cells, x[:, :gdim], ufl_domain, partitioner)
 
         # Create MeshTags for cells
-        local_entities, local_values = _cpp.io.distribute_entity_data(mesh, mesh.topology.dim, cells, cell_values)
+        local_entities, local_values = _cpp.io.distribute_entity_data(
+            mesh._cpp_object, mesh.topology.dim, cells, cell_values)
         mesh.topology.create_connectivity(mesh.topology.dim, 0)
         adj = _cpp.graph.AdjacencyList_int32(local_entities)
         ct = meshtags_from_entities(mesh, mesh.topology.dim, adj, local_values.astype(np.int32, copy=False))
@@ -257,7 +257,7 @@ def model_to_mesh(model: gmsh.model, comm: _MPI.Comm, rank: int, gdim: int = 3,
         # Create MeshTags for facets
         topology = mesh.topology
         if has_facet_data:
-            # Permute facets from MSH to Dolfin-X ordering
+            # Permute facets from MSH to DOLFINx ordering
             # FIXME: This does not work for prism meshes
             if topology.cell_type == CellType.prism or topology.cell_type == CellType.pyramid:
                 raise RuntimeError(f"Unsupported cell type {topology.cell_type}")
@@ -267,7 +267,7 @@ def model_to_mesh(model: gmsh.model, comm: _MPI.Comm, rank: int, gdim: int = 3,
             marked_facets = marked_facets[:, gmsh_facet_perm]
 
             local_entities, local_values = _cpp.io.distribute_entity_data(
-                mesh, mesh.topology.dim - 1, marked_facets, facet_values)
+                mesh._cpp_object, mesh.topology.dim - 1, marked_facets, facet_values)
             mesh.topology.create_connectivity(topology.dim - 1, topology.dim)
             adj = _cpp.graph.AdjacencyList_int32(local_entities)
             ft = meshtags_from_entities(mesh, topology.dim - 1, adj, local_values.astype(np.int32, copy=False))
@@ -301,11 +301,11 @@ def read_from_msh(filename: str, comm: _MPI.Comm, rank: int = 0, gdim: int = 3,
             gmsh.initialize()
             gmsh.model.add("Mesh from file")
             gmsh.merge(filename)
-
-        output = model_to_mesh(gmsh.model, comm, rank, gdim=gdim, partitioner=partitioner)
-        if comm.rank == rank:
+            msh = model_to_mesh(gmsh.model, comm, rank, gdim=gdim, partitioner=partitioner)
             gmsh.finalize()
-        return output
+            return msh
+        else:
+            return model_to_mesh(gmsh.model, comm, rank, gdim=gdim, partitioner=partitioner)
 
     # Map from Gmsh cell type identifier (integer) to DOLFINx cell type
     # and degree http://gmsh.info//doc/texinfo/gmsh.html#MSH-file-format