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