From 9c562ef598a939f815124c984206954f2a2c49ab Mon Sep 17 00:00:00 2001 From: jorgensd Date: Sun, 13 Sep 2026 20:31:02 +0000 Subject: [PATCH 01/11] Add hardening test for compute_dofmap_pos (work item 7 n submesh plan). --- src/io4dolfinx/utils.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/src/io4dolfinx/utils.py b/src/io4dolfinx/utils.py index 17731cc..9e3263b 100644 --- a/src/io4dolfinx/utils.py +++ b/src/io4dolfinx/utils.py @@ -174,9 +174,19 @@ def compute_dofmap_pos( Compute a map from each owned dof in the dofmap to a single cell owned by the process, and the relative position of the dof. + .. note:: + Every owned dof must be reachable from a cell owned by the process. This + holds for a mesh created by :func:`dolfinx.mesh.create_mesh`, but not + necessarily for one created by :func:`dolfinx.mesh.create_submesh`, which + builds its vertex map allowing owner changes: a rank may own a vertex -- + and therefore a dof -- that is only incident to cells it ghosts. Such + dofs cannot be assigned a position and raise rather than silently + yielding an arbitrary one. + :param V: The function space :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the number of owned dofs (unrolled for block size) + :raises RuntimeError: If an owned dof is not contained in any owned cell. """ dofs = V.dofmap.list mesh = V.mesh @@ -184,10 +194,12 @@ def compute_dofmap_pos( dofmap_bs = V.dofmap.bs num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs - local_cell = np.empty( - num_owned_dofs, dtype=np.int32 + # Initialize to -1 rather than leaving uninitialized, so that a dof that is + # never visited below is detectable instead of holding arbitrary memory. + local_cell = np.full( + num_owned_dofs, -1, dtype=np.int32 ) # Local cell index for each dof owned by process - dof_pos = np.empty(num_owned_dofs, dtype=np.int32) # Position in dofmap for said dof + dof_pos = np.full(num_owned_dofs, -1, dtype=np.int32) # Position in dofmap for said dof unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs) markers = unrolled_dofmap < num_owned_dofs @@ -199,6 +211,17 @@ def compute_dofmap_pos( indicator = unrolled_dofmap[markers].reshape(-1) local_cell[indicator] = cell_indicator[markers].reshape(-1) dof_pos[indicator] = local_indices[markers].reshape(-1) + + num_unreachable = int(np.count_nonzero(local_cell < 0)) + if num_unreachable > 0: + raise RuntimeError( + f"{num_unreachable} of {num_owned_dofs} owned dofs are not contained in any" + " cell owned by this process, so their position in the dofmap is undefined." + " This happens on meshes created with `dolfinx.mesh.create_submesh`, which" + " may move vertex ownership to a process that does not own an incident cell." + " Read the function on a standalone mesh (see `io4dolfinx.read_submesh`)" + " rather than on a mesh derived with `create_submesh`." + ) return local_cell, dof_pos From 5c51c53d779391e03cc2cce95db6b428131c2f2e Mon Sep 17 00:00:00 2001 From: jorgensd Date: Sun, 13 Sep 2026 21:05:55 +0000 Subject: [PATCH 02/11] Add initial work. Should really fix that interpolate onto manifold issue in DOLFINx (https://github.com/FEniCS/dolfinx/issues/3619) --- _toc.yml | 1 + docs/submesh.py | 144 +++++++ src/io4dolfinx/__init__.py | 10 + src/io4dolfinx/backends/__init__.py | 48 ++- src/io4dolfinx/backends/adios2/backend.py | 216 +++++++--- src/io4dolfinx/backends/h5py/backend.py | 46 +- src/io4dolfinx/backends/vtkhdf/backend.py | 7 +- src/io4dolfinx/checkpointing.py | 74 ++++ src/io4dolfinx/comm_helpers.py | 93 ++++ src/io4dolfinx/readers.py | 23 + src/io4dolfinx/submesh.py | 491 ++++++++++++++++++++++ tests/test_reconstruct_mesh.py | 60 +++ tests/test_submesh.py | 460 ++++++++++++++++++++ 13 files changed, 1590 insertions(+), 83 deletions(-) create mode 100644 docs/submesh.py create mode 100644 src/io4dolfinx/submesh.py create mode 100644 tests/test_reconstruct_mesh.py create mode 100644 tests/test_submesh.py diff --git a/_toc.yml b/_toc.yml index 21a7ef3..cad3d99 100644 --- a/_toc.yml +++ b/_toc.yml @@ -21,6 +21,7 @@ parts: - file: "docs/partitioned_mesh" - file: "docs/time_dependent_mesh" - file: "docs/meshtags" + - file: "docs/submesh" - caption: Writing and reading functions chapters: diff --git a/docs/submesh.py b/docs/submesh.py new file mode 100644 index 0000000..aa71f98 --- /dev/null +++ b/docs/submesh.py @@ -0,0 +1,144 @@ +# # Checkpointing a submesh +# +# A submesh created with {py:func}`dolfinx.mesh.create_submesh` cannot be +# checkpointed the way an ordinary mesh is, for two reasons that both come from +# how DOLFINx builds it. +# +# First, when the submesh has co-dimension greater than zero its cells are +# *entities* of the parent, and the vertices of an entity are ordered by their +# current global vertex index. Reading the parent back re-partitions it, which +# renumbers the vertices, so the same cell comes back with its vertices in a +# different order and its degrees of freedom in different positions. +# +# Second, `create_submesh` builds its vertex map allowing ownership to move, so +# a process can own a vertex -- and a degree of freedom -- that is only incident +# to cells it ghosts. Such a dof has no position in any cell that process owns, +# which is exactly what the checkpoint reader needs to find it. +# +# So io4dolfinx never reads a function onto a re-derived submesh. It stores the +# submesh as an ordinary, independent mesh, reads it back as one, and then moves +# the data across to a submesh re-derived from the parent when you need the +# entity maps that make mixed-dimensional assembly possible. + +# + +from pathlib import Path + +from mpi4py import MPI + +import dolfinx +import numpy as np + +import io4dolfinx + +mesh = dolfinx.mesh.create_unit_cube( + MPI.COMM_WORLD, 6, 6, 6, ghost_mode=dolfinx.mesh.GhostMode.shared_facet +) +tdim = mesh.topology.dim +# - + +# We take the face $x=0$ as the submesh. It is a two-dimensional mesh embedded +# in three dimensions -- a manifold -- which the checkpoint handles just as well +# as a co-dimension 0 submesh. + +# + +fdim = tdim - 1 +mesh.topology.create_entities(fdim) +facets = dolfinx.mesh.locate_entities(mesh, fdim, lambda x: np.isclose(x[0], 0.0)) +submesh, cell_map, vertex_map, node_map = dolfinx.mesh.create_submesh(mesh, fdim, facets) + +V = dolfinx.fem.functionspace(submesh, ("Lagrange", 2)) +u = dolfinx.fem.Function(V, name="u") +u.interpolate(lambda x: np.sin(3 * x[1]) + x[2]) +# - + +# ## Writing +# +# The parent goes in first. {py:func}`io4dolfinx.write_submesh` then stores the +# submesh under a name of its own, together with a *parent link*: for each stored +# submesh cell, the parent geometry nodes of the entity it came from. That link +# is what lets the submesh be found again in a re-partitioned parent. +# +# The link tags entities of the parent, so it is stored alongside the parent. +# Here everything goes in one file; pass `parent_filename` to keep the submesh in +# a file of its own. + +# + +filename = Path("submesh_checkpoint.bp") +io4dolfinx.write_mesh(filename, mesh) +io4dolfinx.write_submesh(filename, submesh, mesh, cell_map, mesh_name="wall") +io4dolfinx.write_function(filename, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall") +# - + +# ## Reading +# +# Both meshes are read independently. Because the submesh was stored as an +# ordinary mesh, reading a function on it is the ordinary path -- no special +# handling, and the values come back exactly. + +# + +parent = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD) +stored = io4dolfinx.read_mesh(filename, MPI.COMM_WORLD, mesh_name="wall") + +V_stored = dolfinx.fem.functionspace(stored, ("Lagrange", 2)) +u_stored = dolfinx.fem.Function(V_stored, name="u") +io4dolfinx.read_function(filename, u_stored, time=0.0, name="u", mesh_name="wall") +# - + +# If all you wanted was the data, stop here. If you need to assemble a +# mixed-dimensional form over the parent and the submesh, you need a submesh +# DOLFINx recognises as derived from *this* parent, with the entity maps to prove +# it. {py:func}`io4dolfinx.read_submesh` re-derives one. + +# + +checkpoint = io4dolfinx.read_submesh(filename, parent, mesh_name="wall") +V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) +u_sub = dolfinx.fem.Function(V_sub, name="u") +io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) +# - + +# The transfer routes each interpolation point straight to the process and cell +# that can evaluate it, using the cell correspondence the checkpoint already +# records. No geometric point location is involved, so there is no search +# tolerance to tune and nothing that can silently fail to find a point. +# +# The result agrees with the original to machine precision. + +# + +reference = dolfinx.fem.Function(V_sub) +reference.interpolate(lambda x: np.sin(3 * x[1]) + x[2]) +num_owned = V_sub.dofmap.index_map.size_local * V_sub.dofmap.index_map_bs +error = np.max(np.abs(u_sub.x.array[:num_owned] - reference.x.array[:num_owned])) +print(f"Max error after transfer: {parent.comm.allreduce(error, MPI.MAX):.3e}") +# - + +# `checkpoint.cell_map` and `checkpoint.vertex_map` are genuine +# {py:class}`dolfinx.mesh.EntityMap` objects relating the submesh to `parent`, so +# they can be passed to {py:func}`dolfinx.fem.form`. + +# + +import ufl # noqa: E402 + +one = dolfinx.fem.Function(V_sub) +one.x.array[:] = 1.0 +area_form = dolfinx.fem.form(one * ufl.ds(domain=parent), entity_maps=[checkpoint.cell_map]) +area = parent.comm.allreduce(dolfinx.fem.assemble_scalar(area_form), MPI.SUM) +print(f"Area of the submesh assembled over the parent: {area:.3f}") +# - + +# ## What is not supported +# +# H(div) and H(curl) on a submesh of co-dimension greater than zero can be +# **written and read** exactly -- steps 1-3 above -- but cannot be transferred to +# a re-derived submesh. `transfer_submesh_function` raises `NotImplementedError` +# for them. +# +# The reason is worth knowing. DOLFINx cannot reconcile the reference and +# physical value sizes of these families when `tdim < gdim`, and in this +# particular path it does not say so: `interpolate_nonmatching` returns +# successfully with values that are wrong. Measured on a facet submesh of a unit +# cube, the L2 norm of an N1curl field fell from 9.096 to 4.904. The guard exists +# so that this shows up as an error rather than as a plausible-looking result. +# +# There is no sound workaround. The stored submesh is an independent mesh, +# derived from nothing, so a projection -- or any other form assembled between it +# and the re-derived submesh -- would have no basis. diff --git a/src/io4dolfinx/__init__.py b/src/io4dolfinx/__init__.py index e88ca65..b9742c6 100644 --- a/src/io4dolfinx/__init__.py +++ b/src/io4dolfinx/__init__.py @@ -31,6 +31,12 @@ read_point_data, ) from .snapshot import snapshot_checkpoint +from .submesh import ( + SubmeshCheckpoint, + read_submesh, + transfer_submesh_function, + write_submesh, +) from .utils import reconstruct_mesh meta = metadata("io4dolfinx") @@ -65,4 +71,8 @@ "write_cell_data", "write_point_data", "reconstruct_mesh", + "SubmeshCheckpoint", + "write_submesh", + "read_submesh", + "transfer_submesh_function", ] diff --git a/src/io4dolfinx/backends/__init__.py b/src/io4dolfinx/backends/__init__.py index 8f2726d..68c5f2c 100644 --- a/src/io4dolfinx/backends/__init__.py +++ b/src/io4dolfinx/backends/__init__.py @@ -11,10 +11,41 @@ from ..structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData -__all__ = ["FileMode", "IOBackend", "get_backend", "set_default_backend"] +__all__ = [ + "DEFAULT_MESH_NAME", + "FileMode", + "IOBackend", + "get_backend", + "get_mesh_name", + "set_default_backend", +] _DEFAULT_BACKEND = "adios2" +#: Name used for a mesh when the caller does not supply one. A file written +#: before named meshes existed holds exactly this mesh, so backends must store +#: it exactly as they did before (see :func:`get_mesh_name`). +DEFAULT_MESH_NAME = "mesh" + + +def get_mesh_name(backend_args: dict[str, Any] | None) -> str: + """Get the name of the mesh a backend call refers to. + + A checkpoint may hold several meshes -- for instance a mesh and a submesh of + it -- distinguished by name. Backends take the name from + ``backend_args["name"]``. + + Args: + backend_args: Arguments to backend, or None + + Returns: + The mesh name, :data:`DEFAULT_MESH_NAME` if none was given. + """ + if backend_args is None: + return DEFAULT_MESH_NAME + name = backend_args.get("name", DEFAULT_MESH_NAME) + return DEFAULT_MESH_NAME if name is None else str(name) + def set_default_backend(backend: str): """Set the global default backend for io4dolfinx.""" @@ -60,6 +91,18 @@ class FileMode(Enum): # See https://peps.python.org/pep-0544/#modules-as-implementations-of-protocols class IOBackend(Protocol): + """Interface a storage backend must provide. + + .. note:: + A checkpoint may hold more than one mesh, each with its own topology, + geometry, meshtags, cell permutations and functions. Every method that + touches mesh-associated data takes the mesh it refers to from + ``backend_args["name"]``; use :func:`get_mesh_name` to read it, so that + omitting it selects :data:`DEFAULT_MESH_NAME` consistently. A backend + must store that default mesh exactly as it did before named meshes + existed, so that older checkpoints stay readable. + """ + read_mode: ReadMode def get_default_backend_args(self, arguments: dict[str, Any] | None) -> dict[str, Any]: @@ -254,6 +297,9 @@ def read_cell_perms( Read cell permutation from file with given communicator, Split in continuous chunks based on number of cells in the input data. + Cell permutations belong to a mesh, not to a function, so which mesh is + meant is taken from ``backend_args["name"]``. + Args: comm: MPI communicator used in storage filename: Path to file to read from diff --git a/src/io4dolfinx/backends/adios2/backend.py b/src/io4dolfinx/backends/adios2/backend.py index b023700..a5fc077 100644 --- a/src/io4dolfinx/backends/adios2/backend.py +++ b/src/io4dolfinx/backends/adios2/backend.py @@ -11,7 +11,7 @@ from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from ...utils import check_file_exists, compute_local_range -from .. import FileMode, ReadMode +from .. import DEFAULT_MESH_NAME, FileMode, ReadMode, get_mesh_name from .helpers import ( ADIOSFile, adios_to_numpy_dtype, @@ -28,14 +28,70 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: """Get default arguements (sets engine to BP4).""" - args = arguments or {} + args = dict(arguments) if arguments else {} # Copy: do not mutate the caller's dict if "engine" not in args.keys(): args["engine"] = "BP4" if "legacy" not in args.keys(): args["legacy"] = False # Only used for legacy HDF5 meshtags + if "name" not in args.keys(): + args["name"] = DEFAULT_MESH_NAME return args +def mesh_prefix(backend_args: dict[str, Any] | None) -> str: + """ADIOS2 variable-name prefix for the mesh named in ``backend_args``. + + The default mesh gets an *empty* prefix, so that it is stored under exactly + the names used before checkpoints could hold more than one mesh. Files + written by earlier versions therefore stay readable, and a file holding only + the default mesh stays readable by them. Any other mesh is namespaced under + ``"/"``. + + .. note:: + This is io4dolfinx's own checkpoint schema, not ADIOS2's VTX + visualization schema -- no ``vtk.xml`` attribute is written, and these + files are not readable by ParaView's VTX reader whether they hold one + mesh or several. VTX supports only a single mesh per file, so holding + several here costs nothing that was ever available. For a + visualizable file with several meshes use the ``vtkhdf`` backend, whose + ``MultiBlockDataSet`` is designed for it. + + Args: + backend_args: Arguments to backend + + Returns: + The prefix to prepend to every variable and attribute name of that mesh. + """ + name = get_mesh_name(backend_args) + return "" if name == DEFAULT_MESH_NAME else f"{name}/" + + +def _seek_step(adios_file, variable: str, start: int = 0) -> int: + """Advance the reader to the first step at or after ``start`` holding ``variable``. + + With several meshes in one file a step no longer holds every variable, so a + reader cannot assume the one it wants is in step 0. Leaves the file *inside* + the matching step (the caller is responsible for ``EndStep``). + + Args: + adios_file: The open :class:`ADIOSFile` + variable: Name of the variable to seek + start: First step to consider + + Returns: + Index of the step holding ``variable``. + + Raises: + KeyError: If no step holds ``variable``. + """ + for i in range(start, adios_file.file.Steps()): + adios_file.file.BeginStep() + if variable in adios_file.io.AvailableVariables().keys(): + return i + adios_file.file.EndStep() + raise KeyError(f"'{variable}' not found in file") + + def convert_file_mode(mode: FileMode) -> adios2.Mode: # type: ignore[override] match mode: case FileMode.append: @@ -141,6 +197,7 @@ def read_timestamps( adios = adios2.ADIOS(comm) backend_args = get_default_backend_args(backend_args) + prefix = mesh_prefix(backend_args) with ADIOSFile( adios=adios, filename=filename, @@ -148,7 +205,7 @@ def read_timestamps( engine=backend_args["engine"], io_name="TimestepReader", ) as adios_file: - time_name = f"{function_name}_time" + time_name = f"{prefix}{function_name}_time" time_stamps = [] for _ in range(adios_file.file.Steps()): adios_file.file.BeginStep() @@ -188,10 +245,22 @@ def write_mesh( backend_args = get_default_backend_args(backend_args) if "io_name" not in backend_args.keys(): backend_args["io_name"] = "MeshWriter" + prefix = mesh_prefix(backend_args) mode = convert_file_mode(mode) gdim = mesh.local_geometry.shape[1] adios = adios2.ADIOS(comm) + + # Topology is constant in time and written once. In append mode it may + # already be present either because this mesh was written at an earlier time + # step, or because a *different* mesh was written first -- so the test is + # whether this mesh's topology exists, not whether the file is new. + topology_exists = False + if mode == adios2.Mode.Append: + topology_exists = check_variable_exists( + adios, filename, f"{prefix}Topology", engine=backend_args["engine"] + ) + with ADIOSFile( adios=adios, filename=filename, @@ -203,7 +272,7 @@ def write_mesh( adios_file.file.BeginStep() # Write geometry pointvar = adios_file.io.DefineVariable( - "Points", + f"{prefix}Points", mesh.local_geometry, shape=[mesh.num_nodes_global, gdim], start=[mesh.local_geometry_pos[0], 0], @@ -211,16 +280,18 @@ def write_mesh( ) adios_file.file.Put(pointvar, mesh.local_geometry, adios2.Mode.Sync) - if mode == adios2.Mode.Write: - adios_file.io.DefineAttribute("CellType", mesh.cell_type) - adios_file.io.DefineAttribute("Degree", np.array([mesh.degree], dtype=np.int32)) + if not topology_exists: + adios_file.io.DefineAttribute(f"{prefix}CellType", mesh.cell_type) + adios_file.io.DefineAttribute( + f"{prefix}Degree", np.array([mesh.degree], dtype=np.int32) + ) adios_file.io.DefineAttribute( - "LagrangeVariant", np.array([mesh.lagrange_variant], dtype=np.int32) + f"{prefix}LagrangeVariant", np.array([mesh.lagrange_variant], dtype=np.int32) ) # Write topology (on;y on first write as topology is constant) num_dofs_per_cell = mesh.local_topology.shape[1] dvar = adios_file.io.DefineVariable( - "Topology", + f"{prefix}Topology", mesh.local_topology, shape=[mesh.num_cells_global, num_dofs_per_cell], start=[mesh.local_topology_pos[0], 0], @@ -235,7 +306,7 @@ def write_mesh( if mesh.store_partition: assert mesh.partition_range is not None par_data = adios_file.io.DefineVariable( - "PartitioningData", + f"{prefix}PartitioningData", mesh.ownership_array, shape=[mesh.partition_global], start=[mesh.partition_range[0]], @@ -246,7 +317,7 @@ def write_mesh( adios_file.file.Put(par_data, mesh.ownership_array) assert mesh.ownership_offset is not None par_offset = adios_file.io.DefineVariable( - "PartitioningOffset", + f"{prefix}PartitioningOffset", mesh.ownership_offset, shape=[mesh.num_cells_global + 1], start=[mesh.local_topology_pos[0]], @@ -255,15 +326,16 @@ def write_mesh( adios_file.file.Put(par_offset, mesh.ownership_offset) assert mesh.partition_processes is not None adios_file.io.DefineAttribute( - "PartitionProcesses", np.array([mesh.partition_processes], dtype=np.int32) + f"{prefix}PartitionProcesses", + np.array([mesh.partition_processes], dtype=np.int32), ) - if mode == adios2.Mode.Append and mesh.store_partition: + if topology_exists and mesh.store_partition: warnings.warn("Partitioning data is not written in append mode") # Add time step to file t_arr = np.array([time], dtype=np.float64) time_var = adios_file.io.DefineVariable( - "MeshTime", + f"{prefix}MeshTime", t_arr, shape=[1], start=[0], @@ -300,6 +372,7 @@ def read_mesh_data( legacy = backend_args.get("legacy", False) io_name = backend_args.get("io_name", "MeshReader") engine = backend_args["engine"] + prefix = mesh_prefix(backend_args) with ADIOSFile( adios=adios, filename=filename, @@ -307,12 +380,14 @@ def read_mesh_data( engine=engine, io_name=io_name, ) as adios_file: - # Get time independent mesh variables (mesh topology and cell type info) first - adios_file.file.BeginStep() - # Get mesh topology (distributed) - if "Topology" not in adios_file.io.AvailableVariables().keys(): - raise KeyError(f"Mesh topology not found at Topology in {filename}") - topology = adios_file.io.InquireVariable("Topology") + # Get time independent mesh variables (mesh topology and cell type info) + # first. With several meshes in the file this mesh's topology need not be + # in step 0, so seek the step holding it. + try: + step = _seek_step(adios_file, f"{prefix}Topology") + except KeyError as e: + raise KeyError(f"Mesh topology not found at {prefix}Topology in {filename}") from e + topology = adios_file.io.InquireVariable(f"{prefix}Topology") shape = topology.Shape() local_range = compute_local_range(comm, shape[0]) topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]]) @@ -321,31 +396,33 @@ def read_mesh_data( # Check validity of partitioning information if read_from_partition: - if "PartitionProcesses" not in adios_file.io.AvailableAttributes().keys(): + if f"{prefix}PartitionProcesses" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Partitioning information not found in {filename}") - par_num_procs = adios_file.io.InquireAttribute("PartitionProcesses") + par_num_procs = adios_file.io.InquireAttribute(f"{prefix}PartitionProcesses") num_procs = par_num_procs.Data()[0] if num_procs != comm.size: raise ValueError(f"Number of processes in file ({num_procs})!=({comm.size=})") # Get mesh cell type - if "CellType" not in adios_file.io.AvailableAttributes().keys(): - raise KeyError(f"Mesh cell type not found at CellType in {filename}") - celltype = adios_file.io.InquireAttribute("CellType") + if f"{prefix}CellType" not in adios_file.io.AvailableAttributes().keys(): + raise KeyError(f"Mesh cell type not found at {prefix}CellType in {filename}") + celltype = adios_file.io.InquireAttribute(f"{prefix}CellType") cell_type = celltype.DataString()[0] # Get basix info - if "LagrangeVariant" not in adios_file.io.AvailableAttributes().keys(): + if f"{prefix}LagrangeVariant" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Mesh LagrangeVariant not found in {filename}") - lvar = adios_file.io.InquireAttribute("LagrangeVariant").Data()[0] - if "Degree" not in adios_file.io.AvailableAttributes().keys(): + lvar = adios_file.io.InquireAttribute(f"{prefix}LagrangeVariant").Data()[0] + if f"{prefix}Degree" not in adios_file.io.AvailableAttributes().keys(): raise KeyError(f"Mesh degree not found in {filename}") - degree = adios_file.io.InquireAttribute("Degree").Data()[0] + degree = adios_file.io.InquireAttribute(f"{prefix}Degree").Data()[0] if not legacy: - time_name = "MeshTime" - for i in range(adios_file.file.Steps()): - if i > 0: + # Geometry for this time stamp is written in the same step as the + # topology or a later one, so resume the scan from there. + time_name = f"{prefix}MeshTime" + for i in range(step, adios_file.file.Steps()): + if i > step: adios_file.file.BeginStep() if time_name in adios_file.io.AvailableVariables().keys(): arr = adios_file.io.InquireVariable(time_name) @@ -366,9 +443,9 @@ def read_mesh_data( raise KeyError(f"No data associated with {time_name}={time} found in {filename}") # Get mesh geometry - if "Points" not in adios_file.io.AvailableVariables().keys(): - raise KeyError(f"Mesh coordinates not found at Points in {filename}") - geometry = adios_file.io.InquireVariable("Points") + if f"{prefix}Points" not in adios_file.io.AvailableVariables().keys(): + raise KeyError(f"Mesh coordinates not found at {prefix}Points in {filename}") + geometry = adios_file.io.InquireVariable(f"{prefix}Points") x_shape = geometry.Shape() geometry_range = compute_local_range(comm, x_shape[0]) geometry.SetSelection( @@ -390,8 +467,8 @@ def read_mesh_data( adios, comm, filename, - "PartitioningData", - "PartitioningOffset", + f"{prefix}PartitioningData", + f"{prefix}PartitioningOffset", backend_args["engine"], ) else: @@ -424,6 +501,7 @@ def write_meshtags( backend_args = {} if backend_args is None else backend_args io_name = backend_args.get("io_name", "MeshTagWriter") engine = backend_args.get("engine", "BP4") + prefix = mesh_prefix(backend_args) adios = adios2.ADIOS(comm) with ADIOSFile( adios=adios, @@ -436,7 +514,7 @@ def write_meshtags( # Write meshtag topology topology_var = adios_file.io.DefineVariable( - data.name + "_topology", + f"{prefix}{data.name}_topology", data.indices, shape=[data.num_entities_global, data.num_dofs_per_entity], start=[data.local_start, 0], @@ -447,7 +525,7 @@ def write_meshtags( # Write meshtag values vals = np.array(data.values) values_var = adios_file.io.DefineVariable( - data.name + "_values", + f"{prefix}{data.name}_values", vals, shape=[data.num_entities_global], start=[data.local_start], @@ -456,7 +534,9 @@ def write_meshtags( adios_file.file.Put(values_var, vals, adios2.Mode.Sync) # Write meshtag dim - adios_file.io.DefineAttribute(data.name + "_dim", np.array([data.dim], dtype=np.uint8)) + adios_file.io.DefineAttribute( + f"{prefix}{data.name}_dim", np.array([data.dim], dtype=np.uint8) + ) adios_file.file.PerformPuts() adios_file.file.EndStep() @@ -481,6 +561,7 @@ def read_meshtags_data( io_name = backend_args.get("io_name", "MeshTagsReader") engine = backend_args["engine"] legacy = backend_args["legacy"] + prefix = mesh_prefix(backend_args) with ADIOSFile( adios=adios, filename=filename, @@ -490,7 +571,7 @@ def read_meshtags_data( ) as adios_file: if not legacy: # Get mesh cell type - dim_attr_name = f"{name}_dim" + dim_attr_name = f"{prefix}{name}_dim" step = 0 for i in range(adios_file.file.Steps()): adios_file.file.BeginStep() @@ -505,7 +586,7 @@ def read_meshtags_data( dim = int(m_dim.Data()[0]) # Get mesh tags entites - topology_name = f"{name}_topology" + topology_name = f"{prefix}{name}_topology" for i in range(step, adios_file.file.Steps()): if i > step: adios_file.file.BeginStep() @@ -531,7 +612,7 @@ def read_meshtags_data( adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred) # Get mesh tags values - values_name = f"{name}_values" + values_name = f"{prefix}{name}_values" if values_name not in adios_file.io.AvailableVariables().keys(): raise KeyError(f"{values_name} not found") @@ -594,8 +675,11 @@ def read_meshtags_data( adios_file.file.PerformGets() adios_file.file.EndStep() + # Tag values are int32 unless the caller asks for something wider. The + # submesh parent link stores global cell indices, which need int64. + values_dtype = np.dtype(backend_args.get("values_dtype", np.int32)) return MeshTagsData( - name=name, values=tag_values.astype(np.int32), indices=mesh_entities, dim=dim + name=name, values=tag_values.astype(values_dtype), indices=mesh_entities, dim=dim ) @@ -614,6 +698,7 @@ def read_dofmap( Dofmap as an AdjacencyList """ backend_args = {} if backend_args is None else backend_args + prefix = mesh_prefix(backend_args) # Handles legacy io4dolfinx files, modern files, and custom location of dofmap. legacy = backend_args.get("legacy", False) @@ -623,13 +708,13 @@ def read_dofmap( if legacy: dofmap_path = "Dofmap" else: - dofmap_path = f"{name}_dofmap" + dofmap_path = f"{prefix}{name}_dofmap" if (xdofmap_path := backend_args.get("offsets", None)) is None: if legacy: xdofmap_path = "XDofmap" else: - xdofmap_path = f"{name}_XDofmap" + xdofmap_path = f"{prefix}{name}_XDofmap" engine = backend_args.get("engine", "BP4") @@ -665,6 +750,7 @@ def read_dofs( legacy = backend_args.get("legacy", False) engine = backend_args.get("engine", "BP4") io_name = backend_args.get("io_name", f"{name}_FunctionReader") + prefix = mesh_prefix(backend_args) # Check that file contains the function to read adios = adios2.ADIOS(comm) check_file_exists(filename) @@ -680,8 +766,11 @@ def read_dofs( variables = set( sorted( map( - lambda x: x.split("_time")[0], - filter(lambda x: x.endswith("_time"), adios_file.io.AvailableVariables()), + lambda x: x[len(prefix) :].split("_time")[0], + filter( + lambda x: x.startswith(prefix) and x.endswith("_time"), + adios_file.io.AvailableVariables(), + ), ) ) ) @@ -691,9 +780,9 @@ def read_dofs( if legacy: array_path = "Values" else: - array_path = f"{name}_values" + array_path = f"{prefix}{name}_values" - time_name = f"{name}_time" + time_name = f"{prefix}{name}_time" return read_array(adios, filename, array_path, engine, comm, time, time_name, legacy=legacy) @@ -724,9 +813,10 @@ def read_cell_perms( # Open ADIOS engine backend_args = {} if backend_args is None else backend_args engine = backend_args.get("engine", "BP4") + prefix = mesh_prefix(backend_args) cell_perms, _ = read_array( - adios, filename, "CellPermutations", engine=engine, comm=comm, legacy=True + adios, filename, f"{prefix}CellPermutations", engine=engine, comm=comm, legacy=True ) return cell_perms.astype(np.uint32) @@ -780,6 +870,7 @@ def write_function( backend_args = get_default_backend_args(backend_args) engine = backend_args["engine"] io_name = backend_args.get("io_name", "{name}_writer") + prefix = mesh_prefix(backend_args) adios = adios2.ADIOS(comm) cell_permutations_exists = False @@ -787,10 +878,14 @@ def write_function( XDofmap_exists = False if adios_mode == adios2.Mode.Append: cell_permutations_exists = check_variable_exists( - adios, filename, "CellPermutations", engine=engine + adios, filename, f"{prefix}CellPermutations", engine=engine + ) + dofmap_exists = check_variable_exists( + adios, filename, f"{prefix}{u.name}_dofmap", engine=engine + ) + XDofmap_exists = check_variable_exists( + adios, filename, f"{prefix}{u.name}_XDofmap", engine=engine ) - dofmap_exists = check_variable_exists(adios, filename, f"{u.name}_dofmap", engine=engine) - XDofmap_exists = check_variable_exists(adios, filename, f"{u.name}_XDofmap", engine=engine) with ADIOSFile( adios=adios, filename=filename, mode=adios_mode, engine=engine, io_name=io_name, comm=comm @@ -798,9 +893,10 @@ def write_function( adios_file.file.BeginStep() if not cell_permutations_exists: - # Add mesh permutations + # Add mesh permutations. These belong to the mesh, not the function, + # so they are shared by every function on that mesh. pvar = adios_file.io.DefineVariable( - "CellPermutations", + f"{prefix}CellPermutations", u.cell_permutations, shape=[u.num_cells_global], start=[u.local_cell_range[0]], @@ -811,7 +907,7 @@ def write_function( if not dofmap_exists: # Add dofmap dofmap_var = adios_file.io.DefineVariable( - f"{u.name}_dofmap", + f"{prefix}{u.name}_dofmap", u.dofmap_array, shape=[u.global_dofs_in_dofmap], start=[u.dofmap_range[0]], @@ -822,7 +918,7 @@ def write_function( if not XDofmap_exists: # Add XDofmap xdofmap_var = adios_file.io.DefineVariable( - f"{u.name}_XDofmap", + f"{prefix}{u.name}_XDofmap", u.dofmap_offsets, shape=[u.num_cells_global + 1], start=[u.local_cell_range[0]], @@ -831,7 +927,7 @@ def write_function( adios_file.file.Put(xdofmap_var, u.dofmap_offsets) val_var = adios_file.io.DefineVariable( - f"{u.name}_values", + f"{prefix}{u.name}_values", u.values, shape=[u.num_dofs_global], start=[u.dof_range[0]], @@ -842,7 +938,7 @@ def write_function( # Add time step to file t_arr = np.array([time], dtype=np.float64) time_var = adios_file.io.DefineVariable( - f"{u.name}_time", + f"{prefix}{u.name}_time", t_arr, shape=[1], start=[0], diff --git a/src/io4dolfinx/backends/h5py/backend.py b/src/io4dolfinx/backends/h5py/backend.py index aa8e4bc..5b3e8dd 100644 --- a/src/io4dolfinx/backends/h5py/backend.py +++ b/src/io4dolfinx/backends/h5py/backend.py @@ -20,7 +20,7 @@ from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from ...utils import check_file_exists, compute_local_range -from .. import FileMode, ReadMode +from .. import DEFAULT_MESH_NAME, FileMode, ReadMode, get_mesh_name read_mode = ReadMode.parallel @@ -61,7 +61,9 @@ def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None): def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: - args = arguments or {"legacy": False} # If meshtags is read from legacy + args = dict(arguments) if arguments else {} + args.setdefault("legacy", False) # If meshtags is read from legacy + args.setdefault("name", DEFAULT_MESH_NAME) # Which mesh in the file to act on return args @@ -144,7 +146,7 @@ def read_timestamps( The time-stamps """ check_file_exists(filename) - mesh_name = "mesh" + mesh_name = get_mesh_name(backend_args) with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: mesh_directory = h5file[mesh_name] functions = mesh_directory["functions"] @@ -172,7 +174,7 @@ def write_mesh( """ backend_args = get_default_backend_args(backend_args) h5_mode = convert_file_mode(mode) - mesh_name = "mesh" + mesh_name = get_mesh_name(backend_args) with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: if mesh_name in h5file.keys() and h5_mode == "a": mesh_directory = h5file[mesh_name] @@ -254,11 +256,12 @@ def read_mesh_data( """ backend_args = get_default_backend_args(backend_args) + mesh_name = get_mesh_name(backend_args) with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: - if "mesh" not in h5file.keys(): - raise KeyError("Could not find mesh in file") - mesh_group = h5file["mesh"] + if mesh_name not in h5file.keys(): + raise KeyError(f"Could not find mesh '{mesh_name}' in {filename}") + mesh_group = h5file[mesh_name] timestamps = mesh_group.attrs["timestamps"] assert time is not None parent_group = np.flatnonzero(np.isclose(timestamps, time)) @@ -333,10 +336,11 @@ def write_meshtags( """ backend_args = get_default_backend_args(backend_args) + mesh_name = get_mesh_name(backend_args) with h5pyfile(filename, filemode="a", comm=comm, force_serial=False) as h5file: - if "mesh" not in h5file.keys(): - raise KeyError("Could not find mesh in file") - mesh_group = h5file["mesh"] + if mesh_name not in h5file.keys(): + raise KeyError(f"Could not find mesh '{mesh_name}' in file") + mesh_group = h5file[mesh_name] if "tags" not in mesh_group.keys(): tags = mesh_group.create_group("tags") else: @@ -394,14 +398,15 @@ def read_meshtags_data( dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type)) values = mesh["values"] else: - if "mesh" not in h5file.keys(): - raise KeyError("No mesh found") - mesh = h5file["mesh"] + mesh_name = get_mesh_name(backend_args) + if mesh_name not in h5file.keys(): + raise KeyError(f"No mesh '{mesh_name}' found") + mesh = h5file[mesh_name] if "tags" not in mesh.keys(): raise KeyError("Could not find 'tags' in file, are you sure this is a checkpoint?") tags = mesh["tags"] if name not in tags.keys(): - raise KeyError(f"Could not find {name} in '/mesh/tags/' in {filename}") + raise KeyError(f"Could not find {name} in '/{mesh_name}/tags/' in {filename}") tag = tags[name] dim = tag.attrs["dim"] @@ -410,7 +415,8 @@ def read_meshtags_data( num_entities_global = topology.shape[0] topology_range = compute_local_range(comm, num_entities_global) indices = topology[slice(*topology_range), :].astype(np.int64) - vals = values[slice(*topology_range)].astype(np.int32) + values_dtype = np.dtype(backend_args.get("values_dtype", np.int32)) + vals = values[slice(*topology_range)].astype(values_dtype) return MeshTagsData(name=name, values=vals, indices=indices, dim=dim) @@ -433,7 +439,7 @@ def read_dofmap( # If dofmap is read with full path, it is passed through backend_args dofmap_key = backend_args.get("dofmap", None) if dofmap_key is None: - mesh_name = "mesh" # Prepare for multiple meshes + mesh_name = get_mesh_name(backend_args) if mesh_name not in h5file.keys(): raise KeyError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] @@ -490,7 +496,7 @@ def read_dofs( Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: - mesh_name = "mesh" # Prepare for multiple meshes + mesh_name = get_mesh_name(backend_args) if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] @@ -532,7 +538,7 @@ def read_cell_perms( """ with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: - mesh_name = "mesh" # Prepare for multiple meshes + mesh_name = get_mesh_name(backend_args) if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] @@ -562,7 +568,7 @@ def write_function( backend_args: Arguments to backend """ - mesh_name = "mesh" # Prepare for multiple meshes + mesh_name = get_mesh_name(backend_args) backend_args = get_default_backend_args(backend_args) h5_mode = convert_file_mode(mode) with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: @@ -754,7 +760,7 @@ def read_function_names( check_file_exists(filename) backend_args = get_default_backend_args(backend_args) with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: - mesh_name = "mesh" # Prepare for multiple meshes + mesh_name = get_mesh_name(backend_args) if mesh_name not in h5file.keys(): raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") mesh = h5file[mesh_name] diff --git a/src/io4dolfinx/backends/vtkhdf/backend.py b/src/io4dolfinx/backends/vtkhdf/backend.py index e4f1737..6a73080 100644 --- a/src/io4dolfinx/backends/vtkhdf/backend.py +++ b/src/io4dolfinx/backends/vtkhdf/backend.py @@ -16,7 +16,7 @@ from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData from io4dolfinx.utils import check_file_exists, compute_local_range -from .. import FileMode, ReadMode +from .. import DEFAULT_MESH_NAME, FileMode, ReadMode from ..h5py.backend import convert_file_mode, h5pyfile from ..pyvista.backend import _arbitrary_lagrange_vtk, _cell_degree, _first_order_vtk @@ -34,7 +34,10 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any] Returns: Updated backend arguments """ - args = arguments or {"name": "mesh"} + # `arguments or {...}` would drop the default whenever the caller passed any + # other argument, so set it on a copy instead. + args = dict(arguments) if arguments else {} + args.setdefault("name", DEFAULT_MESH_NAME) return args diff --git a/src/io4dolfinx/checkpointing.py b/src/io4dolfinx/checkpointing.py index c261f53..3e6c8be 100644 --- a/src/io4dolfinx/checkpointing.py +++ b/src/io4dolfinx/checkpointing.py @@ -55,6 +55,32 @@ logger = logging.getLogger(__name__) +def _with_mesh_name( + backend_args: dict[str, Any] | None, mesh_name: str | None +) -> dict[str, Any] | None: + """Select which mesh in the file a backend call refers to. + + A checkpoint may hold several meshes -- a mesh and a submesh of it, say -- + and backends take the one to act on from ``backend_args["name"]``. This + folds an explicit ``mesh_name`` in without mutating the caller's dictionary. + ``mesh_name=None`` leaves ``backend_args`` untouched, so an existing + ``backend_args["name"]`` still wins and files written before named meshes + existed stay readable. + + Args: + backend_args: Arguments for the backend, or None + mesh_name: Name of the mesh to act on, or None for the default + + Returns: + Backend arguments naming the requested mesh. + """ + if mesh_name is None: + return backend_args + args = dict(backend_args) if backend_args else {} + args["name"] = mesh_name + return args + + def write_attributes( filename: Path | str, comm: MPI.Comm, @@ -111,6 +137,7 @@ def read_timestamps( function_name: str, backend_args: dict[str, typing.Any] | None = None, backend: str | None = None, + mesh_name: str | None = None, ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] """ Read time-stamps from a checkpoint file. @@ -121,9 +148,13 @@ def read_timestamps( function_name: Name of the function to read time-stamps for backend_args: Arguments for backend, for instance file type. backend: What backend to use for writing. + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. Returns: The time-stamps """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Reading time-stamps from {filename} for function {function_name}") logger.debug(f"Using {backend} backend with arguments {backend_args} to read time-stamps") check_file_exists(filename) @@ -140,6 +171,7 @@ def write_meshtags( backend_args: dict[str, Any] | None = None, backend: str | None = None, on_input_mesh: bool = False, + mesh_name: str | None = None, ): """ Write meshtags associated with input mesh to file. @@ -157,7 +189,11 @@ def write_meshtags( backend: IO backend on_input_mesh: If True, the meshtags are written with the node ordering of the input mesh. + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Writing meshtags to {filename} for meshtag {meshtag_name or meshtags.name}") logger.debug(f"Using {backend} backend with arguments {backend_args} to write meshtags") @@ -219,6 +255,7 @@ def read_meshtags( meshtag_name: str, backend_args: dict[str, Any] | None = None, backend: str | None = None, + mesh_name: str | None = None, ) -> dolfinx.mesh.MeshTags: """ Read meshtags from file and return a :class:`dolfinx.mesh.MeshTags` object. @@ -228,9 +265,13 @@ def read_meshtags( mesh: The mesh associated with the meshtags meshtag_name: The name of the meshtag to read engine: Adios2 Engine + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. Returns: The meshtags """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Reading meshtags from {filename} for meshtag {meshtag_name}") logger.debug(f"Using {backend} backend with arguments {backend_args} to read meshtags") check_file_exists(filename) @@ -260,6 +301,7 @@ def read_function( name: str | None = None, backend_args: dict[str, Any] | None = None, backend: str | None = None, + mesh_name: str | None = None, ): """ Read checkpoint from file and fill it into `u`. @@ -269,7 +311,11 @@ def read_function( u: Function to fill time: Time-stamp associated with checkpoint name: If not provided, `u.name` is used to search through the input file for the function + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug( f"Reading function checkpoint from {filename} for function {name or u.name} at time {time}" ) @@ -410,6 +456,7 @@ def read_mesh( backend_args: dict[str, Any] | None = None, backend: str | None = None, max_facet_to_cell_links: int = 2, + mesh_name: str | None = None, ) -> dolfinx.mesh.Mesh: """ Read an ADIOS2 mesh into DOLFINx. @@ -424,9 +471,13 @@ def read_mesh( backend_args: List of arguments to reader backend max_facet_to_cell_links: Maximum number of cells a facet can be connected to. + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. Returns: The distributed mesh """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Reading mesh from {filename}") logger.debug(f"Using {backend} backend with arguments {backend_args}") logger.debug(f"Time {time} and read_from_partition {read_from_partition}") @@ -521,6 +572,7 @@ def write_mesh( store_partition_info: bool = False, backend_args: dict[str, Any] | None = None, backend: str | None = None, + mesh_name: str | None = None, ): """ Write a mesh to file. @@ -530,7 +582,11 @@ def write_mesh( mesh: The mesh to write to file store_partition_info: Store mesh partitioning (including ghosting) to file + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Writing mesh to {filename}") logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}") mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info) @@ -555,6 +611,7 @@ def write_function( name: str | None = None, backend_args: dict[str, Any] | None = None, backend: str | None = None, + mesh_name: str | None = None, ): """ Write function checkpoint to file. @@ -567,7 +624,11 @@ def write_function( name: Name of function to write. If None, the name of the function is used. backend_args: Arguments to the IO backend. backend: The backend to use + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. """ + backend_args = _with_mesh_name(backend_args, mesh_name) n = u.name if name is None else name logger.debug(f"Writing function checkpoint to {filename} for function {n} at time {time}") logger.debug(f"Using {backend} backend with arguments {backend_args}") @@ -633,6 +694,7 @@ def read_function_names( comm: MPI.Comm, backend_args: dict[str, Any] | None = None, backend: str = "h5py", + mesh_name: str | None = None, ) -> list[str]: """Read all function names from a file. @@ -640,10 +702,14 @@ def read_function_names( filename: Path to file comm: MPI communicator to launch IO on. backend_args: Arguments to backend + mesh_name: Name of the mesh in the file to act on. Defaults to the + mesh written without an explicit name, so existing checkpoints and + call sites are unaffected. Returns: A list of function names. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Reading function names from {filename}") logger.debug(f"Using {backend} backend with arguments {backend_args} to read function names") check_file_exists(filename) @@ -658,6 +724,7 @@ def write_point_data( mode: FileMode, backend_args: dict[str, Any] | None, backend: str = "vtkhdf", + mesh_name: str | None = None, ): """Write function to file by interpolating into geometry nodes. @@ -669,7 +736,10 @@ def write_point_data( mode: Append or write backend_args: The backend arguments backend: Which backend to use. + mesh_name: Name of the mesh in the file to write the data against. + Defaults to the mesh written without an explicit name. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}") V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape))) v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype) @@ -698,6 +768,7 @@ def write_cell_data( mode: FileMode, backend_args: dict[str, Any] | None, backend: str = "vtkhdf", + mesh_name: str | None = None, ): """Write function to file by interpolating into cell midpoints. @@ -708,7 +779,10 @@ def write_cell_data( time: Time stamp mode: Append or write backend_args: The backend arguments + mesh_name: Name of the mesh in the file to write the data against. + Defaults to the mesh written without an explicit name. """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}") V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape)) v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype) diff --git a/src/io4dolfinx/comm_helpers.py b/src/io4dolfinx/comm_helpers.py index 8a4b62e..8f1b52e 100644 --- a/src/io4dolfinx/comm_helpers.py +++ b/src/io4dolfinx/comm_helpers.py @@ -15,6 +15,8 @@ "send_dofmap_and_recv_values", "send_and_recv_cell_perm", "send_dofs_and_recv_values", + "neighbourhood_ranks", + "exchange_to_owners", "numpy_to_mpi", ] @@ -261,3 +263,94 @@ def send_dofs_and_recv_values( values = np.empty_like(inc_values, dtype=input_array.dtype) values[proc_to_local] = inc_values return values + + +def neighbourhood_ranks( + comm: MPI.Intracomm, destinations: npt.NDArray[np.int32] +) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: + """Find the ranks this process exchanges data with. + + A process knows which ranks it must send to, but not which ranks will send + to it. Building a distributed graph over the whole communicator and reading + its neighbours back discovers the incoming side without an all-to-all. + + Both lists are sorted ascending so that every process packs and unpacks + neighbour blocks in the same order. + + Args: + comm: The MPI communicator + destinations: Ranks this process sends to (duplicates allowed) + + Returns: + ``(sources, destinations)``, each sorted and without duplicates. + """ + dest = np.unique(np.asarray(destinations, dtype=np.int32)) + graph = comm.Create_dist_graph([comm.rank], [len(dest)], dest.tolist(), reorder=False) + sources, _, _ = graph.Get_dist_neighbors() + graph.Free() + return np.unique(np.asarray(sources, dtype=np.int32)), dest + + +def exchange_to_owners( + comm: MPI.Intracomm, + owners: npt.NDArray[np.int32], + arrays: list[npt.NDArray], +) -> tuple[ + list[npt.NDArray], + npt.NDArray[np.int32], + npt.NDArray[np.int32], + npt.NDArray[np.int32], + npt.NDArray[np.int32], +]: + """Send each entry of ``arrays`` to the rank naming it in ``owners``. + + Entry ``i`` of every array in ``arrays`` travels together to ``owners[i]``. + Data is packed grouped by destination rank in ascending rank order and is + received grouped by source rank in ascending rank order, preserving the + sender's relative order within each group. Both properties matter to callers + that must pair a reply with the request that produced it. + + Args: + comm: The MPI communicator + owners: Destination rank of each entry; ``-1`` is not allowed + arrays: Arrays to send, each of length ``len(owners)``. The leading axis + is the one distributed; trailing axes travel as a block. + + Returns: + ``(received, sources, send_counts, recv_counts, insert_position)`` where + ``received`` holds the incoming counterpart of each entry of ``arrays``, + ``sources`` the ascending source ranks, ``send_counts`` and + ``recv_counts`` the number of entries exchanged with each destination and + source, and ``insert_position`` the index each local entry was packed at + (so a reply can be unpacked back into local order). + """ + owners = np.asarray(owners, dtype=np.int32) + sources, destinations = neighbourhood_ranks(comm, owners) + + send_counts = np.zeros(len(destinations), dtype=np.int32) + if len(owners) > 0: + present, counts = np.unique(owners, return_counts=True) + send_counts[np.searchsorted(destinations, present)] = counts + insert_position = compute_insert_position(owners, destinations, send_counts) + + recv_counts = np.zeros(len(sources), dtype=np.int32) + forward = comm.Create_dist_graph_adjacent( + sources.tolist(), destinations.tolist(), reorder=False + ) + forward.Neighbor_alltoall(send_counts, recv_counts) + + received = [] + for array in arrays: + array = np.asarray(array) + block = int(np.prod(array.shape[1:], dtype=np.int64)) if array.ndim > 1 else 1 + packed = np.zeros(array.shape, dtype=array.dtype) + packed[insert_position] = array + out = np.zeros((int(recv_counts.sum()), *array.shape[1:]), dtype=array.dtype) + mpi_type = numpy_to_mpi[array.dtype.type] + forward.Neighbor_alltoallv( + [packed.reshape(-1), send_counts * block, mpi_type], + [out.reshape(-1), recv_counts * block, mpi_type], + ) + received.append(out) + forward.Free() + return received, sources, send_counts, recv_counts, insert_position diff --git a/src/io4dolfinx/readers.py b/src/io4dolfinx/readers.py index 0de359a..ef3456d 100644 --- a/src/io4dolfinx/readers.py +++ b/src/io4dolfinx/readers.py @@ -35,6 +35,21 @@ logger = logging.getLogger(__name__) +def _with_mesh_name( + backend_args: dict[str, Any] | None, mesh_name: str | None +) -> dict[str, Any] | None: + """Select which mesh in the file a backend call refers to. + + See :func:`io4dolfinx.checkpointing._with_mesh_name`; duplicated here to keep + :mod:`io4dolfinx.readers` free of a circular import. + """ + if mesh_name is None: + return backend_args + args = dict(backend_args) if backend_args else {} + args["name"] = mesh_name + return args + + def map_dofmap(dofmap: dolfinx.graph.AdjacencyList, bs: int | np.int64) -> npt.NDArray[np.int64]: """ Map xxxyyyzzz to xyzxyz @@ -399,6 +414,7 @@ def read_point_data( time: float | None = None, backend_args: dict[str, Any] | None = None, backend: str = "xdmf", + mesh_name: str | None = None, ) -> dolfinx.fem.Function: """Read data from the nodes of a mesh. @@ -410,12 +426,15 @@ def read_point_data( name: Name of point data mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`. time: Time-step to read from. + mesh_name: Name of the mesh in the file the data belongs to. Defaults + to the mesh written without an explicit name. Returns: A function in the space equivalent to the mesh coordinate element (up to shape). """ + backend_args = _with_mesh_name(backend_args, mesh_name) logger.debug(f"Reading point data from {filename} with name {name} at time {time}") logger.debug(f"Using backend {backend} with arguments {backend_args}") backend_cls = get_backend(backend) @@ -465,6 +484,7 @@ def read_cell_data( time: float | None = None, backend_args: dict[str, Any] | None = None, backend: str = "xdmf", + mesh_name: str | None = None, ) -> dolfinx.fem.Function: """Read data from the nodes of a mesh. @@ -476,10 +496,13 @@ def read_cell_data( name: Name of point data mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`. time: Time-step to read from. + mesh_name: Name of the mesh in the file the data belongs to. Defaults + to the mesh written without an explicit name. Returns: A function in a DG-0 space on the mesh. The cells not found in input is set to zero. """ + backend_args = _with_mesh_name(backend_args, mesh_name) backend_cls = get_backend(backend) diff --git a/src/io4dolfinx/submesh.py b/src/io4dolfinx/submesh.py new file mode 100644 index 0000000..b36d1db --- /dev/null +++ b/src/io4dolfinx/submesh.py @@ -0,0 +1,491 @@ +# Copyright (C) 2026 Jørgen Schartum Dokken +# +# This file is part of io4dolfinx +# +# SPDX-License-Identifier: MIT + +"""Checkpointing of submeshes. + +A :func:`dolfinx.mesh.create_submesh` submesh cannot be checkpointed the way an +ordinary mesh is. Two properties of the DOLFINx construction get in the way: + +1. For a submesh of co-dimension greater than zero, a cell of the submesh is an + *entity* of the parent, and the vertices of an entity are ordered by their + current global vertex index. Re-partitioning the parent renumbers the + vertices, so the same submesh cell comes back with its vertices in a + different order and the degrees of freedom inside it land in different + positions. +2. :func:`dolfinx.mesh.create_submesh` builds its vertex map allowing owner + changes, so a process can own a vertex -- and a dof -- that is only incident + to cells it ghosts. Such a dof has no position in any owned cell, which is + what :func:`io4dolfinx.utils.compute_dofmap_pos` needs. + +So a function is never read onto a re-derived submesh. Instead the submesh is +stored as an ordinary, independent mesh (whose cells keep the stored node order +and whose ownership is the ordinary one), read back as such, and the data is +then transferred to a mesh re-derived from the parent when the caller needs the +entity maps that make mixed-dimensional assembly possible. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from mpi4py import MPI + +import dolfinx +import numpy as np +import numpy.typing as npt +from dolfinx import cpp as _cpp + +from .backends import FileMode, get_backend +from .checkpointing import _with_mesh_name, write_mesh +from .comm_helpers import exchange_to_owners, numpy_to_mpi +from .structures import MeshTagsData +from .utils import check_file_exists, compute_local_range, index_owner + +__all__ = [ + "SubmeshCheckpoint", + "read_submesh", + "transfer_submesh_function", + "write_submesh", +] + +logger = logging.getLogger(__name__) + + +def _link_name(mesh_name: str) -> str: + """Name of the meshtag holding the parent link of the submesh ``mesh_name``. + + The link is a tag on the *parent* mesh, so it is stored in the parent's + namespace and must not collide with a tag a user wrote. + """ + return f"__submesh_link_{mesh_name}" + + +@dataclass +class SubmeshCheckpoint: + """A submesh re-derived from a parent mesh that was read from a checkpoint.""" + + #: The submesh, derived from the parent with :func:`dolfinx.mesh.create_submesh` + submesh: dolfinx.mesh.Mesh + #: Map from cells of :attr:`submesh` to entities of the parent + cell_map: Any + #: Map from vertices of :attr:`submesh` to vertices of the parent + vertex_map: Any + #: Map from geometry nodes of :attr:`submesh` to nodes of the parent + node_map: npt.NDArray[np.int32] + #: Parent entities that make up the submesh (local indices, owned and ghost) + parent_entities: npt.NDArray[np.int32] + #: Index the cell had in the stored submesh, per local cell (owned and ghost) + stored_cells: npt.NDArray[np.int64] + + +def write_submesh( + filename: Path | str, + submesh: dolfinx.mesh.Mesh, + parent: dolfinx.mesh.Mesh, + cell_map: Any, + mesh_name: str, + parent_mesh_name: str | None = None, + parent_filename: Path | str | None = None, + time: float = 0.0, + mode: FileMode = FileMode.append, + backend_args: dict[str, Any] | None = None, + backend: str | None = None, +): + """Write a submesh, and the link that ties it back to its parent. + + The submesh is stored as an ordinary independent mesh under ``mesh_name``, so + :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function` read it + back with no special handling. Alongside it goes a *parent link*: for every + stored submesh cell, the parent geometry nodes of the entity it came from. + That is what lets :func:`read_submesh` find the same entities in a + re-partitioned parent. + + The link tags entities of the *parent*, so it is stored in the parent's + namespace, in ``parent_filename``. Leave that unset to keep everything in one + checkpoint, or point it at the parent's own file to keep the two apart; the + submesh always goes to ``filename``. Either way the parent must already have + been written there with :func:`io4dolfinx.write_mesh`. + + Args: + filename: Path to write the submesh to + submesh: The submesh, as returned by :func:`dolfinx.mesh.create_submesh` + parent: The mesh the submesh was derived from + cell_map: Map from submesh cells to parent entities, as returned by + :func:`dolfinx.mesh.create_submesh` + mesh_name: Name to store the submesh under + parent_mesh_name: Name the parent mesh is stored under + parent_filename: File holding the parent mesh, if not ``filename`` + time: Time stamp associated with the submesh geometry + mode: Whether to create the file or append to it + backend_args: Arguments for the backend + backend: Which backend to use + """ + logger.debug(f"Writing submesh '{mesh_name}' to {filename}") + dim = submesh.topology.dim + + # 1. The submesh itself, as a perfectly ordinary mesh. + write_mesh( + Path(filename), + submesh, + mode=mode, + time=time, + backend_args=backend_args, + backend=backend, + mesh_name=mesh_name, + ) + + # 2. The parent link. Structurally this is a meshtag on the parent: the + # tagged entities are the parent entities that became submesh cells, and + # each value is the index that cell has in the stored submesh. + cell_imap = submesh.topology.index_map(dim) + num_owned = cell_imap.size_local + parent_entities = cell_map.sub_topology_to_topology(np.arange(num_owned, dtype=np.int32), False) + values = np.arange(*cell_imap.local_range, dtype=np.int64) + + parent.topology.create_connectivity(dim, parent.topology.dim) + parent.topology.create_connectivity(0, parent.topology.dim) + entities_to_geometry = dolfinx.cpp.mesh.entities_to_geometry( + parent._cpp_object, dim, parent_entities, False + ) + indices = ( + parent.geometry.index_map() + .local_to_global(entities_to_geometry.reshape(-1)) + .reshape(entities_to_geometry.shape) + ) + + tag_data = MeshTagsData( + values=values, + num_entities_global=cell_imap.size_global, + num_dofs_per_entity=entities_to_geometry.shape[1], + indices=indices, + name=_link_name(mesh_name), + local_start=cell_imap.local_range[0], + dim=dim, + cell_type=submesh.topology.cell_name(), + ) + backend_cls = get_backend(backend) + parent_args = backend_cls.get_default_backend_args( + _with_mesh_name(backend_args, parent_mesh_name) + ) + link_file = filename if parent_filename is None else parent_filename + backend_cls.write_meshtags(link_file, parent.comm, tag_data, backend_args=parent_args) + + +def read_submesh( + filename: Path | str, + parent: dolfinx.mesh.Mesh, + mesh_name: str, + parent_mesh_name: str | None = None, + backend_args: dict[str, Any] | None = None, + backend: str | None = None, +) -> SubmeshCheckpoint: + """Re-derive a stored submesh from a parent mesh that has been read back. + + The result is a genuine :func:`dolfinx.mesh.create_submesh` submesh, so its + entity maps are accepted by :func:`dolfinx.fem.form` for mixed-dimensional + assembly against ``parent``. + + It does **not** carry the checkpointed data. Read that onto the standalone + submesh with :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function`, + then move it across with :func:`transfer_submesh_function`. + + Args: + filename: File holding the parent mesh and the parent link, i.e. whatever + was passed as ``parent_filename`` to :func:`write_submesh` + parent: The parent mesh, as read back from the checkpoint + mesh_name: Name the submesh was stored under + parent_mesh_name: Name the parent mesh is stored under + backend_args: Arguments for the backend + backend: Which backend to use + + Returns: + The re-derived submesh and the maps relating it to ``parent``. + """ + logger.debug(f"Reading submesh '{mesh_name}' from {filename}") + check_file_exists(filename) + backend_cls = get_backend(backend) + parent_args = backend_cls.get_default_backend_args( + _with_mesh_name(backend_args, parent_mesh_name) + ) + # Stored cell indices are global, so they must survive as int64. + parent_args["values_dtype"] = np.int64 + data = backend_cls.read_meshtags_data(filename, parent.comm, _link_name(mesh_name), parent_args) + dim = int(data.dim) + + local_entities, local_values = dolfinx.io.distribute_entity_data( + parent, dim, data.indices, data.values.astype(np.int64) + ) + parent.topology.create_connectivity(dim, 0) + parent.topology.create_connectivity(dim, parent.topology.dim) + adj = dolfinx.graph.adjacencylist(local_entities) + tags = dolfinx.mesh.meshtags_from_entities( + parent, dim, adj, np.asarray(local_values, dtype=np.int64) + ) + + # `distribute_entity_data` only reports entities to processes that hold them; + # a process ghosting an entity may not be among them. Scattering through an + # index map gives every process the owner's value for its ghosts, so the + # submesh is built from a consistent set of entities everywhere. + entity_map = parent.topology.index_map(dim) + marker = dolfinx.la.vector(entity_map, dtype=np.int64) + marker.array[:] = -1 + marker.array[tags.indices] = tags.values + marker.scatter_forward() + entities = np.flatnonzero(marker.array >= 0).astype(np.int32) + + submesh, cell_map, vertex_map, node_map = dolfinx.mesh.create_submesh(parent, dim, entities) + sub_imap = submesh.topology.index_map(submesh.topology.dim) + num_cells = sub_imap.size_local + sub_imap.num_ghosts + parent_entities = cell_map.sub_topology_to_topology(np.arange(num_cells, dtype=np.int32), False) + stored_cells = marker.array[parent_entities].astype(np.int64) + return SubmeshCheckpoint( + submesh=submesh, + cell_map=cell_map, + vertex_map=vertex_map, + node_map=node_map, + parent_entities=parent_entities, + stored_cells=stored_cells, + ) + + +def _point_ownership_data( + comm: MPI.Intracomm, + src_owner: npt.NDArray[np.int32], + points: npt.NDArray[np.floating], + source_cells: npt.NDArray[np.int32], +) -> Any: + """Build a :class:`dolfinx.geometry.PointOwnershipData` from a known routing. + + :func:`dolfinx.fem.create_interpolation_data` would derive the same object by + searching a bounding-box tree for each point, with a padding tolerance and an + extrapolation fallback. Here the owning rank and cell of every point are + already known exactly, so the routing is done directly and no geometric + search -- and no tolerance -- is involved. + + The layout required of the result is dictated by ``fem::interpolate`` and + ``impl::scatter_values``: + + * ``src_owner`` is per interpolation point of the receiving function, in the + order :func:`dolfinx.cpp.fem.interpolation_coords` produced them. + * ``dest_owners`` must be in **contiguous runs by rank**; ``scatter_values`` + reduces it with a non-sorting ``unique``, so interleaved ranks would + silently build the wrong neighbourhood. :func:`exchange_to_owners` groups + by ascending source rank, which satisfies this. + * ``dest_points`` is flat row-major ``(n_recv, 3)`` and ``dest_cells`` indexes + cells of the *sending* function's mesh on the receiving process. + + Args: + comm: The MPI communicator both meshes live on + src_owner: Rank holding the source cell of each point + points: The points, shape ``(num_points, 3)`` + source_cells: Source-mesh local cell of each point, on ``src_owner`` + + Returns: + Interpolation data for :meth:`dolfinx.fem.Function.interpolate_nonmatching`. + """ + (recv_points, recv_cells), sources, _, recv_counts, _ = exchange_to_owners( + comm, src_owner, [points, source_cells] + ) + + # One entry per received point, naming the rank it came from. `sources` is + # ascending and `exchange_to_owners` receives in that order, so the runs are + # contiguous by construction -- which `scatter_values` requires. + dest_owners = np.repeat(sources, recv_counts).astype(np.int32) + + dtype = np.dtype(points.dtype) + cls = getattr(_cpp.geometry, f"PointOwnershipData_float{8 * dtype.itemsize}") + return dolfinx.geometry.PointOwnershipData( + cls( + np.ascontiguousarray(src_owner, dtype=np.int32), + dest_owners, + np.ascontiguousarray(recv_points, dtype=dtype).reshape(-1), + np.ascontiguousarray(recv_cells, dtype=np.int32), + ) + ) + + +def transfer_submesh_function( + u_source: dolfinx.fem.Function, + u_dest: dolfinx.fem.Function, + stored_cells: npt.NDArray[np.int64], +): + """Move data from a standalone submesh onto one re-derived from a parent. + + ``u_source`` lives on the submesh as it was stored -- read with + :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function` -- and + ``u_dest`` on the submesh :func:`read_submesh` derived from the parent. The + two meshes cover the same cells but number their vertices differently and are + partitioned independently, so the transfer cannot be a copy. + + Every cell of ``u_dest``'s mesh has a known counterpart in ``u_source``'s, so + each interpolation point is routed straight to the process and cell that can + evaluate it, with no geometric point location. This is an evaluation followed + by a scatter and a local interpolation: it assembles no form and needs no + entity map between the two meshes, which would have no meaning as neither is + derived from the other. + + Args: + u_source: Function on the submesh as stored + u_dest: Function to fill, on the submesh re-derived from the parent + stored_cells: For each cell of ``u_dest``'s mesh (owned and ghost), the + index that cell had in the stored submesh. From + :attr:`SubmeshCheckpoint.stored_cells`. + + Raises: + NotImplementedError: For H(div)/H(curl) spaces on a manifold submesh, + where DOLFINx's interpolation would return wrong values without + reporting an error. + """ + V_dest = u_dest.function_space + dest_mesh = V_dest.mesh + source_mesh = u_source.function_space.mesh + comm = dest_mesh.comm + assert isinstance(comm, MPI.Intracomm) + + tdim = dest_mesh.topology.dim + if tdim < dest_mesh.geometry.dim and V_dest.element.needs_dof_transformations: + raise NotImplementedError( + "Transferring an H(div)/H(curl) function onto a submesh with" + f" tdim ({tdim}) < gdim ({dest_mesh.geometry.dim}) is not supported." + " DOLFINx cannot reconcile the reference and physical value sizes of" + " these families on a manifold, and the interpolation this transfer" + " relies on does not report that: it returns silently with values that" + " are wrong (measured: the L2 norm of an N1curl field on a facet" + " submesh fell from 9.096 to 4.904). Writing and reading the standalone" + " submesh is exact for these spaces -- use `read_mesh` and" + " `read_function` with the submesh's `mesh_name` and stop there. Only" + " re-deriving from the parent is unavailable." + ) + + sub_imap = dest_mesh.topology.index_map(tdim) + num_cells = sub_imap.size_local + sub_imap.num_ghosts + if len(stored_cells) != num_cells: + raise ValueError( + f"Expected one stored cell index per cell of the destination mesh" + f" ({num_cells}), got {len(stored_cells)}." + ) + cells = np.arange(num_cells, dtype=np.int32) + + # Interpolation points of the destination, grouped per cell. + points = _cpp.fem.interpolation_coords( + V_dest.element._cpp_object, # type: ignore[arg-type] + dest_mesh.geometry._cpp_object, # type: ignore[arg-type] + cells, + ) + points = np.ascontiguousarray(points.T) # (num_points, 3) + points_per_cell = 0 if num_cells == 0 else points.shape[0] // num_cells + + # Which process holds each stored cell in the source mesh, and where. + source_tdim = source_mesh.topology.dim + source_imap = source_mesh.topology.index_map(source_tdim) + owner_rank, owner_cell = _lookup_stored_cells( + comm, + np.asarray(source_mesh.topology.original_cell_index[: source_imap.size_local]), + np.asarray(stored_cells, dtype=np.int64), + source_imap.size_global, + ) + if num_cells > 0 and (owner_rank < 0).any(): + missing = np.unique(stored_cells[owner_rank < 0]) + raise RuntimeError( + f"{len(missing)} cells of the destination submesh have no counterpart in" + f" the stored submesh (for instance {missing[:5]}). The two must describe" + " the same submesh of the same parent." + ) + + interpolation_data = _point_ownership_data( + comm, + np.repeat(owner_rank, points_per_cell), + points, + np.repeat(owner_cell, points_per_cell), + ) + u_dest.interpolate_nonmatching(u_source, cells, interpolation_data=interpolation_data) + u_dest.x.scatter_forward() + + +def _lookup_stored_cells( + comm: MPI.Intracomm, + owned_stored: npt.NDArray[np.int64], + queried_stored: npt.NDArray[np.int64], + num_cells_global: int, +) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: + """Find which process holds each stored cell, and its local index there. + + The two meshes are partitioned independently, so neither side knows where the + other put a given cell. Both sides talk to a third: the process that + :func:`index_owner` assigns the cell in the equal-split layout acts as a + directory. Owners publish into it, queriers read out of it, each in one + neighbourhood exchange, so no process ever holds a global table. + + Args: + comm: The MPI communicator + owned_stored: Stored index of each cell this process owns in the source + mesh; position in the array is the local cell index + queried_stored: Stored indices this process wants to locate + num_cells_global: Number of cells in the stored mesh + + Returns: + ``(rank, local_cell)`` aligned with ``queried_stored``; ``-1`` where the + stored cell was never published. + """ + directory_range = compute_local_range(comm, num_cells_global) + directory_size = int(directory_range[1] - directory_range[0]) + directory_rank = np.full(directory_size, -1, dtype=np.int32) + directory_cell = np.full(directory_size, -1, dtype=np.int32) + + # Publish: each owner tells the directory where it keeps the cell. + owned_stored = np.asarray(owned_stored, dtype=np.int64) + local_cells = np.arange(len(owned_stored), dtype=np.int32) + publish_to = ( + index_owner(comm, owned_stored, num_cells_global) + if len(owned_stored) + else np.empty(0, dtype=np.int32) + ) + (pub_keys, pub_cells), pub_sources, _, pub_counts, _ = exchange_to_owners( + comm, publish_to, [owned_stored, local_cells] + ) + slots = (pub_keys - directory_range[0]).astype(np.int64) + directory_rank[slots] = np.repeat(pub_sources, pub_counts) + directory_cell[slots] = pub_cells + + # Query: ask the directory, and unpack the reply into the order asked. + queried_stored = np.asarray(queried_stored, dtype=np.int64) + query_to = ( + index_owner(comm, queried_stored, num_cells_global) + if len(queried_stored) + else np.empty(0, dtype=np.int32) + ) + (inc_keys,), q_sources, q_send_counts, q_recv_counts, q_insert = exchange_to_owners( + comm, query_to, [queried_stored] + ) + inc_slots = (inc_keys - directory_range[0]).astype(np.int64) + reply_rank = directory_rank[inc_slots] + reply_cell = directory_cell[inc_slots] + + # Send the answers back along the reverse of the query graph. + q_destinations = np.unique(query_to) + reverse = comm.Create_dist_graph_adjacent( + q_destinations.tolist(), q_sources.tolist(), reorder=False + ) + packed_rank = np.zeros(len(queried_stored), dtype=np.int32) + packed_cell = np.zeros(len(queried_stored), dtype=np.int32) + reverse.Neighbor_alltoallv( + [reply_rank, q_recv_counts, numpy_to_mpi[np.int32]], + [packed_rank, q_send_counts, numpy_to_mpi[np.int32]], + ) + reverse.Neighbor_alltoallv( + [reply_cell, q_recv_counts, numpy_to_mpi[np.int32]], + [packed_cell, q_send_counts, numpy_to_mpi[np.int32]], + ) + reverse.Free() + + # `q_insert[i]` is where local entry `i` was packed into the outgoing buffer, + # and the reply comes back in that same packed order, so read the answers + # back out by gathering -- scattering here would apply the inverse + # permutation and silently mix up which cell each answer belongs to. + return packed_rank[q_insert], packed_cell[q_insert] diff --git a/tests/test_reconstruct_mesh.py b/tests/test_reconstruct_mesh.py new file mode 100644 index 0000000..f1d64f2 --- /dev/null +++ b/tests/test_reconstruct_mesh.py @@ -0,0 +1,60 @@ +from mpi4py import MPI + +import basix.ufl +import dolfinx +import numpy as np +import numpy.typing as npt +import pytest +import ufl + +from io4dolfinx import reconstruct_mesh + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("degree", [1, 2, 3]) +@pytest.mark.parametrize( + "cell_type", [dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron] +) +def test_reconstruct(cell_type: dolfinx.mesh.CellType, degree: int, dtype: npt.DTypeLike): + def f(x): + return (0.1 * x[0] ** degree + x[1] ** (degree - 1), x[0], 0.1 * x[2] ** degree) + + el = basix.ufl.element( + "Lagrange", dolfinx.mesh.to_string(cell_type), degree, shape=(3,), dtype=dtype + ) + mesh = dolfinx.mesh.create_unit_cube( + MPI.COMM_WORLD, 10, 10, 10, dtype=dtype, cell_type=cell_type + ) + V = dolfinx.fem.functionspace(mesh, el) + u = dolfinx.fem.Function(V, dtype=dtype) + u.interpolate(f) + + new_mesh = reconstruct_mesh(mesh, degree) + V2 = dolfinx.fem.functionspace(new_mesh, el) + v = dolfinx.fem.Function(V2, dtype=dtype) + cell_map = mesh.topology.index_map(mesh.topology.dim) + num_cells_local = cell_map.size_local + cell_map.num_ghosts + cells = np.arange(num_cells_local, dtype=np.int32) + v.interpolate(u, cells, cells) + + assert new_mesh.geometry.cmap.degree == degree + + f_ex = ufl.as_vector(f(ufl.SpatialCoordinate(new_mesh))) + + local_vol = dolfinx.fem.assemble_scalar( + dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, dtype(1.0)) * ufl.dx, dtype=dtype) + ) + vol = new_mesh.comm.allreduce(local_vol, op=MPI.SUM) + tol = 4e3 * degree * np.finfo(dtype).eps + assert np.isclose(vol, 1, atol=tol, rtol=tol) + + local_surf = dolfinx.fem.assemble_scalar( + dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, dtype(1.0)) * ufl.ds, dtype=dtype) + ) + surf = new_mesh.comm.allreduce(local_surf, op=MPI.SUM) + assert np.isclose(surf, 6, atol=tol, rtol=tol) + + diff = dolfinx.fem.form(ufl.inner(v - f_ex, v - f_ex) * ufl.dx, dtype=dtype) + local_diff = dolfinx.fem.assemble_scalar(diff) + diff_glob = new_mesh.comm.allreduce(local_diff, op=MPI.SUM) + assert np.sqrt(np.isclose(diff_glob, 0.0, atol=tol, rtol=tol)) diff --git a/tests/test_submesh.py b/tests/test_submesh.py new file mode 100644 index 0000000..06aaf63 --- /dev/null +++ b/tests/test_submesh.py @@ -0,0 +1,460 @@ +from mpi4py import MPI + +import dolfinx +import numpy as np +import pytest +import ufl + +import io4dolfinx + +SUFFIX = {"adios2": ".bp", "h5py": ".h5"} + + +def _reference(x, vector: bool, gdim: int): + """A function that is exactly representable in every space tested below.""" + if vector: + return np.vstack([np.sin(1.3 * x[0] + 0.2), x[1], 0.3 + 0.0 * x[2]])[:gdim] + return np.sin(1.3 * x[0] + 0.2) + x[1] - 0.5 * x[2] + + +def _make_submesh(comm, codim, n=4): + """A cube, and a submesh of co-dimension ``codim`` of it.""" + mesh = dolfinx.mesh.create_unit_cube( + comm, n, n, n, ghost_mode=dolfinx.mesh.GhostMode.shared_facet + ) + dim = mesh.topology.dim - codim + mesh.topology.create_entities(dim) + if codim == 0: + + def locator(x): + return x[0] <= 0.5 + 1e-12 + else: + + def locator(x): + return np.isclose(x[0], 0.0) + + entities = dolfinx.mesh.locate_entities(mesh, dim, locator) + submesh, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, dim, entities) + return mesh, submesh, cell_map + + +def _max_error(u, reference): + """Largest difference over dofs owned by any process. + + A process may own no dofs at all -- a small submesh spread over many ranks -- + so the local maximum is only taken when there is something to take it over. + The reduction itself still has to run on every process. + """ + V = u.function_space + num_owned = V.dofmap.index_map.size_local * V.dofmap.index_map_bs + local = ( + float(np.max(np.abs(u.x.array[:num_owned] - reference.x.array[:num_owned]))) + if num_owned + else 0.0 + ) + return V.mesh.comm.allreduce(local, MPI.MAX) + + +@pytest.mark.parametrize("codim", [0, 1]) +@pytest.mark.parametrize( + "family,degree", + [("Lagrange", 1), ("Lagrange", 2), ("Lagrange", 3), ("Lagrange", 4), ("DG", 0), ("DG", 4)], +) +@pytest.mark.parametrize("same_file", [True, False]) +def test_submesh_roundtrip(tmp_path, backend, codim, family, degree, same_file): + """Write a submesh and a function on it, read both back, reattach to the parent.""" + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + parent_file = folder / f"parent{SUFFIX[backend]}" + sub_file = parent_file if same_file else folder / f"sub{SUFFIX[backend]}" + + mesh, submesh, cell_map = _make_submesh(comm, codim) + gdim = submesh.geometry.dim + + def f(x): + return _reference(x, vector=False, gdim=gdim) + + V = dolfinx.fem.functionspace(submesh, (family, degree)) + u = dolfinx.fem.Function(V, name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(parent_file, mesh, backend=backend) + io4dolfinx.write_submesh( + sub_file, + submesh, + mesh, + cell_map, + mesh_name="wall", + parent_filename=None if same_file else parent_file, + mode=io4dolfinx.FileMode.append if same_file else io4dolfinx.FileMode.write, + backend=backend, + ) + io4dolfinx.write_function( + sub_file, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map, u + + # Steps 1-3: the submesh is an ordinary mesh, so this is the ordinary path. + parent = io4dolfinx.read_mesh(parent_file, comm, backend=backend) + stored = io4dolfinx.read_mesh(sub_file, comm, mesh_name="wall", backend=backend) + V_stored = dolfinx.fem.functionspace(stored, (family, degree)) + u_stored = dolfinx.fem.Function(V_stored, name="u") + io4dolfinx.read_function( + sub_file, u_stored, time=0.0, name="u", mesh_name="wall", backend=backend + ) + reference = dolfinx.fem.Function(V_stored) + reference.interpolate(f) + assert _max_error(u_stored, reference) < 1e-13 + + # Steps 4-5: re-derive from the parent and move the data across. + checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, (family, degree)) + u_sub = dolfinx.fem.Function(V_sub) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + + reference_sub = dolfinx.fem.Function(V_sub) + reference_sub.interpolate(f) + assert _max_error(u_sub, reference_sub) < 1e-13 + + # An L2 error assembled on the re-derived submesh alone: one mesh, no maps. + error = ufl.inner(u_sub - reference_sub, u_sub - reference_sub) * ufl.dx + l2 = np.sqrt(comm.allreduce(dolfinx.fem.assemble_scalar(dolfinx.fem.form(error)), MPI.SUM)) + assert l2 < 1e-13 + + +@pytest.mark.parametrize("family", ["N1curl", "RT"]) +@pytest.mark.parametrize("degree", [1, 2]) +def test_submesh_roundtrip_codim0_vector(tmp_path, backend, family, degree): + """H(div)/H(curl) on a co-dimension 0 submesh, where dof transformations apply.""" + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + parent_file = folder / f"parent{SUFFIX[backend]}" + + mesh, submesh, cell_map = _make_submesh(comm, codim=0) + gdim = submesh.geometry.dim + + def f(x): + return _reference(x, vector=True, gdim=gdim) + + V = dolfinx.fem.functionspace(submesh, (family, degree)) + u = dolfinx.fem.Function(V, name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(parent_file, mesh, backend=backend) + io4dolfinx.write_submesh( + parent_file, submesh, mesh, cell_map, mesh_name="wall", backend=backend + ) + io4dolfinx.write_function( + parent_file, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(parent_file, comm, backend=backend) + stored = io4dolfinx.read_mesh(parent_file, comm, mesh_name="wall", backend=backend) + V_stored = dolfinx.fem.functionspace(stored, (family, degree)) + u_stored = dolfinx.fem.Function(V_stored, name="u") + io4dolfinx.read_function( + parent_file, u_stored, time=0.0, name="u", mesh_name="wall", backend=backend + ) + reference = dolfinx.fem.Function(V_stored) + reference.interpolate(f) + assert _max_error(u_stored, reference) < 1e-13 + + checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, (family, degree)) + u_sub = dolfinx.fem.Function(V_sub) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + reference_sub = dolfinx.fem.Function(V_sub) + reference_sub.interpolate(f) + assert _max_error(u_sub, reference_sub) < 1e-13 + + +@pytest.mark.parametrize("family", ["N1curl", "RT"]) +def test_manifold_hcurl_roundtrips_but_transfer_raises(tmp_path, backend, family): + """H(div)/H(curl) on a manifold submesh: storable, but not reattachable. + + Writing and reading the standalone submesh is exact for these spaces, and + that is checked here so the capability is not lost by accident. + + The transfer, though, is refused. DOLFINx cannot reconcile the reference and + physical value sizes of these families when ``tdim < gdim``, and its + interpolation does not report the problem -- it returns values that are + simply wrong (the L2 norm of an N1curl field measured 9.096 before and 4.904 + after). A loud refusal is the only safe behaviour. + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + parent_file = folder / f"parent{SUFFIX[backend]}" + + mesh, submesh, cell_map = _make_submesh(comm, codim=1) + V = dolfinx.fem.functionspace(submesh, (family, 1)) + u = dolfinx.fem.Function(V, name="u") + # These spaces cannot be interpolated into on a manifold, so fill by global + # dof index instead: reproducible no matter how the mesh is partitioned. + imap = V.dofmap.index_map + num_owned = imap.size_local * V.dofmap.index_map_bs + global_index = imap.local_to_global(np.arange(imap.size_local, dtype=np.int32)) + u.x.array[:num_owned] = np.sin(0.37 * global_index.astype(np.float64) + 0.11) + u.x.scatter_forward() + norm_written = comm.allreduce( + dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(u, u) * ufl.dx)), MPI.SUM + ) + + io4dolfinx.write_mesh(parent_file, mesh, backend=backend) + io4dolfinx.write_submesh( + parent_file, submesh, mesh, cell_map, mesh_name="wall", backend=backend + ) + io4dolfinx.write_function( + parent_file, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map + + parent = io4dolfinx.read_mesh(parent_file, comm, backend=backend) + stored = io4dolfinx.read_mesh(parent_file, comm, mesh_name="wall", backend=backend) + u_stored = dolfinx.fem.Function(dolfinx.fem.functionspace(stored, (family, 1)), name="u") + io4dolfinx.read_function( + parent_file, u_stored, time=0.0, name="u", mesh_name="wall", backend=backend + ) + + # Steps 1-3 are exact even here. + norm_read = comm.allreduce( + dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(u_stored, u_stored) * ufl.dx)), + MPI.SUM, + ) + assert np.isclose(norm_read, norm_written, rtol=0, atol=1e-12) + + checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) + u_sub = dolfinx.fem.Function(dolfinx.fem.functionspace(checkpoint.submesh, (family, 1))) + with pytest.raises(NotImplementedError, match="not supported"): + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + + +@pytest.mark.parametrize("codim", [0, 1]) +def test_read_submesh_supports_mixed_dimensional_assembly(tmp_path, backend, codim): + """The re-derived submesh carries entity maps usable against its parent. + + This is the whole reason to re-derive rather than stop at the standalone + submesh, and it is the only cross-mesh assembly that is well founded here: + the submesh really is derived from this parent. + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + parent_file = folder / f"parent{SUFFIX[backend]}" + + mesh, submesh, cell_map = _make_submesh(comm, codim) + io4dolfinx.write_mesh(parent_file, mesh, backend=backend) + io4dolfinx.write_submesh( + parent_file, submesh, mesh, cell_map, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map + + parent = io4dolfinx.read_mesh(parent_file, comm, backend=backend) + checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) + + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 1)) + u_sub = dolfinx.fem.Function(V_sub) + u_sub.x.array[:] = 1.0 + + dim = checkpoint.submesh.topology.dim + measure = ufl.dx if codim == 0 else ufl.ds + form = dolfinx.fem.form(u_sub * measure(domain=parent), entity_maps=[checkpoint.cell_map]) + area = comm.allreduce(dolfinx.fem.assemble_scalar(form), MPI.SUM) + expected = 0.5 if (codim == 0 and dim == 3) else 1.0 + assert np.isclose(area, expected) + + +def test_submesh_with_empty_ranks(tmp_path, backend): + """A submesh small enough that some process owns none of it.""" + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + parent_file = folder / f"parent{SUFFIX[backend]}" + + mesh = dolfinx.mesh.create_unit_cube( + comm, 6, 6, 6, ghost_mode=dolfinx.mesh.GhostMode.shared_facet + ) + dim = mesh.topology.dim - 1 + mesh.topology.create_entities(dim) + # A single sliver of the x=0 face, so on several ranks this is empty. + entities = dolfinx.mesh.locate_entities( + mesh, dim, lambda x: np.isclose(x[0], 0.0) & (x[1] <= 1.0 / 6.0 + 1e-12) + ) + submesh, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, dim, entities) + + def f(x): + return x[1] - 0.5 * x[2] + + V = dolfinx.fem.functionspace(submesh, ("Lagrange", 2)) + u = dolfinx.fem.Function(V, name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(parent_file, mesh, backend=backend) + io4dolfinx.write_submesh( + parent_file, submesh, mesh, cell_map, mesh_name="strip", backend=backend + ) + io4dolfinx.write_function( + parent_file, + u, + time=0.0, + mode=io4dolfinx.FileMode.append, + mesh_name="strip", + backend=backend, + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(parent_file, comm, backend=backend) + stored = io4dolfinx.read_mesh(parent_file, comm, mesh_name="strip", backend=backend) + V_stored = dolfinx.fem.functionspace(stored, ("Lagrange", 2)) + u_stored = dolfinx.fem.Function(V_stored, name="u") + io4dolfinx.read_function( + parent_file, u_stored, time=0.0, name="u", mesh_name="strip", backend=backend + ) + + checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="strip", backend=backend) + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) + u_sub = dolfinx.fem.Function(V_sub) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + + reference = dolfinx.fem.Function(V_sub) + reference.interpolate(f) + assert _max_error(u_sub, reference) < 1e-13 + + +def test_named_meshes_do_not_collide(tmp_path, backend): + """Two meshes in one file keep their own topology, geometry and functions.""" + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"two{SUFFIX[backend]}" + + square = dolfinx.mesh.create_unit_square(comm, 4, 4) + cube = dolfinx.mesh.create_unit_cube(comm, 3, 3, 3) + io4dolfinx.write_mesh(path, square, backend=backend) + io4dolfinx.write_mesh( + path, cube, mesh_name="cube", mode=io4dolfinx.FileMode.append, backend=backend + ) + + u_square = dolfinx.fem.Function(dolfinx.fem.functionspace(square, ("Lagrange", 1)), name="u") + u_square.interpolate(lambda x: x[0] + 2.0 * x[1]) + u_cube = dolfinx.fem.Function(dolfinx.fem.functionspace(cube, ("Lagrange", 1)), name="u") + u_cube.interpolate(lambda x: 3.0 * x[2]) + io4dolfinx.write_function(path, u_square, time=0.0, backend=backend) + io4dolfinx.write_function(path, u_cube, time=0.0, mesh_name="cube", backend=backend) + del square, cube, u_square, u_cube + + read_square = io4dolfinx.read_mesh(path, comm, backend=backend) + read_cube = io4dolfinx.read_mesh(path, comm, mesh_name="cube", backend=backend) + assert read_square.topology.dim == 2 + assert read_cube.topology.dim == 3 + + v_square = dolfinx.fem.Function( + dolfinx.fem.functionspace(read_square, ("Lagrange", 1)), name="u" + ) + io4dolfinx.read_function(path, v_square, time=0.0, name="u", backend=backend) + ref_square = dolfinx.fem.Function(v_square.function_space) + ref_square.interpolate(lambda x: x[0] + 2.0 * x[1]) + assert _max_error(v_square, ref_square) < 1e-13 + + v_cube = dolfinx.fem.Function(dolfinx.fem.functionspace(read_cube, ("Lagrange", 1)), name="u") + io4dolfinx.read_function(path, v_cube, time=0.0, name="u", mesh_name="cube", backend=backend) + ref_cube = dolfinx.fem.Function(v_cube.function_space) + ref_cube.interpolate(lambda x: 3.0 * x[2]) + assert _max_error(v_cube, ref_cube) < 1e-13 + + +@pytest.mark.parametrize("kind", ["point", "cell"]) +def test_submesh_point_and_cell_data(tmp_path, kind): + """Visualisation data attached to a named submesh in a multi-block file. + + ``read_point_data`` and ``read_cell_data`` rebuild their space from the + mesh's own coordinate element, so they work on a submesh once the mesh name + reaches the backend. + """ + pytest.importorskip("h5py") + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / "vis.vtkhdf" + + mesh = dolfinx.mesh.create_unit_cube( + comm, 4, 4, 4, ghost_mode=dolfinx.mesh.GhostMode.shared_facet + ) + fdim = mesh.topology.dim - 1 + mesh.topology.create_entities(fdim) + facets = dolfinx.mesh.locate_entities(mesh, fdim, lambda x: np.isclose(x[0], 0.0)) + submesh = dolfinx.mesh.create_submesh(mesh, fdim, facets)[0] + + def f(x): + return np.sin(2 * x[1]) + x[2] + + # Point data goes through the coordinate element, cell data through DG-0. + element = ("Lagrange", 1) if kind == "point" else ("DG", 0) + u = dolfinx.fem.Function(dolfinx.fem.functionspace(submesh, element), name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(path, mesh, backend="vtkhdf") + io4dolfinx.write_mesh( + path, submesh, mesh_name="wall", mode=io4dolfinx.FileMode.append, backend="vtkhdf" + ) + writer = io4dolfinx.write_point_data if kind == "point" else io4dolfinx.write_cell_data + writer( + path, + u, + time=0.0, + mode=io4dolfinx.FileMode.append, + backend_args=None, + backend="vtkhdf", + mesh_name="wall", + ) + del mesh, submesh, u + + stored = io4dolfinx.read_mesh(path, comm, mesh_name="wall", backend="vtkhdf", time=0.0) + reader = io4dolfinx.read_point_data if kind == "point" else io4dolfinx.read_cell_data + read_back = reader(path, "u", stored, time=0.0, backend="vtkhdf", mesh_name="wall") + + reference = dolfinx.fem.Function(read_back.function_space) + reference.interpolate(f) + assert _max_error(read_back, reference) < 1e-13 + + +def test_submesh_with_named_parent(tmp_path, backend): + """Neither mesh has to use the default name.""" + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"named{SUFFIX[backend]}" + + mesh, submesh, cell_map = _make_submesh(comm, codim=1) + + def f(x): + return np.sin(2 * x[1]) + x[2] + + u = dolfinx.fem.Function(dolfinx.fem.functionspace(submesh, ("Lagrange", 2)), name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(path, mesh, mesh_name="bulk", backend=backend) + io4dolfinx.write_submesh( + path, + submesh, + mesh, + cell_map, + mesh_name="wall", + parent_mesh_name="bulk", + backend=backend, + ) + io4dolfinx.write_function( + path, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(path, comm, mesh_name="bulk", backend=backend) + stored = io4dolfinx.read_mesh(path, comm, mesh_name="wall", backend=backend) + u_stored = dolfinx.fem.Function(dolfinx.fem.functionspace(stored, ("Lagrange", 2)), name="u") + io4dolfinx.read_function(path, u_stored, time=0.0, name="u", mesh_name="wall", backend=backend) + + checkpoint = io4dolfinx.read_submesh( + path, parent, mesh_name="wall", parent_mesh_name="bulk", backend=backend + ) + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) + u_sub = dolfinx.fem.Function(V_sub) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + + reference = dolfinx.fem.Function(V_sub) + reference.interpolate(f) + assert _max_error(u_sub, reference) < 1e-13 From 55104026d28f08795fc779466762433567524d75 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Mon, 14 Sep 2026 05:03:34 +0000 Subject: [PATCH 03/11] Add higher order tests for read-write point ata. --- src/io4dolfinx/backends/vtkhdf/backend.py | 53 +++++++- tests/test_submesh.py | 148 ++++++++++++++++++++++ 2 files changed, 198 insertions(+), 3 deletions(-) diff --git a/src/io4dolfinx/backends/vtkhdf/backend.py b/src/io4dolfinx/backends/vtkhdf/backend.py index 6a73080..c66d377 100644 --- a/src/io4dolfinx/backends/vtkhdf/backend.py +++ b/src/io4dolfinx/backends/vtkhdf/backend.py @@ -23,6 +23,15 @@ read_mode = ReadMode.parallel _vtk_hdf_version = np.array([2, 1], dtype=np.int32) +#: Top-level group for metadata io4dolfinx needs but VTKHDF has no slot for. +#: The specification reserves everything outside ``/VTKHDF`` for exactly this: +#: "Top-level groups outside of /VTKHDF do not contain any information related to +#: the VTK data model and are outside of the scope of this specification. They can +#: be useful to store meta-information that could be read and written by custom +#: VTKHDF implementation." +_metadata_group = "/io4dolfinx" +#: Basix Lagrange variant of a mesh's coordinate element, per mesh name. +_lagrange_variant_attr = "LagrangeVariant" def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]: @@ -127,6 +136,26 @@ def _get_time_index(hdf: h5py.Group, time: float | str, filename: str | Path) -> return pos[0] +def _read_lagrange_variant(h5file, name: str) -> int: + """Lagrange variant of mesh ``name``, or the VTKHDF default. + + A file written by anything other than io4dolfinx -- or by io4dolfinx before + this metadata existed -- holds VTKHDF's equispaced higher-order cells. + + Args: + h5file: The open HDF5 file + name: Name of the mesh + + Returns: + The basix Lagrange variant as an integer. + """ + equispaced = int(basix.LagrangeVariant.equispaced) + metadata = h5file.get(_metadata_group) + if metadata is None or name not in metadata: + return equispaced + return int(metadata[name].attrs.get(_lagrange_variant_attr, equispaced)) + + def read_mesh_data( filename: Path | str, comm: MPI.Comm, @@ -152,6 +181,7 @@ def read_mesh_data( raise RuntimeError("Cannot read partition data with VTKHDF") with h5pyfile(filename, "r", comm=comm) as h5file: hdf = _get_vtk_group(h5file, backend_args["name"]) + lagrange_variant = _read_lagrange_variant(h5file, backend_args["name"]) if time is None: num_cells_global = hdf["Types"].size local_cell_range = compute_local_range(comm, num_cells_global) @@ -214,11 +244,17 @@ def read_mesh_data( # NOTE: Currently we limit ourselfs to a single celltype, as it makes life easier, # other things have to change in `MeshReadData` to support this. num_nodes_per_cell = offset[1:] - offset[:-1] - unique_cells = find_all_unique_cell_types(MPI.COMM_WORLD, cell_types_local, num_nodes_per_cell) + unique_cells = find_all_unique_cell_types(comm, cell_types_local, num_nodes_per_cell) if unique_cells.shape[0] > 1: raise NotImplementedError("io4dolfinx does not support mixed celltype grids") - topology = topology.reshape(-1, num_nodes_per_cell[0]) + if unique_cells.shape[0] == 0: + name = backend_args["name"] + raise ValueError(f"Grid '{name}' in {filename} has no cells on any process.") + # Take the width from the globally agreed cell type rather than from this + # process's own cells: a mesh with fewer cells than processes leaves some of + # them with nothing local to measure. cell_type, number_of_nodes = unique_cells[0] + topology = topology.reshape(-1, number_of_nodes) gtype = backend_args.get("dtype", points_local.dtype) if cell_type in _first_order_vtk.keys(): ct = _first_order_vtk[cell_type] @@ -230,7 +266,9 @@ def read_mesh_data( raise ValueError(f"Unknown VTK cell type {cell_type} in {filename}") perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(ct), number_of_nodes) topology = topology[:, perm] - lvar = int(basix.LagrangeVariant.equispaced) + # Equispaced is the VTKHDF default; a file io4dolfinx wrote says which + # variant its nodes were actually placed with. + lvar = int(lagrange_variant) return ReadMeshData( cells=topology, cell_type=ct, x=points_local.astype(gtype), lvar=lvar, degree=degree ) @@ -536,6 +574,15 @@ def write_mesh( if name not in mesh_assembly.keys(): mesh_assembly[name] = h5py.SoftLink(f"/VTKHDF/{name}") + # VTKHDF defines its higher-order cells as equispaced and has nowhere to + # record a Lagrange variant, so it goes in our own top-level group, which + # the specification reserves for exactly this. Without it a degree>1 mesh + # reads back with its coordinates intact but its cells curving + # differently between them. + metadata = _create_group(h5file, _metadata_group, h5_mode) + mesh_metadata = _create_group(metadata, name, h5_mode) + mesh_metadata.attrs[_lagrange_variant_attr] = np.int32(mesh.lagrange_variant) + # Write time dependent points number_of_points = _create_dataset( mesh_group, diff --git a/tests/test_submesh.py b/tests/test_submesh.py index 06aaf63..b98cfa5 100644 --- a/tests/test_submesh.py +++ b/tests/test_submesh.py @@ -1,5 +1,6 @@ from mpi4py import MPI +import basix import dolfinx import numpy as np import pytest @@ -458,3 +459,150 @@ def f(x): reference = dolfinx.fem.Function(V_sub) reference.interpolate(f) assert _max_error(u_sub, reference) < 1e-13 + + +def _curved_mesh(comm, dim, degree): + """A disk (dim 2) or ball (dim 3) with degree-``degree`` geometry. + + Built by lifting a straight-sided box to a higher-order coordinate element + and then pushing the nodes onto the circle/sphere, so cell edges are genuine + curves rather than chords -- which is the point: it puts the interior + geometry nodes somewhere a first-order mesh could not represent. + """ + # Degree > 2 has no default node placement, so name the variant. + variant = basix.LagrangeVariant.gll_isaac if degree > 2 else basix.LagrangeVariant.unset + if dim == 3: + # At least three divisions per axis, so that a cut offset from the + # symmetry plane still has whole cells on one side of it. + base = dolfinx.mesh.create_box( + comm, + [np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0])], + [3, 3, 3], + ghost_mode=dolfinx.mesh.GhostMode.shared_facet, + ) + else: + base = dolfinx.mesh.create_rectangle( + comm, + [np.array([-1.0, -1.0]), np.array([1.0, 1.0])], + [4, 4], + ghost_mode=dolfinx.mesh.GhostMode.shared_facet, + ) + cmap = dolfinx.fem.coordinate_element(base.topology.cell_type, degree, int(variant)) + mesh = dolfinx.fem.interpolate_geometry(base, cmap) + + # Elliptical grid mapping: sends the box onto the ball, smoothly. + x = mesh.geometry.x + if dim == 3: + a, b, c = x[:, 0].copy(), x[:, 1].copy(), x[:, 2].copy() + x[:, 0] = a * np.sqrt(1 - b**2 / 2 - c**2 / 2 + b**2 * c**2 / 3) + x[:, 1] = b * np.sqrt(1 - c**2 / 2 - a**2 / 2 + c**2 * a**2 / 3) + x[:, 2] = c * np.sqrt(1 - a**2 / 2 - b**2 / 2 + a**2 * b**2 / 3) + else: + a, b = x[:, 0].copy(), x[:, 1].copy() + x[:, 0] = a * np.sqrt(1 - b**2 / 2) + x[:, 1] = b * np.sqrt(1 - a**2 / 2) + return mesh + + +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("degree", [2, 4]) +@pytest.mark.parametrize("codim", [0, 1]) +def test_point_data_on_curved_submesh(tmp_path, dim, degree, codim): + """Point data on a submesh of a curved, higher-order mesh. + + Point data lives on the geometry nodes, so a degree-4 mesh puts most of it on + interior and edge nodes rather than vertices. The codim-1 case is the curved + boundary itself -- a manifold submesh whose cells are genuinely curved. + """ + pytest.importorskip("h5py") + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / "curved.vtkhdf" + + mesh = _curved_mesh(comm, dim, degree) + tdim = mesh.topology.dim + sub_dim = tdim - codim + mesh.topology.create_entities(sub_dim) + if codim == 0: + # Offset from the symmetry plane on purpose. A cut that grazes node + # positions makes `locate_entities` disagree between a cell's owner and + # the ranks ghosting it on a higher-order mesh, and `create_submesh` + # rejects the result with "Index owner change detected". Measured at + # np=3 on a degree-2 cube: 3 mismatches at `x[0] <= 0.0`, none at -0.1, + # with or without the curving perturbation. + entities = dolfinx.mesh.locate_entities(mesh, sub_dim, lambda x: x[0] <= -0.1) + else: + mesh.topology.create_connectivity(sub_dim, tdim) + entities = dolfinx.mesh.exterior_facet_indices(mesh.topology) + submesh = dolfinx.mesh.create_submesh(mesh, sub_dim, entities)[0] + assert submesh.geometry.cmaps[0].degree == degree + + def f(x): + return np.sin(1.7 * x[0]) + 0.5 * x[1] - 0.3 * x[tdim - 1] + + V = io4dolfinx.readers.create_geometry_function_space(submesh, 1) + u = dolfinx.fem.Function(V, name="u") + u.interpolate(f) + nodes_written = submesh.geometry.index_map().size_global + variant_written = submesh.geometry.cmaps[0].variant + + io4dolfinx.write_mesh(path, submesh, mesh_name="curved", backend="vtkhdf") + io4dolfinx.write_point_data( + path, + u, + time=0.0, + mode=io4dolfinx.FileMode.append, + backend_args=None, + backend="vtkhdf", + mesh_name="curved", + ) + del mesh, submesh, u + + stored = io4dolfinx.read_mesh(path, comm, mesh_name="curved", backend="vtkhdf", time=0.0) + # The higher-order geometry must survive whole: degree, node placement and + # node count, not just the coordinates. + assert stored.geometry.cmaps[0].degree == degree + assert stored.geometry.cmaps[0].variant == variant_written + assert stored.geometry.index_map().size_global == nodes_written + + read_back = io4dolfinx.read_point_data( + path, "u", stored, time=0.0, backend="vtkhdf", mesh_name="curved" + ) + reference = dolfinx.fem.Function(read_back.function_space) + reference.interpolate(f) + assert _max_error(read_back, reference) < 1e-13 + + +@pytest.mark.parametrize("degree", [2, 4]) +@pytest.mark.parametrize("store", ["adios2", "h5py", "vtkhdf"]) +def test_lagrange_variant_survives_mesh_roundtrip(tmp_path, store, degree): + """A higher-order mesh keeps its node placement, not just its node count. + + Two coordinate elements of the same degree but different Lagrange variants + put their nodes at different reference positions, so a checkpoint that drops + the variant restores a mesh whose cells curve differently between the same + node coordinates -- and does so silently, since every coordinate still + round-trips exactly. + + ``vtkhdf`` is included deliberately. VTKHDF defines its higher-order cells as + equispaced and has no slot for a variant, so io4dolfinx records one in an + extra attribute; this is what pins that down. + """ + pytest.importorskip("h5py") + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + suffix = {"adios2": ".bp", "h5py": ".h5", "vtkhdf": ".vtkhdf"}[store] + path = folder / f"variant{suffix}" + + base = dolfinx.mesh.create_rectangle( + comm, [np.array([-1.0, -1.0]), np.array([1.0, 1.0])], [4, 4] + ) + variant = basix.LagrangeVariant.gll_isaac if degree > 2 else basix.LagrangeVariant.unset + cmap = dolfinx.fem.coordinate_element(base.topology.cell_type, degree, int(variant)) + mesh = dolfinx.fem.interpolate_geometry(base, cmap) + written = (mesh.geometry.cmaps[0].degree, mesh.geometry.cmaps[0].variant) + + io4dolfinx.write_mesh(path, mesh, backend=store) + read_kwargs = {"time": 0.0} if store == "vtkhdf" else {} + read_back = io4dolfinx.read_mesh(path, comm, backend=store, **read_kwargs) + assert (read_back.geometry.cmaps[0].degree, read_back.geometry.cmaps[0].variant) == written From 40f2c562d1487af1f01b2598de0b5af087ba67cc Mon Sep 17 00:00:00 2001 From: jorgensd Date: Tue, 22 Sep 2026 20:04:06 +0000 Subject: [PATCH 04/11] Minor type fixes --- src/io4dolfinx/utils.py | 3 ++- tests/test_reconstruct_mesh.py | 12 ++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/io4dolfinx/utils.py b/src/io4dolfinx/utils.py index 17184ef..ed5fdcb 100644 --- a/src/io4dolfinx/utils.py +++ b/src/io4dolfinx/utils.py @@ -299,10 +299,11 @@ def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(new_c_el)) else: # Use the new interpolate_geometry function + assert len(mesh.geometry.cmaps) == 1, "Only single geometry maps are supported" cmap = dolfinx.fem.coordinate_element( mesh.topology.cell_type, coordinate_element_degree, dtype=mesh.geometry.x.dtype, - variant=mesh.geometry.cmap.variant, + variant=mesh.geometry.cmaps[0].variant, ) return dolfinx.fem.interpolate_geometry(mesh, cmap) diff --git a/tests/test_reconstruct_mesh.py b/tests/test_reconstruct_mesh.py index f1d64f2..c2a0364 100644 --- a/tests/test_reconstruct_mesh.py +++ b/tests/test_reconstruct_mesh.py @@ -1,3 +1,5 @@ +from typing import cast + from mpi4py import MPI import basix.ufl @@ -37,19 +39,21 @@ def f(x): cells = np.arange(num_cells_local, dtype=np.int32) v.interpolate(u, cells, cells) - assert new_mesh.geometry.cmap.degree == degree + assert new_mesh.geometry.cmaps[0].degree == degree f_ex = ufl.as_vector(f(ufl.SpatialCoordinate(new_mesh))) + one = np.array(1.0, dtype=dtype) local_vol = dolfinx.fem.assemble_scalar( - dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, dtype(1.0)) * ufl.dx, dtype=dtype) + dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, one) * ufl.dx, dtype=dtype) ) vol = new_mesh.comm.allreduce(local_vol, op=MPI.SUM) - tol = 4e3 * degree * np.finfo(dtype).eps + float_dtype = cast(type[np.floating], dtype) + tol = 4e3 * degree * np.finfo(float_dtype).eps assert np.isclose(vol, 1, atol=tol, rtol=tol) local_surf = dolfinx.fem.assemble_scalar( - dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, dtype(1.0)) * ufl.ds, dtype=dtype) + dolfinx.fem.form(dolfinx.fem.Constant(new_mesh, one) * ufl.ds, dtype=dtype) ) surf = new_mesh.comm.allreduce(local_surf, op=MPI.SUM) assert np.isclose(surf, 6, atol=tol, rtol=tol) From 309ce3c9370801315b2b054d7fabe2d1832f8ec1 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Tue, 22 Sep 2026 20:33:04 +0000 Subject: [PATCH 05/11] Start rewriting docs --- docs/submesh.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/docs/submesh.py b/docs/submesh.py index aa71f98..017f185 100644 --- a/docs/submesh.py +++ b/docs/submesh.py @@ -10,9 +10,10 @@ # renumbers the vertices, so the same cell comes back with its vertices in a # different order and its degrees of freedom in different positions. # -# Second, `create_submesh` builds its vertex map allowing ownership to move, so -# a process can own a vertex -- and a degree of freedom -- that is only incident -# to cells it ghosts. Such a dof has no position in any cell that process owns, +# Second, {py:func}`create_submesh` builds its vertex +# map allowing ownership to move, so a process can own a vertex and a degree of +# freedom that is only incident to the cells it ghosts. +# Such a dof has no position in any cell that process owns, # which is exactly what the checkpoint reader needs to find it. # # So io4dolfinx never reads a function onto a re-derived submesh. It stores the @@ -96,12 +97,18 @@ io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) # - +# The {py:class}`checkpoint` stores the +# {py:attr}`cell_map` and +# {py:attr}`vertex_map` and +# {py:attr}`node_map` that are +# similar to the ones created by {py:func}`dolfinx.mesh.create_submesh`. +# There are also some extra attributes that are stored that are used +# for the transfer of functions from the stored submesh to the re-derived submesh. + # The transfer routes each interpolation point straight to the process and cell # that can evaluate it, using the cell correspondence the checkpoint already # records. No geometric point location is involved, so there is no search # tolerance to tune and nothing that can silently fail to find a point. -# -# The result agrees with the original to machine precision. # + reference = dolfinx.fem.Function(V_sub) From 0d7aa0d38e3607689908666a423b8cb231b307b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:33:28 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/submesh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/submesh.py b/docs/submesh.py index 017f185..15e3b58 100644 --- a/docs/submesh.py +++ b/docs/submesh.py @@ -97,7 +97,7 @@ io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) # - -# The {py:class}`checkpoint` stores the +# The {py:class}`checkpoint` stores the # {py:attr}`cell_map` and # {py:attr}`vertex_map` and # {py:attr}`node_map` that are From 575cb414e4ae865e0c450f64b79330a4bde100fa Mon Sep 17 00:00:00 2001 From: jorgensd Date: Tue, 22 Sep 2026 21:19:47 +0000 Subject: [PATCH 07/11] Rewrites and avoiding hanging in parallel. --- docs/submesh.py | 69 ++++--- src/io4dolfinx/backends/__init__.py | 2 +- src/io4dolfinx/backends/adios2/backend.py | 2 +- src/io4dolfinx/backends/exodus/backend.py | 2 +- src/io4dolfinx/backends/h5py/backend.py | 2 +- src/io4dolfinx/original_checkpoint.py | 4 +- src/io4dolfinx/readers.py | 6 +- src/io4dolfinx/submesh.py | 212 +++++++++++++--------- src/io4dolfinx/writers.py | 2 +- tests/test_submesh.py | 200 +++++++++++++++++++- 10 files changed, 375 insertions(+), 126 deletions(-) diff --git a/docs/submesh.py b/docs/submesh.py index 15e3b58..2a2ac47 100644 --- a/docs/submesh.py +++ b/docs/submesh.py @@ -55,11 +55,13 @@ # ## Writing # # The parent goes in first. {py:func}`io4dolfinx.write_submesh` then stores the -# submesh under a name of its own, together with a *parent link*: for each stored -# submesh cell, the parent geometry nodes of the entity it came from. That link -# is what lets the submesh be found again in a re-partitioned parent. +# submesh under a name of its own, together with its *post codes*: each parent +# entity that became a submesh cell is tagged, by its parent geometry nodes, with +# the index that cell has in the stored submesh. Addressing the entities that way +# is what lets them be found again in a re-partitioned parent, and the values are +# what the data is later routed by. # -# The link tags entities of the parent, so it is stored alongside the parent. +# The post codes tag entities of the parent, so they are stored alongside the parent. # Here everything goes in one file; pass `parent_filename` to keep the submesh in # a file of its own. @@ -94,7 +96,7 @@ checkpoint = io4dolfinx.read_submesh(filename, parent, mesh_name="wall") V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) u_sub = dolfinx.fem.Function(V_sub, name="u") -io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) +io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) # - # The {py:class}`checkpoint` stores the @@ -102,13 +104,36 @@ # {py:attr}`vertex_map` and # {py:attr}`node_map` that are # similar to the ones created by {py:func}`dolfinx.mesh.create_submesh`. -# There are also some extra attributes that are stored that are used -# for the transfer of functions from the stored submesh to the re-derived submesh. - -# The transfer routes each interpolation point straight to the process and cell -# that can evaluate it, using the cell correspondence the checkpoint already -# records. No geometric point location is involved, so there is no search -# tolerance to tune and nothing that can silently fail to find a point. +# It also carries the +# {py:attr}`post_code` that the transfer +# above was given, which is worth a word of its own. +# +# ## The post office +# +# The stored submesh and the re-derived one are partitioned independently, so +# neither knows where the other put a given cell. What they agree on is a cell's +# *post code*: the index it had in the stored submesh. The stored submesh knows +# it as its ordinary `topology.original_cell_index`, being a mesh like any other. +# The re-derived one cannot -- {py:func}`dolfinx.mesh.create_submesh` does not +# give a submesh an original cell index, and for a submesh of co-dimension +# greater than zero there is none to give, since the parent's input data numbers +# cells and not facets. So the post code is what +# {py:func}`io4dolfinx.write_submesh` tags onto the parent and +# {py:func}`io4dolfinx.read_submesh` hands back. +# +# Resolving one works like posting a letter. No process holds a global table of +# who owns what. Instead the post code itself decides which process acts as its +# *post office*, by a rule every process applies identically +# ({py:func}`io4dolfinx.utils.index_owner`, an equal split of the global cell range). +# Each process publishes to the post office of every cell it owns, saying where it +# keeps that cell; each process then asks the post office of every cell it wants. +# Two neighbourhood exchanges, and the memory per process stays proportional to +# the cells it actually touches. +# +# The reply; a rank and a local cell, is what makes the transfer a routing +# rather than a search. Each interpolation point goes straight to the process and +# cell that can evaluate it, so no geometric point location is involved, there is +# no search tolerance to tune, and nothing can silently fail to find a point. # + reference = dolfinx.fem.Function(V_sub) @@ -118,10 +143,6 @@ print(f"Max error after transfer: {parent.comm.allreduce(error, MPI.MAX):.3e}") # - -# `checkpoint.cell_map` and `checkpoint.vertex_map` are genuine -# {py:class}`dolfinx.mesh.EntityMap` objects relating the submesh to `parent`, so -# they can be passed to {py:func}`dolfinx.fem.form`. - # + import ufl # noqa: E402 @@ -139,12 +160,16 @@ # a re-derived submesh. `transfer_submesh_function` raises `NotImplementedError` # for them. # -# The reason is worth knowing. DOLFINx cannot reconcile the reference and -# physical value sizes of these families when `tdim < gdim`, and in this -# particular path it does not say so: `interpolate_nonmatching` returns -# successfully with values that are wrong. Measured on a facet submesh of a unit -# cube, the L2 norm of an N1curl field fell from 9.096 to 4.904. The guard exists -# so that this shows up as an error rather than as a plausible-looking result. +# The reason is worth knowing. A finite element space on a manifold does not +# carry its value shape in *physical* space, so DOLFINx cannot reconcile the +# reference and physical value sizes of these families when `tdim < gdim` +# ([FEniCS/dolfinx#3619](https://github.com/FEniCS/dolfinx/issues/3619), open at +# the time of writing). Elsewhere that surfaces as an outright error: +# *"Interpolation: elements have different value dimensions"*, but in this +# particular path it does not: `interpolate_nonmatching` returns successfully +# with values that are wrong. The guard exists so that this +# shows up as an error rather than as a plausible-looking result, and it can be +# lifted once that issue is fixed. # # There is no sound workaround. The stored submesh is an independent mesh, # derived from nothing, so a projection -- or any other form assembled between it diff --git a/src/io4dolfinx/backends/__init__.py b/src/io4dolfinx/backends/__init__.py index 68c5f2c..2f03a63 100644 --- a/src/io4dolfinx/backends/__init__.py +++ b/src/io4dolfinx/backends/__init__.py @@ -264,7 +264,7 @@ def read_dofmap( backend_args: Arguments to backend Returns: - Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` + Dofmap as an :class:`dolfinx.graph.AdjacencyList` """ def read_dofs( diff --git a/src/io4dolfinx/backends/adios2/backend.py b/src/io4dolfinx/backends/adios2/backend.py index a5fc077..7e791d5 100644 --- a/src/io4dolfinx/backends/adios2/backend.py +++ b/src/io4dolfinx/backends/adios2/backend.py @@ -676,7 +676,7 @@ def read_meshtags_data( adios_file.file.EndStep() # Tag values are int32 unless the caller asks for something wider. The - # submesh parent link stores global cell indices, which need int64. + # submesh post codes are global cell indices, which need int64. values_dtype = np.dtype(backend_args.get("values_dtype", np.int32)) return MeshTagsData( name=name, values=tag_values.astype(values_dtype), indices=mesh_entities, dim=dim diff --git a/src/io4dolfinx/backends/exodus/backend.py b/src/io4dolfinx/backends/exodus/backend.py index 06a7e5e..24a5188 100644 --- a/src/io4dolfinx/backends/exodus/backend.py +++ b/src/io4dolfinx/backends/exodus/backend.py @@ -519,7 +519,7 @@ def read_dofmap( backend_args: Arguments to backend Returns: - Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` + Dofmap as an :class:`dolfinx.graph.AdjacencyList` """ raise NotImplementedError("The Exodus backend cannot read dofmap.") diff --git a/src/io4dolfinx/backends/h5py/backend.py b/src/io4dolfinx/backends/h5py/backend.py index 5b3e8dd..f9cc83b 100644 --- a/src/io4dolfinx/backends/h5py/backend.py +++ b/src/io4dolfinx/backends/h5py/backend.py @@ -432,7 +432,7 @@ def read_dofmap( backend_args: Arguments to backend Returns: - Dofmap as an {py:class}`dolfinx.graph.AdjacencyList` + Dofmap as an :class:`dolfinx.graph.AdjacencyList` """ backend_args = {} if backend_args is None else backend_args with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: diff --git a/src/io4dolfinx/original_checkpoint.py b/src/io4dolfinx/original_checkpoint.py index 9c2be5f..aa565f2 100644 --- a/src/io4dolfinx/original_checkpoint.py +++ b/src/io4dolfinx/original_checkpoint.py @@ -355,7 +355,7 @@ def write_function_on_input_mesh( Write function checkpoint (to be read with the input mesh). Note: - Requires backend to implement {py:class}`io4dolfinx.backends.write_function`. + Requires backend to implement :meth:`io4dolfinx.backends.IOBackend.write_function`. Args: filename: The filename to write to @@ -398,7 +398,7 @@ def write_mesh_input_order( Write mesh to checkpoint file in original input ordering. Note: - Requires backend to implement {py:class}`io4dolfinx.backends.write_mesh`. + Requires backend to implement :meth:`io4dolfinx.backends.IOBackend.write_mesh`. Args: filename: The filename to write to diff --git a/src/io4dolfinx/readers.py b/src/io4dolfinx/readers.py index 834e0fe..7b6fa8a 100644 --- a/src/io4dolfinx/readers.py +++ b/src/io4dolfinx/readers.py @@ -419,7 +419,7 @@ def read_point_data( """Read data from the nodes of a mesh. Note: - Backend has to implement {py:class}`io4dolfinx.backends.read_cell_data`. + Backend has to implement :meth:`io4dolfinx.backends.IOBackend.read_point_data`. Args: filename: Path to file @@ -486,10 +486,10 @@ def read_cell_data( backend: str = "xdmf", mesh_name: str | None = None, ) -> dolfinx.fem.Function: - """Read data from the nodes of a mesh. + """Read data from the cells of a mesh. Note: - Backend has to implement {py:class}`io4dolfinx.backends.read_cell_data`. + Backend has to implement :meth:`io4dolfinx.backends.IOBackend.read_cell_data`. Args: filename: Path to file diff --git a/src/io4dolfinx/submesh.py b/src/io4dolfinx/submesh.py index b36d1db..a8b3a15 100644 --- a/src/io4dolfinx/submesh.py +++ b/src/io4dolfinx/submesh.py @@ -57,13 +57,22 @@ logger = logging.getLogger(__name__) -def _link_name(mesh_name: str) -> str: - """Name of the meshtag holding the parent link of the submesh ``mesh_name``. +def _post_code_name(mesh_name: str) -> str: + """Name of the meshtag holding the post codes of the submesh ``mesh_name``. - The link is a tag on the *parent* mesh, so it is stored in the parent's + The post codes tag the *parent* mesh, so they are stored in the parent's namespace and must not collide with a tag a user wrote. """ - return f"__submesh_link_{mesh_name}" + return f"__submesh_post_code_{mesh_name}" + + +def _parent_entities( + submesh: dolfinx.mesh.Mesh, cell_map: dolfinx.mesh.EntityMap +) -> npt.NDArray[np.int32]: + """Parent entity of every cell of ``submesh``, owned and ghost.""" + imap = submesh.topology.index_map(submesh.topology.dim) + cells = np.arange(imap.size_local + imap.num_ghosts, dtype=np.int32) + return cell_map.sub_topology_to_topology(cells, False) @dataclass @@ -73,15 +82,22 @@ class SubmeshCheckpoint: #: The submesh, derived from the parent with :func:`dolfinx.mesh.create_submesh` submesh: dolfinx.mesh.Mesh #: Map from cells of :attr:`submesh` to entities of the parent - cell_map: Any + cell_map: dolfinx.mesh.EntityMap #: Map from vertices of :attr:`submesh` to vertices of the parent - vertex_map: Any + vertex_map: dolfinx.mesh.EntityMap #: Map from geometry nodes of :attr:`submesh` to nodes of the parent node_map: npt.NDArray[np.int32] - #: Parent entities that make up the submesh (local indices, owned and ghost) - parent_entities: npt.NDArray[np.int32] - #: Index the cell had in the stored submesh, per local cell (owned and ghost) - stored_cells: npt.NDArray[np.int64] + + #: Index the cell had in the stored submesh, per local cell (owned and ghost). + #: This is the key both submeshes agree on: :func:`transfer_submesh_function` + #: derives from it the post office holding each cell, and so routes the data + #: from the stored submesh to this one. + post_code: npt.NDArray[np.int64] + + @property + def parent_entities(self) -> npt.NDArray[np.int32]: + """Parent entities of the submesh cells, aligned with :attr:`post_code`.""" + return _parent_entities(self.submesh, self.cell_map) def write_submesh( @@ -97,16 +113,18 @@ def write_submesh( backend_args: dict[str, Any] | None = None, backend: str | None = None, ): - """Write a submesh, and the link that ties it back to its parent. + """Write a submesh, and the post codes that tie it back to its parent. The submesh is stored as an ordinary independent mesh under ``mesh_name``, so :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function` read it - back with no special handling. Alongside it goes a *parent link*: for every - stored submesh cell, the parent geometry nodes of the entity it came from. - That is what lets :func:`read_submesh` find the same entities in a - re-partitioned parent. - - The link tags entities of the *parent*, so it is stored in the parent's + back with no special handling. Alongside it goes a tag of *post codes*: every + parent entity that became a submesh cell is given the index that cell has in + the stored submesh. The entities are addressed by their parent geometry + nodes, which is what lets :func:`read_submesh` find them again in a + re-partitioned parent; the values come back as + :attr:`SubmeshCheckpoint.post_code`. + + The post codes tag entities of the *parent*, so they are stored in the parent's namespace, in ``parent_filename``. Leave that unset to keep everything in one checkpoint, or point it at the parent's own file to keep the two apart; the submesh always goes to ``filename``. Either way the parent must already have @@ -140,13 +158,13 @@ def write_submesh( mesh_name=mesh_name, ) - # 2. The parent link. Structurally this is a meshtag on the parent: the + # 2. The post codes. Structurally this is a meshtag on the parent: the # tagged entities are the parent entities that became submesh cells, and - # each value is the index that cell has in the stored submesh. + # each value is that cell's post code, its index in the stored submesh. cell_imap = submesh.topology.index_map(dim) num_owned = cell_imap.size_local parent_entities = cell_map.sub_topology_to_topology(np.arange(num_owned, dtype=np.int32), False) - values = np.arange(*cell_imap.local_range, dtype=np.int64) + post_code = np.arange(*cell_imap.local_range, dtype=np.int64) parent.topology.create_connectivity(dim, parent.topology.dim) parent.topology.create_connectivity(0, parent.topology.dim) @@ -160,11 +178,11 @@ def write_submesh( ) tag_data = MeshTagsData( - values=values, + values=post_code, num_entities_global=cell_imap.size_global, num_dofs_per_entity=entities_to_geometry.shape[1], indices=indices, - name=_link_name(mesh_name), + name=_post_code_name(mesh_name), local_start=cell_imap.local_range[0], dim=dim, cell_type=submesh.topology.cell_name(), @@ -173,8 +191,8 @@ def write_submesh( parent_args = backend_cls.get_default_backend_args( _with_mesh_name(backend_args, parent_mesh_name) ) - link_file = filename if parent_filename is None else parent_filename - backend_cls.write_meshtags(link_file, parent.comm, tag_data, backend_args=parent_args) + post_code_file = filename if parent_filename is None else parent_filename + backend_cls.write_meshtags(post_code_file, parent.comm, tag_data, backend_args=parent_args) def read_submesh( @@ -196,7 +214,7 @@ def read_submesh( then move it across with :func:`transfer_submesh_function`. Args: - filename: File holding the parent mesh and the parent link, i.e. whatever + filename: File holding the parent mesh and the post codes, i.e. whatever was passed as ``parent_filename`` to :func:`write_submesh` parent: The parent mesh, as read back from the checkpoint mesh_name: Name the submesh was stored under @@ -213,9 +231,11 @@ def read_submesh( parent_args = backend_cls.get_default_backend_args( _with_mesh_name(backend_args, parent_mesh_name) ) - # Stored cell indices are global, so they must survive as int64. + # Post codes are global cell indices, so they must survive as int64. parent_args["values_dtype"] = np.int64 - data = backend_cls.read_meshtags_data(filename, parent.comm, _link_name(mesh_name), parent_args) + data = backend_cls.read_meshtags_data( + filename, parent.comm, _post_code_name(mesh_name), parent_args + ) dim = int(data.dim) local_entities, local_values = dolfinx.io.distribute_entity_data( @@ -240,17 +260,13 @@ def read_submesh( entities = np.flatnonzero(marker.array >= 0).astype(np.int32) submesh, cell_map, vertex_map, node_map = dolfinx.mesh.create_submesh(parent, dim, entities) - sub_imap = submesh.topology.index_map(submesh.topology.dim) - num_cells = sub_imap.size_local + sub_imap.num_ghosts - parent_entities = cell_map.sub_topology_to_topology(np.arange(num_cells, dtype=np.int32), False) - stored_cells = marker.array[parent_entities].astype(np.int64) + post_code = marker.array[_parent_entities(submesh, cell_map)].astype(np.int64) return SubmeshCheckpoint( submesh=submesh, cell_map=cell_map, vertex_map=vertex_map, node_map=node_map, - parent_entities=parent_entities, - stored_cells=stored_cells, + post_code=post_code, ) @@ -265,11 +281,11 @@ def _point_ownership_data( :func:`dolfinx.fem.create_interpolation_data` would derive the same object by searching a bounding-box tree for each point, with a padding tolerance and an extrapolation fallback. Here the owning rank and cell of every point are - already known exactly, so the routing is done directly and no geometric - search -- and no tolerance -- is involved. + already known exactly, so the routing is done directly. - The layout required of the result is dictated by ``fem::interpolate`` and - ``impl::scatter_values``: + The layout required of the returned + :class:`dolfinx.geometry.PointOwnershipData` is dictated by + ``fem::interpolate`` and ``impl::scatter_values``: * ``src_owner`` is per interpolation point of the receiving function, in the order :func:`dolfinx.cpp.fem.interpolation_coords` produced them. @@ -313,7 +329,7 @@ def _point_ownership_data( def transfer_submesh_function( u_source: dolfinx.fem.Function, u_dest: dolfinx.fem.Function, - stored_cells: npt.NDArray[np.int64], + post_code: npt.NDArray[np.int64], ): """Move data from a standalone submesh onto one re-derived from a parent. @@ -333,9 +349,9 @@ def transfer_submesh_function( Args: u_source: Function on the submesh as stored u_dest: Function to fill, on the submesh re-derived from the parent - stored_cells: For each cell of ``u_dest``'s mesh (owned and ghost), the - index that cell had in the stored submesh. From - :attr:`SubmeshCheckpoint.stored_cells`. + post_code: For each cell of ``u_dest``'s mesh (owned and ghost), the + index that cell had in the stored submesh, which is what the cell is + routed by. From :attr:`SubmeshCheckpoint.post_code`. Raises: NotImplementedError: For H(div)/H(curl) spaces on a manifold submesh, @@ -354,21 +370,27 @@ def transfer_submesh_function( "Transferring an H(div)/H(curl) function onto a submesh with" f" tdim ({tdim}) < gdim ({dest_mesh.geometry.dim}) is not supported." " DOLFINx cannot reconcile the reference and physical value sizes of" - " these families on a manifold, and the interpolation this transfer" - " relies on does not report that: it returns silently with values that" - " are wrong (measured: the L2 norm of an N1curl field on a facet" - " submesh fell from 9.096 to 4.904). Writing and reading the standalone" - " submesh is exact for these spaces -- use `read_mesh` and" + " these families on a manifold (FEniCS/dolfinx#3619," + " https://github.com/FEniCS/dolfinx/issues/3619), and the interpolation" + " this transfer relies on does not report that: it returns silently" + " with values that are wrong (measured: the L2 norm of an N1curl field" + " on a facet submesh fell from 9.096 to 4.904). Writing and reading the" + " standalone submesh is exact for these spaces -- use `read_mesh` and" " `read_function` with the submesh's `mesh_name` and stop there. Only" " re-deriving from the parent is unavailable." ) sub_imap = dest_mesh.topology.index_map(tdim) num_cells = sub_imap.size_local + sub_imap.num_ghosts - if len(stored_cells) != num_cells: + # Both checks below are collective. What follows them is a neighbourhood + # exchange, so a process that raised on its own would leave the others + # waiting at it: every process has to reach the same verdict, even one whose + # own data is fine. + if comm.allreduce(len(post_code) != num_cells, MPI.LOR): raise ValueError( - f"Expected one stored cell index per cell of the destination mesh" - f" ({num_cells}), got {len(stored_cells)}." + "Expected one stored cell index per cell of the destination mesh;" + " the counts disagree on at least one process." + f" Here: {num_cells} cells, {len(post_code)} indices." ) cells = np.arange(num_cells, dtype=np.int32) @@ -384,18 +406,22 @@ def transfer_submesh_function( # Which process holds each stored cell in the source mesh, and where. source_tdim = source_mesh.topology.dim source_imap = source_mesh.topology.index_map(source_tdim) - owner_rank, owner_cell = _lookup_stored_cells( + owner_rank, owner_cell = _lookup_post_code( comm, np.asarray(source_mesh.topology.original_cell_index[: source_imap.size_local]), - np.asarray(stored_cells, dtype=np.int64), + np.asarray(post_code, dtype=np.int64), source_imap.size_global, ) - if num_cells > 0 and (owner_rank < 0).any(): - missing = np.unique(stored_cells[owner_rank < 0]) + # Counting owned cells only makes this an exact global tally: a ghost carries + # the same post code as its owner and so resolves the same way, and would + # otherwise be counted once per process holding it. + missing = np.flatnonzero(owner_rank[: sub_imap.size_local] < 0) + num_missing = comm.allreduce(len(missing), MPI.SUM) + if num_missing > 0: raise RuntimeError( - f"{len(missing)} cells of the destination submesh have no counterpart in" - f" the stored submesh (for instance {missing[:5]}). The two must describe" - " the same submesh of the same parent." + f"{num_missing} cells of the destination submesh have no counterpart" + " in the stored submesh. The two must describe the same submesh of the" + f" same parent. Post codes missing here: {post_code[missing][:5]}." ) interpolation_data = _point_ownership_data( @@ -408,30 +434,46 @@ def transfer_submesh_function( u_dest.x.scatter_forward() -def _lookup_stored_cells( +def _lookup_post_code( comm: MPI.Intracomm, - owned_stored: npt.NDArray[np.int64], - queried_stored: npt.NDArray[np.int64], + owned_post_code: npt.NDArray[np.int64], + queried_post_code: npt.NDArray[np.int64], num_cells_global: int, ) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: - """Find which process holds each stored cell, and its local index there. + """Locate each post code in the *stored* submesh: which process, which cell. + + Which side is which follows from what each submesh is. The stored submesh is + the standalone one, independent of the parent, and it is where the data + lives, so it is the side that publishes: it says where each of its cells is + held. The submesh re-derived from the parent is the side that needs that + information, so it is the side that queries. + + The two are partitioned independently, so neither knows where the other put a + given cell. Both talk to a third: the process that :func:`index_owner` + assigns the post code in the equal-split layout acts as a directory. Owners + publish into it, queriers read out of it, each in one neighbourhood exchange, + so no process ever holds a global table. - The two meshes are partitioned independently, so neither side knows where the - other put a given cell. Both sides talk to a third: the process that - :func:`index_owner` assigns the cell in the equal-split layout acts as a - directory. Owners publish into it, queriers read out of it, each in one - neighbourhood exchange, so no process ever holds a global table. + Every process is on both sides, it holds part of each submesh, so it + publishes for the stored cells it owns and queries for the re-derived cells + it holds, which is why the two arrays below come from different meshes. Args: - comm: The MPI communicator - owned_stored: Stored index of each cell this process owns in the source - mesh; position in the array is the local cell index - queried_stored: Stored indices this process wants to locate - num_cells_global: Number of cells in the stored mesh + comm: The MPI communicator both submeshes live on + owned_post_code: Post code of each cell this process owns in the + **stored** submesh, which for a mesh read back from a checkpoint is + its ordinary ``topology.original_cell_index``. The position in the + array is that cell's local index in the stored submesh, and that + index is what gets published. + queried_post_code: Post codes this process wants to locate, one per cell + of the **re-derived** submesh, owned and ghost + num_cells_global: Number of cells in the stored submesh, globally Returns: - ``(rank, local_cell)`` aligned with ``queried_stored``; ``-1`` where the - stored cell was never published. + ``(rank, local_cell)`` aligned with ``queried_post_code``: the process + that owns that cell in the **stored** submesh, and the cell's local index + in *that* process's stored submesh -- not in the re-derived one. Both are + ``-1`` where the post code was never published. """ directory_range = compute_local_range(comm, num_cells_global) directory_size = int(directory_range[1] - directory_range[0]) @@ -439,29 +481,21 @@ def _lookup_stored_cells( directory_cell = np.full(directory_size, -1, dtype=np.int32) # Publish: each owner tells the directory where it keeps the cell. - owned_stored = np.asarray(owned_stored, dtype=np.int64) - local_cells = np.arange(len(owned_stored), dtype=np.int32) - publish_to = ( - index_owner(comm, owned_stored, num_cells_global) - if len(owned_stored) - else np.empty(0, dtype=np.int32) - ) + owned_post_code = np.asarray(owned_post_code, dtype=np.int64) + local_cells = np.arange(len(owned_post_code), dtype=np.int32) + publish_to = index_owner(comm, owned_post_code, num_cells_global) (pub_keys, pub_cells), pub_sources, _, pub_counts, _ = exchange_to_owners( - comm, publish_to, [owned_stored, local_cells] + comm, publish_to, [owned_post_code, local_cells] ) slots = (pub_keys - directory_range[0]).astype(np.int64) directory_rank[slots] = np.repeat(pub_sources, pub_counts) directory_cell[slots] = pub_cells # Query: ask the directory, and unpack the reply into the order asked. - queried_stored = np.asarray(queried_stored, dtype=np.int64) - query_to = ( - index_owner(comm, queried_stored, num_cells_global) - if len(queried_stored) - else np.empty(0, dtype=np.int32) - ) + queried_post_code = np.asarray(queried_post_code, dtype=np.int64) + query_to = index_owner(comm, queried_post_code, num_cells_global) (inc_keys,), q_sources, q_send_counts, q_recv_counts, q_insert = exchange_to_owners( - comm, query_to, [queried_stored] + comm, query_to, [queried_post_code] ) inc_slots = (inc_keys - directory_range[0]).astype(np.int64) reply_rank = directory_rank[inc_slots] @@ -472,8 +506,8 @@ def _lookup_stored_cells( reverse = comm.Create_dist_graph_adjacent( q_destinations.tolist(), q_sources.tolist(), reorder=False ) - packed_rank = np.zeros(len(queried_stored), dtype=np.int32) - packed_cell = np.zeros(len(queried_stored), dtype=np.int32) + packed_rank = np.zeros(len(queried_post_code), dtype=np.int32) + packed_cell = np.zeros(len(queried_post_code), dtype=np.int32) reverse.Neighbor_alltoallv( [reply_rank, q_recv_counts, numpy_to_mpi[np.int32]], [packed_rank, q_send_counts, numpy_to_mpi[np.int32]], diff --git a/src/io4dolfinx/writers.py b/src/io4dolfinx/writers.py index 8ad5b58..58759b7 100644 --- a/src/io4dolfinx/writers.py +++ b/src/io4dolfinx/writers.py @@ -21,7 +21,7 @@ def prepare_meshdata_for_storage(mesh: dolfinx.mesh.Mesh, store_partition_info: bool) -> MeshData: """ Helper function for extracting the required data from a distributed - {py:class}`dolfinx.mesh.Mesh`. + :class:`dolfinx.mesh.Mesh`. Args: mesh: The mesh diff --git a/tests/test_submesh.py b/tests/test_submesh.py index b98cfa5..f018324 100644 --- a/tests/test_submesh.py +++ b/tests/test_submesh.py @@ -111,7 +111,7 @@ def f(x): checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) V_sub = dolfinx.fem.functionspace(checkpoint.submesh, (family, degree)) u_sub = dolfinx.fem.Function(V_sub) - io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) reference_sub = dolfinx.fem.Function(V_sub) reference_sub.interpolate(f) @@ -164,7 +164,7 @@ def f(x): checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) V_sub = dolfinx.fem.functionspace(checkpoint.submesh, (family, degree)) u_sub = dolfinx.fem.Function(V_sub) - io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) reference_sub = dolfinx.fem.Function(V_sub) reference_sub.interpolate(f) assert _max_error(u_sub, reference_sub) < 1e-13 @@ -227,7 +227,7 @@ def test_manifold_hcurl_roundtrips_but_transfer_raises(tmp_path, backend, family checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="wall", backend=backend) u_sub = dolfinx.fem.Function(dolfinx.fem.functionspace(checkpoint.submesh, (family, 1))) with pytest.raises(NotImplementedError, match="not supported"): - io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) @pytest.mark.parametrize("codim", [0, 1]) @@ -264,6 +264,113 @@ def test_read_submesh_supports_mixed_dimensional_assembly(tmp_path, backend, cod assert np.isclose(area, expected) +@pytest.mark.parametrize("codim", [0, 1]) +@pytest.mark.parametrize("degree", [0, 1, 2]) +def test_transfer_of_discontinuous_field(tmp_path, backend, codim, degree): + """A DG field that genuinely jumps inside the submesh survives the transfer. + + The DG spaces tested elsewhere carry a smooth function, which cannot tell + whether a value came from the right cell. Here the field is discontinuous + across a plane of the mesh, and for degree > 0 the interpolation points sit + on the vertices and facets that straddle it -- the places where taking the + value from the neighbouring cell would go unnoticed by a continuous field. + Each point is evaluated in the cell its post code names, so the side is never + in doubt. + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"jump{SUFFIX[backend]}" + + def jumpy(x): + # Discontinuous across x = 0.25, which is a plane of the 4x4x4 mesh. + return np.where(x[0] < 0.25, 3.0, -1.0) + + element = ("DG", degree) + mesh, submesh, cell_map = _make_submesh(comm, codim) + u = dolfinx.fem.Function(dolfinx.fem.functionspace(submesh, element), name="u") + u.interpolate(jumpy) + + io4dolfinx.write_mesh(path, mesh, backend=backend) + io4dolfinx.write_submesh(path, submesh, mesh, cell_map, mesh_name="wall", backend=backend) + io4dolfinx.write_function( + path, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="wall", backend=backend + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(path, comm, backend=backend) + stored = io4dolfinx.read_mesh(path, comm, mesh_name="wall", backend=backend) + u_stored = dolfinx.fem.Function(dolfinx.fem.functionspace(stored, element), name="u") + io4dolfinx.read_function(path, u_stored, time=0.0, name="u", mesh_name="wall", backend=backend) + + checkpoint = io4dolfinx.read_submesh(path, parent, mesh_name="wall", backend=backend) + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, element) + u_sub = dolfinx.fem.Function(V_sub, name="u") + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) + + reference = dolfinx.fem.Function(V_sub) + reference.interpolate(jumpy) + assert _max_error(u_sub, reference) < 1e-13 + + +def test_interface_submesh_assembly(tmp_path, backend): + """A facet submesh on an interface, integrated against a parent field. + + Only quantities that do not depend on which side is ``"+"`` are asserted: + the restrictions follow the facet-to-cell connectivity, so their assignment + is not a property of the mesh and no checkpoint preserves it. It is not + stable without IO either -- the same form on a freshly built mesh changes + sign between one process and three. ``abs`` of the jump is invariant; a form + that needs the *signed* jump has to pin the restrictions itself, e.g. by + ordering the integration entities by cell marker + (``scifem.compute_interface_data``). + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"iface{SUFFIX[backend]}" + + mesh = dolfinx.mesh.create_unit_cube( + comm, 4, 4, 4, ghost_mode=dolfinx.mesh.GhostMode.shared_facet + ) + tdim = mesh.topology.dim + mesh.topology.create_entities(tdim - 1) + facets = dolfinx.mesh.locate_entities(mesh, tdim - 1, lambda x: np.isclose(x[0], 0.5)) + interface, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, tdim - 1, facets) + + io4dolfinx.write_mesh(path, mesh, backend=backend) + io4dolfinx.write_submesh( + path, interface, mesh, cell_map, mesh_name="interface", backend=backend + ) + del mesh, interface, cell_map + + parent = io4dolfinx.read_mesh(path, comm, backend=backend) + checkpoint = io4dolfinx.read_submesh(path, parent, mesh_name="interface", backend=backend) + + # A parent field with a jump of exactly 1 across the interface. + k = dolfinx.fem.Function(dolfinx.fem.functionspace(parent, ("DG", 0))) + midpoints = parent.geometry.x[parent.geometry.dofmaps[0]].mean(axis=1) + k.x.array[: len(midpoints)] = np.where(midpoints[:, 0] < 0.5, 1.0, 2.0) + k.x.scatter_forward() + + ptdim = parent.topology.dim + parent.topology.create_entities(ptdim - 1) + parent.topology.create_connectivity(ptdim - 1, ptdim) + marked = dolfinx.mesh.locate_entities(parent, ptdim - 1, lambda x: np.isclose(x[0], 0.5)) + tags = dolfinx.mesh.meshtags(parent, ptdim - 1, marked, np.ones(len(marked), dtype=np.int32)) + dS = ufl.Measure("dS", domain=parent, subdomain_data=tags, subdomain_id=1) + + w = dolfinx.fem.Function(dolfinx.fem.functionspace(checkpoint.submesh, ("DG", 0))) + w.x.array[:] = 1.0 + + for expr, expected in [ + (w("+") * dS, 1.0), # area of the interface + (ufl.avg(k) * w("+") * dS, 1.5), # restriction-symmetric + (abs(ufl.jump(k)) * w("+") * dS, 1.0), # restriction-invariant + ]: + form = dolfinx.fem.form(expr, entity_maps=[checkpoint.cell_map]) + value = comm.allreduce(dolfinx.fem.assemble_scalar(form), MPI.SUM) + assert np.isclose(value, expected), (value, expected) + + def test_submesh_with_empty_ranks(tmp_path, backend): """A submesh small enough that some process owns none of it.""" comm = MPI.COMM_WORLD @@ -313,7 +420,7 @@ def f(x): checkpoint = io4dolfinx.read_submesh(parent_file, parent, mesh_name="strip", backend=backend) V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) u_sub = dolfinx.fem.Function(V_sub) - io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) reference = dolfinx.fem.Function(V_sub) reference.interpolate(f) @@ -361,6 +468,89 @@ def test_named_meshes_do_not_collide(tmp_path, backend): assert _max_error(v_cube, ref_cube) < 1e-13 +def test_multiple_submeshes_of_one_parent(tmp_path, backend): + """Several submeshes of the same parent share one file without colliding. + + Two of them have the same co-dimension, so their post codes tag the same + dimension of the parent, and one name is a prefix of the other -- which is + what would break if tags were matched by prefix rather than by name. + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"many{SUFFIX[backend]}" + + def f(x): + return np.sin(1.3 * x[0] + 0.2) + x[1] - 0.5 * x[2] + + mesh = dolfinx.mesh.create_unit_cube( + comm, 4, 4, 4, ghost_mode=dolfinx.mesh.GhostMode.shared_facet + ) + tdim = mesh.topology.dim + mesh.topology.create_entities(tdim - 1) + + # A cell submesh and two facet submeshes, on opposite faces. + subs = {} + for name, dim, locator in [ + ("half", tdim, lambda x: x[0] <= 0.5 + 1e-12), + ("wall", tdim - 1, lambda x: np.isclose(x[0], 0.0)), + ("wall2", tdim - 1, lambda x: np.isclose(x[0], 1.0)), + ]: + entities = dolfinx.mesh.locate_entities(mesh, dim, locator) + submesh, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, dim, entities) + subs[name] = (dim, submesh, cell_map) + + io4dolfinx.write_mesh(path, mesh, backend=backend) + for name, (_, submesh, cell_map) in subs.items(): + u = dolfinx.fem.Function(dolfinx.fem.functionspace(submesh, ("Lagrange", 2)), name="u") + u.interpolate(f) + io4dolfinx.write_submesh( + path, + submesh, + mesh, + cell_map, + mesh_name=name, + mode=io4dolfinx.FileMode.append, + backend=backend, + ) + io4dolfinx.write_function( + path, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name=name, backend=backend + ) + expected = { + name: (dim, submesh.topology.index_map(dim).size_global) + for name, (dim, submesh, _) in subs.items() + } + del mesh, subs + + parent = io4dolfinx.read_mesh(path, comm, backend=backend) + for name, (dim, num_cells_global) in expected.items(): + stored = io4dolfinx.read_mesh(path, comm, mesh_name=name, backend=backend) + assert stored.topology.dim == dim + u_stored = dolfinx.fem.Function( + dolfinx.fem.functionspace(stored, ("Lagrange", 2)), name="u" + ) + io4dolfinx.read_function( + path, u_stored, time=0.0, name="u", mesh_name=name, backend=backend + ) + + checkpoint = io4dolfinx.read_submesh(path, parent, mesh_name=name, backend=backend) + assert checkpoint.submesh.topology.dim == dim + assert checkpoint.submesh.topology.index_map(dim).size_global == num_cells_global + + # Each submesh must come back as itself: the two facet submeshes would be + # indistinguishable by size alone. + if name == "wall": + assert np.all(checkpoint.submesh.geometry.x[:, 0] < 1e-12) + elif name == "wall2": + assert np.all(checkpoint.submesh.geometry.x[:, 0] > 1.0 - 1e-12) + + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) + u_sub = dolfinx.fem.Function(V_sub, name="u") + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) + reference = dolfinx.fem.Function(V_sub) + reference.interpolate(f) + assert _max_error(u_sub, reference) < 1e-13 + + @pytest.mark.parametrize("kind", ["point", "cell"]) def test_submesh_point_and_cell_data(tmp_path, kind): """Visualisation data attached to a named submesh in a multi-block file. @@ -454,7 +644,7 @@ def f(x): ) V_sub = dolfinx.fem.functionspace(checkpoint.submesh, ("Lagrange", 2)) u_sub = dolfinx.fem.Function(V_sub) - io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.stored_cells) + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) reference = dolfinx.fem.Function(V_sub) reference.interpolate(f) From 6cfff4dfa8a4d444a0e2d80a7aec89b8d76b9324 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Tue, 22 Sep 2026 21:24:11 +0000 Subject: [PATCH 08/11] Another global check --- src/io4dolfinx/submesh.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/io4dolfinx/submesh.py b/src/io4dolfinx/submesh.py index a8b3a15..e17117e 100644 --- a/src/io4dolfinx/submesh.py +++ b/src/io4dolfinx/submesh.py @@ -392,6 +392,7 @@ def transfer_submesh_function( " the counts disagree on at least one process." f" Here: {num_cells} cells, {len(post_code)} indices." ) + post_code = np.asarray(post_code, dtype=np.int64) cells = np.arange(num_cells, dtype=np.int32) # Interpolation points of the destination, grouped per cell. @@ -406,11 +407,22 @@ def transfer_submesh_function( # Which process holds each stored cell in the source mesh, and where. source_tdim = source_mesh.topology.dim source_imap = source_mesh.topology.index_map(source_tdim) + # `index_owner` asserts that its indices are in range, and that assert is + # local: a process tripping it would abort while every other process waited + # for it in the exchange below. So the range is checked collectively first. + num_stored = source_imap.size_global + out_of_range = post_code[(post_code < 0) | (post_code >= num_stored)] + if comm.allreduce(len(out_of_range), MPI.SUM) > 0: + raise ValueError( + f"Post codes must index cells of the stored submesh, which has" + f" {num_stored} of them. Out of range here: {out_of_range[:5]}." + ) + owner_rank, owner_cell = _lookup_post_code( comm, np.asarray(source_mesh.topology.original_cell_index[: source_imap.size_local]), - np.asarray(post_code, dtype=np.int64), - source_imap.size_global, + post_code, + num_stored, ) # Counting owned cells only makes this an exact global tally: a ghost carries # the same post code as its owner and so resolves the same way, and would From dc795d6dd408221ec41b35bd96042a85cad8d7e0 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Tue, 22 Sep 2026 21:34:45 +0000 Subject: [PATCH 09/11] Add extra test. Higher order curved parent mesh and submesh read. --- tests/test_submesh.py | 126 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 126 insertions(+) diff --git a/tests/test_submesh.py b/tests/test_submesh.py index f018324..3cbcdff 100644 --- a/tests/test_submesh.py +++ b/tests/test_submesh.py @@ -763,6 +763,132 @@ def f(x): assert _max_error(read_back, reference) < 1e-13 +@pytest.mark.parametrize("degree", [2, 4]) +def test_transfer_on_curved_higher_order_parent(tmp_path, backend, degree): + """The whole chain on a parent whose geometry is genuinely curved. + + Every other transfer test uses an affine parent, where a cell's higher-order + nodes are collinear and a degree-4 mesh is indistinguishable from a degree-1 + one. Here the parent is a ball, so the submesh is the curved sphere surface: + its cells are curves, and the geometry nodes that carry that curvature have + to survive the parent's round-trip, ``create_submesh`` on the re-derived + parent, and the interpolation the transfer performs. + """ + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / f"curved_parent{SUFFIX[backend]}" + + mesh = _curved_mesh(comm, 3, degree) + tdim = mesh.topology.dim + mesh.topology.create_entities(tdim - 1) + mesh.topology.create_connectivity(tdim - 1, tdim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + submesh, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, tdim - 1, facets) + assert submesh.geometry.cmaps[0].degree == degree + + def f(x): + return np.sin(1.7 * x[0]) + 0.5 * x[1] - 0.3 * x[2] + + element = ("Lagrange", degree) + u = dolfinx.fem.Function(dolfinx.fem.functionspace(submesh, element), name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(path, mesh, backend=backend) + io4dolfinx.write_submesh(path, submesh, mesh, cell_map, mesh_name="shell", backend=backend) + io4dolfinx.write_function( + path, u, time=0.0, mode=io4dolfinx.FileMode.append, mesh_name="shell", backend=backend + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(path, comm, backend=backend) + assert parent.geometry.cmaps[0].degree == degree + stored = io4dolfinx.read_mesh(path, comm, mesh_name="shell", backend=backend) + u_stored = dolfinx.fem.Function(dolfinx.fem.functionspace(stored, element), name="u") + io4dolfinx.read_function(path, u_stored, time=0.0, name="u", mesh_name="shell", backend=backend) + + checkpoint = io4dolfinx.read_submesh(path, parent, mesh_name="shell", backend=backend) + assert checkpoint.submesh.geometry.cmaps[0].degree == degree + V_sub = dolfinx.fem.functionspace(checkpoint.submesh, element) + u_sub = dolfinx.fem.Function(V_sub, name="u") + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) + + reference = dolfinx.fem.Function(V_sub) + reference.interpolate(f) + assert _max_error(u_sub, reference) < 1e-12 + + +@pytest.mark.parametrize("degree", [1, 4]) +@pytest.mark.parametrize("kind", ["point", "cell"]) +def test_vtkhdf_visualisation_data_can_be_transferred(tmp_path, kind, degree): + """The vtkhdf route: read the stored submesh's data, then transfer it. + + vtkhdf cannot make function checkpoints, so ``read_function`` is unavailable, + but ``read_point_data`` and ``read_cell_data`` return ordinary functions and + feed ``transfer_submesh_function`` like any other source. Point data follows + the mesh's coordinate element, so on a degree-4 mesh this transfers a + degree-4 field, not a linear one. + """ + pytest.importorskip("h5py") + comm = MPI.COMM_WORLD + folder = comm.bcast(tmp_path, root=0) + path = folder / "vis_transfer.vtkhdf" + + mesh = _curved_mesh(comm, 3, degree) + tdim = mesh.topology.dim + mesh.topology.create_entities(tdim - 1) + mesh.topology.create_connectivity(tdim - 1, tdim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + submesh, cell_map, _, _ = dolfinx.mesh.create_submesh(mesh, tdim - 1, facets) + + def f(x): + return np.sin(1.7 * x[0]) + 0.5 * x[1] - 0.3 * x[2] + + def space(m): + if kind == "point": + return io4dolfinx.readers.create_geometry_function_space(m, 1) + return dolfinx.fem.functionspace(m, ("DG", 0)) + + u = dolfinx.fem.Function(space(submesh), name="u") + u.interpolate(f) + + io4dolfinx.write_mesh(path, mesh, backend="vtkhdf") + io4dolfinx.write_submesh( + path, + submesh, + mesh, + cell_map, + mesh_name="shell", + mode=io4dolfinx.FileMode.append, + backend="vtkhdf", + ) + writer = io4dolfinx.write_point_data if kind == "point" else io4dolfinx.write_cell_data + writer( + path, + u, + time=0.0, + mode=io4dolfinx.FileMode.append, + backend_args=None, + backend="vtkhdf", + mesh_name="shell", + ) + del mesh, submesh, cell_map, u + + parent = io4dolfinx.read_mesh(path, comm, backend="vtkhdf") + stored = io4dolfinx.read_mesh(path, comm, mesh_name="shell", backend="vtkhdf") + reader = io4dolfinx.read_point_data if kind == "point" else io4dolfinx.read_cell_data + u_stored = reader(path, "u", stored, time=0.0, backend="vtkhdf", mesh_name="shell") + if kind == "point": + assert u_stored.function_space.ufl_element().embedded_superdegree == degree + + checkpoint = io4dolfinx.read_submesh(path, parent, mesh_name="shell", backend="vtkhdf") + u_sub = dolfinx.fem.Function(space(checkpoint.submesh), name="u") + io4dolfinx.transfer_submesh_function(u_stored, u_sub, checkpoint.post_code) + + reference = dolfinx.fem.Function(u_sub.function_space) + reference.interpolate(f) + assert _max_error(u_sub, reference) < 1e-12 + + @pytest.mark.parametrize("degree", [2, 4]) @pytest.mark.parametrize("store", ["adios2", "h5py", "vtkhdf"]) def test_lagrange_variant_survives_mesh_roundtrip(tmp_path, store, degree): From f6cc88787b2526037d8076495ea82e51aede6ccd Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 23 Sep 2026 06:48:23 +0000 Subject: [PATCH 10/11] Use alltoall helper to avoid issues with openmpi. --- src/io4dolfinx/comm_helpers.py | 45 ++++++++++++++++++++++----- src/io4dolfinx/original_checkpoint.py | 12 +++---- src/io4dolfinx/readers.py | 4 +-- tests/test_original_checkpoint.py | 2 +- 4 files changed, 45 insertions(+), 18 deletions(-) diff --git a/src/io4dolfinx/comm_helpers.py b/src/io4dolfinx/comm_helpers.py index 8f1b52e..1036b9f 100644 --- a/src/io4dolfinx/comm_helpers.py +++ b/src/io4dolfinx/comm_helpers.py @@ -18,6 +18,7 @@ "neighbourhood_ranks", "exchange_to_owners", "numpy_to_mpi", + "all_to_all", ] numpy_to_mpi = { @@ -30,6 +31,38 @@ } +def all_to_all(comm, send_data, recv_data): + """ + Exchange a single item with each neighbor in a distributed graph communicator. + + Note: + The count is passed explicitly, and is 1 on every process. MPI-4.1 9.6.2 requires + the type signature of ``sendcount``/``sendtype`` at a process to equal that of + ``recvcount``/``recvtype`` at *any other* process in the communicator, not just at + its neighbors, so the count must be identical on every process whatever its degree. + Left implicit, mpi4py derives it as ``buffer size // degree`` and falls back to the + whole buffer when the degree is zero, making it rank-local: 1 where the degree is + nonzero, 0 where it is zero. Such a call is erroneous; Open MPI rejects it with + ``MPI_ERR_TRUNCATE`` while MPICH happens to accept it. See + https://github.com/open-mpi/ompi/issues/14452 for the discussion. + + ``all_to_allv`` is not affected: the vector variant is only required to match + pairwise along each edge, so per-process counts may legitimately differ there. + """ + dtype = numpy_to_mpi[send_data.dtype.type] + assert recv_data.dtype == send_data.dtype, ( + f"Data types do not match, {recv_data.dtype} != {send_data.dtype}" + ) + indegree, outdegree, _ = comm.Get_dist_neighbors_count() + assert (d_size := send_data.size) == outdegree, ( + f"Number of send data {d_size} does not match number of destinations {outdegree}" + ) + assert (d_size := recv_data.size) == indegree, ( + f"Number of recv data {d_size} does not match number of sources {indegree}" + ) + comm.Neighbor_alltoall([send_data, 1, dtype], [recv_data, 1, dtype]) + + def send_dofmap_and_recv_values( comm: MPI.Comm, source_ranks: npt.NDArray[np.int32], @@ -84,7 +117,7 @@ def send_dofmap_and_recv_values( mesh_to_data_comm = comm.Create_dist_graph_adjacent( source_ranks.tolist(), dest_ranks.tolist(), reorder=False ) - mesh_to_data_comm.Neighbor_alltoall(dest_size, recv_size) + all_to_all(mesh_to_data_comm, dest_size, recv_size) # Prepare data-structures for receiving total_incoming = sum(recv_size) @@ -166,7 +199,7 @@ def send_and_recv_cell_perm( # Send sizes to create data structures for receiving from NeighAlltoAllv recv_size = np.zeros_like(source, dtype=np.int32) - mesh_to_data.Neighbor_alltoall(dest_size, recv_size) + all_to_all(mesh_to_data, dest_size, recv_size) # Prepare data-structures for receiving total_incoming = sum(recv_size) @@ -233,11 +266,7 @@ def send_dofs_and_recv_values( # Send sizes to create data structures for receiving from NeighAlltoAllv recv_size = np.zeros_like(source, dtype=np.int32) - recv_size.resize(max(len(recv_size), 1)) # Minimal resize to work with ompi - dest_size.resize(max(len(dest_size), 1)) # Mininal resize to work with ompi - dofmap_to_values.Neighbor_alltoall(dest_size, recv_size) - dest_size.resize(len(dest)) - recv_size.resize(len(source)) + all_to_all(dofmap_to_values, dest_size, recv_size) # Send input dofs to processes holding input array inc_dofs = np.zeros(sum(recv_size), dtype=np.int64) @@ -337,7 +366,7 @@ def exchange_to_owners( forward = comm.Create_dist_graph_adjacent( sources.tolist(), destinations.tolist(), reorder=False ) - forward.Neighbor_alltoall(send_counts, recv_counts) + all_to_all(forward, send_counts, recv_counts) received = [] for array in arrays: diff --git a/src/io4dolfinx/original_checkpoint.py b/src/io4dolfinx/original_checkpoint.py index aa565f2..9307f7f 100644 --- a/src/io4dolfinx/original_checkpoint.py +++ b/src/io4dolfinx/original_checkpoint.py @@ -15,6 +15,7 @@ import dolfinx import numpy as np +from .comm_helpers import all_to_all from . import compat from .backends import FileMode, get_backend from .comm_helpers import numpy_to_mpi @@ -64,12 +65,9 @@ def create_original_mesh_data(mesh: dolfinx.mesh.Mesh) -> MeshData: # Compute number of recieving cells recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32) - if len(send_cells_per_proc) == 0: - send_cells_per_proc = np.zeros(1, dtype=np.int32) - if len(recv_cells_per_proc) == 0: - recv_cells_per_proc = np.zeros(1, dtype=np.int32) send_cells_per_proc = send_cells_per_proc.astype(np.int32) - cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc) + all_to_all(cell_to_output_comm, send_cells_per_proc, recv_cells_per_proc) + assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0] # Pack and send cell indices (used for mapping topology dofmap later) cell_insert_position = compute_insert_position( @@ -150,7 +148,7 @@ def create_original_mesh_data(mesh: dolfinx.mesh.Mesh) -> MeshData: # Send and recieve geometry sizes send_coordinate_sizes = (send_nodes_per_proc * 3).astype(np.int32) recv_coordinate_sizes = np.zeros_like(node_sources, dtype=np.int32) - geometry_to_owner_comm.Neighbor_alltoall(send_coordinate_sizes, recv_coordinate_sizes) + all_to_all(geometry_to_owner_comm, send_coordinate_sizes, recv_coordinate_sizes) # Send node coordinates recv_coordinates = np.empty(recv_coordinate_sizes.sum(), dtype=mesh.geometry.x.dtype) @@ -246,7 +244,7 @@ def create_function_data_on_original_mesh( # Compute number of recieving cells recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32) send_cells_per_proc = send_cells_per_proc.astype(np.int32) - cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc) + all_to_all(cell_to_output_comm, send_cells_per_proc, recv_cells_per_proc) assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0] # Pack and send cell indices (used for mapping topology dofmap later) diff --git a/src/io4dolfinx/readers.py b/src/io4dolfinx/readers.py index 7b6fa8a..a82b2f2 100644 --- a/src/io4dolfinx/readers.py +++ b/src/io4dolfinx/readers.py @@ -22,7 +22,7 @@ from . import compat from .backends import ReadMode, get_backend -from .comm_helpers import send_dofs_and_recv_values +from .comm_helpers import send_dofs_and_recv_values, all_to_all from .utils import ( check_file_exists, compute_dofmap_pos, @@ -100,7 +100,7 @@ def send_cells_and_receive_dofmap_index( source_ranks.tolist(), dest_ranks.tolist(), reorder=False ) # Send sizes to create data structures for receiving from NeighAlltoAllv - mesh_to_data_comm.Neighbor_alltoall(dest_size, recv_size) + all_to_all(mesh_to_data_comm, dest_size, recv_size) # Sort output for sending and fill send data out_cells = np.zeros(len(output_owners), dtype=np.int64) diff --git a/tests/test_original_checkpoint.py b/tests/test_original_checkpoint.py index 59370e4..da734ec 100644 --- a/tests/test_original_checkpoint.py +++ b/tests/test_original_checkpoint.py @@ -704,7 +704,7 @@ def read_xdmf_and_write_distributed(fname): query = cluster[:].apply_async(read_xdmf_and_write_distributed, fname) query.wait() - assert query.successful(), query.error + assert query.successful(), sum(str(e) for e in query.error) hash, f_dtype = query.result()[0] if write_mesh: mesh_fname = hash From 096744a9f531d2b28bdabe8d6e7854523857b66a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 06:48:34 +0000 Subject: [PATCH 11/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/io4dolfinx/original_checkpoint.py | 3 +-- src/io4dolfinx/readers.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/io4dolfinx/original_checkpoint.py b/src/io4dolfinx/original_checkpoint.py index 9307f7f..518df11 100644 --- a/src/io4dolfinx/original_checkpoint.py +++ b/src/io4dolfinx/original_checkpoint.py @@ -15,10 +15,9 @@ import dolfinx import numpy as np -from .comm_helpers import all_to_all from . import compat from .backends import FileMode, get_backend -from .comm_helpers import numpy_to_mpi +from .comm_helpers import all_to_all, numpy_to_mpi from .structures import FunctionData, MeshData from .utils import ( compute_insert_position, diff --git a/src/io4dolfinx/readers.py b/src/io4dolfinx/readers.py index a82b2f2..85408dc 100644 --- a/src/io4dolfinx/readers.py +++ b/src/io4dolfinx/readers.py @@ -22,7 +22,7 @@ from . import compat from .backends import ReadMode, get_backend -from .comm_helpers import send_dofs_and_recv_values, all_to_all +from .comm_helpers import all_to_all, send_dofs_and_recv_values from .utils import ( check_file_exists, compute_dofmap_pos,