Add mesh movement function that simplifies movement for any ufl expression

src/scifem/bcs.py

@@ -104,9 +104,11 @@ def interpolate_function_onto_facet_dofs(
         ] = values_per_entity.reshape(-1)
 
     qh = dolfinx.fem.Function(Q)
-    qh._cpp_object.interpolate(
-        values.reshape(-1, domain.geometry.dim).T.copy(), all_connected_cells
-    )
+    if hasattr(qh._cpp_object, "interpolate_f"):
+        interpolate_func = qh._cpp_object.interpolate_f
+    else:
+        interpolate_func = qh._cpp_object.interpolate
+    interpolate_func(values.reshape(-1, domain.geometry.dim).T.copy(), all_connected_cells)
     qh.x.scatter_forward()
 
     return qh

src/scifem/interpolation.py

@@ -386,5 +386,9 @@ def interpolate_to_surface_submesh(
     else:
         shaped_data = data.reshape(expr.value_size, -1)
 
-    u_surface._cpp_object.interpolate(shaped_data, submesh_facets)
+    if hasattr(u_surface._cpp_object, "interpolate_f"):
+        interpolate_func = u_surface._cpp_object.interpolate_f
+    else:
+        interpolate_func = u_surface._cpp_object.interpolate
+    interpolate_func(shaped_data, submesh_facets)
     u_surface.x.scatter_forward()

src/scifem/mesh.py

@@ -8,7 +8,8 @@
 import numpy.typing as npt
 from typing import Protocol
 from packaging.version import Version
-
+import basix
+import ufl
 
 __all__ = [
     "create_entity_markers",
@@ -19,6 +20,8 @@
     "find_interface",
     "compute_interface_data",
     "compute_subdomain_exterior_facets",
+    "create_geometry_function_space",
+    "move",
 ]
 
 if typing.TYPE_CHECKING:
@@ -396,3 +399,85 @@ def compute_interface_data(
     if True in switch:
         ordered_idata[switch, :] = ordered_idata[switch][:, [2, 3, 0, 1]]
     return ordered_idata
+
+
+def create_geometry_function_space(
+    mesh: dolfinx.mesh.Mesh, N: int | None = None
+) -> dolfinx.fem.FunctionSpace:
+    """
+    Reconstruct a vector space with N components using
+    the geometry dofmap to ensure a 1-1 mapping between mesh nodes and DOFs.
+
+    Args:
+        mesh: The mesh to create the function space on.
+        N: The number of components. If not provided the geometrical dimension is chosen
+
+    """
+    geom_imap = mesh.geometry.index_map()
+    geom_dofmap = mesh.geometry.dofmap
+    ufl_domain = mesh.ufl_domain()
+    assert ufl_domain is not None
+    sub_el = ufl_domain.ufl_coordinate_element().sub_elements[0]
+    adj_list = dolfinx.cpp.graph.AdjacencyList_int32(geom_dofmap)
+
+    value_shape: tuple[int, ...]
+    if N is None:
+        ufl_el = basix.ufl.blocked_element(sub_el, shape=(mesh.geometry.dim,))
+        value_shape = (mesh.geometry.dim,)
+        N = value_shape[0]
+    elif N == 1:
+        ufl_el = sub_el
+        value_shape = ()
+    else:
+        ufl_el = basix.ufl.blocked_element(sub_el, shape=(N,))
+        value_shape = (N,)
+
+    if ufl_el.dtype == np.float32:
+        _fe_constructor = dolfinx.cpp.fem.FiniteElement_float32
+        _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float32
+    elif ufl_el.dtype == np.float64:
+        _fe_constructor = dolfinx.cpp.fem.FiniteElement_float64
+        _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float64
+    else:
+        raise RuntimeError(f"Unsupported type {ufl_el.dtype}")
+    try:
+        cpp_el = _fe_constructor(ufl_el.basix_element._e, block_shape=value_shape, symmetric=False)
+    except TypeError:
+        cpp_el = _fe_constructor(ufl_el.basix_element._e, block_size=N, symmetric=False)
+    dof_layout = dolfinx.cpp.fem.create_element_dof_layout(cpp_el, [])
+    cpp_dofmap = dolfinx.cpp.fem.DofMap(dof_layout, geom_imap, N, adj_list, N)
+
+    # Create function space
+    try:
+        cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap)
+    except TypeError:
+        cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap, value_shape=value_shape)
+
+    return dolfinx.fem.FunctionSpace(mesh, ufl_el, cpp_space)
+
+
+def move(
+    mesh: dolfinx.mesh.Mesh,
+    u: dolfinx.fem.Function
+    | ufl.core.expr.Expr
+    | typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.inexact]],
+):
+    """
+    Move the geometry nodes of a mesh given by the movement of a function u.
+
+    Args:
+        mesh: The mesh to move
+        u: The displacement as a :py:class:dolfinx.fem.Function`,
+            :py:class:`ufl.core.expr.Expr` or a lambda function.
+    """
+
+    V_geom = create_geometry_function_space(mesh)
+    u_geom = dolfinx.fem.Function(V_geom, dtype=mesh.geometry.x.dtype)
+    if isinstance(u, dolfinx.fem.Function):
+        u_geom.interpolate(u)
+    elif isinstance(u, ufl.core.expr.Expr):
+        u_compiled = dolfinx.fem.Expression(u, V_geom.element.interpolation_points)
+        u_geom.interpolate(u_compiled)
+    else:
+        u_geom.interpolate(u)
+    mesh.geometry.x[:, : mesh.geometry.dim] += u_geom.x.array[:].reshape(-1, mesh.geometry.dim)

