Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/h5py/backend.py: 90%

1""" 

2H5py interface to io4dolfinx 

3 

4SPDX License identifier: MIT 

5 

6Copyright: Jørgen S. Dokken, Henrik N.T. Finsberg, Simula Research Laboratory 

7""" 

8 

9import contextlib 

10from pathlib import Path 

11from typing import Any 

12 

13from mpi4py import MPI 

14 

15import dolfinx 

16import h5py 

17import numpy as np 

18import numpy.typing as npt 

19from dolfinx.graph import adjacencylist 

20 

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

22from ...utils import check_file_exists, compute_local_range 

23from .. import DEFAULT_MESH_NAME, FileMode, ReadMode, get_mesh_name 

24 

25read_mode = ReadMode.parallel 

26 

27# try: 

28# except ModuleNotFoundError: 

29# raise ModuleNotFoundError("This backend requires h5py to be installed.") 

30 

31 

32@contextlib.contextmanager 

33def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None): 

34 """Context manager for opening an HDF5 file with h5py. 

35 

36 Args: 

37 h5name: The name of the HDF5 file. 

38 filemode: The file mode. 

39 force_serial: Force serial access to the file. 

40 comm: The MPI communicator 

41 

42 """ 

43 if comm is None: 

44 comm = MPI.COMM_WORLD 

45 

46 if h5py.h5.get_config().mpi and not force_serial: 

47 h5file = h5py.File(h5name, filemode, driver="mpio", comm=comm) 

48 else: 

49 if comm.size > 1 and not force_serial: 

50 raise ValueError( 

51 f"h5py is not installed with MPI support, while using {comm.size} processes.", 

52 "If you really want to do this, turn on the `force_serial` flag.", 

53 ) 

54 h5file = h5py.File(h5name, filemode) 

55 

56 try: 

57 yield h5file 

58 finally: 

59 if h5file.id: 

60 h5file.close() 

61 

62 

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

64 args = dict(arguments) if arguments else {} 

65 args.setdefault("legacy", False) # If meshtags is read from legacy 

66 args.setdefault("name", DEFAULT_MESH_NAME) # Which mesh in the file to act on 

67 return args 

68 

69 

70def convert_file_mode(mode: FileMode) -> str: 

71 match mode: 

72 case FileMode.append: 

73 return "a" 

74 case FileMode.read: 

75 return "r" 

76 case FileMode.write: 

77 return "w" 

78 case _: 

79 raise NotImplementedError(f"File mode {mode} not implemented") 

80 

81 

82def write_attributes( 

83 filename: Path | str, 

84 comm: MPI.Comm, 

85 name: str, 

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

87 backend_args: dict[str, Any] | None = None, 

88): 

89 """Write attributes to file using H5PY. 

90 

91 Args: 

92 filename: Path to file to write to 

93 comm: MPI communicator used in storage 

94 name: Name of the attributes 

95 attributes: Dictionary of attributes to write to file 

96 engine: ADIOS2 engine to use 

97 """ 

98 filemode = "a" if Path(filename).exists() else "w" 

99 with h5pyfile(filename, filemode=filemode, comm=comm, force_serial=False) as h5file: 

100 if name not in h5file.keys(): 

101 h5file.create_group(name, track_order=True) 

102 group = h5file[name] 

103 

104 for key, val in attributes.items(): 

105 group.attrs[key] = val 

106 

107 

108def read_attributes( 

109 filename: Path | str, 

110 comm: MPI.Comm, 

111 name: str, 

112 backend_args: dict[str, Any] | None = None, 

113) -> dict[str, Any]: 

114 """Read attributes from file using H5PY. 

115 

116 Args: 

117 filename: Path to file to read from 

118 comm: MPI communicator used in storage 

119 name: Name of the attributes 

120 Returns: 

121 The attributes 

122 """ 

123 check_file_exists(filename) 

124 output_attrs = {} 

125 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

126 for key, val in h5file[name].attrs.items(): 

127 output_attrs[key] = val 

128 return output_attrs 

129 

130 

