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 read_component_major_rows( 

64 dataset: h5py.Dataset, num_entities_global: int, local_range: list[int] 

65) -> np.ndarray: 

66 """Read rows ``local_range`` of a 2D array stored component by component. 

67 

68 The array is a one-dimensional dataset holding all first components, then all 

69 second components and so on, i.e. ``x0 x1 ... xn y0 y1 ... yn z0 z1 ... zn`` 

70 rather than ``x0 y0 z0 x1 y1 z1 ...`` (``MED_NO_INTERLACE`` in the MED format). 

71 The rows are one block per component, read as a single hyperslab. 

72 

73 Args: 

74 dataset: One-dimensional dataset 

75 num_entities_global: Number of entities (rows) in the whole dataset 

76 local_range: Range ``[start, end)`` of rows to read on this process 

77 

78 Returns: 

79 Array of shape ``(end - start, num_components)``, where 

80 ``num_components = dataset.size // num_entities_global`` 

81 """ 

82 if num_entities_global <= 0 or dataset.size == 0 or dataset.size % num_entities_global != 0: 

83 raise ValueError( 

84 f"Size {dataset.size} of {dataset.name} is not a positive multiple of " 

85 f"{num_entities_global} entities." 

86 ) 

87 num_components = dataset.size // num_entities_global 

88 num_rows = local_range[1] - local_range[0] 

89 if num_rows == 0: 

90 # h5py.MultiBlockSlice rejects empty blocks 

91 return np.empty((0, num_components), dtype=dataset.dtype) 

92 blocks = h5py.MultiBlockSlice( 

93 start=local_range[0], stride=num_entities_global, count=num_components, block=num_rows 

94 ) 

95 return np.ascontiguousarray(dataset[blocks].reshape(num_components, num_rows).T) 

96 

97 

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

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

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

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

102 return args 

103 

104 

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

106 match mode: 

107 case FileMode.append: 

108 return "a" 

109 case FileMode.read: 

110 return "r" 

111 case FileMode.write: 

112 return "w" 

113 case _: 

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

115 

116 

117def write_attributes( 

118 filename: Path | str, 

119 comm: MPI.Comm, 

120 name: str, 

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

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

123): 

124 """Write attributes to file using H5PY. 

125 

126 Args: 

127 filename: Path to file to write to 

128 comm: MPI communicator used in storage 

129 name: Name of the attributes 

130 attributes: Dictionary of attributes to write to file 

131 engine: ADIOS2 engine to use 

132 """ 

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

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

135 if name not in h5file.keys(): 

136 h5file.create_group(name, track_order=True) 

137 group = h5file[name] 

138 

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

140 group.attrs[key] = val 

141 

142 

143def read_attributes( 

144 filename: Path | str, 

145 comm: MPI.Comm, 

146 name: str, 

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

148) -> dict[str, Any]: 

149 """Read attributes from file using H5PY. 

150 

151 Args: 

152 filename: Path to file to read from 

153 comm: MPI communicator used in storage 

154 name: Name of the attributes 

155 Returns: 

156 The attributes 

157 """ 

158 check_file_exists(filename) 

159 output_attrs = {} 

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

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

162 output_attrs[key] = val 

163 return output_attrs 

164 

165 

