Coverage for  / dolfinx-env / lib / python3.12 / site-packages / io4dolfinx / backends / adios2 / helpers.py: 82%

1""" 

2Helpers reading/writing data with ADIOS2 

3""" 

4 

5from __future__ import annotations 

6 

7import shutil 

8from contextlib import contextmanager 

9from pathlib import Path 

10from typing import NamedTuple 

11 

12from mpi4py import MPI 

13 

14import adios2 

15import dolfinx.cpp.graph 

16import dolfinx.graph 

17import numpy as np 

18import numpy.typing as npt 

19 

20from io4dolfinx.utils import compute_local_range, valid_function_types 

21 

22 

23def resolve_adios_scope(adios2): 

24 scope = adios2.bindings if hasattr(adios2, "bindings") else adios2 

25 if not scope.is_built_with_mpi: 

26 raise ImportError("ADIOS2 must be built with MPI support") 

27 return scope 

28 

29 

30adios2 = resolve_adios_scope(adios2) 

31 

32 

33__all__ = [ 

34 "AdiosFile", 

35 "ADIOSFile", 

36 "check_variable_exists", 

37 "read_array", 

38 "read_adjacency_list", 

39 "adios_to_numpy_dtype", 

40] 

41 

42adios_to_numpy_dtype = { 

43 "float": np.float32, 

44 "double": np.float64, 

45 "float complex": np.complex64, 

46 "double complex": np.complex128, 

47 "uint32_t": np.uint32, 

48} 

49 

50 

51class AdiosFile(NamedTuple): 

52 io: adios2.IO 

53 file: adios2.Engine 

54 

55 

56@contextmanager 

57def ADIOSFile( 

58 adios: adios2.ADIOS, 

59 filename: Path | str, 

60 engine: str, 

61 mode: adios2.Mode, 

62 io_name: str, 

63 comm: MPI.Intracomm | None = None, 

64): 

65 io = adios.DeclareIO(io_name) 

66 io.SetEngine(engine) 

67 # ADIOS2 sometimes struggles with existing files/folders it should overwrite 

68 if mode == adios2.Mode.Write: 

69 filename = Path(filename) 

70 if filename.exists() and comm is not None and comm.rank == 0: 

71 if filename.is_dir(): 

72 shutil.rmtree(filename) 

73 else: 

74 filename.unlink() 

75 if comm is not None: 

76 comm.Barrier() 

77 

78 file = io.Open(str(filename), mode) 

79 try: 

80 yield AdiosFile(io=io, file=file) 

81 finally: 

82 file.Close() 

83 adios.RemoveIO(io_name) 

84 

85 

86def check_variable_exists( 

87 adios: adios2.ADIOS, 

88 filename: Path | str, 

89 variable: str, 

90 engine: str, 

91) -> bool: 

92 io_name = f"{variable}_reader" 

93 

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

95 return False 

96 

97 variable_found = False 

98 with ADIOSFile( 

99 adios=adios, 

100 engine=engine, 

101 filename=filename, 

102 mode=adios2.Mode.Read, 

103 io_name=io_name, 

104 ) as adios_file: 

105 # Find step that has cell permutation 

106 for _ in range(adios_file.file.Steps()): 

107 adios_file.file.BeginStep() 

108 if variable in adios_file.io.AvailableVariables().keys(): 

109 variable_found = True 

110 break 

111 adios_file.file.EndStep() 

112 

113 # Not sure if this is needed, but just in case 

114 if variable in adios_file.io.AvailableVariables().keys(): 

115 variable_found = True 

116 return variable_found 

117 

118 

119def read_adjacency_list( 

120 adios: adios2.ADIOS, 

121 comm: MPI.Intracomm, 

122 filename: Path | str, 

123 data_name: str, 

124 offsets_name: str, 

125 engine: str, 

126) -> dolfinx.graph.AdjacencyList: 

127 """ 

128 Read an adjacency-list from an ADIOS file with given communicator. 

129 The adjancency list is split in to a flat array (data) and its corresponding offset. 

130 

131 Args: 

132 adios: The ADIOS instance 

133 comm: The MPI communicator used to read the data 

134 filename: Path to input file 

135 data_name: Name of variable containing the indices of the adjacencylist 

136 dofmap_offsets: Name of variable containing offsets of the adjacencylist 

137 engine: Type of ADIOS engine to use for reading data 

138 

139 Returns: 

140 The local part of dofmap from input dofs 

141 

142 .. note:: 

143 No MPI communication is done during this call 

144 """ 

145 

146 # Open ADIOS engine 

147 io_name = f"{data_name=}_reader" 

148 

149 with ADIOSFile( 

150 adios=adios, 

151 engine=engine, 

152 filename=filename, 

153 mode=adios2.Mode.Read, 

154 io_name=io_name, 

155 ) as adios_file: 

156 # First find step with dofmap offsets, to be able to read 

157 # in a full row of the dofmap 

158 for _ in range(adios_file.file.Steps()): 

159 adios_file.file.BeginStep() 

160 if offsets_name in adios_file.io.AvailableVariables().keys(): 

161 break 

162 adios_file.file.EndStep() 

163 if offsets_name not in adios_file.io.AvailableVariables().keys(): 

164 raise KeyError(f"Dof offsets not found at '{offsets_name}' in {filename}") 

165 

166 # Get global shape of dofmap-offset, and read in data with an overlap 

167 d_offsets = adios_file.io.InquireVariable(offsets_name) 

168 shape = d_offsets.Shape() 

169 num_nodes = shape[0] - 1 

170 local_range = compute_local_range(comm, num_nodes) 

171 

172 # As the offsets are one longer than the number of cells, we need to read in with an overlap 

173 if len(shape) == 1: 