131def read_timestamps( 

132 filename: Path | str, 

133 comm: MPI.Comm, 

134 function_name: str, 

135 backend_args: dict[str, Any] | None = None, 

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

137 """Read time-stamps from a checkpoint file. 

138 

139 Args: 

140 comm: MPI communicator 

141 filename: Path to file 

142 function_name: Name of the function to read time-stamps for 

143 backend_args: Arguments for backend, for instance file type. 

144 

145 Returns: 

146 The time-stamps 

147 """ 

148 check_file_exists(filename) 

149 mesh_name = get_mesh_name(backend_args) 

150 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

151 mesh_directory = h5file[mesh_name] 

152 functions = mesh_directory["functions"] 

153 u = functions[function_name] 

154 timestamps = u.attrs["timestamps"] 

155 return timestamps 

156 

157 

158def write_mesh( 

159 filename: Path | str, 

160 comm: MPI.Comm, 

161 mesh: MeshData, 

162 backend_args: dict[str, Any] | None = None, 

163 mode: FileMode = FileMode.write, 

164 time: float = 0.0, 

165): 

166 """Write a mesh to file using H5PY 

167 

168 Args: 

169 comm: MPI communicator used in storage 

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

171 filename: Path to file to write to. 

172 mode: Mode to use (write or append) 

173 time: Time stamp 

174 """ 

175 backend_args = get_default_backend_args(backend_args) 

176 h5_mode = convert_file_mode(mode) 

177 mesh_name = get_mesh_name(backend_args) 

178 with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: 

179 if mesh_name in h5file.keys() and h5_mode == "a": 

180 mesh_directory = h5file[mesh_name] 

181 timestamps = mesh_directory.attrs["timestamps"] 

182 if np.isclose(time, timestamps).any(): 

183 raise ValueError("Mesh has already been stored at time={time_stamp}.") 

184 else: 

185 mesh_directory.attrs["timestamps"] = np.append( 

186 mesh_directory.attrs["timestamps"], time 

187 ) 

188 idx = len(mesh_directory.attrs["timestamps"]) - 1 

189 write_topology = False 

190 else: 

191 mesh_directory = h5file.create_group(mesh_name) 

192 mesh_directory.attrs["timestamps"] = np.array([time], dtype=np.float64) 

193 idx = 0 

194 write_topology = True 

195 

196 geometry_group = mesh_directory.create_group(f"{idx}") 

197 

198 # Write geometry data 

199 gdim = mesh.local_geometry.shape[1] 

200 geometry_dataset = geometry_group.create_dataset( 

201 "Points", 

202 [mesh.num_nodes_global, gdim], 

203 dtype=mesh.local_geometry.dtype, 

204 ) 

205 geometry_dataset[slice(*mesh.local_geometry_pos), :] = mesh.local_geometry 

206 

207 # Write static partitioning data 

208 if "PartitioningData" not in mesh_directory.keys() and mesh.store_partition: 

209 assert mesh.partition_range is not None 

210 assert mesh.ownership_array is not None 

211 par_dataset = mesh_directory.create_dataset( 

212 "PartitioningData", [mesh.partition_global], dtype=mesh.ownership_array.dtype 

213 ) 

214 par_dataset[slice(*mesh.partition_range)] = mesh.ownership_array 

215 

216 if "PartitioningOffset" not in mesh_directory.keys() and mesh.store_partition: 

217 assert mesh.ownership_offset is not None 

218 par_dataset = mesh_directory.create_dataset( 

219 "PartitioningOffset", [mesh.num_cells_global + 1], dtype=np.int64 

220 ) 

221 par_dataset[mesh.local_topology_pos[0] : mesh.local_topology_pos[1] + 1] = ( 

222 mesh.ownership_offset 

223 ) 

224 

225 if "PartitionProcesses" not in mesh_directory.attrs.keys() and mesh.store_partition: 

226 mesh_directory.attrs["PartitionProcesses"] = mesh.partition_processes 

227 

228 # Write static data 

229 if write_topology: 

230 mesh_directory.attrs["CellType"] = mesh.cell_type 

231 mesh_directory.attrs["Degree"] = mesh.degree 

232 mesh_directory.attrs["LagrangeVariant"] = mesh.lagrange_variant 

233 num_dofs_per_cell = mesh.local_topology.shape[1] 

234 topology_dataset = mesh_directory.create_dataset( 

235 "Topology", [mesh.num_cells_global, num_dofs_per_cell], dtype=np.int64 

236 ) 

237 topology_dataset[slice(*mesh.local_topology_pos), :] = mesh.local_topology 

238 

239 

240def read_mesh_data( 

241 filename: Path | str, 

242 comm: MPI.Comm, 

243 time: str | float | None = 0.0, 

244 read_from_partition: bool = False, 

245 backend_args: dict[str, Any] | None = None, 

246) -> ReadMeshData: 

247 """Read mesh data from h5py based checkpoint files. 

248 

249 Args: 

250 filename: Path to input file 

251 comm: The MPI communciator to distribute the mesh over 

252 time: Time stamp associated with mesh 

253 read_from_partition: Read mesh with partition from file 

254 Returns: 

255 The mesh topology, geometry, UFL domain and partition function 

256 """ 

257 

258 backend_args = get_default_backend_args(backend_args) 

259 mesh_name = get_mesh_name(backend_args) 

260 

261 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

262 if mesh_name not in h5file.keys(): 

263 raise KeyError(f"Could not find mesh '{mesh_name}' in {filename}") 

264 mesh_group = h5file[mesh_name] 

265 timestamps = mesh_group.attrs["timestamps"] 

266 assert time is not None 

267 parent_group = np.flatnonzero(np.isclose(timestamps, time)) 

268 if len(parent_group) != 1: 

269 raise RuntimeError( 

270 f"Time step {time} not found in file, available steps is {timestamps}" 

271 ) 

272 time_group = f"{parent_group[0]:d}" 

273 

274 # Get mesh topology (distributed) 

275 topology = mesh_group["Topology"] 

276 local_range = compute_local_range(comm, topology.shape[0]) 

277 mesh_topology = topology[slice(*local_range), :] 

278 

279 cell_type = mesh_group.attrs["CellType"] 

280 lvar = mesh_group.attrs["LagrangeVariant"] 

281 degree = mesh_group.attrs["Degree"] 

282 

283 geometry_group = mesh_group[time_group] 

284 geometry_dataset = geometry_group["Points"] 

285 x_shape = geometry_dataset.shape 

286 

287 geometry_range = compute_local_range(comm, x_shape[0]) 

288 mesh_geometry = geometry_dataset[slice(*geometry_range), :] 

289 

290 # Check validity of partitioning information 

291 if read_from_partition: 

292 if "PartitionProcesses" not in mesh_group.attrs.keys(): 

293 raise KeyError(f"Partitioning information not found in {filename}") 

294 par_keys = ("PartitioningData", "PartitioningOffset") 

295 for key in par_keys: 

296 if key not in mesh_group.keys(): 

297 raise KeyError(f"Partitioning information not found in {filename}") 

298 

299 par_num_procs = mesh_group.attrs["PartitionProcesses"] 

300 if par_num_procs != comm.size: 

301 raise ValueError(f"Number of processes in file ({par_num_procs})!=({comm.size=})") 

302 

303 # First read in offsets based on the number of cells [0, num_cells_local] 

304 par_offsets = mesh_group["PartitioningOffset"][local_range[0] : local_range[1] + 1] 

305 # Then read the data based of offsets 

306 par_data = mesh_group["PartitioningData"][par_offsets[0] : par_offsets[-1]] 

307 # Then make offsets local 

308 par_offsets[:] -= par_offsets[0] 

309 partition_graph = adjacencylist(par_data, par_offsets.astype(np.int32)) 

310 else: 

311 partition_graph = None 

312 

313 return ReadMeshData( 

314 cells=mesh_topology, 

315 cell_type=cell_type, 

316 x=mesh_geometry, 

317 degree=degree, 

318 lvar=lvar, 

319 partition_graph=partition_graph, 

320 ) 

321 

322 

323def write_meshtags( 

324 filename: str | Path, 

325 comm: MPI.Comm, 

326 data: MeshTagsData, 

327 backend_args: dict[str, Any] | None = None, 

328): 

329 """Write mesh tags to file. 

330 

331 Args: 

332 filename: Path to file to write to 

333 comm: MPI communicator used in storage 

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

335 backend_args: Arguments to backend 

336 """ 

337 backend_args = get_default_backend_args(backend_args) 

338 

339 mesh_name = get_mesh_name(backend_args) 

340 with h5pyfile(filename, filemode="a", comm=comm, force_serial=False) as h5file: 

341 if mesh_name not in h5file.keys(): 

342 raise KeyError(f"Could not find mesh '{mesh_name}' in file") 

343 mesh_group = h5file[mesh_name] 

344 if "tags" not in mesh_group.keys(): 

345 tags = mesh_group.create_group("tags") 

346 else: 

347 tags = mesh_group["tags"] 

348 if data.name in tags.keys(): 

349 raise KeyError(f"MeshTags with {data.name=} already exists in this file") 

350 tag = tags.create_group(data.name) 

351 

352 # Add topology 

353 topology = tag.create_dataset( 

354 "Topology", shape=[data.num_entities_global, data.num_dofs_per_entity], dtype=np.int64 

355 ) 

356 assert data.local_start is not None 

357 topology[data.local_start : data.local_start + len(data.indices), :] = data.indices 

358 

359 # Add cell_type attribute 

360 tag.attrs["CellType"] = data.cell_type 

361 

362 # Add values 

363 values = tag.create_dataset( 

364 "Values", shape=[data.num_entities_global], dtype=data.values.dtype 

365 ) 

366 values[data.local_start : data.local_start + len(data.indices)] = data.values 

367 

368 # Add dimension 

369 tag.attrs["dim"] = data.dim 

370 

371 

372def read_meshtags_data( 

373 filename: str | Path, comm: MPI.Comm, name: str, backend_args: dict[str, Any] | None = None 

374) -> MeshTagsData: 

375 """Read mesh tags from file. 

376 

377 Args: 

378 filename: Path to file to read from 

379 comm: MPI communicator used in storage 

380 name: Name of the mesh tags to read 

381 backend_args: Arguments to backend. If "legacy_dolfin" is supplied as argument 

382 the HDF5 file is assumed to have been made with DOLFIN 

383 

384 Returns: 

385 Internal data structure for the mesh tags read from file 

386 """ 

387 backend_args = get_default_backend_args(backend_args) 

388 legacy = backend_args["legacy"] 

389 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

390 if legacy: 

391 if name not in h5file.keys(): 

392 raise RuntimeError(f"MeshTag {name} not found in {filename}.") 

393 mesh = h5file[name] 

394 topology = mesh["topology"] 

395 cell_type = topology.attrs["celltype"] 

396 if isinstance(cell_type, np.bytes_): 

397 cell_type = cell_type.decode("utf-8") 

398 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type)) 