166def read_timestamps( 

167 filename: Path | str, 

168 comm: MPI.Comm, 

169 function_name: str, 

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

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

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

173 

174 Args: 

175 comm: MPI communicator 

176 filename: Path to file 

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

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

179 

180 Returns: 

181 The time-stamps 

182 """ 

183 check_file_exists(filename) 

184 mesh_name = get_mesh_name(backend_args) 

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

186 mesh_directory = h5file[mesh_name] 

187 functions = mesh_directory["functions"] 

188 u = functions[function_name] 

189 timestamps = u.attrs["timestamps"] 

190 return timestamps 

191 

192 

193def write_mesh( 

194 filename: Path | str, 

195 comm: MPI.Comm, 

196 mesh: MeshData, 

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

198 mode: FileMode = FileMode.write, 

199 time: float = 0.0, 

200): 

201 """Write a mesh to file using H5PY 

202 

203 Args: 

204 comm: MPI communicator used in storage 

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

206 filename: Path to file to write to. 

207 mode: Mode to use (write or append) 

208 time: Time stamp 

209 """ 

210 backend_args = get_default_backend_args(backend_args) 

211 h5_mode = convert_file_mode(mode) 

212 mesh_name = get_mesh_name(backend_args) 

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

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

215 mesh_directory = h5file[mesh_name] 

216 timestamps = mesh_directory.attrs["timestamps"] 

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

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

219 else: 

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

221 mesh_directory.attrs["timestamps"], time 

222 ) 

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

224 write_topology = False 

225 else: 

226 mesh_directory = h5file.create_group(mesh_name) 

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

228 idx = 0 

229 write_topology = True 

230 

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

232 

233 # Write geometry data 

234 gdim = mesh.local_geometry.shape[1] 

235 geometry_dataset = geometry_group.create_dataset( 

236 "Points", 

237 [mesh.num_nodes_global, gdim], 

238 dtype=mesh.local_geometry.dtype, 

239 ) 

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

241 

242 # Write static partitioning data 

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

244 assert mesh.partition_range is not None 

245 assert mesh.ownership_array is not None 

246 par_dataset = mesh_directory.create_dataset( 

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

248 ) 

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

250 

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

252 assert mesh.ownership_offset is not None 

253 par_dataset = mesh_directory.create_dataset( 

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

255 ) 

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

257 mesh.ownership_offset 

258 ) 

259 

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

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

262 

263 # Write static data 

264 if write_topology: 

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

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

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

268 num_dofs_per_cell = mesh.local_topology.shape[1] 

269 topology_dataset = mesh_directory.create_dataset( 

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

271 ) 

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

273 

274 

275def read_mesh_data( 

276 filename: Path | str, 

277 comm: MPI.Comm, 

278 time: str | float | None = 0.0, 

279 read_from_partition: bool = False, 

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

281) -> ReadMeshData: 

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

283 

284 Args: 

285 filename: Path to input file 

286 comm: The MPI communciator to distribute the mesh over 

287 time: Time stamp associated with mesh 

288 read_from_partition: Read mesh with partition from file 

289 Returns: 

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

291 """ 

292 

293 backend_args = get_default_backend_args(backend_args) 

294 mesh_name = get_mesh_name(backend_args) 

295 

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

297 if mesh_name not in h5file.keys(): 

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

299 mesh_group = h5file[mesh_name] 

300 timestamps = mesh_group.attrs["timestamps"] 

301 assert time is not None 

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

303 if len(parent_group) != 1: 

304 raise RuntimeError( 

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

306 ) 

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

308 

309 # Get mesh topology (distributed) 

310 topology = mesh_group["Topology"] 

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

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

313 

314 cell_type = mesh_group.attrs["CellType"] 

315 lvar = mesh_group.attrs["LagrangeVariant"] 

316 degree = mesh_group.attrs["Degree"] 

317 

318 geometry_group = mesh_group[time_group] 

319 geometry_dataset = geometry_group["Points"] 

320 x_shape = geometry_dataset.shape 

321 

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

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

324 

325 # Check validity of partitioning information 

326 if read_from_partition: 

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

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

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

330 for key in par_keys: 

331 if key not in mesh_group.keys(): 

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

333 

334 par_num_procs = mesh_group.attrs["PartitionProcesses"] 

335 if par_num_procs != comm.size: 

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

337 

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

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

340 # Then read the data based of offsets 

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

342 # Then make offsets local 

343 par_offsets[:] -= par_offsets[0] 

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

345 else: 

346 partition_graph = None 

347 

348 return ReadMeshData( 

349 cells=mesh_topology, 

350 cell_type=cell_type, 

351 x=mesh_geometry, 

352 degree=degree, 

353 lvar=lvar, 

354 partition_graph=partition_graph, 

355 ) 

356 

357 

358def write_meshtags( 

359 filename: str | Path, 

360 comm: MPI.Comm, 

361 data: MeshTagsData, 

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

363): 

364 """Write mesh tags to file. 

365 

366 Args: 

367 filename: Path to file to write to 

368 comm: MPI communicator used in storage 

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

370 backend_args: Arguments to backend 

371 """ 

372 backend_args = get_default_backend_args(backend_args) 

373 

374 mesh_name = get_mesh_name(backend_args) 

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

376 if mesh_name not in h5file.keys(): 

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

378 mesh_group = h5file[mesh_name] 

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

380 tags = mesh_group.create_group("tags") 

381 else: 

382 tags = mesh_group["tags"] 

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

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

385 tag = tags.create_group(data.name) 

386 

387 # Add topology 

388 topology = tag.create_dataset( 

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

390 ) 

391 assert data.local_start is not None 

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

393 

394 # Add cell_type attribute 

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

396 

397 # Add values 

398 values = tag.create_dataset( 

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

400 ) 

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

402 

403 # Add dimension 

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

405 

406 

407def read_meshtags_data( 

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

409) -> MeshTagsData: 

410 """Read mesh tags from file. 

411 

412 Args: 

413 filename: Path to file to read from 

414 comm: MPI communicator used in storage 

415 name: Name of the mesh tags to read 

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

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

418 

419 Returns: 

420 Internal data structure for the mesh tags read from file 

421 """ 

422 backend_args = get_default_backend_args(backend_args) 

423 legacy = backend_args["legacy"] 

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

425 if legacy: 

426 if name not in h5file.keys(): 

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

428 mesh = h5file[name] 

429 topology = mesh["topology"] 

430 cell_type = topology.attrs["celltype"] 

431 if isinstance(cell_type, np.bytes_): 

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

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

434 values = mesh["values"] 

435 else: 

436 mesh_name = get_mesh_name(backend_args) 

437 if mesh_name not in h5file.keys(): 

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

439 mesh = h5file[mesh_name] 

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

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

442 tags = mesh["tags"] 

443 if name not in tags.keys(): 

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

445 tag = tags[name] 

446 

447 dim = tag.attrs["dim"] 

448 topology = tag["Topology"] 

449 values = tag["Values"] 

450 num_entities_global = topology.shape[0] 

451 topology_range = compute_local_range(comm, num_entities_global) 

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

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

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

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

456 

457 

458def read_dofmap( 

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

460) -> dolfinx.graph.AdjacencyList: 

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

462 

463 Args: 

464 filename: Path to file to read from 

465 comm: MPI communicator used in storage 

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

467 backend_args: Arguments to backend 

468 

469 Returns: 

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

471 """ 

472 backend_args = {} if backend_args is None else backend_args 

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

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

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

476 if dofmap_key is None: 

477 mesh_name = get_mesh_name(backend_args) 

478 if mesh_name not in h5file.keys(): 

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

480 mesh = h5file[mesh_name] 

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

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

483 functions = mesh["functions"] 

484 if name not in functions.keys(): 

485 raise KeyError( 

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

487 ) 

488 function = functions[name] 

489 offset_key = "dofmap_offsets" 

490 dofmap_key = "dofmap" 

491 offsets = function[offset_key] 

492 dofmap = function[dofmap_key] 

493 else: 

494 offset_key = backend_args["offsets"] 

495 dofmap = h5file[dofmap_key] 

496 offsets = h5file[offset_key] 

497 

498 num_cells = offsets.shape[0] - 1 

499 local_range = compute_local_range(comm, num_cells) 

500 

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

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

503 

504 # Then read the data based of offsets 

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

506 

507 # Then make offsets local 

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

509 return adjacencylist(dofmap_data, loc_offsets) 

510 

511 

512def read_dofs( 

513 filename: str | Path, 

514 comm: MPI.Comm, 

515 name: str, 

516 time: float, 

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

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

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

520 

521 Args: 

522 filename: Path to file to read from 

523 comm: MPI communicator used in storage 

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

525 time: Time stamp associated with the function to read 

526 backend_args: Arguments to backend 

527 

528 Returns: 

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

530 and the global starting point on the process. 

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

532 """ 

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

534 mesh_name = get_mesh_name(backend_args) 

535 if mesh_name not in h5file.keys(): 

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

537 mesh = h5file[mesh_name] 

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

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

540 functions = mesh["functions"] 

541 if name not in functions.keys(): 

542 raise RuntimeError( 

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

544 ) 

545 function = functions[name] 

546 timestamps = function.attrs["timestamps"] 

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

548 if len(idx) != 1: 

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

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

551 data_group = u_t["array"] 

552 num_dofs_global = data_group.shape[0] 

553 local_range = compute_local_range(comm, num_dofs_global) 

554 local_array = data_group[slice(*local_range)] 

555 return local_array, local_range[0] 

556 

557 

558def read_cell_perms( 

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

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

561 """ 

562 Read cell permutation from file with given communicator, 

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

564 

565 Args: 

566 comm: MPI communicator used in storage 

567 filename: Path to file to read from 

568 backend_args: Arguments to backend 

569 

570 Returns: 

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

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

573 """ 

574 

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

576 mesh_name = get_mesh_name(backend_args) 

577 if mesh_name not in h5file.keys(): 

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

579 mesh = h5file[mesh_name] 

580 data_group = mesh["CellPermutations"] 

581 num_cells_global = data_group.shape[0] 

582 local_range = compute_local_range(comm, num_cells_global) 

583 local_array = data_group[slice(*local_range)] 

584 return local_array 

585 

586 

587def write_function( 

588 filename: str | Path, 

589 comm: MPI.Comm, 

590 u: FunctionData, 

591 time: float, 

592 mode: FileMode, 

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

594): 

595 """Write a function to file. 

596 

597 Args: 

598 comm: MPI communicator used in storage 

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

600 filename: Path to file to write to 

601 time: Time stamp associated with function 

602 mode: File-mode to store the function 

603 backend_args: Arguments to backend 

604 """ 

605 

606 mesh_name = get_mesh_name(backend_args) 

607 backend_args = get_default_backend_args(backend_args) 

608 h5_mode = convert_file_mode(mode) 

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

610 cell_permutations_exist = False 

611 dofmap_exists = False 

612 dofmap_offsets_exists = False 

613 if h5_mode == "a": 

614 if mesh_name not in h5file.keys(): 

615 mesh = h5file.create_group(mesh_name) 

616 else: 

617 mesh = h5file[mesh_name] 

618 

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

620 

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

622 functions = mesh.create_group("functions") 

623 else: 

624 functions = mesh["functions"] 

625 

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

627 function = functions.create_group(u.name) 

628 else: 

629 function = functions[u.name] 

630 

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

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

633 

634 if not cell_permutations_exist: 

635 cell_perms = mesh.create_dataset( 

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

637 ) 

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

639 

640 if not dofmap_exists: 

641 dofmap = function.create_dataset( 

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

643 ) 

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

645 

646 if not dofmap_offsets_exists: 

647 dofmap_offsets = function.create_dataset( 

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

649 ) 

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

651 

652 # Write timestamp 

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

654 timestamps = function.attrs["timestamps"] 

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

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

657 else: 

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

659 else: 

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

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

662 

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

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

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

666 

667 

668def read_legacy_mesh( 

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

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

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

672 legacy DOLFIN HDF5-file. 

673 

674 Args: 

675 filename: Path to file to read from 

676 comm: MPI communicator used in storage 

677 group: Group in HDF5 file where mesh is stored 

678 

679 Returns: 

680 Tuple containing: 

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

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

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

684 """ 

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

686 if group not in h5file.keys(): 

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

688 mesh = h5file[group] 

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

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

691 

692 topology = mesh["topology"] 

693 shape = topology.shape 

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

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

696 

697 # Get mesh cell type 

698 cell_type = None 

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

700 cell_type = topology.attrs["celltype"] 

701 if isinstance(cell_type, np.bytes_): 

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

703 

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

705 if geometry_key in mesh.keys(): 

706 break 

707 else: 

708 raise KeyError( 

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

710 + f" in {filename}." 

711 ) 

712 geometry = mesh[geometry_key] 

713 g_shape = geometry.shape 

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

715 mesh_geometry = geometry[slice(*local_g_range)] 

716 

717 return mesh_topology, mesh_geometry, cell_type 

718 

719 

720def read_hdf5_array( 

721 comm: MPI.Comm, 

722 filename: Path | str, 

723 group: str, 

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

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

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

727 

728 Args: 

729 comm: MPI communicator used in storage 

730 filename: Path to file to read from 

731 group: Group in HDF5 file where array is stored 

732 backend_args: Arguments to backend 

733 

734 Returns: 

735 Tuple containing: 

736 - Numpy array read from file 

737 - Global starting point on the process. 

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

739 """ 

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

741 data = h5file[group] 

742 shape = data.shape[0] 

743 local_range = compute_local_range(comm, shape) 

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

745 return out_data, local_range[0] 

746 

747 

748def snapshot_checkpoint( 

749 filename: Path | str, 

750 mode: FileMode, 

751 u: dolfinx.fem.Function, 

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

753): 

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

755 

756 Args: 

757 filename: Path to file to read from 

758 mode: File-mode to store the function 

759 u: dolfinx function to create a snapshot checkpoint for 

760 backend_args: Arguments to backend 

761 """ 

762 comm = u.function_space.mesh.comm 

763 dofmap = u.function_space.dofmap 

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

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

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