It’s probably a lot easier to interpolate the displacement field or an absolute coordinate change to the mesh’s coordinate element:
import dolfinx
from mpi4py import MPI
import numpy as np
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4)
Vex = mesh.ufl_domain().ufl_coordinate_element()
V = dolfinx.fem.FunctionSpace(mesh, Vex)
u = dolfinx.fem.Function(V)
u.interpolate(lambda x: np.stack((x[1]**2, x[0]**2)))
mesh.geometry.x[:,:mesh.geometry.dim] = u.x.array.reshape((-1, mesh.geometry.dim))
Serial:
Parallel: