Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/utils.py: 99%
76 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-23 13:30 +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 "codim_names",
33]
35#: Codimension of the entities named ``"cell"``, ``"facet"``, ``"ridge"`` and ``"peak"``,
36#: as in :func:`dolfinx.io.gmsh.model_to_mesh`
37codim_names = {"cell": 0, "facet": 1, "ridge": 2, "peak": 3}
40def check_file_exists(filename: Path | str):
41 """Check if file exists."""
42 if not Path(filename).exists():
43 raise FileNotFoundError(f"{filename} not found")
46valid_function_types = np.float32 | np.float64 | np.complex64 | np.complex128
47valid_real_types = np.float32 | np.float64
50def element_signature(V):
51 return V.element.signature
54def compute_insert_position(
55 data_owner: npt.NDArray[np.int32],
56 destination_ranks: npt.NDArray[np.int32],
57 out_size: npt.NDArray[np.int32],
58) -> npt.NDArray[np.int32]:
59 """
60 Giving a list of ranks, compute the local insert position for each rank in a list
61 sorted by destination ranks. This function is used for packing data from a
62 given process to its destination processes.
64 Example:
66 .. highlight:: python
67 .. code-block:: python
69 data_owner = [0, 1, 1, 0, 2, 3]
70 destination_ranks = [2,0,3,1]
71 out_size = [1, 2, 1, 2]
72 insert_position = compute_insert_position(data_owner, destination_ranks, out_size)
74 Insert position is then ``[1, 4, 5, 2, 0, 3]``
75 """
76 process_pos_indicator = data_owner.reshape(-1, 1) == destination_ranks
78 # Compute offsets for insertion based on input size
79 send_offsets = np.zeros(len(out_size) + 1, dtype=np.intc)
80 send_offsets[1:] = np.cumsum(out_size)
81 assert send_offsets[-1] == len(data_owner)
83 # Compute local insert index on each process
84 proc_row, proc_col = np.nonzero(process_pos_indicator)
85 cum_pos = np.cumsum(process_pos_indicator, axis=0)
86 insert_position = cum_pos[proc_row, proc_col] - 1
88 # Add process offset for each local index
89 insert_position += send_offsets[proc_col]
90 return insert_position
93def unroll_insert_position(
94 insert_position: npt.NDArray[np.int32], block_size: int
95) -> npt.NDArray[np.int32]:
96 """
97 Unroll insert position by a block size
99 Example:
102 .. highlight:: python
103 .. code-block:: python
105 insert_position = [1, 4, 5, 2, 0, 3]
106 unrolled_ip = unroll_insert_position(insert_position, 3)
108 where ``unrolled_ip = [3, 4 ,5, 12, 13, 14, 15, 16, 17, 6, 7, 8, 0, 1, 2, 9, 10, 11]``
109 """
110 unrolled_ip = np.repeat(insert_position, block_size) * block_size
111 unrolled_ip += np.tile(np.arange(block_size), len(insert_position))
112 return unrolled_ip
115def compute_local_range(comm: MPI.Comm, N: int | np.int64):
116 """
117 Divide a set of `N` objects into `M` partitions, where `M` is
118 the size of the MPI communicator `comm`.
120 NOTE: If N is not divisible by the number of ranks, the first `r`
121 processes gets an extra value
123 Returns the local range of values
124 """
125 rank = comm.rank
126 size = comm.size
127 n = N // size
128 r = N % size
129 # First r processes has one extra value
130 if rank < r:
131 return [rank * (n + 1), (rank + 1) * (n + 1)]
132 else:
133 return [rank * n + r, (rank + 1) * n + r]
136def index_owner(
137 comm: MPI.Comm, indices: npt.NDArray[np.int64], N: int | np.int64
138) -> npt.NDArray[np.int32]:
139 """
140 Find which rank (local to comm) which owns an `index`, given that
141 data of size `N` has been split equally among the ranks.
143 NOTE: If `N` is not divisible by the number of ranks, the first `r`
144 processes gets an extra value.
145 """
146 size = comm.size
147 assert (indices < N).all()
148 n = N // size
149 r = N % size
151 owner = np.empty_like(indices, dtype=np.int32)
152 inc_remainder = indices < (n + 1) * r
153 owner[inc_remainder] = indices[inc_remainder] // (n + 1)
154 owner[~inc_remainder] = r + (indices[~inc_remainder] - r * (n + 1)) // n
155 return owner
158def unroll_dofmap(dofs: npt.NDArray[np.int32], bs: int) -> npt.NDArray[np.int32]:
159 """
160 Given a two-dimensional dofmap of size `(num_cells, num_dofs_per_cell)`
161 Expand the dofmap by its block size such that the resulting array
162 is of size `(num_cells, bs*num_dofs_per_cell)`
163 """
164 num_cells, num_dofs_per_cell = dofs.shape
165 unrolled_dofmap = np.repeat(dofs, bs).reshape(num_cells, num_dofs_per_cell * bs) * bs
166 unrolled_dofmap += np.tile(np.arange(bs), num_dofs_per_cell)
167 return unrolled_dofmap
170def compute_dofmap_pos(
171 V: dolfinx.fem.FunctionSpace,
172) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
173 """
174 Compute a map from each owned dof in the dofmap to a single cell owned by the
175 process, and the relative position of the dof.
177 .. note::
178 Every owned dof must be reachable from a cell owned by the process. This
179 holds for a mesh created by :func:`dolfinx.mesh.create_mesh`, but not
180 necessarily for one created by :func:`dolfinx.mesh.create_submesh`, which
181 builds its vertex map allowing owner changes: a rank may own a vertex --
182 and therefore a dof -- that is only incident to cells it ghosts. Such
183 dofs cannot be assigned a position and raise rather than silently
184 yielding an arbitrary one.
186 :param V: The function space
187 :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the
188 number of owned dofs (unrolled for block size)
189 :raises RuntimeError: If an owned dof is not contained in any owned cell.
190 """
191 dofs = V.dofmap.list
192 mesh = V.mesh
193 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
194 dofmap_bs = V.dofmap.bs
195 num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
197 # Initialize to -1 rather than leaving uninitialized, so that a dof that is
198 # never visited below is detectable instead of holding arbitrary memory.
199 local_cell = np.full(
200 num_owned_dofs, -1, dtype=np.int32
201 ) # Local cell index for each dof owned by process
202 dof_pos = np.full(num_owned_dofs, -1, dtype=np.int32) # Position in dofmap for said dof
204 unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs)
205 markers = unrolled_dofmap < num_owned_dofs
206 local_indices = np.broadcast_to(np.arange(markers.shape[1]), markers.shape)
207 cell_indicator = np.broadcast_to(
208 np.arange(num_owned_cells, dtype=np.int32).reshape(-1, 1),
209 (num_owned_cells, markers.shape[1]),
210 )
211 indicator = unrolled_dofmap[markers].reshape(-1)
212 local_cell[indicator] = cell_indicator[markers].reshape(-1)
213 dof_pos[indicator] = local_indices[markers].reshape(-1)
215 num_unreachable = int(np.count_nonzero(local_cell < 0))
216 if num_unreachable > 0:
217 raise RuntimeError(
218 f"{num_unreachable} of {num_owned_dofs} owned dofs are not contained in any"
219 " cell owned by this process, so their position in the dofmap is undefined."
220 " This happens on meshes created with `dolfinx.mesh.create_submesh`, which"
221 " may move vertex ownership to a process that does not own an incident cell."
222 " Read the function on a standalone mesh (see `io4dolfinx.read_submesh`)"
223 " rather than on a mesh derived with `create_submesh`."
224 )
225 return local_cell, dof_pos
228def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh:
229 """
230 Make a copy of a mesh and potentially change the element of the coordinate element.
232 Note:
233 The topology is shared with the original mesh but the geometry is reconstructed.
235 Args:
236 mesh: Mesh to reconstruct
237 coordinate_element_degree: Degree to use for coordinate element
239 Returns:
240 The new mesh
242 """
243 warnings.warn(
244 "reconstruct_mesh is deprecated and will be removed in a future release. "
245 + "Use dolfinx.fem.interpolate_geometry instead, available from DOLFINx>=0.11, "
246 + "or fenicsx_compat.reconstruct_mesh.",
247 category=DeprecationWarning,
248 stacklevel=2,
249 )
250 return fenicsx_compat.reconstruct_mesh(mesh, coordinate_element_degree)