Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/__init__.py: 88%

65 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-21 11:46 +0000

1from enum import Enum 

2from importlib import import_module 

3from pathlib import Path 

4from typing import Any, Protocol 

5 

6from mpi4py import MPI 

7 

8import dolfinx 

9import numpy as np 

10import numpy.typing as npt 

11 

12from ..structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData 

13 

14__all__ = [ 

15 "DEFAULT_MESH_NAME", 

16 "FileMode", 

17 "IOBackend", 

18 "get_backend", 

19 "get_mesh_name", 

20 "set_default_backend", 

21] 

22 

23_DEFAULT_BACKEND = "adios2" 

24 

25#: Name used for a mesh when the caller does not supply one. A file written 

26#: before named meshes existed holds exactly this mesh, so backends must store 

27#: it exactly as they did before (see :func:`get_mesh_name`). 

28DEFAULT_MESH_NAME = "mesh" 

29 

30 

31def get_mesh_name(backend_args: dict[str, Any] | None) -> str: 

32 """Get the name of the mesh a backend call refers to. 

33 

34 A checkpoint may hold several meshes -- for instance a mesh and a submesh of 

35 it -- distinguished by name. Backends take the name from 

36 ``backend_args["name"]``. 

37 

38 Args: 

39 backend_args: Arguments to backend, or None 

40 

41 Returns: 

42 The mesh name, :data:`DEFAULT_MESH_NAME` if none was given. 

43 """ 

44 if backend_args is None: 

45 return DEFAULT_MESH_NAME 

46 name = backend_args.get("name", DEFAULT_MESH_NAME) 

47 return DEFAULT_MESH_NAME if name is None else str(name) 

48 

49 

50def set_default_backend(backend: str): 

51 """Set the global default backend for io4dolfinx.""" 

52 global _DEFAULT_BACKEND 

53 if backend not in BUILTIN_BAKENDS: 

54 try: 

55 get_backend(backend) 

56 except ImportError as e: 

57 raise ImportError(f"Backend {backend} not found.") from e 

58 _DEFAULT_BACKEND = backend 

59 

60 

61def get_default_backend(): 

62 """Get the global default backend for io4dolfinx.""" 

63 return _DEFAULT_BACKEND 

64 

65 

66class ReadMode(Enum): 

67 serial = 10 # This means that all data is read in on root rank 

68 

69 # Total number of data P, num processes = i + 1. 

70 # All processes reads at least `P // (i+1)` items 

71 # The first j=P%(i+1) processes reads `P // (i+1) + 1` items 

72 # ```python 

73 # def compute_partitioning(P, J): 

74 # min_num = P // J 

75 # num_per_proc = np.full(J, min_num) 

76 # rem = P % J 

77 # num_per_proc[:int(rem)] += 1 

78 # assert(sum(num_per_proc)) == P 

79 # return num_per_proc 

80 # ``` 

81 parallel = 20 

82 

83 

84class FileMode(Enum): 

85 """Filen mode used for opening files.""" 

86 

87 append = 10 #: Append data to file 

88 write = 20 #: Write data to file 

89 read = 30 #: Read data from file 

90 

91 

92# See https://peps.python.org/pep-0544/#modules-as-implementations-of-protocols 

93class IOBackend(Protocol): 

94 """Interface a storage backend must provide. 

95 

96 .. note:: 

97 A checkpoint may hold more than one mesh, each with its own topology, 

98 geometry, meshtags, cell permutations and functions. Every method that 

99 touches mesh-associated data takes the mesh it refers to from 

100 ``backend_args["name"]``; use :func:`get_mesh_name` to read it, so that 

101 omitting it selects :data:`DEFAULT_MESH_NAME` consistently. A backend 

102 must store that default mesh exactly as it did before named meshes 

103 existed, so that older checkpoints stay readable. 

104 """ 

105 

106 read_mode: ReadMode 

107 

108 def get_default_backend_args(self, arguments: dict[str, Any] | None) -> dict[str, Any]: 

109 """Get default backend arguments given a set of input arguments. 

110 

111 Args: 

112 arguments: Input backend arguments 

113 

114 Returns: 

115 Updated backend arguments 

116 """ 