399 values = mesh["values"] 

400 else: 

401 mesh_name = get_mesh_name(backend_args) 

402 if mesh_name not in h5file.keys(): 

403 raise KeyError(f"No mesh '{mesh_name}' found") 

404 mesh = h5file[mesh_name] 

405 if "tags" not in mesh.keys(): 

406 raise KeyError("Could not find 'tags' in file, are you sure this is a checkpoint?") 

407 tags = mesh["tags"] 

408 if name not in tags.keys(): 

409 raise KeyError(f"Could not find {name} in '/{mesh_name}/tags/' in {filename}") 

410 tag = tags[name] 

411 

412 dim = tag.attrs["dim"] 

413 topology = tag["Topology"] 

414 values = tag["Values"] 

415 num_entities_global = topology.shape[0] 

416 topology_range = compute_local_range(comm, num_entities_global) 

417 indices = topology[slice(*topology_range), :].astype(np.int64) 

418 values_dtype = np.dtype(backend_args.get("values_dtype", np.int32)) 

419 vals = values[slice(*topology_range)].astype(values_dtype) 

420 return MeshTagsData(name=name, values=vals, indices=indices, dim=dim) 

421 

422 

423def read_dofmap( 

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

425) -> dolfinx.graph.AdjacencyList: 

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

427 

