Merge pull request #24 from scientificcomputing/dokken/reconstruct-mesh

Add mesh reconstruction command for curving/uncurving

Author: jorgensd

src/io4dolfinx/__init__.py

@@ -31,6 +31,7 @@
     read_point_data,
 )
 from .snapshot import snapshot_checkpoint
+from .utils import reconstruct_mesh
 
 meta = metadata("io4dolfinx")
 __version__ = meta["Version"]
@@ -62,4 +63,5 @@
     "get_backend",
     "write_cell_data",
     "write_point_data",
+    "reconstruct_mesh",
 ]

src/io4dolfinx/utils.py

@@ -14,9 +14,11 @@
 
 from mpi4py import MPI
 
+import basix.ufl
 import dolfinx
 import numpy as np
 import numpy.typing as npt
+import ufl
 from packaging.version import Version
 
 __all__ = [
@@ -27,6 +29,7 @@
     "unroll_dofmap",
     "compute_insert_position",
     "unroll_insert_position",
+    "reconstruct_mesh",
 ]
 
 
@@ -196,3 +199,62 @@ def compute_dofmap_pos(
     local_cell[indicator] = cell_indicator[markers].reshape(-1)
     dof_pos[indicator] = local_indices[markers].reshape(-1)
     return local_cell, dof_pos
+
+
+def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh:
+    """
+    Make a copy of a mesh and potentially change the element of the coordinate element.
+
+    Note:
+        The topology is shared with the original mesh but the geometry is reconstructed.
+
+    Args:
+        mesh: Mesh to reconstruct
+        coordinate_element_degree: Degree to use for coordinate element
+
+    Returns:
+        The new mesh
+
+    """
+    # Extract cell properties
+    ud = mesh.ufl_domain()
+    assert ud is not None
+    c_el = ud.ufl_coordinate_element()
+    family = c_el.family_name
+    lvar = c_el.lagrange_variant
+    ct = c_el.cell_type
+
+    # Create new UFL element
+    new_c_el = basix.ufl.element(
+        family,
+        ct,
+        coordinate_element_degree,
+        shape=(mesh.geometry.dim,),
+        lagrange_variant=lvar,
+        dtype=mesh.geometry.x.dtype,
+    )
+
+    # Extract new node coordinates
+    V_tmp = dolfinx.fem.functionspace(mesh, new_c_el)
+    gdim = mesh.geometry.dim
+    x = V_tmp.tabulate_dof_coordinates()[:, :gdim]
+
+    # Create new geoemtry
+    geom_imap = V_tmp.dofmap.index_map
+    geom_dofmap = V_tmp.dofmap.list
+    num_nodes_local = geom_imap.size_local + geom_imap.num_ghosts
+    original_input_indices = geom_imap.local_to_global(np.arange(num_nodes_local, dtype=np.int32))
+    coordinate_element = dolfinx.fem.coordinate_element(
+        mesh.topology.cell_type, coordinate_element_degree, lvar, dtype=mesh.geometry.x.dtype
+    )
+    # Could use create_geometry here when things are fixed
+    geom = dolfinx.mesh.Geometry(
+        type(mesh.geometry._cpp_object)(
+            geom_imap, geom_dofmap, coordinate_element._cpp_object, x, original_input_indices
+        )
+    )
+
+    # Create new mesh
+    new_top = mesh.topology
+    cpp_mesh = type(mesh._cpp_object)(mesh.comm, new_top._cpp_object, geom._cpp_object)
+    return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(new_c_el))

tests/test_mesh.py

@@ -0,0 +1,56 @@
+from mpi4py import MPI
+
+import dolfinx
+import numpy as np
+import pytest
+import ufl
+
+from io4dolfinx import reconstruct_mesh
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("degree", [2, 3])
+@pytest.mark.parametrize("R", [0.1, 1, 10])
+def test_curve_mesh(degree, dtype, R):
+    N = 8
+    mesh = dolfinx.mesh.create_rectangle(
+        MPI.COMM_WORLD,
+        [[-1, -1], [1, 1]],
+        [N, N],
+        diagonal=dolfinx.mesh.DiagonalType.crossed,
+        dtype=dtype,
+    )
+    org_area = dolfinx.fem.form(1 * ufl.dx(domain=mesh), dtype=dtype)
+
+    curved_mesh = reconstruct_mesh(mesh, degree)
+
+    def transform(x):
+        u = R * x[0] * np.sqrt(1 - (x[1] ** 2 / (2)))
+        v = R * x[1] * np.sqrt(1 - (x[0] ** 2 / (2)))
+        return np.asarray([u, v])
+
+    curved_mesh.geometry.x[:, : curved_mesh.geometry.dim] = transform(curved_mesh.geometry.x.T).T
+
+    area = dolfinx.fem.form(1 * ufl.dx(domain=curved_mesh), dtype=dtype)
+    circumference = dolfinx.fem.form(1 * ufl.ds(domain=curved_mesh), dtype=dtype)
+
+    computed_area = curved_mesh.comm.allreduce(dolfinx.fem.assemble_scalar(area), op=MPI.SUM)
+    computed_circumference = curved_mesh.comm.allreduce(
+        dolfinx.fem.assemble_scalar(circumference), op=MPI.SUM
+    )
+
+    tol = 10 * np.finfo(dtype).eps
+    assert np.isclose(computed_area, np.pi * R**2, atol=tol)
+    assert np.isclose(computed_circumference, 2 * np.pi * R, atol=tol)
+
+    linear_mesh = reconstruct_mesh(curved_mesh, 1)
+    linear_area = dolfinx.fem.form(1 * ufl.dx(domain=linear_mesh), dtype=dtype)
+
+    recovered_area = linear_mesh.comm.allreduce(
+        dolfinx.fem.assemble_scalar(linear_area), op=MPI.SUM
+    )
+
+    # Curve original mesh
+    mesh.geometry.x[:, : mesh.geometry.dim] = transform(mesh.geometry.x.T).T
+    ref_area = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(org_area), op=MPI.SUM)
+    assert np.isclose(recovered_area, ref_area, atol=tol)