Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/utils.py: 99%
75 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:46 +0000
1# Copyright (C) 2023 Jørgen Schartum Dokken
2#
3# This file is part of io4dolfinx
4#
5# SPDX-License-Identifier: MIT
7"""
8Vectorized numpy operations used internally in io4dolfinx
9"""
11from __future__ import annotations
13import warnings
14from pathlib import Path
16from mpi4py import MPI
18import dolfinx
19import fenicsx_compat
20import numpy as np
21import numpy.typing as npt
23__all__ = [
24 "check_file_exists",
25 "compute_local_range",
26 "index_owner",
27 "compute_dofmap_pos",
28 "unroll_dofmap",
29 "compute_insert_position",
30 "unroll_insert_position",
31 "reconstruct_mesh",
32]
35def check_file_exists(filename: Path | str):
36 """Check if file exists."""
37 if not Path(filename).exists():
38 raise FileNotFoundError(f"{filename} not found")
41valid_function_types = np.float32 | np.float64 | np.complex64 | np.complex128
42valid_real_types = np.float32 | np.float64
45def element_signature(V):
46 return V.element.signature
49def compute_insert_position(
50 data_owner: npt.NDArray[np.int32],
51 destination_ranks: npt.NDArray[np.int32],
52 out_size: npt.NDArray[np.int32],
53) -> npt.NDArray[np.int32]:
54 """
55 Giving a list of ranks, compute the local insert position for each rank in a list
56 sorted by destination ranks. This function is used for packing data from a
57 given process to its destination processes.
59 Example:
61 .. highlight:: python
62 .. code-block:: python
64 data_owner = [0, 1, 1, 0, 2, 3]
65 destination_ranks = [2,0,3,1]
66 out_size = [1, 2, 1, 2]
67 insert_position = compute_insert_position(data_owner, destination_ranks, out_size)
69 Insert position is then ``[1, 4, 5, 2, 0, 3]``
70 """
71 process_pos_indicator = data_owner.reshape(-1, 1) == destination_ranks
73 # Compute offsets for insertion based on input size
74 send_offsets = np.zeros(len(out_size) + 1, dtype=np.intc)
75 send_offsets[1:] = np.cumsum(out_size)
76 assert send_offsets[-1] == len(data_owner)
78 # Compute local insert index on each process
79 proc_row, proc_col = np.nonzero(process_pos_indicator)
80 cum_pos = np.cumsum(process_pos_indicator, axis=0)
81 insert_position = cum_pos[proc_row, proc_col] - 1
83 # Add process offset for each local index
84 insert_position += send_offsets[proc_col]
85 return insert_position
88def unroll_insert_position(
89 insert_position: npt.NDArray[np.int32], block_size: int
90) -> npt.NDArray[np.int32]:
91 """
92 Unroll insert position by a block size
94 Example:
97 .. highlight:: python
98 .. code-block:: python
100 insert_position = [1, 4, 5, 2, 0, 3]
101 unrolled_ip = unroll_insert_position(insert_position, 3)
103 where ``unrolled_ip = [3, 4 ,5, 12, 13, 14, 15, 16, 17, 6, 7, 8, 0, 1, 2, 9, 10, 11]``
104 """
105 unrolled_ip = np.repeat(insert_position, block_size) * block_size
106 unrolled_ip += np.tile(np.arange(block_size), len(insert_position))
107 return unrolled_ip
110def compute_local_range(comm: MPI.Comm, N: int | np.int64):
111 """
112 Divide a set of `N` objects into `M` partitions, where `M` is
113 the size of the MPI communicator `comm`.
115 NOTE: If N is not divisible by the number of ranks, the first `r`
116 processes gets an extra value
118 Returns the local range of values
119 """
120 rank = comm.rank
121 size = comm.size
122 n = N // size
123 r = N % size
124 # First r processes has one extra value
125 if rank < r:
126 return [rank * (n + 1), (rank + 1) * (n + 1)]
127 else:
128 return [rank * n + r, (rank + 1) * n + r]
131def index_owner(
132 comm: MPI.Comm, indices: npt.NDArray[np.int64], N: int | np.int64
133) -> npt.NDArray[np.int32]:
134 """
135 Find which rank (local to comm) which owns an `index`, given that
136 data of size `N` has been split equally among the ranks.
138 NOTE: If `N` is not divisible by the number of ranks, the first `r`
139 processes gets an extra value.
140 """
141 size = comm.size
142 assert (indices < N).all()
143 n = N // size
144 r = N % size
146 owner = np.empty_like(indices, dtype=np.int32)
147 inc_remainder = indices < (n + 1) * r
148 owner[inc_remainder] = indices[inc_remainder] // (n + 1)
149 owner[~inc_remainder] = r + (indices[~inc_remainder] - r * (n + 1)) // n
150 return owner
153def unroll_dofmap(dofs: npt.NDArray[np.int32], bs: int) -> npt.NDArray[np.int32]:
154 """
155 Given a two-dimensional dofmap of size `(num_cells, num_dofs_per_cell)`
156 Expand the dofmap by its block size such that the resulting array
157 is of size `(num_cells, bs*num_dofs_per_cell)`
158 """
159 num_cells, num_dofs_per_cell = dofs.shape
160 unrolled_dofmap = np.repeat(dofs, bs).reshape(num_cells, num_dofs_per_cell * bs) * bs
161 unrolled_dofmap += np.tile(np.arange(bs), num_dofs_per_cell)
162 return unrolled_dofmap
165def compute_dofmap_pos(
166 V: dolfinx.fem.FunctionSpace,
167) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
168 """
169 Compute a map from each owned dof in the dofmap to a single cell owned by the
170 process, and the relative position of the dof.
172 .. note::
173 Every owned dof must be reachable from a cell owned by the process. This
174 holds for a mesh created by :func:`dolfinx.mesh.create_mesh`, but not
175 necessarily for one created by :func:`dolfinx.mesh.create_submesh`, which
176 builds its vertex map allowing owner changes: a rank may own a vertex --
177 and therefore a dof -- that is only incident to cells it ghosts. Such
178 dofs cannot be assigned a position and raise rather than silently
179 yielding an arbitrary one.
181 :param V: The function space
182 :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the
183 number of owned dofs (unrolled for block size)
184 :raises RuntimeError: If an owned dof is not contained in any owned cell.
185 """
186 dofs = V.dofmap.list
187 mesh = V.mesh
188 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
189 dofmap_bs = V.dofmap.bs
190 num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
192 # Initialize to -1 rather than leaving uninitialized, so that a dof that is
193 # never visited below is detectable instead of holding arbitrary memory.
194 local_cell = np.full(
195 num_owned_dofs, -1, dtype=np.int32
196 ) # Local cell index for each dof owned by process
197 dof_pos = np.full(num_owned_dofs, -1, dtype=np.int32) # Position in dofmap for said dof
199 unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs)
200 markers = unrolled_dofmap < num_owned_dofs
201 local_indices = np.broadcast_to(np.arange(markers.shape[1]), markers.shape)
202 cell_indicator = np.broadcast_to(
203 np.arange(num_owned_cells, dtype=np.int32).reshape(-1, 1),
204 (num_owned_cells, markers.shape[1]),
205 )
206 indicator = unrolled_dofmap[markers].reshape(-1)
207 local_cell[indicator] = cell_indicator[markers].reshape(-1)
208 dof_pos[indicator] = local_indices[markers].reshape(-1)
210 num_unreachable = int(np.count_nonzero(local_cell < 0))
211 if num_unreachable > 0:
212 raise RuntimeError(
213 f"{num_unreachable} of {num_owned_dofs} owned dofs are not contained in any"
214 " cell owned by this process, so their position in the dofmap is undefined."
215 " This happens on meshes created with `dolfinx.mesh.create_submesh`, which"
216 " may move vertex ownership to a process that does not own an incident cell."
217 " Read the function on a standalone mesh (see `io4dolfinx.read_submesh`)"
218 " rather than on a mesh derived with `create_submesh`."
219 )
220 return local_cell, dof_pos
223def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh:
224 """
225 Make a copy of a mesh and potentially change the element of the coordinate element.
227 Note:
228 The topology is shared with the original mesh but the geometry is reconstructed.
230 Args:
231 mesh: Mesh to reconstruct
232 coordinate_element_degree: Degree to use for coordinate element
234 Returns:
235 The new mesh
237 """
238 warnings.warn(
239 "reconstruct_mesh is deprecated and will be removed in a future release. "
240 + "Use dolfinx.fem.interpolate_geometry instead, available from DOLFINx>=0.11, "
241 + "or fenicsx_compat.reconstruct_mesh.",
242 category=DeprecationWarning,
243 stacklevel=2,
244 )
245 return fenicsx_compat.reconstruct_mesh(mesh, coordinate_element_degree)