428 Args: 

429 filename: Path to file to read from 

430 comm: MPI communicator used in storage 

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

432 backend_args: Arguments to backend 

433 

434 Returns: 

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

436 """ 

437 backend_args = {} if backend_args is None else backend_args 

438 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

439 # If dofmap is read with full path, it is passed through backend_args 

440 dofmap_key = backend_args.get("dofmap", None) 

441 if dofmap_key is None: 

442 mesh_name = get_mesh_name(backend_args) 

443 if mesh_name not in h5file.keys(): 

444 raise KeyError(f"No mesh '{mesh_name}' found in {filename}") 

445 mesh = h5file[mesh_name] 

446 if "functions" not in mesh.keys(): 

447 raise KeyError(f"No functions stored in '{mesh_name}' in {filename}") 

448 functions = mesh["functions"] 

449 if name not in functions.keys(): 

450 raise KeyError( 

451 f"No function with name '{name}' on '{mesh_name}' stored in {filename}" 

452 ) 

453 function = functions[name] 

454 offset_key = "dofmap_offsets" 

455 dofmap_key = "dofmap" 

456 offsets = function[offset_key] 

457 dofmap = function[dofmap_key] 

458 else: 

459 offset_key = backend_args["offsets"] 

460 dofmap = h5file[dofmap_key] 

461 offsets = h5file[offset_key] 

462 

463 num_cells = offsets.shape[0] - 1 

464 local_range = compute_local_range(comm, num_cells) 

465 

466 # First read in offsets based on the number of cells [0, num_cells_local] 

467 glob_offsets = offsets[local_range[0] : local_range[1] + 1].flatten().astype(np.int64) 

468 

469 # Then read the data based of offsets 

470 dofmap_data = dofmap[glob_offsets[0] : glob_offsets[-1]].flatten() 

471 

472 # Then make offsets local 

473 loc_offsets = (glob_offsets - glob_offsets[0]).astype(np.int32) 

474 return adjacencylist(dofmap_data, loc_offsets) 

475 

476 

477def read_dofs( 

478 filename: str | Path, 

479 comm: MPI.Comm, 

480 name: str, 

481 time: float, 

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

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

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

485 

486 Args: 

487 filename: Path to file to read from 

488 comm: MPI communicator used in storage 

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

490 time: Time stamp associated with the function to read 

491 backend_args: Arguments to backend 

492 

493 Returns: 

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

495 and the global starting point on the process. 

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

497 """ 

