Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/original_checkpoint.py: 99%
192 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 18:20 +0000
1# Copyright (C) 2024 Jørgen Schartum Dokken
2#
3# This file is part of io4dolfinx
4#
5# SPDX-License-Identifier: MIT
7from __future__ import annotations
9import logging
10import typing
11from pathlib import Path
13from mpi4py import MPI
15import dolfinx
16import numpy as np
18from . import compat
19from .backends import FileMode, get_backend
20from .comm_helpers import numpy_to_mpi
21from .structures import FunctionData, MeshData
22from .utils import (
23 compute_insert_position,
24 compute_local_range,
25 index_owner,
26 unroll_dofmap,
27 unroll_insert_position,
28)
30__all__ = ["write_function_on_input_mesh", "write_mesh_input_order"]
31logger = logging.getLogger(__name__)
34def create_original_mesh_data(mesh: dolfinx.mesh.Mesh) -> MeshData:
35 """
36 Store data locally on output process
37 """
39 # 1. Send cell indices owned by current process to the process which owned its input
41 # Get the input cell index for cells owned by this process
42 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
43 original_cell_index = mesh.topology.original_cell_index[:num_owned_cells]
45 # Compute owner of cells on this process based on the original cell index
46 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
47 output_cell_owner = index_owner(mesh.comm, original_cell_index, num_cells_global)
48 local_cell_range = compute_local_range(mesh.comm, num_cells_global)
50 # Compute outgoing edges from current process to outputting process
51 # Computes the number of cells sent to each process at the same time
52 cell_destinations, _send_cells_per_proc = np.unique(output_cell_owner, return_counts=True)
53 send_cells_per_proc = _send_cells_per_proc.astype(np.int32)
54 del _send_cells_per_proc
55 assert isinstance(mesh.comm, MPI.Intracomm)
56 cell_to_output_comm = mesh.comm.Create_dist_graph(
57 [mesh.comm.rank],
58 [len(cell_destinations)],
59 cell_destinations.tolist(),
60 reorder=False,
61 )
62 cell_sources, cell_dests, _ = cell_to_output_comm.Get_dist_neighbors()
63 assert np.allclose(cell_dests, cell_destinations)
65 # Compute number of recieving cells
66 recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32)
67 if len(send_cells_per_proc) == 0:
68 send_cells_per_proc = np.zeros(1, dtype=np.int32)
69 if len(recv_cells_per_proc) == 0:
70 recv_cells_per_proc = np.zeros(1, dtype=np.int32)
71 send_cells_per_proc = send_cells_per_proc.astype(np.int32)
72 cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc)
73 assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0]
74 # Pack and send cell indices (used for mapping topology dofmap later)
75 cell_insert_position = compute_insert_position(
76 output_cell_owner, cell_destinations, send_cells_per_proc
77 )
78 send_cells = np.empty_like(cell_insert_position, dtype=np.int64)
79 send_cells[cell_insert_position] = original_cell_index
80 recv_cells = np.empty(recv_cells_per_proc.sum(), dtype=np.int64)
81 send_cells_msg = [send_cells, send_cells_per_proc, MPI.INT64_T]
82 recv_cells_msg = [recv_cells, recv_cells_per_proc, MPI.INT64_T]
83 cell_to_output_comm.Neighbor_alltoallv(send_cells_msg, recv_cells_msg)
84 del send_cells_msg, recv_cells_msg, send_cells
86 # Map received cells to the local index
87 local_cell_index = recv_cells - local_cell_range[0]
89 # 2. Create dofmap based on original geometry indices and re-order in the same order as original
90 # cell indices on output process
92 # Get original node index for all nodes (including ghosts) and convert dofmap to these indices
93 original_node_index = mesh.geometry.input_global_indices
94 _, num_nodes_per_cell = compat.dofmap(mesh).shape
95 local_geometry_dofmap = compat.dofmap(mesh)[:num_owned_cells, :]
96 global_geometry_dofmap = original_node_index[local_geometry_dofmap.reshape(-1)]
98 # Unroll insert position for geometry dofmap
99 dofmap_insert_position = unroll_insert_position(cell_insert_position, num_nodes_per_cell)
101 # Create and commmnicate connecitivity in original geometry indices
102 send_geometry_dofmap = np.empty_like(dofmap_insert_position, dtype=np.int64)
103 send_geometry_dofmap[dofmap_insert_position] = global_geometry_dofmap
104 del global_geometry_dofmap
105 send_sizes_dofmap = send_cells_per_proc * num_nodes_per_cell
106 recv_sizes_dofmap = recv_cells_per_proc * num_nodes_per_cell
107 recv_geometry_dofmap = np.empty(recv_sizes_dofmap.sum(), dtype=np.int64)
108 send_geometry_dofmap_msg = [send_geometry_dofmap, send_sizes_dofmap, MPI.INT64_T]
109 recv_geometry_dofmap_msg = [recv_geometry_dofmap, recv_sizes_dofmap, MPI.INT64_T]
110 cell_to_output_comm.Neighbor_alltoallv(send_geometry_dofmap_msg, recv_geometry_dofmap_msg)
111 del send_geometry_dofmap_msg, recv_geometry_dofmap_msg
113 # Reshape dofmap and sort by original cell index
114 recv_dofmap = recv_geometry_dofmap.reshape(-1, num_nodes_per_cell)
115 sorted_recv_dofmap = np.empty_like(recv_dofmap)
116 sorted_recv_dofmap[local_cell_index] = recv_dofmap
118 # 3. Move geometry coordinates to input process
119 # Compute outgoing edges from current process and create neighbourhood communicator
120 # Also create number of outgoing cells at the same time
121 num_owned_nodes = mesh.geometry.index_map().size_local
122 num_nodes_global = mesh.geometry.index_map().size_global
123 output_node_owner = index_owner(
124 mesh.comm, original_node_index[:num_owned_nodes], num_nodes_global
125 )
127 node_destinations, _send_nodes_per_proc = np.unique(output_node_owner, return_counts=True)
128 send_nodes_per_proc = _send_nodes_per_proc.astype(np.int32)
129 del _send_nodes_per_proc
131 geometry_to_owner_comm = mesh.comm.Create_dist_graph(
132 [mesh.comm.rank],
133 [len(node_destinations)],
134 node_destinations.tolist(),
135 reorder=False,
136 )
138 node_sources, node_dests, _ = geometry_to_owner_comm.Get_dist_neighbors()
139 assert np.allclose(node_dests, node_destinations)
141 # Compute send node insert positions
142 send_nodes_position = compute_insert_position(
143 output_node_owner, node_destinations, send_nodes_per_proc
144 )
145 unrolled_nodes_positiion = unroll_insert_position(send_nodes_position, 3)
147 send_coordinates = np.empty_like(unrolled_nodes_positiion, dtype=mesh.geometry.x.dtype)
148 send_coordinates[unrolled_nodes_positiion] = mesh.geometry.x[:num_owned_nodes, :].reshape(-1)
150 # Send and recieve geometry sizes
151 send_coordinate_sizes = (send_nodes_per_proc * 3).astype(np.int32)
152 recv_coordinate_sizes = np.zeros_like(node_sources, dtype=np.int32)
153 geometry_to_owner_comm.Neighbor_alltoall(send_coordinate_sizes, recv_coordinate_sizes)
155 # Send node coordinates
156 recv_coordinates = np.empty(recv_coordinate_sizes.sum(), dtype=mesh.geometry.x.dtype)
157 mpi_type = numpy_to_mpi[recv_coordinates.dtype.type]
158 send_coord_msg = [send_coordinates, send_coordinate_sizes, mpi_type]
159 recv_coord_msg = [recv_coordinates, recv_coordinate_sizes, mpi_type]
160 geometry_to_owner_comm.Neighbor_alltoallv(send_coord_msg, recv_coord_msg)
161 del send_coord_msg, recv_coord_msg
163 # Send node ordering for reordering the coordinates on output process
164 send_nodes = np.empty(num_owned_nodes, dtype=np.int64)
165 send_nodes[send_nodes_position] = original_node_index[:num_owned_nodes]
167 recv_indices = np.empty(recv_coordinate_sizes.sum() // 3, dtype=np.int64)
168 send_nodes_msg = [send_nodes, send_nodes_per_proc, MPI.INT64_T]
169 recv_nodes_msg = [recv_indices, recv_coordinate_sizes // 3, MPI.INT64_T]
170 geometry_to_owner_comm.Neighbor_alltoallv(send_nodes_msg, recv_nodes_msg)
172 # Compute local ording of received nodes
173 local_node_range = compute_local_range(mesh.comm, num_nodes_global)
174 recv_indices -= local_node_range[0]
176 # Sort geometry based on input index and strip to gdim
177 gdim = mesh.geometry.dim
178 recv_nodes = recv_coordinates.reshape(-1, 3)
179 _geometry = np.empty(recv_nodes.shape, dtype=mesh.geometry.x.dtype)
180 _geometry[recv_indices, :] = recv_nodes
181 geometry = _geometry[:, :gdim].copy()
182 del _geometry, recv_nodes
184 assert local_node_range[1] - local_node_range[0] == geometry.shape[0]
185 cmap = compat.cmap(mesh)
187 cell_to_output_comm.Free()
188 geometry_to_owner_comm.Free()
190 # NOTE: Could in theory store partitioning information, but would not work nicely
191 # as one would need to read this data rather than the xdmffile.
192 # NOTE: Local geometry type hint skip is only required on DOLFINX<0.10 where
193 # proper `dolfinx.mesh.Geometry` wrapper doesn't exist
194 return MeshData(
195 local_geometry=geometry, # type: ignore[arg-type]
196 local_geometry_pos=local_node_range,
197 num_nodes_global=num_nodes_global,
198 local_topology=sorted_recv_dofmap,
199 local_topology_pos=local_cell_range,
200 num_cells_global=num_cells_global,
201 cell_type=mesh.topology.cell_name(),
202 degree=cmap.degree,
203 lagrange_variant=cmap.variant,
204 store_partition=False,
205 partition_processes=None,
206 ownership_array=None,
207 ownership_offset=None,
208 partition_range=None,
209 partition_global=None,
210 )
213def create_function_data_on_original_mesh(
214 u: dolfinx.fem.Function, name: typing.Optional[str] = None
215) -> FunctionData:
216 """
217 Create data object to save with ADIOS2
218 """
219 mesh = u.function_space.mesh
221 # Compute what cells owned by current process should be sent to what output process
222 # FIXME: Cache this
223 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
224 original_cell_index = mesh.topology.original_cell_index[:num_owned_cells]
226 # Compute owner of cells on this process based on the original cell index
227 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
228 output_cell_owner = index_owner(mesh.comm, original_cell_index, num_cells_global)
229 local_cell_range = compute_local_range(mesh.comm, num_cells_global)
231 # Compute outgoing edges from current process to outputting process
232 # Computes the number of cells sent to each process at the same time
233 cell_destinations, _send_cells_per_proc = np.unique(output_cell_owner, return_counts=True)
234 send_cells_per_proc = _send_cells_per_proc.astype(np.int32)
235 del _send_cells_per_proc
236 assert isinstance(mesh.comm, MPI.Intracomm)
237 cell_to_output_comm = mesh.comm.Create_dist_graph(
238 [mesh.comm.rank],
239 [len(cell_destinations)],
240 cell_destinations.tolist(),
241 reorder=False,
242 )
243 cell_sources, cell_dests, _ = cell_to_output_comm.Get_dist_neighbors()
244 assert np.allclose(cell_dests, cell_destinations)
246 # Compute number of recieving cells
247 recv_cells_per_proc = np.zeros_like(cell_sources, dtype=np.int32)
248 send_cells_per_proc = send_cells_per_proc.astype(np.int32)
249 cell_to_output_comm.Neighbor_alltoall(send_cells_per_proc, recv_cells_per_proc)
250 assert recv_cells_per_proc.sum() == local_cell_range[1] - local_cell_range[0]
252 # Pack and send cell indices (used for mapping topology dofmap later)
253 cell_insert_position = compute_insert_position(
254 output_cell_owner, cell_destinations, send_cells_per_proc
255 )
256 send_cells = np.empty_like(cell_insert_position, dtype=np.int64)
257 send_cells[cell_insert_position] = original_cell_index
258 recv_cells = np.empty(recv_cells_per_proc.sum(), dtype=np.int64)
259 send_cells_msg = [send_cells, send_cells_per_proc, MPI.INT64_T]
260 recv_cells_msg = [recv_cells, recv_cells_per_proc, MPI.INT64_T]
261 cell_to_output_comm.Neighbor_alltoallv(send_cells_msg, recv_cells_msg)
262 del send_cells_msg, recv_cells_msg
264 # Map received cells to the local index
265 local_cell_index = recv_cells - local_cell_range[0]
267 # Pack and send cell permutation info
268 mesh.topology.create_entity_permutations()
269 cell_permutation_info = mesh.topology.get_cell_permutation_info()[:num_owned_cells]
270 send_perm = np.empty_like(send_cells, dtype=np.uint32)
271 send_perm[cell_insert_position] = cell_permutation_info
272 recv_perm = np.empty_like(recv_cells, dtype=np.uint32)
273 send_perm_msg = [send_perm, send_cells_per_proc, MPI.UINT32_T]
274 recv_perm_msg = [recv_perm, recv_cells_per_proc, MPI.UINT32_T]
275 cell_to_output_comm.Neighbor_alltoallv(send_perm_msg, recv_perm_msg)
276 cell_permutation_info = np.empty_like(recv_perm)
277 cell_permutation_info[local_cell_index] = recv_perm
279 # 2. Extract function data (array is the same, keeping global indices from DOLFINx)
280 # Dofmap is moved by the original cell index similar to the mesh geometry dofmap
281 dofmap = u.function_space.dofmap
282 dmap = dofmap.list
283 num_dofs_per_cell = dmap.shape[1]
284 dofmap_bs = dofmap.bs
285 index_map_bs = dofmap.index_map_bs
287 # Unroll dofmap for block size
288 unrolled_dofmap = unroll_dofmap(dofmap.list[:num_owned_cells, :], dofmap_bs)
289 dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1)
290 dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1)
292 # Convert imap index to global index
293 imap_global = dofmap.index_map.local_to_global(dmap_loc)
294 dofmap_global = (imap_global * index_map_bs + dmap_rem).reshape(unrolled_dofmap.shape)
295 assert len(dofmap_global.shape) >= 2
296 num_dofs_per_cell = dofmap_global.shape[1]
297 dofmap_insert_position = unroll_insert_position(cell_insert_position, num_dofs_per_cell)
299 # Create and send array for global dofmap
300 send_function_dofmap = np.empty(len(dofmap_insert_position), dtype=np.int64)
301 send_function_dofmap[dofmap_insert_position] = dofmap_global.reshape(-1)
302 send_sizes_dofmap = send_cells_per_proc * num_dofs_per_cell
303 recv_size_dofmap = recv_cells_per_proc * num_dofs_per_cell
304 recv_function_dofmap = np.empty(recv_size_dofmap.sum(), dtype=np.int64)
305 cell_to_output_comm.Neighbor_alltoallv(
306 [send_function_dofmap, send_sizes_dofmap, MPI.INT64_T],
307 [recv_function_dofmap, recv_size_dofmap, MPI.INT64_T],
308 )
310 shaped_dofmap = recv_function_dofmap.reshape(
311 local_cell_range[1] - local_cell_range[0], num_dofs_per_cell
312 ).copy()
313 _final_dofmap = np.empty_like(shaped_dofmap)
314 _final_dofmap[local_cell_index] = shaped_dofmap
315 final_dofmap = _final_dofmap.reshape(-1)
317 # Get offsets of dofmap
318 num_cells_local = local_cell_range[1] - local_cell_range[0]
319 num_dofs_local_dmap = num_cells_local * num_dofs_per_cell
320 dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap)
321 local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64)
322 local_dofmap_offsets[:] *= num_dofs_per_cell
323 local_dofmap_offsets[:] += dofmap_imap.local_range[0]
325 num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs
326 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs
327 local_range = np.asarray(dofmap.index_map.local_range, dtype=np.int64) * dofmap.index_map_bs
328 func_name = name if name is not None else u.name
329 cell_to_output_comm.Free()
330 return FunctionData(
331 cell_permutations=cell_permutation_info,
332 local_cell_range=local_cell_range,
333 num_cells_global=num_cells_global,
334 dofmap_array=final_dofmap,
335 dofmap_offsets=local_dofmap_offsets,
336 values=u.x.array[:num_dofs_local].copy(),
337 dof_range=local_range,
338 num_dofs_global=num_dofs_global,
339 dofmap_range=dofmap_imap.local_range,
340 global_dofs_in_dofmap=dofmap_imap.size_global,
341 name=func_name,
342 )
345def write_function_on_input_mesh(
346 filename: Path | str,
347 u: dolfinx.fem.Function,
348 time: float = 0.0,
349 name: typing.Optional[str] = None,
350 mode: FileMode = FileMode.append,
351 backend_args: dict[str, typing.Any] | None = None,
352 backend: str | None = None,
353):
354 """
355 Write function checkpoint (to be read with the input mesh).
357 Note:
358 Requires backend to implement {py:class}`io4dolfinx.backends.write_function`.
360 Args:
361 filename: The filename to write to
362 u: The function to checkpoint
363 time: Time-stamp associated with function at current write step
364 mode: The mode to use (write or append)
365 name: Name of function. If None, the name of the function is used.
366 backend_args: Arguments to backend
367 backend: Choice of backend module
368 """
369 logger.debug(
370 f"Writing function on input mesh to {filename} at time {time} with name {name or u.name}"
371 )
372 logger.debug(f"Using backend {backend} with arguments {backend_args} and mode {mode}")
373 mesh = u.function_space.mesh
374 function_data = create_function_data_on_original_mesh(u, name)
375 fname = Path(filename)
377 backend_cls = get_backend(backend)
378 backend_args = backend_cls.get_default_backend_args(backend_args)
379 backend_cls.write_function(
380 fname,
381 mesh.comm,
382 function_data,
383 time=time,
384 mode=mode,
385 backend_args=backend_args,
386 )
389def write_mesh_input_order(
390 filename: Path | str,
391 mesh: dolfinx.mesh.Mesh,
392 time: float = 0.0,
393 mode: FileMode = FileMode.write,
394 backend: str | None = None,
395 backend_args: dict[str, typing.Any] | None = None,
396):
397 """
398 Write mesh to checkpoint file in original input ordering.
400 Note:
401 Requires backend to implement {py:class}`io4dolfinx.backends.write_mesh`.
403 Args:
404 filename: The filename to write to
405 mesh: Mesh to checkpoint
406 time: Time-stamp associated with function at current write step
407 mode: The mode to use (write or append)
408 name: Name of function. If None, the name of the function is used.
409 backend_args: Arguments to backend
410 backend: Choice of backend module
411 """
412 logger.debug(f"Writing mesh in input order to {filename} at time {time}")
413 logger.debug(f"Using backend {backend} with arguments {backend_args} and mode {mode}")
414 mesh_data = create_original_mesh_data(mesh)
415 fname = Path(filename)
417 backend_cls = get_backend(backend)
418 backend_args = backend_cls.get_default_backend_args(backend_args)
419 backend_cls.write_mesh(
420 fname,
421 mesh.comm,
422 mesh_data,
423 backend_args=backend_args,
424 mode=mode,
425 time=time,
426 )