Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/submesh.py: 97%
137 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1# Copyright (C) 2026 Jørgen Schartum Dokken
2#
3# This file is part of io4dolfinx
4#
5# SPDX-License-Identifier: MIT
7"""Checkpointing of submeshes.
9A :func:`dolfinx.mesh.create_submesh` submesh cannot be checkpointed the way an
10ordinary mesh is. Two properties of the DOLFINx construction get in the way:
121. For a submesh of co-dimension greater than zero, a cell of the submesh is an
13 *entity* of the parent, and the vertices of an entity are ordered by their
14 current global vertex index. Re-partitioning the parent renumbers the
15 vertices, so the same submesh cell comes back with its vertices in a
16 different order and the degrees of freedom inside it land in different
17 positions.
182. :func:`dolfinx.mesh.create_submesh` builds its vertex map allowing owner
19 changes, so a process can own a vertex -- and a dof -- that is only incident
20 to cells it ghosts. Such a dof has no position in any owned cell, which is
21 what :func:`io4dolfinx.utils.compute_dofmap_pos` needs.
23So a function is never read onto a re-derived submesh. Instead the submesh is
24stored as an ordinary, independent mesh (whose cells keep the stored node order
25and whose ownership is the ordinary one), read back as such, and the data is
26then transferred to a mesh re-derived from the parent when the caller needs the
27entity maps that make mixed-dimensional assembly possible.
28"""
30from __future__ import annotations
32import logging
33from dataclasses import dataclass
34from pathlib import Path
35from typing import Any
37from mpi4py import MPI
39import dolfinx
40import numpy as np
41import numpy.typing as npt
42from dolfinx import cpp as _cpp
44from .backends import FileMode, get_backend
45from .checkpointing import _with_mesh_name, write_mesh
46from .comm_helpers import exchange_to_owners, numpy_to_mpi
47from .structures import MeshTagsData
48from .utils import check_file_exists, compute_local_range, index_owner
50__all__ = [
51 "SubmeshCheckpoint",
52 "read_submesh",
53 "transfer_submesh_function",
54 "write_submesh",
55]
57logger = logging.getLogger(__name__)
60def _post_code_name(mesh_name: str) -> str:
61 """Name of the meshtag holding the post codes of the submesh ``mesh_name``.
63 The post codes tag the *parent* mesh, so they are stored in the parent's
64 namespace and must not collide with a tag a user wrote.
65 """
66 return f"__submesh_post_code_{mesh_name}"
69def _parent_entities(
70 submesh: dolfinx.mesh.Mesh, cell_map: dolfinx.mesh.EntityMap
71) -> npt.NDArray[np.int32]:
72 """Parent entity of every cell of ``submesh``, owned and ghost."""
73 imap = submesh.topology.index_map(submesh.topology.dim)
74 cells = np.arange(imap.size_local + imap.num_ghosts, dtype=np.int32)
75 return cell_map.sub_topology_to_topology(cells, False)
78@dataclass
79class SubmeshCheckpoint:
80 """A submesh re-derived from a parent mesh that was read from a checkpoint."""
82 #: The submesh, derived from the parent with :func:`dolfinx.mesh.create_submesh`
83 submesh: dolfinx.mesh.Mesh
84 #: Map from cells of :attr:`submesh` to entities of the parent
85 cell_map: dolfinx.mesh.EntityMap
86 #: Map from vertices of :attr:`submesh` to vertices of the parent
87 vertex_map: dolfinx.mesh.EntityMap
88 #: Map from geometry nodes of :attr:`submesh` to nodes of the parent
89 node_map: npt.NDArray[np.int32]
91 #: Index the cell had in the stored submesh, per local cell (owned and ghost).
92 #: This is the key both submeshes agree on: :func:`transfer_submesh_function`
93 #: derives from it the post office holding each cell, and so routes the data
94 #: from the stored submesh to this one.
95 post_code: npt.NDArray[np.int64]
97 @property
98 def parent_entities(self) -> npt.NDArray[np.int32]:
99 """Parent entities of the submesh cells, aligned with :attr:`post_code`."""
100 return _parent_entities(self.submesh, self.cell_map)
103def write_submesh(
104 filename: Path | str,
105 submesh: dolfinx.mesh.Mesh,
106 parent: dolfinx.mesh.Mesh,
107 cell_map: Any,
108 mesh_name: str,
109 parent_mesh_name: str | None = None,
110 parent_filename: Path | str | None = None,
111 time: float = 0.0,
112 mode: FileMode = FileMode.append,
113 backend_args: dict[str, Any] | None = None,
114 backend: str | None = None,
115):
116 """Write a submesh, and the post codes that tie it back to its parent.
118 The submesh is stored as an ordinary independent mesh under ``mesh_name``, so
119 :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function` read it
120 back with no special handling. Alongside it goes a tag of *post codes*: every
121 parent entity that became a submesh cell is given the index that cell has in
122 the stored submesh. The entities are addressed by their parent geometry
123 nodes, which is what lets :func:`read_submesh` find them again in a
124 re-partitioned parent; the values come back as
125 :attr:`SubmeshCheckpoint.post_code`.
127 The post codes tag entities of the *parent*, so they are stored in the parent's
128 namespace, in ``parent_filename``. Leave that unset to keep everything in one
129 checkpoint, or point it at the parent's own file to keep the two apart; the
130 submesh always goes to ``filename``. Either way the parent must already have
131 been written there with :func:`io4dolfinx.write_mesh`.
133 Args:
134 filename: Path to write the submesh to
135 submesh: The submesh, as returned by :func:`dolfinx.mesh.create_submesh`
136 parent: The mesh the submesh was derived from
137 cell_map: Map from submesh cells to parent entities, as returned by
138 :func:`dolfinx.mesh.create_submesh`
139 mesh_name: Name to store the submesh under
140 parent_mesh_name: Name the parent mesh is stored under
141 parent_filename: File holding the parent mesh, if not ``filename``
142 time: Time stamp associated with the submesh geometry
143 mode: Whether to create the file or append to it
144 backend_args: Arguments for the backend
145 backend: Which backend to use
146 """
147 logger.debug(f"Writing submesh '{mesh_name}' to {filename}")
148 dim = submesh.topology.dim
150 # 1. The submesh itself, as a perfectly ordinary mesh.
151 write_mesh(
152 Path(filename),
153 submesh,
154 mode=mode,
155 time=time,
156 backend_args=backend_args,
157 backend=backend,
158 mesh_name=mesh_name,
159 )
161 # 2. The post codes. Structurally this is a meshtag on the parent: the
162 # tagged entities are the parent entities that became submesh cells, and
163 # each value is that cell's post code, its index in the stored submesh.
164 cell_imap = submesh.topology.index_map(dim)
165 num_owned = cell_imap.size_local
166 parent_entities = cell_map.sub_topology_to_topology(np.arange(num_owned, dtype=np.int32), False)
167 post_code = np.arange(*cell_imap.local_range, dtype=np.int64)
169 parent.topology.create_connectivity(dim, parent.topology.dim)
170 parent.topology.create_connectivity(0, parent.topology.dim)
171 entities_to_geometry = dolfinx.cpp.mesh.entities_to_geometry(
172 parent._cpp_object, dim, parent_entities, False
173 )
174 indices = (
175 parent.geometry.index_map()
176 .local_to_global(entities_to_geometry.reshape(-1))
177 .reshape(entities_to_geometry.shape)
178 )
180 tag_data = MeshTagsData(
181 values=post_code,
182 num_entities_global=cell_imap.size_global,
183 num_dofs_per_entity=entities_to_geometry.shape[1],
184 indices=indices,
185 name=_post_code_name(mesh_name),
186 local_start=cell_imap.local_range[0],
187 dim=dim,
188 cell_type=submesh.topology.cell_name(),
189 )
190 backend_cls = get_backend(backend)
191 parent_args = backend_cls.get_default_backend_args(
192 _with_mesh_name(backend_args, parent_mesh_name)
193 )
194 post_code_file = filename if parent_filename is None else parent_filename
195 backend_cls.write_meshtags(post_code_file, parent.comm, tag_data, backend_args=parent_args)
198def read_submesh(
199 filename: Path | str,
200 parent: dolfinx.mesh.Mesh,
201 mesh_name: str,
202 parent_mesh_name: str | None = None,
203 backend_args: dict[str, Any] | None = None,
204 backend: str | None = None,
205) -> SubmeshCheckpoint:
206 """Re-derive a stored submesh from a parent mesh that has been read back.
208 The result is a genuine :func:`dolfinx.mesh.create_submesh` submesh, so its
209 entity maps are accepted by :func:`dolfinx.fem.form` for mixed-dimensional
210 assembly against ``parent``.
212 It does **not** carry the checkpointed data. Read that onto the standalone
213 submesh with :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function`,
214 then move it across with :func:`transfer_submesh_function`.
216 Args:
217 filename: File holding the parent mesh and the post codes, i.e. whatever
218 was passed as ``parent_filename`` to :func:`write_submesh`
219 parent: The parent mesh, as read back from the checkpoint
220 mesh_name: Name the submesh was stored under
221 parent_mesh_name: Name the parent mesh is stored under
222 backend_args: Arguments for the backend
223 backend: Which backend to use
225 Returns:
226 The re-derived submesh and the maps relating it to ``parent``.
227 """
228 logger.debug(f"Reading submesh '{mesh_name}' from {filename}")
229 check_file_exists(filename)
230 backend_cls = get_backend(backend)
231 parent_args = backend_cls.get_default_backend_args(
232 _with_mesh_name(backend_args, parent_mesh_name)
233 )
234 # Post codes are global cell indices, so they must survive as int64.
235 parent_args["values_dtype"] = np.int64
236 data = backend_cls.read_meshtags_data(
237 filename, parent.comm, _post_code_name(mesh_name), parent_args
238 )
239 dim = int(data.dim)
241 local_entities, local_values = dolfinx.io.distribute_entity_data(
242 parent, dim, data.indices, data.values.astype(np.int64)
243 )
244 parent.topology.create_connectivity(dim, 0)
245 parent.topology.create_connectivity(dim, parent.topology.dim)
246 adj = dolfinx.graph.adjacencylist(local_entities)
247 tags = dolfinx.mesh.meshtags_from_entities(
248 parent, dim, adj, np.asarray(local_values, dtype=np.int64)
249 )
251 # `distribute_entity_data` only reports entities to processes that hold them;
252 # a process ghosting an entity may not be among them. Scattering through an
253 # index map gives every process the owner's value for its ghosts, so the
254 # submesh is built from a consistent set of entities everywhere.
255 entity_map = parent.topology.index_map(dim)
256 marker = dolfinx.la.vector(entity_map, dtype=np.int64)
257 marker.array[:] = -1
258 marker.array[tags.indices] = tags.values
259 marker.scatter_forward()
260 entities = np.flatnonzero(marker.array >= 0).astype(np.int32)
262 submesh, cell_map, vertex_map, node_map = dolfinx.mesh.create_submesh(parent, dim, entities)
263 post_code = marker.array[_parent_entities(submesh, cell_map)].astype(np.int64)
264 return SubmeshCheckpoint(
265 submesh=submesh,
266 cell_map=cell_map,
267 vertex_map=vertex_map,
268 node_map=node_map,
269 post_code=post_code,
270 )
273def _point_ownership_data(
274 comm: MPI.Intracomm,
275 src_owner: npt.NDArray[np.int32],
276 points: npt.NDArray[np.floating],
277 source_cells: npt.NDArray[np.int32],
278) -> Any:
279 """Build a :class:`dolfinx.geometry.PointOwnershipData` from a known routing.
281 :func:`dolfinx.fem.create_interpolation_data` would derive the same object by
282 searching a bounding-box tree for each point, with a padding tolerance and an
283 extrapolation fallback. Here the owning rank and cell of every point are
284 already known exactly, so the routing is done directly.
286 The layout required of the returned
287 :class:`dolfinx.geometry.PointOwnershipData` is dictated by
288 ``fem::interpolate`` and ``impl::scatter_values``:
290 * ``src_owner`` is per interpolation point of the receiving function, in the
291 order :func:`dolfinx.cpp.fem.interpolation_coords` produced them.
292 * ``dest_owners`` must be in **contiguous runs by rank**; ``scatter_values``
293 reduces it with a non-sorting ``unique``, so interleaved ranks would
294 silently build the wrong neighbourhood. :func:`exchange_to_owners` groups
295 by ascending source rank, which satisfies this.
296 * ``dest_points`` is flat row-major ``(n_recv, 3)`` and ``dest_cells`` indexes
297 cells of the *sending* function's mesh on the receiving process.
299 Args:
300 comm: The MPI communicator both meshes live on
301 src_owner: Rank holding the source cell of each point
302 points: The points, shape ``(num_points, 3)``
303 source_cells: Source-mesh local cell of each point, on ``src_owner``
305 Returns:
306 Interpolation data for :meth:`dolfinx.fem.Function.interpolate_nonmatching`.
307 """
308 (recv_points, recv_cells), sources, _, recv_counts, _ = exchange_to_owners(
309 comm, src_owner, [points, source_cells]
310 )
312 # One entry per received point, naming the rank it came from. `sources` is
313 # ascending and `exchange_to_owners` receives in that order, so the runs are
314 # contiguous by construction -- which `scatter_values` requires.
315 dest_owners = np.repeat(sources, recv_counts).astype(np.int32)
317 dtype = np.dtype(points.dtype)
318 cls = getattr(_cpp.geometry, f"PointOwnershipData_float{8 * dtype.itemsize}")
319 return dolfinx.geometry.PointOwnershipData(
320 cls(
321 np.ascontiguousarray(src_owner, dtype=np.int32),
322 dest_owners,
323 np.ascontiguousarray(recv_points, dtype=dtype).reshape(-1),
324 np.ascontiguousarray(recv_cells, dtype=np.int32),
325 )
326 )
329def transfer_submesh_function(
330 u_source: dolfinx.fem.Function,
331 u_dest: dolfinx.fem.Function,
332 post_code: npt.NDArray[np.int64],
333):
334 """Move data from a standalone submesh onto one re-derived from a parent.
336 ``u_source`` lives on the submesh as it was stored -- read with
337 :func:`io4dolfinx.read_mesh` and :func:`io4dolfinx.read_function` -- and
338 ``u_dest`` on the submesh :func:`read_submesh` derived from the parent. The
339 two meshes cover the same cells but number their vertices differently and are
340 partitioned independently, so the transfer cannot be a copy.
342 Every cell of ``u_dest``'s mesh has a known counterpart in ``u_source``'s, so
343 each interpolation point is routed straight to the process and cell that can
344 evaluate it, with no geometric point location. This is an evaluation followed
345 by a scatter and a local interpolation: it assembles no form and needs no
346 entity map between the two meshes, which would have no meaning as neither is
347 derived from the other.
349 Args:
350 u_source: Function on the submesh as stored
351 u_dest: Function to fill, on the submesh re-derived from the parent
352 post_code: For each cell of ``u_dest``'s mesh (owned and ghost), the
353 index that cell had in the stored submesh, which is what the cell is
354 routed by. From :attr:`SubmeshCheckpoint.post_code`.
356 Raises:
357 NotImplementedError: For H(div)/H(curl) spaces on a manifold submesh
358 with a DOLFINx lacking ``FiniteElement.physical_base_value_size``,
359 whose interpolation would return wrong values without reporting an
360 error.
361 """
362 V_dest = u_dest.function_space
363 dest_mesh = V_dest.mesh
364 source_mesh = u_source.function_space.mesh
365 comm = dest_mesh.comm
366 assert isinstance(comm, MPI.Intracomm)
368 tdim = dest_mesh.topology.dim
369 supports_manifold_transform = hasattr(dolfinx.fem.FiniteElement, "physical_base_value_size")
370 manifold_transformations = (
371 tdim < dest_mesh.geometry.dim and V_dest.element.needs_dof_transformations
372 )
374 if manifold_transformations and not supports_manifold_transform:
375 raise NotImplementedError(
376 "Transferring an H(div)/H(curl) function onto a submesh with"
377 f" tdim ({tdim}) < gdim ({dest_mesh.geometry.dim}) is not supported."
378 " This DOLFINx cannot reconcile the reference and physical value sizes"
379 " of these families on a manifold (FEniCS/dolfinx#3619,"
380 " https://github.com/FEniCS/dolfinx/issues/3619, fixed on the main"
381 " branch by FEniCS/dolfinx#4511), and the interpolation"
382 " this transfer relies on does not report that: it returns silently"
383 " with values that are wrong (measured: the L2 norm of an N1curl field"
384 " on a facet submesh fell from 9.096 to 4.904). Writing and reading the"
385 " standalone submesh is exact for these spaces -- use `read_mesh` and"
386 " `read_function` with the submesh's `mesh_name` and stop there. Only"
387 " re-deriving from the parent is unavailable."
388 )
390 sub_imap = dest_mesh.topology.index_map(tdim)
391 num_cells = sub_imap.size_local + sub_imap.num_ghosts
392 # Both checks below are collective. What follows them is a neighbourhood
393 # exchange, so a process that raised on its own would leave the others
394 # waiting at it: every process has to reach the same verdict, even one whose
395 # own data is fine.
396 if comm.allreduce(len(post_code) != num_cells, MPI.LOR):
397 raise ValueError(
398 "Expected one stored cell index per cell of the destination mesh;"
399 " the counts disagree on at least one process."
400 f" Here: {num_cells} cells, {len(post_code)} indices."
401 )
402 post_code = np.asarray(post_code, dtype=np.int64)
403 cells = np.arange(num_cells, dtype=np.int32)
405 # Interpolation points of the destination, grouped per cell.
406 points = _cpp.fem.interpolation_coords(
407 V_dest.element._cpp_object, # type: ignore[arg-type]
408 dest_mesh.geometry._cpp_object, # type: ignore[arg-type]
409 cells,
410 )
411 points = np.ascontiguousarray(points.T) # (num_points, 3)
412 points_per_cell = 0 if num_cells == 0 else points.shape[0] // num_cells
414 # Which process holds each stored cell in the source mesh, and where.
415 source_tdim = source_mesh.topology.dim
416 source_imap = source_mesh.topology.index_map(source_tdim)
417 # `index_owner` asserts that its indices are in range, and that assert is
418 # local: a process tripping it would abort while every other process waited
419 # for it in the exchange below. So the range is checked collectively first.
420 num_stored = source_imap.size_global
421 out_of_range = post_code[(post_code < 0) | (post_code >= num_stored)]
422 if comm.allreduce(len(out_of_range), MPI.SUM) > 0:
423 raise ValueError(
424 f"Post codes must index cells of the stored submesh, which has"
425 f" {num_stored} of them. Out of range here: {out_of_range[:5]}."
426 )
428 owner_rank, owner_cell = _lookup_post_code(
429 comm,
430 np.asarray(source_mesh.topology.original_cell_index[: source_imap.size_local]),
431 post_code,
432 num_stored,
433 )
434 # Counting owned cells only makes this an exact global tally: a ghost carries
435 # the same post code as its owner and so resolves the same way, and would
436 # otherwise be counted once per process holding it.
437 missing = np.flatnonzero(owner_rank[: sub_imap.size_local] < 0)
438 num_missing = comm.allreduce(len(missing), MPI.SUM)
439 if num_missing > 0:
440 raise RuntimeError(
441 f"{num_missing} cells of the destination submesh have no counterpart"
442 " in the stored submesh. The two must describe the same submesh of the"
443 f" same parent. Post codes missing here: {post_code[missing][:5]}."
444 )
446 interpolation_data = _point_ownership_data(
447 comm,
448 np.repeat(owner_rank, points_per_cell),
449 points,
450 np.repeat(owner_cell, points_per_cell),
451 )
452 u_dest.interpolate_nonmatching(u_source, cells, interpolation_data=interpolation_data)
453 u_dest.x.scatter_forward()
456def _lookup_post_code(
457 comm: MPI.Intracomm,
458 owned_post_code: npt.NDArray[np.int64],
459 queried_post_code: npt.NDArray[np.int64],
460 num_cells_global: int,
461) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
462 """Locate each post code in the *stored* submesh: which process, which cell.
464 Which side is which follows from what each submesh is. The stored submesh is
465 the standalone one, independent of the parent, and it is where the data
466 lives, so it is the side that publishes: it says where each of its cells is
467 held. The submesh re-derived from the parent is the side that needs that
468 information, so it is the side that queries.
470 The two are partitioned independently, so neither knows where the other put a
471 given cell. Both talk to a third: the process that :func:`index_owner`
472 assigns the post code in the equal-split layout acts as a directory. Owners
473 publish into it, queriers read out of it, each in one neighbourhood exchange,
474 so no process ever holds a global table.
476 Every process is on both sides, it holds part of each submesh, so it
477 publishes for the stored cells it owns and queries for the re-derived cells
478 it holds, which is why the two arrays below come from different meshes.
480 Args:
481 comm: The MPI communicator both submeshes live on
482 owned_post_code: Post code of each cell this process owns in the
483 **stored** submesh, which for a mesh read back from a checkpoint is
484 its ordinary ``topology.original_cell_index``. The position in the
485 array is that cell's local index in the stored submesh, and that
486 index is what gets published.
487 queried_post_code: Post codes this process wants to locate, one per cell
488 of the **re-derived** submesh, owned and ghost
489 num_cells_global: Number of cells in the stored submesh, globally
491 Returns:
492 ``(rank, local_cell)`` aligned with ``queried_post_code``: the process
493 that owns that cell in the **stored** submesh, and the cell's local index
494 in *that* process's stored submesh -- not in the re-derived one. Both are
495 ``-1`` where the post code was never published.
496 """
497 directory_range = compute_local_range(comm, num_cells_global)
498 directory_size = int(directory_range[1] - directory_range[0])
499 directory_rank = np.full(directory_size, -1, dtype=np.int32)
500 directory_cell = np.full(directory_size, -1, dtype=np.int32)
502 # Publish: each owner tells the directory where it keeps the cell.
503 owned_post_code = np.asarray(owned_post_code, dtype=np.int64)
504 local_cells = np.arange(len(owned_post_code), dtype=np.int32)
505 publish_to = index_owner(comm, owned_post_code, num_cells_global)
506 (pub_keys, pub_cells), pub_sources, _, pub_counts, _ = exchange_to_owners(
507 comm, publish_to, [owned_post_code, local_cells]
508 )
509 slots = (pub_keys - directory_range[0]).astype(np.int64)
510 directory_rank[slots] = np.repeat(pub_sources, pub_counts)
511 directory_cell[slots] = pub_cells
513 # Query: ask the directory, and unpack the reply into the order asked.
514 queried_post_code = np.asarray(queried_post_code, dtype=np.int64)
515 query_to = index_owner(comm, queried_post_code, num_cells_global)
516 (inc_keys,), q_sources, q_send_counts, q_recv_counts, q_insert = exchange_to_owners(
517 comm, query_to, [queried_post_code]
518 )
519 inc_slots = (inc_keys - directory_range[0]).astype(np.int64)
520 reply_rank = directory_rank[inc_slots]
521 reply_cell = directory_cell[inc_slots]
523 # Send the answers back along the reverse of the query graph.
524 q_destinations = np.unique(query_to)
525 reverse = comm.Create_dist_graph_adjacent(
526 q_destinations.tolist(), q_sources.tolist(), reorder=False
527 )
528 packed_rank = np.zeros(len(queried_post_code), dtype=np.int32)
529 packed_cell = np.zeros(len(queried_post_code), dtype=np.int32)
530 reverse.Neighbor_alltoallv(
531 [reply_rank, q_recv_counts, numpy_to_mpi[np.int32]],
532 [packed_rank, q_send_counts, numpy_to_mpi[np.int32]],
533 )
534 reverse.Neighbor_alltoallv(
535 [reply_cell, q_recv_counts, numpy_to_mpi[np.int32]],
536 [packed_cell, q_send_counts, numpy_to_mpi[np.int32]],
537 )
538 reverse.Free()
540 # `q_insert[i]` is where local entry `i` was packed into the outgoing buffer,
541 # and the reply comes back in that same packed order, so read the answers
542 # back out by gathering -- scattering here would apply the inverse
543 # permutation and silently mix up which cell each answer belongs to.
544 return packed_rank[q_insert], packed_cell[q_insert]