174 d_offsets.SetSelection([[local_range[0]], [local_range[1] + 1 - local_range[0]]]) 

175 in_offsets = np.empty( 

176 local_range[1] + 1 - local_range[0], 

177 dtype=d_offsets.Type().strip("_t"), 

178 ) 

179 else: 

180 d_offsets.SetSelection( 

181 [ 

182 [local_range[0], 0], 

183 [local_range[1] + 1 - local_range[0], shape[1]], 

184 ] 

185 ) 

186 in_offsets = np.empty( 

187 (local_range[1] + 1 - local_range[0], shape[1]), 

188 dtype=d_offsets.Type().strip("_t"), 

189 ) 

190 

191 adios_file.file.Get(d_offsets, in_offsets, adios2.Mode.Sync) 

192 in_offsets = in_offsets.squeeze() 

193 

194 # Assuming dofmap is saved in stame step 

195 # Get the relevant part of the dofmap 

196 if data_name not in adios_file.io.AvailableVariables().keys(): 

197 raise KeyError(f"Dofs not found at {data_name} in {filename}") 

198 cell_dofs = adios_file.io.InquireVariable(data_name) 

199 if len(shape) == 1: 

200 cell_dofs.SetSelection([[in_offsets[0]], [in_offsets[-1] - in_offsets[0]]]) 

201 in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t")) 

202 else: 

203 cell_dofs.SetSelection([[in_offsets[0], 0], [in_offsets[-1] - in_offsets[0], shape[1]]]) 

204 in_dofmap = np.empty( 

205 (in_offsets[-1] - in_offsets[0], shape[1]), 

206 dtype=cell_dofs.Type().strip("_t"), 

207 ) 

208 assert shape[1] == 1 

209 

210 in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t")) 

211 adios_file.file.Get(cell_dofs, in_dofmap, adios2.Mode.Sync) 

212 in_offsets -= in_offsets[0] 

213 adios_file.file.EndStep() 

214 

215 # Return local dofmap 

216 return dolfinx.graph.adjacencylist(in_dofmap, in_offsets.astype(np.int32)) 

217 

218 

219def read_array( 

220 adios: adios2.ADIOS, 

221 filename: Path | str, 

222 array_name: str, 

223 engine: str, 

224 comm: MPI.Intracomm, 

225 time: float = 0.0, 

226 time_name: str = "", 

227 legacy: bool = False, 

228) -> tuple[npt.NDArray[valid_function_types], int]: 

229 """ 

230 Read an array from file, return the global starting position of the local array 

231 

232 Args: 

233 adios: The ADIOS instance 

234 filename: Path to file to read array from 

235 array_name: Name of array in file 

236 engine: Name of engine to use to read file 

237 comm: MPI communicator used for reading the data 

238 time_name: Name of time variable for modern checkpoints 

239 legacy: If True ignore time_name and read the first available step 

240 Returns: 

241 Local part of array and its global starting position 

242 """ 

243 

244 with ADIOSFile( 

245 adios=adios, 

246 engine=engine, 

247 filename=filename, 

248 mode=adios2.Mode.Read, 

249 io_name="ArrayReader", 

250 ) as adios_file: 

251 # Get time-stamp from first available step 

252 if legacy: 

253 for i in range(adios_file.file.Steps()): 

254 adios_file.file.BeginStep() 

255 if array_name in adios_file.io.AvailableVariables().keys(): 

256 break 

257 adios_file.file.EndStep() 

258 if array_name not in adios_file.io.AvailableVariables().keys(): 

259 raise KeyError(f"No array found at {array_name}") 

260 else: 

261 for i in range(adios_file.file.Steps()): 

262 adios_file.file.BeginStep() 

263 if time_name in adios_file.io.AvailableVariables().keys(): 

264 arr = adios_file.io.InquireVariable(time_name) 

265 time_shape = arr.Shape() 

266 arr.SetSelection([[0], [time_shape[0]]]) 

267 times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()]) 

268 adios_file.file.Get(arr, times, adios2.Mode.Sync) 

269 if times[0] == time: 

270 break 

271 if i == adios_file.file.Steps() - 1: 

272 raise KeyError( 

273 f"No data associated with {time_name}={time} found in {filename}" 

274 ) 

275 

276 adios_file.file.EndStep() 

277 

278 if time_name not in adios_file.io.AvailableVariables().keys(): 

279 raise KeyError(f"No data associated with {time_name}={time} found in {filename}") 

280 

281 if array_name not in adios_file.io.AvailableVariables().keys(): 

282 raise KeyError(f"No array found at {time=} for {array_name}") 

283 

284 arr = adios_file.io.InquireVariable(array_name) 

285 arr_shape = arr.Shape() 

286 # TODO: Should we always pick the first element? 

287 assert len(arr_shape) >= 1 

288 arr_range = compute_local_range(comm, arr_shape[0]) 

289 

290 if len(arr_shape) == 1: 

291 arr.SetSelection([[arr_range[0]], [arr_range[1] - arr_range[0]]]) 

292 vals = np.empty(arr_range[1] - arr_range[0], dtype=adios_to_numpy_dtype[arr.Type()]) 

293 else: 

294 arr.SetSelection([[arr_range[0], 0], [arr_range[1] - arr_range[0], arr_shape[1]]]) 

295 vals = np.empty( 

296 (arr_range[1] - arr_range[0], arr_shape[1]), 

297 dtype=adios_to_numpy_dtype[arr.Type()], 

298 ) 

299 assert arr_shape[1] == 1 

300 

301 adios_file.file.Get(arr, vals, adios2.Mode.Sync) 

302 adios_file.file.EndStep() 

303 

304 return vals.reshape(-1), arr_range[0]