498 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

499 mesh_name = get_mesh_name(backend_args) 

500 if mesh_name not in h5file.keys(): 

501 raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") 

502 mesh = h5file[mesh_name] 

503 if "functions" not in mesh.keys(): 

504 raise RuntimeError(f"No functions stored in '{mesh_name}' in {filename}") 

505 functions = mesh["functions"] 

506 if name not in functions.keys(): 

507 raise RuntimeError( 

508 f"No function with name '{name}' on '{mesh_name}' stored in {filename}" 

509 ) 

510 function = functions[name] 

511 timestamps = function.attrs["timestamps"] 

512 idx = np.flatnonzero(np.isclose(timestamps, time)) 

513 if len(idx) != 1: 

514 raise RuntimeError("Could not find {name}(t={time}) on grid {mesh_name} in {filename}.") 

515 u_t = function[f"{idx[0]:d}"] 

516 data_group = u_t["array"] 

517 num_dofs_global = data_group.shape[0] 

518 local_range = compute_local_range(comm, num_dofs_global) 

519 local_array = data_group[slice(*local_range)] 

520 return local_array, local_range[0] 

521 

522 

523def read_cell_perms( 

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

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

526 """ 

527 Read cell permutation from file with given communicator, 

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

529 

530 Args: 

531 comm: MPI communicator used in storage 

532 filename: Path to file to read from 

533 backend_args: Arguments to backend 

534 

535 Returns: 

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

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

538 """ 

539 

540 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

541 mesh_name = get_mesh_name(backend_args) 

542 if mesh_name not in h5file.keys(): 

543 raise RuntimeError(f"No mesh '{mesh_name}' found in {filename}") 

544 mesh = h5file[mesh_name] 

545 data_group = mesh["CellPermutations"] 

546 num_cells_global = data_group.shape[0] 

547 local_range = compute_local_range(comm, num_cells_global) 

548 local_array = data_group[slice(*local_range)] 

549 return local_array 

550 

551 

552def write_function( 

553 filename: str | Path, 

554 comm: MPI.Comm, 

555 u: FunctionData, 

556 time: float, 

557 mode: FileMode, 

558 backend_args: dict[str, Any] | None = None, 

559): 

560 """Write a function to file. 

561 

562 Args: 

563 comm: MPI communicator used in storage 

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

565 filename: Path to file to write to 

566 time: Time stamp associated with function 

567 mode: File-mode to store the function 

568 backend_args: Arguments to backend 

569 """ 

570 

571 mesh_name = get_mesh_name(backend_args) 

572 backend_args = get_default_backend_args(backend_args) 

573 h5_mode = convert_file_mode(mode) 

574 with h5pyfile(filename, filemode=h5_mode, comm=comm, force_serial=False) as h5file: 

575 cell_permutations_exist = False 

576 dofmap_exists = False 

577 dofmap_offsets_exists = False 

578 if h5_mode == "a": 

579 if mesh_name not in h5file.keys(): 

580 mesh = h5file.create_group(mesh_name) 

581 else: 

582 mesh = h5file[mesh_name] 

583 

584 cell_permutations_exist = "CellPermutations" in mesh.keys() 

585 

586 if "functions" not in mesh.keys(): 

587 functions = mesh.create_group("functions") 

588 else: 

589 functions = mesh["functions"] 

590 

591 if u.name not in functions.keys(): 

592 function = functions.create_group(u.name) 

593 else: 

594 function = functions[u.name] 

595 

596 dofmap_exists = "dofmap" in function.keys() 

597 dofmap_offsets_exists = "dofmap_offsets" in function.keys() 

598 

599 if not cell_permutations_exist: 

600 cell_perms = mesh.create_dataset( 

601 "CellPermutations", shape=[u.num_cells_global], dtype=np.uint32 

602 ) 

603 cell_perms[slice(*u.local_cell_range)] = u.cell_permutations 

604 

605 if not dofmap_exists: 

606 dofmap = function.create_dataset( 

607 "dofmap", shape=[u.global_dofs_in_dofmap], dtype=np.int64 

608 ) 

609 dofmap[slice(*u.dofmap_range)] = u.dofmap_array 

610 

611 if not dofmap_offsets_exists: 

612 dofmap_offsets = function.create_dataset( 

613 "dofmap_offsets", shape=[u.num_cells_global + 1], dtype=np.int64 

614 ) 

615 dofmap_offsets[u.local_cell_range[0] : u.local_cell_range[1] + 1] = u.dofmap_offsets 

616 

617 # Write timestamp 

618 if "timestamps" in function.attrs.keys(): 

619 timestamps = function.attrs["timestamps"] 

620 if np.isclose(time, timestamps).any(): 

621 raise RuntimeError("FUnction has already been stored at time={time_stamp}.") 

622 else: 

623 function.attrs["timestamps"] = np.append(function.attrs["timestamps"], time) 

624 else: 

625 function.attrs["timestamps"] = np.array([time]) 

626 idx = len(function.attrs["timestamps"]) - 1 

627 

628 data_group = function.create_group(f"{idx:d}") 

629 array = data_group.create_dataset("array", shape=[u.num_dofs_global], dtype=u.values.dtype) 

630 array[slice(*u.dof_range)] = u.values 

631 

632 

633def read_legacy_mesh( 

634 filename: Path | str, comm: MPI.Comm, group: str 

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

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

637 legacy DOLFIN HDF5-file. 

638 

639 Args: 

640 filename: Path to file to read from 

641 comm: MPI communicator used in storage 

642 group: Group in HDF5 file where mesh is stored 

643 

644 Returns: 

645 Tuple containing: 

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

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

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

649 """ 

650 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

651 if group not in h5file.keys(): 

652 raise KeyError(f"Could not find {group} in {filename}.") 

653 mesh = h5file[group] 

654 if "topology" not in mesh.keys(): 

655 raise KeyError(f"Could not find '{group}/topology' in {filename}.") 

656 

657 topology = mesh["topology"] 

658 shape = topology.shape 

659 local_range = compute_local_range(comm, shape[0]) 

660 mesh_topology = topology[slice(*local_range)].astype(np.int64) 

661 

662 # Get mesh cell type 

663 cell_type = None 

664 if "celltype" in topology.attrs.keys(): 

665 cell_type = topology.attrs["celltype"] 

666 if isinstance(cell_type, np.bytes_): 

667 cell_type = cell_type.decode("utf-8") 

668 

669 for geometry_key in ["geometry", "coordinates"]: 

670 if geometry_key in mesh.keys(): 

671 break 

672 else: 

673 raise KeyError( 

674 "Could not find geometry in '/mesh/geometry' or '/mesh/coordinates'" 

675 + f" in {filename}." 

676 ) 

677 geometry = mesh[geometry_key] 

678 g_shape = geometry.shape 

679 local_g_range = compute_local_range(comm, g_shape[0]) 

680 mesh_geometry = geometry[slice(*local_g_range)] 

681 

682 return mesh_topology, mesh_geometry, cell_type 

683 

684 

685def read_hdf5_array( 

686 comm: MPI.Comm, 

687 filename: Path | str, 

688 group: str, 

689 backend_args: dict[str, Any] | None = None, 

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

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

692 

693 Args: 

694 comm: MPI communicator used in storage 

695 filename: Path to file to read from 

696 group: Group in HDF5 file where array is stored 

697 backend_args: Arguments to backend 

698 

699 Returns: 

700 Tuple containing: 

701 - Numpy array read from file 

702 - Global starting point on the process. 

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

704 """ 

705 with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file: 

706 data = h5file[group] 

707 shape = data.shape[0] 

708 local_range = compute_local_range(comm, shape) 

709 out_data = data[slice(*local_range)].flatten() 

710 return out_data, local_range[0] 

711 

712 

713def snapshot_checkpoint( 

714 filename: Path | str, 

715 mode: FileMode, 

716 u: dolfinx.fem.Function, 

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

718): 

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

720 

721 Args: 

722 filename: Path to file to read from 

723 mode: File-mode to store the function 

724 u: dolfinx function to create a snapshot checkpoint for 

725 backend_args: Arguments to backend 

726 """ 

727 comm = u.function_space.mesh.comm 

728 dofmap = u.function_space.dofmap 

729 local_range = np.array(dofmap.index_map.local_range) * dofmap.index_map_bs 

730 num_dofs_local = local_range[1] - local_range[0] 

731 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs 

732 h5mode = convert_file_mode(mode) 

733 if h5mode == "w": 

734 with h5pyfile(filename, filemode=h5mode, comm=comm, force_serial=False) as h5file: 

735 local_dofs = u.x.array[:num_dofs_local].copy() 

736 data = h5file.create_group("snapshot") 

737 dataset = data.create_dataset("dofs", shape=num_dofs_global, dtype=local_dofs.dtype) 

738 dataset[slice(*local_range)] = local_dofs 

739 elif h5mode == "r": 

740 with h5pyfile(filename, filemode=h5mode, comm=comm, force_serial=False) as h5file: 

741 data = h5file["snapshot"]["dofs"] 

742 assert data.shape[0] == num_dofs_global 

743 u.x.array[:num_dofs_local] = data[slice(*local_range)] 

744 u.x.scatter_forward() 

745 

746 

747def read_function_names( 

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

749) -> list[str]: 

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

751 

752 Args: 

753 filename: Path to file 

754 comm: MPI communicator to launch IO on. 

755 backend_args: Arguments to backend 

756 

757 Returns: 

758 A list of function names. 

759 """ 

760 check_file_exists(filename)