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

16import numpy as np 

17import numpy.typing as npt 

18 

19from ...utils import compute_local_range, valid_function_types 

20 

21 

22def resolve_adios_scope(adios2): 

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

24 if not scope.is_built_with_mpi: 

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

26 return scope 

27 

28 

29adios2 = resolve_adios_scope(adios2) 

30 

31 

32__all__ = [ 

33 "AdiosFile", 

34 "ADIOSFile", 

35 "check_variable_exists", 

36 "read_array", 

37 "read_adjacency_list", 

38 "adios_to_numpy_dtype", 

39] 

40 

41adios_to_numpy_dtype = { 

42 "float": np.float32, 

43 "double": np.float64, 

44 "float complex": np.complex64, 

45 "double complex": np.complex128, 

46 "uint32_t": np.uint32, 

47} 

48 

49 

50class AdiosFile(NamedTuple): 

51 io: adios2.IO 

52 file: adios2.Engine 

53 

54 

55@contextmanager 

56def ADIOSFile( 

57 adios: adios2.ADIOS, 

58 filename: Path | str, 

59 engine: str, 

60 mode: adios2.Mode, 

61 io_name: str, 

62 comm: MPI.Comm | None = None, 

63): 

64 io = adios.DeclareIO(io_name) 

65 io.SetEngine(engine) 

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

67 if mode == adios2.Mode.Write: 

68 filename = Path(filename) 

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

70 if filename.is_dir(): 

71 shutil.rmtree(filename) 

72 else: 

73 filename.unlink() 

74 if comm is not None: 

75 comm.Barrier() 

76 

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

78 try: 

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

80 finally: 

81 file.Close() 

82 adios.RemoveIO(io_name) 

83 

84 

85def check_variable_exists( 

86 adios: adios2.ADIOS, 

87 filename: Path | str, 

88 variable: str, 

89 engine: str, 

90) -> bool: 

91 io_name = f"{variable}_reader" 

92 

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

94 return False 

95 

96 variable_found = False 

97 with ADIOSFile( 

98 adios=adios, 

99 engine=engine, 

100 filename=filename, 

101 mode=adios2.Mode.Read, 

102 io_name=io_name, 

103 ) as adios_file: 

104 # Find step that has cell permutation 

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

106 adios_file.file.BeginStep() 

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

108 variable_found = True 

109 break 

110 adios_file.file.EndStep() 

111 

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

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

114 variable_found = True 

115 return variable_found 

116 

117 

118def read_adjacency_list( 

119 adios: adios2.ADIOS, 

120 comm: MPI.Comm, 

121 filename: Path | str, 

122 data_name: str, 

123 offsets_name: str, 

124 engine: str, 

125) -> dolfinx.graph.AdjacencyList: 

126 """ 

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

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

129 

130 Args: 

131 adios: The ADIOS instance 

132 comm: The MPI communicator used to read the data 

133 filename: Path to input file 

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

135 dofmap_offsets: Name of variable containing offsets of the adjacencylist 

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

137 

138 Returns: 

139 The local part of dofmap from input dofs 

140 

141 .. note:: 

142 No MPI communication is done during this call 

143 """ 

144 

145 # Open ADIOS engine 

146 io_name = f"{data_name=}_reader" 

147 

148 with ADIOSFile( 

149 adios=adios, 

150 engine=engine, 

151 filename=filename, 

152 mode=adios2.Mode.Read, 

153 io_name=io_name, 

154 ) as adios_file: 

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

156 # in a full row of the dofmap 

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

158 adios_file.file.BeginStep() 

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

160 break 

161 adios_file.file.EndStep() 

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

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

164 

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

166 d_offsets = adios_file.io.InquireVariable(offsets_name) 

167 shape = d_offsets.Shape() 

168 num_nodes = shape[0] - 1 

169 local_range = compute_local_range(comm, num_nodes) 

170 

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

172 if len(shape) == 1: 

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

174 in_offsets = np.empty( 

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

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

177 ) 

178 else: 

179 d_offsets.SetSelection( 

180 [ 

181 [local_range[0], 0], 

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

183 ] 

184 ) 

185 in_offsets = np.empty( 

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

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

188 ) 

189 

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

191 in_offsets = in_offsets.squeeze() 

192 

193 # Assuming dofmap is saved in stame step 

194 # Get the relevant part of the dofmap 

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

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

197 cell_dofs = adios_file.io.InquireVariable(data_name) 

198 if len(shape) == 1: 

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

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

201 else: 

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

203 in_dofmap = np.empty( 

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

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

206 ) 

207 assert shape[1] == 1 

208 

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

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

211 in_offsets -= in_offsets[0] 

212 adios_file.file.EndStep() 

213 

214 # Return local dofmap 

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

216 

217 

218def read_array( 

219 adios: adios2.ADIOS, 

220 filename: Path | str, 

221 array_name: str, 

222 engine: str, 

223 comm: MPI.Comm, 

224 time: float = 0.0, 

225 time_name: str = "", 

226 legacy: bool = False, 

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

228 """ 

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

230 

231 Args: 

232 adios: The ADIOS instance 

233 filename: Path to file to read array from 

234 array_name: Name of array in file 

235 engine: Name of engine to use to read file 

236 comm: MPI communicator used for reading the data 

237 time_name: Name of time variable for modern checkpoints 

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

239 Returns: 

240 Local part of array and its global starting position 

241 """ 

242 

243 with ADIOSFile( 

244 adios=adios, 

245 engine=engine, 

246 filename=filename, 

247 mode=adios2.Mode.Read, 

248 io_name="ArrayReader", 

249 ) as adios_file: 

250 # Get time-stamp from first available step 

251 if legacy: 

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

253 adios_file.file.BeginStep() 

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

255 break 

256 adios_file.file.EndStep() 

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

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

259 else: 

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

261 adios_file.file.BeginStep() 

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

263 arr = adios_file.io.InquireVariable(time_name) 

264 time_shape = arr.Shape() 

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

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

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

268 if times[0] == time: 

269 break 

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

271 raise KeyError( 

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

273 ) 

274 

275 adios_file.file.EndStep() 

276 

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

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

279 

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

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

282 

283 arr = adios_file.io.InquireVariable(array_name) 

284 arr_shape = arr.Shape() 

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

286 assert len(arr_shape) >= 1 

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

288 

289 if len(arr_shape) == 1: 

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

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

292 else: 

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

294 vals = np.empty( 

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

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

297 ) 

298 assert arr_shape[1] == 1 

299 

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

301 adios_file.file.EndStep() 

302 

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