Add compat layer with dofmaps and cmaps

Author: finsberg

examples/transfer_tags_to_submesh.py

@@ -82,7 +82,7 @@ def plot_mesh(mesh: dolfinx.mesh.Mesh, values=None):
     plotter = pyvista.Plotter()
     V_linear = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
     linear_grid = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(V_linear))
-    if compat.get_cmap(mesh).degree > 1:
+    if compat.cmap(mesh).degree > 1:
         ugrid = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(mesh))
         if values is not None:
             ugrid.cell_data["Marker"] = values

src/scifem/compat.py

@@ -1,11 +1,25 @@
 """Layer for small backward compatibility wrappers for DOLFINx"""
 
+import numpy.typing as npt
+import numpy as np
 import dolfinx
 
 
-def get_cmap(mesh: dolfinx.mesh.Mesh) -> dolfinx.fem.CoordinateElement:
+def cmap(mesh: dolfinx.mesh.Mesh) -> dolfinx.fem.CoordinateElement:
     """Get the basix Cmap for the mesh."""
+    if hasattr(mesh.geometry, "cmaps"):
+        return mesh.geometry.cmaps[0]
     if callable(mesh.geometry.cmap):
         return mesh.geometry.cmap()
     else:
         return mesh.geometry.cmap
+
+
+def dofmap(mesh: dolfinx.mesh.Mesh) -> npt.NDArray[np.int32]:
+    """Get the dofmap for the geometry."""
+    if hasattr(mesh.geometry, "dofmaps"):
+        return mesh.geometry.dofmaps[0]
+    if callable(mesh.geometry.dofmap):
+        return mesh.geometry.dofmap()
+    else:
+        return mesh.geometry.dofmap

src/scifem/eval.py

@@ -10,7 +10,7 @@
 import ufl
 from scipy.optimize import minimize
 from ufl.algorithms.signature import compute_expression_signature
-from .compat import get_cmap
+from . import compat
 
 
 T = typing.TypeVar("T", int, float)
@@ -155,9 +155,9 @@ def find_cell_extrema(
         x_ref = x0
 
     _cell = np.array([cell], dtype=np.int32)
-    mesh_nodes = mesh.geometry.x[mesh.geometry.dofmap[cell], : mesh.geometry.dim]
+    mesh_nodes = mesh.geometry.x[compat.dofmap(mesh)[cell], : mesh.geometry.dim]
     _x_p = np.zeros(3)
-    cmap = get_cmap(mesh)
+    cmap = compat.cmap(mesh)
 
     def eval_J(x_ref):
         # Evaluating basis functions through {py:func}`dolfinx.fem.Function.eval`

src/scifem/geometry.py

@@ -3,6 +3,7 @@
 import numpy.typing as npt
 import warnings
 from scifem._scifem import closest_point_projection_float64, closest_point_projection_float32
+from . import compat
 
 
 def project_onto_simplex(
@@ -162,7 +163,7 @@ def _closest_point_projection(
     element = mesh.ufl_domain().ufl_coordinate_element().sub_elements[0]
     tdim = mesh.topology.dim
     # Get the coordinates of the nodes for the specified cell
-    node_coords = mesh.geometry.x[mesh.geometry.dofmap[cells]][:, :, : mesh.geometry.dim]
+    node_coords = mesh.geometry.x[compat.dofmap(mesh)[cells]][:, :, : mesh.geometry.dim]
     target_points = target_points.reshape(-1, 3)
 
     # Constraints and Bounds

src/scifem/mesh.py

@@ -11,6 +11,8 @@
 import basix
 import ufl
 
+from . import compat
+
 __all__ = [
     "create_entity_markers",
     "transfer_meshtags_to_submesh",
@@ -413,7 +415,7 @@ def create_geometry_function_space(
 
     """
     geom_imap = mesh.geometry.index_map()
-    geom_dofmap = mesh.geometry.dofmap
+    geom_dofmap = compat.dofmap(mesh)
     ufl_domain = mesh.ufl_domain()
     assert ufl_domain is not None
     sub_el = ufl_domain.ufl_coordinate_element().sub_elements[0]

src/scifem/point_source.py

@@ -23,7 +23,7 @@
 import numpy.typing as npt
 import ufl
 
-from .compat import get_cmap
+from . import compat
 from .utils import unroll_dofmap
 
 __all__ = ["PointSource"]
@@ -122,11 +122,11 @@ def compute_cell_contributions(self):
         mesh = self._function_space.mesh
         # Pull owning points back to reference cell
         mesh_nodes = mesh.geometry.x
-        cmap = get_cmap(mesh)
+        cmap = compat.cmap(mesh)
 
         ref_x = np.zeros((len(self._cells), mesh.topology.dim), dtype=mesh.geometry.x.dtype)
         for i, (point, cell) in enumerate(zip(self._points, self._cells)):
-            geom_dofs = mesh.geometry.dofmap[cell]
+            geom_dofs = compat.dofmap(mesh)[cell]
             ref_x[i] = cmap.pull_back(point.reshape(-1, 3), mesh_nodes[geom_dofs])
 
         # Create expression evaluating a trial function (i.e. just the basis function)

tests/test_geometry.py

@@ -3,7 +3,7 @@
 import numpy as np
 import numpy.typing as npt
 from scipy.optimize import minimize
-from scifem import closest_point_projection
+from scifem import closest_point_projection, compat
 from scifem.geometry import project_onto_simplex, _closest_point_projection
 import ufl
 import basix.ufl
@@ -35,7 +35,7 @@ def scipy_project_point_to_element(
     element = mesh.ufl_domain().ufl_coordinate_element().sub_elements[0]
 
     # Get the coordinates of the nodes for the specified cell
-    node_coords = mesh.geometry.x[mesh.geometry.dofmap[cells]][:, :, : mesh.geometry.dim]
+    node_coords = mesh.geometry.x[compat.dofmap(mesh)[cells]][:, :, : mesh.geometry.dim]
     target_points = target_points.reshape(-1, 3)
     # cmap = mesh.geometry.cmap
 

tests/test_point_source.py

@@ -88,10 +88,11 @@ def test_outside():
     cells = dolfinx.geometry.compute_colliding_cells(mesh, cell_candidates, point)
     # Only use evaluate for points on current processor
     # BUG: In DOLFINx 0.8.0, links(0) yields a too long array
+
     if cells.offsets[-1] > 0:
         cell = cells.links(0)[0]
-        geom_dofs = mesh.geometry.dofmap[cell]
-        ref_x = compat.get_cmap(mesh).pull_back(point.reshape(-1, 3), mesh.geometry.x[geom_dofs])
+        geom_dofs = compat.dofmap(mesh)[cell]
+        ref_x = compat.cmap(mesh).pull_back(point.reshape(-1, 3), mesh.geometry.x[geom_dofs])
         ref_values = V.ufl_element().tabulate(0, ref_x).flatten()
         b_nonzero = np.flatnonzero(b.x.array)
         dofs = V.dofmap.cell_dofs(cell)