Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/utils.py: 75%

1# Copyright (C) 2023 Jørgen Schartum Dokken 

2# 

3# This file is part of io4dolfinx 

4# 

5# SPDX-License-Identifier: MIT 

6 

7""" 

8Vectorized numpy operations used internally in io4dolfinx 

9""" 

10 

11from __future__ import annotations 

12 

13import warnings 

14from pathlib import Path 

15 

16from mpi4py import MPI 

17 

18import basix.ufl 

19import dolfinx 

20import numpy as np 

21import numpy.typing as npt 

22import ufl 

23from packaging.version import Version 

24 

25__all__ = [ 

26 "check_file_exists", 

27 "compute_local_range", 

28 "index_owner", 

29 "compute_dofmap_pos", 

30 "unroll_dofmap", 

31 "compute_insert_position", 

32 "unroll_insert_position", 

33 "reconstruct_mesh", 

34] 

35 

36 

37def check_file_exists(filename: Path | str): 

38 """Check if file exists.""" 

39 if not Path(filename).exists(): 

40 raise FileNotFoundError(f"{filename} not found") 

41 

42 

43valid_function_types = np.float32 | np.float64 | np.complex64 | np.complex128 

44valid_real_types = np.float32 | np.float64 

45 

46 

47def element_signature(V): 

48 if Version(dolfinx.__version__) > Version("0.9.0"): 

49 return V.element.signature 

50 else: 

51 return V.element.signature() 

52 

53 

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. 

63 

64 Example: 

65 

66 .. highlight:: python 

67 .. code-block:: python 

68 

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) 

73 

74 Insert position is then ``[1, 4, 5, 2, 0, 3]`` 

75 """ 

76 process_pos_indicator = data_owner.reshape(-1, 1) == destination_ranks 

77 

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) 

82 

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 

87 

88 # Add process offset for each local index 

89 insert_position += send_offsets[proc_col] 

90 return insert_position 

91 

92 

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 

98 

99 Example: 

100 

101 

102 .. highlight:: python 

103 .. code-block:: python 

104 

105 insert_position = [1, 4, 5, 2, 0, 3] 

106 unrolled_ip = unroll_insert_position(insert_position, 3) 

107 

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 

113 

114 

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`. 

119 

120 NOTE: If N is not divisible by the number of ranks, the first `r` 

121 processes gets an extra value 

122 

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] 

134 

135 

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. 

142 

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 

150 

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 

156 

157 

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 

168 

169 

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. 

176 

177 :param V: The function space 

178 :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the 

179 number of owned dofs (unrolled for block size) 

180 """ 

181 dofs = V.dofmap.list 

182 mesh = V.mesh 

183 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local 

184 dofmap_bs = V.dofmap.bs 

185 num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs 

186 

187 local_cell = np.empty( 

188 num_owned_dofs, dtype=np.int32 

189 ) # Local cell index for each dof owned by process 

190 dof_pos = np.empty(num_owned_dofs, dtype=np.int32) # Position in dofmap for said dof 

191 

192 unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs) 

193 markers = unrolled_dofmap < num_owned_dofs 

194 local_indices = np.broadcast_to(np.arange(markers.shape[1]), markers.shape) 

195 cell_indicator = np.broadcast_to( 

196 np.arange(num_owned_cells, dtype=np.int32).reshape(-1, 1), 

197 (num_owned_cells, markers.shape[1]), 

198 ) 

199 indicator = unrolled_dofmap[markers].reshape(-1) 

200 local_cell[indicator] = cell_indicator[markers].reshape(-1) 

201 dof_pos[indicator] = local_indices[markers].reshape(-1) 

202 return local_cell, dof_pos 

203 

204 

205def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh: 

206 """ 

207 Make a copy of a mesh and potentially change the element of the coordinate element. 

208 

209 Note: 

210 The topology is shared with the original mesh but the geometry is reconstructed. 

211 

212 Args: 

213 mesh: Mesh to reconstruct 

214 coordinate_element_degree: Degree to use for coordinate element 

215 

216 Returns: 

217 The new mesh 

218 

219 """ 

220 warnings.warn( 

221 "reconstruct_mesh is deprecated and will be removed in a future release. " 

222 + "Use dolfinx.fem.interpolate_geometry instead, available from DOLFINx>=0.11.", 

223 category=DeprecationWarning, 

224 stacklevel=2, 

225 ) 

226 if not hasattr(dolfinx.fem, "interpolate_geometry"): 

227 # Extract cell properties 

228 ud = mesh.ufl_domain() 

229 assert ud is not None 

230 c_el = ud.ufl_coordinate_element() 

231 family = c_el.family_name 

232 lvar = c_el.lagrange_variant 

233 ct = c_el.cell_type 

234 

235 # Create new UFL element 

236 new_c_el = basix.ufl.element( 

237 family, 

238 ct, 

239 coordinate_element_degree, 

240 shape=(mesh.geometry.dim,), 

241 lagrange_variant=lvar, 

242 dtype=mesh.geometry.x.dtype, 

243 ) 

244 

245 # Extract new node coordinates 

246 V_tmp = dolfinx.fem.functionspace(mesh, new_c_el) 

247 gdim = mesh.geometry.dim 

248 x = V_tmp.tabulate_dof_coordinates()[:, :gdim] 

249 

250 # Create new geoemtry 

251 geom_imap = V_tmp.dofmap.index_map 

252 geom_dofmap = V_tmp.dofmap.list 

253 num_nodes_local = geom_imap.size_local + geom_imap.num_ghosts 

254 original_input_indices = geom_imap.local_to_global( 

255 np.arange(num_nodes_local, dtype=np.int32) 

256 ) 

257 coordinate_element = dolfinx.fem.coordinate_element( 

258 mesh.topology.cell_type, coordinate_element_degree, lvar, dtype=mesh.geometry.x.dtype 

259 ) 

260 # Could use create_geometry here when things are fixed 

261 geom = dolfinx.mesh.Geometry( 

262 type(mesh.geometry._cpp_object)( 

263 geom_imap, 

264 geom_dofmap, 

265 coordinate_element._cpp_object, # type: ignore[arg-type] 

266 x, 

267 original_input_indices, 

268 ) 

269 ) 

270 

271 # Create new mesh 

272 new_top = mesh.topology 

273 cpp_mesh = type(mesh._cpp_object)(mesh.comm, new_top._cpp_object, geom._cpp_object) # type: ignore[arg-type] 

274 return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(new_c_el)) 

275 else: 

276 # Use the new interpolate_geometry function 

277 cmap = dolfinx.fem.coordinate_element( 

278 mesh.topology.cell_type, 

279 coordinate_element_degree, 

280 dtype=mesh.geometry.x.dtype, 

281 variant=mesh.geometry.cmap.variant, 

282 ) 

283 return dolfinx.fem.interpolate_geometry(mesh, cmap)