Writing a mesh checkpoint

Writing a mesh checkpoint#

In this example, we will demonstrate how to write a mesh checkpoint to disk.

We start by creating a simple unit-square mesh.

import logging
from pathlib import Path

from mpi4py import MPI

import dolfinx
import ipyparallel as ipp

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)

Note that when a mesh is created in DOLFINx, we send in an MPI communicator. The communicator is used to partition (distribute) the mesh across the available processes. This means that each process only have access to a sub-set of cells and nodes of the mesh. We can inspect these with the following commands:

def print_mesh_info(mesh: dolfinx.mesh.Mesh):
    cell_map = mesh.topology.index_map(mesh.topology.dim)
    node_map = mesh.geometry.index_map()
    print(
        f"Rank {mesh.comm.rank}: number of owned cells {cell_map.size_local}",
        f", number of ghosted cells {cell_map.num_ghosts}\n",
        f"Number of owned nodes {node_map.size_local}",
        f", number of ghosted nodes {node_map.num_ghosts}",
    )

print_mesh_info(mesh)

Rank 0: number of owned cells 200 , number of ghosted cells 0
 Number of owned nodes 121 , number of ghosted nodes 0

Create a distributed mesh#

Next, we can use IPython parallel to inspect a partitioned mesh. We create a convenience function for creating a mesh that shares cells on the boundary between two processes if ghosted=True.

def create_distributed_mesh(ghosted: bool, N: int = 10):
    """
    Create a distributed mesh with N x N cells. Share cells on process boundaries
    if ghosted is set to True
    """
    from mpi4py import MPI

    import dolfinx

    ghost_mode = dolfinx.mesh.GhostMode.shared_facet if ghosted else dolfinx.mesh.GhostMode.none
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, N, N, ghost_mode=ghost_mode)
    print(f"{ghost_mode=}")
    print_mesh_info(mesh)

Next we start up a new cluster with three engines. As we defined print_mesh_info locally on this process, we need to push it to all engines.

with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster:
    # Push print_mesh_info to all engines
    cluster[:].push({"print_mesh_info": print_mesh_info})

    # Create mesh with ghosted cells
    query_true = cluster[:].apply_async(create_distributed_mesh, True)
    query_true.wait()
    assert query_true.successful(), query_true.error
    print("".join(query_true.stdout))
    # Create mesh without ghosted cells
    query_false = cluster[:].apply_async(create_distributed_mesh, False)
    query_false.wait()
    assert query_false.successful(), query_false.error
    print("".join(query_false.stdout))

ghost_mode=GhostMode.shared_facet
Rank 0: number of owned cells 66 , number of ghosted cells 9
 Number of owned nodes 40 , number of ghosted nodes 13
ghost_mode=GhostMode.shared_facet
Rank 1: number of owned cells 67 , number of ghosted cells 14
 Number of owned nodes 37 , number of ghosted nodes 20
ghost_mode=GhostMode.shared_facet
Rank 2: number of owned cells 67 , number of ghosted cells 11
 Number of owned nodes 44 , number of ghosted nodes 13

ghost_mode=GhostMode.none
Rank 0: number of owned cells 66 , number of ghosted cells 0
 Number of owned nodes 40 , number of ghosted nodes 6
ghost_mode=GhostMode.none
Rank 1: number of owned cells 67 , number of ghosted cells 0
 Number of owned nodes 37 , number of ghosted nodes 9
ghost_mode=GhostMode.none
Rank 2: number of owned cells 67 , number of ghosted cells 0
 Number of owned nodes 44 , number of ghosted nodes 4

Writing a mesh checkpoint#

The input data to a mesh is:

  • A geometry: the set of points in R^D that are part of each cell

  • A two-dimensional connectivity array: A list that indicates which nodes of the geometry is part of each cell

  • A reference element: Used for push data back and forth from the reference element and computing Jacobians We now use adios4dolfinx to write a mesh to file.

def write_mesh(filename: Path):
    import subprocess

    from mpi4py import MPI

    import dolfinx

    import adios4dolfinx

    # Create a simple unit square mesh
    mesh = dolfinx.mesh.create_unit_square(
        MPI.COMM_WORLD, 10, 10, cell_type=dolfinx.mesh.CellType.quadrilateral
    )

    # Write mesh checkpoint
    adios4dolfinx.write_mesh(filename, mesh, engine="BP4")

    # Inspect checkpoint on rank 0 with `bpls`
    if mesh.comm.rank == 0:
        output = subprocess.run(["bpls", "-a", "-l", str(filename.absolute())], capture_output=True)
        print(output.stdout.decode("utf-8"))

mesh_file = Path("mesh.bp")

with ipp.Cluster(engines="mpi", n=2, log_level=logging.ERROR) as cluster:
    # Write mesh to file
    query = cluster[:].apply_async(write_mesh, mesh_file)
    query.wait()
    assert query.successful(), query.error
    print("".join(query.stdout))

  string   CellType         attr   = "quadrilateral"
  int32_t  Degree           attr   = 1
  int32_t  LagrangeVariant  attr   = 2
  double   MeshTime         {1} = 0 / 0
  double   Points           {121, 2} = 0 / 1
  int64_t  Topology         {100, 4} = 0 / 120

We observe that we have stored all the data needed to re-create the mesh in the file mesh.bp. We can therefore read it (to any number of processes) with adios4dolfinx.read_mesh

def read_mesh(filename: Path):
    from mpi4py import MPI

    import dolfinx

    import adios4dolfinx

    mesh = adios4dolfinx.read_mesh(
        filename, comm=MPI.COMM_WORLD, engine="BP4", ghost_mode=dolfinx.mesh.GhostMode.none
    )
    print_mesh_info(mesh)

Reading mesh checkpoints (N-to-M)#

We can now read the checkpoint on a different number of processes than we wrote it on.

with ipp.Cluster(engines="mpi", n=4, log_level=logging.ERROR) as cluster:
    # Write mesh to file
    cluster[:].push({"print_mesh_info": print_mesh_info})
    query = cluster[:].apply_async(read_mesh, mesh_file)
    query.wait()
    assert query.successful(), query.error
    print("".join(query.stdout))

Rank 0: number of owned cells 25 , number of ghosted cells 0
 Number of owned nodes 30 , number of ghosted nodes 6
Rank 1: number of owned cells 25 , number of ghosted cells 0
 Number of owned nodes 31 , number of ghosted nodes 5
Rank 2: number of owned cells 25 , number of ghosted cells 0
 Number of owned nodes 27 , number of ghosted nodes 9
Rank 3: number of owned cells 25 , number of ghosted cells 0
 Number of owned nodes 33 , number of ghosted nodes 3