117 

118 def write_attributes( 

119 self, 

120 filename: Path | str, 

121 comm: MPI.Comm, 

122 name: str, 

123 attributes: dict[str, np.ndarray], 

124 backend_args: dict[str, Any] | None, 

125 ): 

126 """Write attributes to file. 

127 

128 Args: 

129 filename: Path to file to write to 

130 comm: MPI communicator used in storage 

131 name: Name of the attribute group 

132 attributes: Dictionary of attributes to write 

133 backend_args: Arguments to backend 

134 """ 

135 

136 def read_attributes( 

137 self, 

138 filename: Path | str, 

139 comm: MPI.Comm, 

140 name: str, 

141 backend_args: dict[str, Any] | None, 

142 ) -> dict[str, Any]: 

143 """Read attributes from file. 

144 

145 Args: 

146 filename: Path to file to read from 

147 comm: MPI communicator used in storage 

148 name: Name of the attribute group 

149 backend_args: Arguments to backend 

150 

151 Returns: 

152 Dictionary of attributes read from file 

153 """ 

154 

155 def read_timestamps( 

156 self, 

157 filename: Path | str, 

158 comm: MPI.Comm, 

159 function_name: str, 

160 backend_args: dict[str, Any] | None, 

161 ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var] 

162 """Read timestamps from file. 

163 

164 Args: 

165 filename: Path to file to read from 

166 comm: MPI communicator used in storage 

167 function_name: Name of the function to read timestamps for 

168 backend_args: Arguments to backend 

169 

170 Returns: 

171 Numpy array of timestamps read from file 

172 """ 

173 

174 def write_mesh( 

175 self, 

176 filename: Path | str, 

177 comm: MPI.Comm, 

178 mesh: MeshData, 

179 backend_args: dict[str, Any] | None, 

180 mode: FileMode, 

181 time: float, 

182 ): 

183 """ 

184 Write a mesh to file. 

185 

186 Args: 

187 comm: MPI communicator used in storage 

188 mesh: Internal data structure for the mesh data to save to file 

189 filename: Path to file to write to 

190 backend_args: Arguments to backend 

191 mode: File-mode to store the mesh 

192 time: Time stamp associated with the mesh 

193 """ 

194 

195 def write_meshtags( 

196 self, 

197 filename: str | Path, 

198 comm: MPI.Comm, 

199 data: MeshTagsData, 

200 backend_args: dict[str, Any] | None, 

201 ): 

202 """Write mesh tags to file. 

203 

204 Args: 

205 filename: Path to file to write to 

206 comm: MPI communicator used in storage 

207 data: Internal data structure for the mesh tags to save to file 

208 backend_args: Arguments to backend 

209 """ 

210 

211 def read_mesh_data( 

212 self, 

213 filename: Path | str, 

214 comm: MPI.Comm, 

215 time: str | float | None, 

216 read_from_partition: bool, 

217 backend_args: dict[str, Any] | None, 

218 ) -> ReadMeshData: 

219 """Read mesh data from file. 

220 

221 Args: 

222 filename: Path to file to read from 

223 comm: MPI communicator used in storage 

224 time: Time stamp associated with the mesh to read 

225 read_from_partition: Whether to read partition information 

226 backend_args: Arguments to backend 

227 

228 Returns: 

229 Internal data structure for the mesh data read from file 

230 """ 

231 

232 def read_meshtags_data( 

233 self, 

234 filename: str | Path, 

235 comm: MPI.Comm, 

236 name: str, 

237 backend_args: dict[str, Any] | None, 

238 ) -> MeshTagsData: 

239 """Read mesh tags from file. 

240 

241 Args: 

242 filename: Path to file to read from 

243 comm: MPI communicator used in storage 

244 name: Name of the mesh tags to read 

245 backend_args: Arguments to backend 

246 

247 Returns: 

248 Internal data structure for the mesh tags read from file 

249 """ 

250 

251 def read_dofmap( 

252 self, 

253 filename: str | Path, 

254 comm: MPI.Comm, 

255 name: str, 

256 backend_args: dict[str, Any] | None, 

257 ) -> dolfinx.graph.AdjacencyList: 

