Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/comm_helpers.py: 100%
112 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 18:20 +0000
1"""
2Helpers for sending and receiving values for checkpointing
3"""
5from __future__ import annotations
7from mpi4py import MPI
9import numpy as np
10import numpy.typing as npt
12from .utils import compute_insert_position, compute_local_range, valid_function_types
14__all__ = [
15 "send_dofmap_and_recv_values",
16 "send_and_recv_cell_perm",
17 "send_dofs_and_recv_values",
18 "numpy_to_mpi",
19]
21numpy_to_mpi = {
22 np.float64: MPI.DOUBLE,
23 np.float32: MPI.FLOAT,
24 np.complex64: MPI.COMPLEX,
25 np.complex128: MPI.DOUBLE_COMPLEX,
26 np.int64: MPI.INT64_T,
27 np.int32: MPI.INT32_T,
28}
31def send_dofmap_and_recv_values(
32 comm: MPI.Comm,
33 source_ranks: npt.NDArray[np.int32],
34 dest_ranks: npt.NDArray[np.int32],
35 output_owners: npt.NDArray[np.int32],
36 dest_size: npt.NDArray[np.int32],
37 input_cells: npt.NDArray[np.int64],
38 dofmap_pos: npt.NDArray[np.int32],
39 num_cells_global: int | np.int64,
40 values: npt.NDArray[valid_function_types],
41 dofmap_offsets: npt.NDArray[np.int32],
42) -> npt.NDArray[valid_function_types]:
43 """
44 Given a set of positions in input dofmap, give the global input index of this dofmap entry
45 in input file.
47 Args:
48 comm: The MPI communicator to create the Neighbourhood-communicator from
49 source_ranks: Ranks that will send dofmap indices to current process
50 dest_ranks: Ranks that will receive dofmap indices from current process
51 output_owners: The owners of each dofmap entry on this process. The unique set of
52 these entries should be the same as the dest_ranks.
53 dest_size: The number of entries sent to each owner
54 input_cells: A cell associated with the degree of freedom sent (global index).
55 dofmap_pos: The local position in the dofmap. I.e.
56 `dof = dofmap.links(input_cells)[dofmap_pos]`
57 num_cells_global: Number of global cells
58 values: Values currently held by this process. These are
59 ordered (num_cells_local, num_dofs_per_cell), flattened row-major.
60 dofmap_offsets: Local dofmap offsets to access the correct `values`.
62 Returns:
63 Values corresponding to the dofs owned by this process.
64 """
65 insert_position = compute_insert_position(output_owners, dest_ranks, dest_size)
67 # Pack the cells and dofmap position for all dofs this process is distributing
68 out_cells = np.zeros(len(output_owners), dtype=np.int64)
69 out_cells[insert_position] = input_cells
70 out_pos = np.zeros(len(output_owners), dtype=np.int32)
71 out_pos[insert_position] = dofmap_pos
73 # Compute map from the data index sent to each process and the local
74 # number on the current process
75 proc_to_dof = np.zeros_like(input_cells, dtype=np.int32)
76 proc_to_dof[insert_position] = np.arange(len(input_cells), dtype=np.int32)
77 del insert_position
79 # Send sizes to create data structures for receiving from NeighAlltoAllv
80 recv_size = np.zeros(len(source_ranks), dtype=np.int32)
81 assert isinstance(comm, MPI.Intracomm)
82 mesh_to_data_comm = comm.Create_dist_graph_adjacent(
83 source_ranks.tolist(), dest_ranks.tolist(), reorder=False
84 )
85 mesh_to_data_comm.Neighbor_alltoall(dest_size, recv_size)
87 # Prepare data-structures for receiving
88 total_incoming = sum(recv_size)
89 inc_cells = np.zeros(total_incoming, dtype=np.int64)
90 inc_pos = np.zeros(total_incoming, dtype=np.intc)
92 # Compute incoming offset
93 inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc)
94 inc_offsets[1:] = np.cumsum(recv_size)
96 # Send data
97 s_msg = [out_cells, dest_size, MPI.INT64_T]
98 r_msg = [inc_cells, recv_size, MPI.INT64_T]
99 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
101 s_msg = [out_pos, dest_size, MPI.INT32_T]
102 r_msg = [inc_pos, recv_size, MPI.INT32_T]
103 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
104 mesh_to_data_comm.Free()
106 local_input_range = compute_local_range(comm, num_cells_global)
107 values_to_distribute = np.zeros_like(inc_pos, dtype=values.dtype)
109 # Map values based on input cells and dofmap
110 local_cells = inc_cells - local_input_range[0]
111 values_to_distribute = values[dofmap_offsets[local_cells] + inc_pos]
113 # Send input dofs back to owning process
114 data_to_mesh_comm = comm.Create_dist_graph_adjacent(
115 dest_ranks.tolist(), source_ranks.tolist(), reorder=False
116 )
118 incoming_global_dofs = np.zeros(sum(dest_size), dtype=values.dtype)
119 s_msg = [values_to_distribute, recv_size, numpy_to_mpi[values.dtype.type]]
120 r_msg = [incoming_global_dofs, dest_size, numpy_to_mpi[values.dtype.type]]
121 data_to_mesh_comm.Neighbor_alltoallv(s_msg, r_msg)
123 # Sort incoming global dofs as they were inputted
124 assert len(incoming_global_dofs) == len(input_cells)
125 sorted_global_dofs = np.zeros_like(incoming_global_dofs, dtype=values.dtype)
126 sorted_global_dofs[proc_to_dof] = incoming_global_dofs
128 data_to_mesh_comm.Free()
129 return sorted_global_dofs
132def send_and_recv_cell_perm(
133 cells: npt.NDArray[np.int64],
134 perms: npt.NDArray[np.uint32],
135 cell_owners: npt.NDArray[np.int32],
136 comm: MPI.Comm,
137) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.uint32]]:
138 """
139 Send global cell index and permutation to corresponding entry in `dest_ranks`.
141 Args:
142 cells: The global input index of the cell
143 perms: The corresponding cell permutation of the cell
144 cell_owners: The rank to send the i-th entry of cells and perms to
145 comm: Rank of comm to generate neighbourhood communicator from
146 """
147 dest_ranks, _dest_size = np.unique(cell_owners, return_counts=True)
148 dest_size = _dest_size.astype(np.int32)
149 del _dest_size
150 assert isinstance(comm, MPI.Intracomm)
151 mesh_to_data = comm.Create_dist_graph(
152 [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False
153 )
154 source, dest, _ = mesh_to_data.Get_dist_neighbors()
155 assert np.allclose(dest, dest_ranks)
156 insert_position = compute_insert_position(cell_owners, dest_ranks.astype(np.int32), dest_size)
158 # Pack cells and permutations for sending
159 out_cells = np.zeros_like(cells, dtype=np.int64)
160 out_perm = np.zeros_like(perms, dtype=np.uint32)
161 out_cells[insert_position] = cells
162 out_perm[insert_position] = perms
163 del insert_position
165 # Send sizes to create data structures for receiving from NeighAlltoAllv
166 recv_size = np.zeros_like(source, dtype=np.int32)
167 mesh_to_data.Neighbor_alltoall(dest_size, recv_size)
169 # Prepare data-structures for receiving
170 total_incoming = sum(recv_size)
171 inc_cells = np.zeros(total_incoming, dtype=np.int64)
172 inc_perm = np.zeros(total_incoming, dtype=np.uint32)
174 # Compute incoming offset
175 inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc)
176 inc_offsets[1:] = np.cumsum(recv_size)
178 # Send data
179 s_msg = [out_cells, dest_size, MPI.INT64_T]
180 r_msg = [inc_cells, recv_size, MPI.INT64_T]
181 mesh_to_data.Neighbor_alltoallv(s_msg, r_msg)
183 s_msg = [out_perm, dest_size, MPI.UINT32_T]
184 r_msg = [inc_perm, recv_size, MPI.UINT32_T]
185 mesh_to_data.Neighbor_alltoallv(s_msg, r_msg)
186 mesh_to_data.Free()
187 return inc_cells, inc_perm
190def send_dofs_and_recv_values(
191 input_dofmap: npt.NDArray[np.int64],
192 dofmap_owners: npt.NDArray[np.int32],
193 comm: MPI.Comm,
194 input_array: npt.NDArray[valid_function_types],
195 array_start: int,
196):
197 """
198 Send a set of dofs (global index) to the process holding the DOF values to retrieve them.
200 Args:
201 input_dofmap: List of dofs (global index) that this process wants values for
202 dofmap_owners: The process currently holding the values this process want to get.
203 comm: MPI communicator
204 input_array: Values for dofs
205 array_start: The global starting index of `input_array`.
206 """
207 dest_ranks, _dest_size = np.unique(dofmap_owners, return_counts=True)
208 dest_size = _dest_size.astype(np.int32)
209 del _dest_size
211 assert isinstance(comm, MPI.Intracomm)
212 dofmap_to_values = comm.Create_dist_graph(
213 [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False
214 )
216 source, dest, _ = dofmap_to_values.Get_dist_neighbors()
217 assert np.allclose(dest_ranks, dest)
218 # Compute amount of data to send to each process
220 insert_position = compute_insert_position(dofmap_owners, dest_ranks, dest_size)
222 # Pack dofs for sending
223 out_dofs = np.zeros(len(dofmap_owners), dtype=np.int64)
224 out_dofs[insert_position] = input_dofmap
226 # Compute map from the data index sent to each process and the local number on
227 # the current process
228 proc_to_local = np.zeros_like(input_dofmap, dtype=np.int32)
229 proc_to_local[insert_position] = np.arange(len(input_dofmap), dtype=np.int32)
230 del insert_position
232 # Send sizes to create data structures for receiving from NeighAlltoAllv
233 recv_size = np.zeros_like(source, dtype=np.int32)
234 recv_size.resize(max(len(recv_size), 1)) # Minimal resize to work with ompi
235 dest_size.resize(max(len(dest_size), 1)) # Mininal resize to work with ompi
236 dofmap_to_values.Neighbor_alltoall(dest_size, recv_size)
237 dest_size.resize(len(dest))
238 recv_size.resize(len(source))
240 # Send input dofs to processes holding input array
241 inc_dofs = np.zeros(sum(recv_size), dtype=np.int64)
242 s_msg = [out_dofs, dest_size, MPI.INT64_T]
243 r_msg = [inc_dofs, recv_size, MPI.INT64_T]
244 dofmap_to_values.Neighbor_alltoallv(s_msg, r_msg)
245 dofmap_to_values.Free()
247 # Send back appropriate input values
248 if len(input_array) > 0:
249 sending_values = input_array[inc_dofs - array_start]
250 else:
251 sending_values = np.zeros(0, dtype=input_array.dtype)
253 values_to_dofmap = comm.Create_dist_graph_adjacent(dest, source, reorder=False)
254 inc_values = np.zeros_like(out_dofs, dtype=input_array.dtype)
255 s_msg_rev = [sending_values, recv_size, numpy_to_mpi[input_array.dtype.type]]
256 r_msg_rev = [inc_values, dest_size, numpy_to_mpi[input_array.dtype.type]]
257 values_to_dofmap.Neighbor_alltoallv(s_msg_rev, r_msg_rev)
258 values_to_dofmap.Free()
260 # Sort inputs according to local dof number (input process)
261 values = np.empty_like(inc_values, dtype=input_array.dtype)
262 values[proc_to_local] = inc_values
263 return values