tests/test_mesh_movement.py

@@ -0,0 +1,145 @@
+import pytest
+import numpy as np
+from mpi4py import MPI
+import dolfinx
+import ufl
+from dolfinx.fem import Function, functionspace
+from scifem.mesh import move
+
+
+@pytest.fixture
+def mesh_fixture():
+    """
+    Creates a fresh unit square mesh for each test.
+    """
+    comm = MPI.COMM_WORLD
+    # Create a simple 2D unit square mesh
+    mesh = dolfinx.mesh.create_unit_square(comm, 5, 5)
+    return mesh
+
+
+def test_move_with_dolfinx_function(mesh_fixture):
+    """
+    Test moving the mesh with a dolfinx.fem.Function.
+    """
+    mesh = mesh_fixture
+    gdim = mesh.geometry.dim
+
+    # 1. Create a displacement function in a Vector FunctionSpace
+    V = functionspace(mesh, ("Lagrange", 2, (gdim,)))
+    displacement = Function(V)
+
+    # Define a simple constant displacement: u = (0.1, 0.2)
+    disp_val = np.array([0.1, 0.2])
+
+    # Using interpolate to set constant value for test clarity
+    displacement.interpolate(lambda x: np.full((gdim, x.shape[1]), disp_val[:, None]))
+
+    # 2. Store original coords to compare later
+    original_coords = mesh.geometry.x.copy()
+
+    # 3. Apply move
+    move(mesh, displacement)
+
+    # 4. Verify
+    tol = 10 * np.finfo(mesh.geometry.x.dtype).eps
+
+    expected_coords = original_coords[:, : mesh.geometry.dim] + disp_val[None, : mesh.geometry.dim]
+    np.testing.assert_allclose(
+        mesh.geometry.x[:, : mesh.geometry.dim],
+        expected_coords,
+        rtol=tol,
+        atol=tol,
+        err_msg="Mesh did not move correctly using dolfinx.fem.Function",
+    )
+
+
+def test_move_with_ufl_expression(mesh_fixture):
+    """
+    Test moving the mesh with a UFL Expression.
+    """
+    mesh = mesh_fixture
+
+    # 1. Define UFL expression
+    x = ufl.SpatialCoordinate(mesh)
+    # Shift nodes radially: u = 0.1 * x
+    ufl_expr = 0.1 * x
+
+    # 2. Store original coords
+    original_coords = mesh.geometry.x.copy()
+
+    # 3. Apply move
+    move(mesh, ufl_expr)
+
+    # 4. Verify
+    # The expected new position is x_new = x_old + 0.1 * x_old = 1.1 * x_old
+    expected_coords = original_coords * 1.1
+    tol = 10 * np.finfo(mesh.geometry.x.dtype).eps
+    np.testing.assert_allclose(
+        mesh.geometry.x,
+        expected_coords,
+        rtol=tol,
+        atol=tol,
+        err_msg="Mesh did not move correctly using UFL Expression",
+    )
+
+
+def test_move_with_python_callable(mesh_fixture):
+    """
+    Test moving the mesh with a generic Python callable.
+    """
+    mesh = mesh_fixture
+
+    # 1. Define python callable
+    # Input x has shape (3, N) or (gdim, N)
+    # Output must match shape
+    def shear_displacement(x):
+        # Shear mapping: dx = 0.5 * y, dy = 0
+        vals = np.zeros((mesh.geometry.dim, x.shape[1]), dtype=mesh.geometry.x.dtype)
+        vals[0] = 0.5 * x[1]  # Displacement in x depends on y
+        vals[1] = 0.0  # No displacement in y
+        return vals
+
+    # 2. Store original coords
+    original_coords = mesh.geometry.x.copy()
+
+    # 3. Apply move
+    move(mesh, shear_displacement)
+
+    # 4. Verify
+    # Manually calculate expected shear
+    expected_coords = original_coords.copy()
+    expected_coords[:, 0] += 0.5 * original_coords[:, 1]
+
+    tol = 10 * np.finfo(mesh.geometry.x.dtype).eps
+    np.testing.assert_allclose(
+        mesh.geometry.x,
+        expected_coords,
+        rtol=tol,
+        atol=tol,
+        err_msg="Mesh did not move correctly using Python callable",
+    )
+
+
+def test_move_accumulates_changes(mesh_fixture):
+    """
+    Ensure that calling move twice accumulates the displacement.
+    """
+    mesh = mesh_fixture
+    # Move all coords by 0.1
+    u = lambda x: np.full((mesh.geometry.dim, x.shape[1]), 0.1, dtype=mesh.geometry.x.dtype)
+    original = mesh.geometry.x.copy()
+
+    move(mesh, u)
+    move(mesh, u)
+
+    expected = original
+    expected[:, : mesh.geometry.dim] += 0.2
+    tol = 10 * np.finfo(mesh.geometry.x.dtype).eps
+    np.testing.assert_allclose(
+        mesh.geometry.x,
+        expected,
+        rtol=tol,
+        atol=tol,
+        err_msg="Multiple calls to move did not accumulate correctly",
+    )