258 """Read the dofmap of a function with a given name. 

259 

260 Args: 

261 filename: Path to file to read from 

262 comm: MPI communicator used in storage 

263 name: Name of the function to read the dofmap for 

264 backend_args: Arguments to backend 

265 

266 Returns: 

267 Dofmap as an :class:`dolfinx.graph.AdjacencyList` 

268 """ 

269 

270 def read_dofs( 

271 self, 

272 filename: str | Path, 

273 comm: MPI.Comm, 

274 name: str, 

275 time: float, 

276 backend_args: dict[str, Any] | None, 

277 ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]: 

278 """Read the dofs (values) of a function with a given name from a given timestep. 

279 

280 Args: 

281 filename: Path to file to read from 

282 comm: MPI communicator used in storage 

283 name: Name of the function to read the dofs for 

284 time: Time stamp associated with the function to read 

285 backend_args: Arguments to backend 

286 

287 Returns: 

288 Contiguous sequence of degrees of freedom (with respect to input data) 

289 and the global starting point on the process. 

290 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. 

291 """ 

292 

293 def read_cell_perms( 

294 self, comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None 

295 ) -> npt.NDArray[np.uint32]: 

296 """ 

297 Read cell permutation from file with given communicator, 

298 Split in continuous chunks based on number of cells in the input data. 

299 

300 Cell permutations belong to a mesh, not to a function, so which mesh is 

301 meant is taken from ``backend_args["name"]``. 

302 

303 Args: 

304 comm: MPI communicator used in storage 

305 filename: Path to file to read from 

306 backend_args: Arguments to backend 

307 

308 Returns: 

309 Contiguous sequence of permutations (with respect to input data) 

310 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. 

311 """ 

312 

313 def write_function( 

314 self, 

315 filename: Path, 

316 comm: MPI.Comm, 

317 u: FunctionData, 

318 time: float, 

319 mode: FileMode, 

320 backend_args: dict[str, Any] | None, 

321 ): 

322 """Write a function to file. 

323 

324 Args: 

325 comm: MPI communicator used in storage 

326 u: Internal data structure for the function data to save to file 

327 filename: Path to file to write to 

328 time: Time stamp associated with function 

329 mode: File-mode to store the function 

330 backend_args: Arguments to backend 

331 """ 

332 

333 def read_legacy_mesh( 

334 self, filename: Path | str, comm: MPI.Comm, group: str 

335 ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]: 

336 """Read in the mesh topology, geometry and (optionally) cell type from a 

337 legacy DOLFIN HDF5-file. 

338 

339 Args: 

340 filename: Path to file to read from 

341 comm: MPI communicator used in storage 

342 group: Group in HDF5 file where mesh is stored 

343 

344 Returns: 

345 Tuple containing: 

346 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices 

347 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates 

348 - Cell type as a string (e.g. "tetrahedron") or None if not found 

349 """ 

350 

351 def snapshot_checkpoint( 

352 self, 

353 filename: Path | str, 

354 mode: FileMode, 

355 u: dolfinx.fem.Function, 

356 backend_args: dict[str, Any] | None, 

357 ): 

358 """Create a snapshot checkpoint of a dolfinx function. 

359 

360 Args: 

361 filename: Path to file to read from 

362 mode: File-mode to store the function 

363 u: dolfinx function to create a snapshot checkpoint for 

364 backend_args: Arguments to backend 

365 """ 

366 

367 def read_hdf5_array( 

368 self, 

369 comm: MPI.Comm, 

370 filename: Path | str, 

371 group: str, 

372 backend_args: dict[str, Any] | None, 

373 ) -> tuple[np.ndarray, int]: 

374 """Read an array from an HDF5 file. 

375 

376 Args: 

377 comm: MPI communicator used in storage 

378 filename: Path to file to read from 

379 group: Group in HDF5 file where array is stored 

380 backend_args: Arguments to backend 

381 

382 Returns: 

383 Tuple containing: 

384 - Numpy array read from file 

385 - Global starting point on the process. 

386 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc. 

387 """ 

388 

389 def read_point_data( 

390 self, 

391 filename: Path | str, 

392 name: str, 

393 comm: MPI.Comm, 

394 time: str | float | None, 

395 backend_args: dict[str, Any] | None, 

396 ) -> tuple[np.ndarray, int]: 

397 """Read data from the nodes of a mesh. 

398 

399 Args: 

400 filename: Path to file 

401 name: Name of point data 

402 comm: Communicator to launch IO on. 

403 time: The time stamp 

404 backend_args: The backend arguments 

405 Returns: 

406 Data local to process (contiguous, no mpi comm) and local start range 

407 """ 

408 ... 

409 

410 def read_function_names( 

411 self, filename: Path | str, comm: MPI.Comm, backend_args: dict[str, Any] | None 

412 ) -> list[str]: 

413 """Read all function names from a file. 

414 

415 Args: 

416 filename: Path to file 

417 comm: MPI communicator to launch IO on. 

418 backend_args: Arguments to backend 

419 

420 Returns: 

421 A list of function names. 

422 """ 

423 ... 

424 

425 def read_cell_data( 

426 self, 

427 filename: Path | str, 

428 name: str, 

429 comm: MPI.Comm, 

430 time: str | float | None, 

431 backend_args: dict[str, Any] | None, 

432 ) -> tuple[npt.NDArray[np.int64], np.ndarray]: 

433 """Read data from the cells of a mesh. 

434 

435 Args: 

436 filename: Path to file 

437 name: Name of point data 

438 comm: Communicator to launch IO on. 

439 time: The time stamp 

440 backend_args: The backend arguments 

441 Returns: 

442 A tuple (topology, dofs) where topology contains the 

443 vertex indices of the cells, dofs the degrees of 

444 freedom within that cell. 

445 """ 

446 ... 

447 

448 def write_data( 

449 self, 

450 filename: Path | str, 

451 array_data: ArrayData, 

452 comm: MPI.Comm, 

453 time: str | float | None, 

454 mode: FileMode, 

455 backend_args: dict[str, Any] | None, 

456 ): 

457 """Write a 2D-array to file. 

458 

459 

460 Args: 

461 filename: Path to file 

462 array_data: Data to write to file. 

463 comm: The MPI communicator to open the writer with. 

464 time: The time stamp 

465 mode: Append or write 

466 backend_args: The backend arguments 

467 """ 

468 ... 

469 

470 

471def get_backend(backend: str | None = None) -> IOBackend: 

472 """Get backend class from backend name. 

473 

474 Args: 

475 backend: Name of the backend to get 

476 

477 Returns: 

478 Backend class 

479 """ 

480 if backend is None: 

481 backend = _DEFAULT_BACKEND 

482 

483 if backend == "h5py": 

484 from .h5py import backend as H5PYInterface 

485 

486 return H5PYInterface 

487 elif backend == "adios2": 

488 from .adios2 import backend as ADIOS2Interface 

489 

490 return ADIOS2Interface 

491 elif backend == "pyvista": 

492 from .pyvista import backend as PYVISTAInterface 

493 

494 return PYVISTAInterface 

495 elif backend == "xdmf": 

496 from .xdmf import backend as XDMFInterface 

497 

498 return XDMFInterface 

499 elif backend == "vtkhdf": 

500 from .vtkhdf import backend as VTKDHFInterface 

501 

502 return VTKDHFInterface 

503 elif backend == "exodus": 

504 from .exodus import backend as EXODUSInterface 

505 

506 return EXODUSInterface 

507 else: 

508 return import_module(backend) 

509 

510 

511BUILTIN_BAKENDS = ("h5py", "adios2", "pyvista", "xdmf", "vtkhdf", "exodus") 

512 

513 

514def list_builtin_backends() -> list[str]: 

515 """List available builtin backends. 

516 

517 Returns: 

518 List of available backends 

519 """ 

520 lst = [] 

521 for backend in BUILTIN_BAKENDS: 

522 try: 

523 get_backend(backend) 

524 except ImportError: 

525 continue 

526 lst.append(backend) 

527 return lst