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

488 statements  

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

1""" 

2Module that can read the VTKHDF format using h5py. 

3""" 

4 

5from pathlib import Path 

6from typing import Any 

7 

8from mpi4py import MPI 

9 

10import basix 

11import dolfinx 

12import fenicsx_compat 

13import h5py 

14import numpy as np 

15import numpy.typing as npt 

16 

17from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData 

18from io4dolfinx.utils import check_file_exists, compute_local_range 

19 

20from .. import DEFAULT_MESH_NAME, FileMode, ReadMode 

21from ..h5py.backend import convert_file_mode, h5pyfile 

22from ..pyvista.backend import _arbitrary_lagrange_vtk, _cell_degree, _first_order_vtk 

23 

24read_mode = ReadMode.parallel 

25 

26_vtk_hdf_version = np.array([2, 1], dtype=np.int32) 

27#: Top-level group for metadata io4dolfinx needs but VTKHDF has no slot for. 

28#: The specification reserves everything outside ``/VTKHDF`` for exactly this: 

29#: "Top-level groups outside of /VTKHDF do not contain any information related to 

30#: the VTK data model and are outside of the scope of this specification. They can 

31#: be useful to store meta-information that could be read and written by custom 

32#: VTKHDF implementation." 

33_metadata_group = "/io4dolfinx" 

34#: Basix Lagrange variant of a mesh's coordinate element, per mesh name. 

35_lagrange_variant_attr = "LagrangeVariant" 

36 

37 

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

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

40 

41 Args: 

42 arguments: Input backend arguments 

43 

44 Returns: 

45 Updated backend arguments 

46 """ 

47 # `arguments or {...}` would drop the default whenever the caller passed any 

48 # other argument, so set it on a copy instead. 

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

50 args.setdefault("name", DEFAULT_MESH_NAME) 

51 return args 

52 

53 

54def find_all_unique_cell_types(comm, cell_types, num_nodes): 

55 """ 

56 Given a set of cell types and number of nodes per cell, find all unique cell types 

57 across all ranks. 

58 

59 Args: 

60 comm: MPI communicator 

61 cell_types: Local cell types 

62 num_nodes: Number of nodes per cell 

63 

64 

65 Returns: 

66 A 2D array where each row corresponds to a cell type (vtk int) 

67 and the number of nodes. 

68 """ 

69 # Combine cell_types, num_nodes as tuple 

70 c_hash = np.zeros((2, len(cell_types)), dtype=np.int32) 

71 c_hash[0] = cell_types 

72 c_hash[1] = num_nodes 

73 indexes = np.unique(c_hash.T, axis=0, return_index=True)[1] 

74 local_unique_cells = c_hash.T[indexes] 

75 

76 all_cell_types = np.vstack(comm.allgather(local_unique_cells)) 

77 indexes = np.unique(all_cell_types, axis=0, return_index=True)[1] 

78 all_unique_cell_types = all_cell_types[indexes] 

79 return all_unique_cell_types 

80 

81 

82def _decode_bytes_if_needed(value: bytes | str) -> str: 

83 """Decode bytes to string if necessary (for h5py compatibility)""" 

84 if isinstance(value, bytes): 

85 return value.decode("utf-8") 

86 return value 

87 

88 

89def _get_vtk_group(h5file, name: str) -> h5py.Group: 

90 """ 

91 Navigates the VTKHDF group hierarchy to find the specific first 

92 UnstructuredGrid with a specific name. 

93 Handles both MultiBlockDataSet and direct UnstructuredGrid types. 

94 """ 

95 hdf = h5file["VTKHDF"] 

96 file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) 

97 

98 if file_type == "MultiBlockDataSet": 

99 ass = hdf["Assembly"] 

100 

101 def visitor(path): 

102 mesh_group = path.rsplit("/", 1) 

103 mesh_name = mesh_group[0] if len(mesh_group) == 1 else mesh_group[1] 

104 if mesh_name == name: 

105 obj = ass.get(path) 

106 # Check attributes carefully 

107 if "Type" in obj.attrs.keys(): 

108 attr_type = obj.attrs["Type"] 

109 if isinstance(attr_type, bytes): 

110 attr_type = attr_type.decode("utf-8") 

111 if attr_type == "UnstructuredGrid": 

112 return path 

113 return None 

114 

115 mesh_node = ass.visit_links(visitor) 

116 if mesh_node is None: 

117 raise RuntimeError(f"Could not find unique mesh named '{name}' in Assembly.") 

118 

119 return ass[mesh_node] 

120 

121 elif file_type == "UnstructuredGrid": 

122 return hdf 

123 else: 

124 raise RuntimeError(f"Not supported file type {file_type}") 

125 

126 

127def _get_time_index(hdf: h5py.Group, time: float | str, filename: str | Path) -> int: 

128 """Finds the index of a specific time stamp.""" 

129 if "Steps" not in hdf.keys(): 

130 raise RuntimeError(f"No timestepping information found in {filename}.") 

131 stamps = hdf["Steps"]["Values"][:] 

132 pos = np.flatnonzero(np.isclose(stamps, time)) 

133 if len(pos) == 0: 

134 raise RuntimeError(f"Could not find mesh at t={time} in {filename}.") 

135 elif len(pos) > 1: 

136 raise RuntimeError(f"Multiple time steps for mesh at t={time} in {filename}") 

137 return pos[0] 

138 

139 

140def _read_lagrange_variant(h5file, name: str) -> int: 

141 """Lagrange variant of mesh ``name``, or the VTKHDF default. 

142 

143 A file written by anything other than io4dolfinx -- or by io4dolfinx before 

144 this metadata existed -- holds VTKHDF's equispaced higher-order cells. 

145 

146 Args: 

147 h5file: The open HDF5 file 

148 name: Name of the mesh 

149 

150 Returns: 

151 The basix Lagrange variant as an integer. 

152 """ 

153 equispaced = int(basix.LagrangeVariant.equispaced) 

154 metadata = h5file.get(_metadata_group) 

155 if metadata is None or name not in metadata: 

156 return equispaced 

157 return int(metadata[name].attrs.get(_lagrange_variant_attr, equispaced)) 

158 

159 

160def read_mesh_data( 

161 filename: Path | str, 

162 comm: MPI.Comm, 

163 time: str | float | None, 

164 read_from_partition: bool, 

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

166) -> ReadMeshData: 

167 """Read mesh data from file. 

168 

169 Args: 

170 filename: Path to file to read from 

171 comm: MPI communicator used in storage 

172 time: Time stamp associated with the mesh to read 

173 read_from_partition: Whether to read partition information 

174 backend_args: Arguments to backend 

175 

176 Returns: 

177 Internal data structure for the mesh data read from file 

178 """ 

179 backend_args = get_default_backend_args(backend_args) 

180 check_file_exists(filename) 

181 if read_from_partition: 

182 raise RuntimeError("Cannot read partition data with VTKHDF") 

183 with h5pyfile(filename, "r", comm=comm) as h5file: 

184 hdf = _get_vtk_group(h5file, backend_args["name"]) 

185 lagrange_variant = _read_lagrange_variant(h5file, backend_args["name"]) 

186 if time is None: 

187 num_cells_global = hdf["Types"].size 

188 local_cell_range = compute_local_range(comm, num_cells_global) 

189 cell_types_local = hdf["Types"][slice(*local_cell_range)] 

190 

191 num_points_global = hdf["NumberOfPoints"][0] 

192 local_point_range = compute_local_range(comm, num_points_global) 

193 points_local = hdf["Points"][slice(*local_point_range)] 

194 

195 # Connectivity read 

196 offsets = hdf["Offsets"] 

197 local_connectivity_offset = offsets[local_cell_range[0] : local_cell_range[1] + 1] 

198 topology = hdf["Connectivity"][ 

199 local_connectivity_offset[0] : local_connectivity_offset[-1] 

200 ] 

201 offset = local_connectivity_offset - local_connectivity_offset[0] 

202 else: 

203 time_index = _get_time_index(hdf, time, filename) 

204 

205 stamps = hdf["Steps"]["Values"][:] 

206 # Get number of points 

207 point_node = hdf["Points"] 

208 step_start = hdf["Steps"]["PointOffsets"][time_index] 

209 

210 # NOTE: currently, it doesn't seem like we follow: 

211 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata 

212 # As only one num_points is stored irregardless of time steps added. 

213 if hdf["NumberOfPoints"].shape[0] != len(stamps): 

214 num_pts = hdf["NumberOfPoints"][0] 

215 else: 

216 num_pts = hdf["NumberOfPoints"][time_index] 

217 lr = compute_local_range(comm, num_pts) 

218 points_local = point_node[step_start + lr[0] : step_start + lr[1]] 

219 

220 # Get cell-types in step 

221 cell_start = hdf["Steps"]["CellOffsets"][time_index] 

222 if hdf["NumberOfCells"].shape[0] != len(stamps): 

223 num_cells = hdf["NumberOfCells"][0] 

224 else: 

225 num_cells = hdf["NumberOfCells"][time_index] 

226 local_cell_range = compute_local_range(comm, num_cells) 

227 cell_types_local = hdf["Types"][ 

228 cell_start + local_cell_range[0] : cell_start + local_cell_range[1] 

229 ] 

230 

231 # Get connectivity in step 

232 connectivity_start = hdf["Steps"]["ConnectivityIdOffsets"][time_index] 

233 # Connectivity read 

234 offsets = hdf["Offsets"] 

235 local_connectivity_offset = offsets[ 

236 connectivity_start + local_cell_range[0] : connectivity_start 

237 + local_cell_range[1] 

238 + 1 

239 ] 

240 topology = hdf["Connectivity"][ 

241 local_connectivity_offset[0] : local_connectivity_offset[-1] 

242 ] 

243 offset = local_connectivity_offset - local_connectivity_offset[0] 

244 

245 # NOTE: Currently we limit ourselfs to a single celltype, as it makes life easier, 

246 # other things have to change in `MeshReadData` to support this. 

247 num_nodes_per_cell = offset[1:] - offset[:-1] 

248 unique_cells = find_all_unique_cell_types(comm, cell_types_local, num_nodes_per_cell) 

249 if unique_cells.shape[0] > 1: 

250 raise NotImplementedError("io4dolfinx does not support mixed celltype grids") 

251 if unique_cells.shape[0] == 0: 

252 name = backend_args["name"] 

253 raise ValueError(f"Grid '{name}' in {filename} has no cells on any process.") 

254 # Take the width from the globally agreed cell type rather than from this 

255 # process's own cells: a mesh with fewer cells than processes leaves some of 

256 # them with nothing local to measure. 

257 cell_type, number_of_nodes = unique_cells[0] 

258 topology = topology.reshape(-1, number_of_nodes) 

259 gtype = backend_args.get("dtype", points_local.dtype) 

260 if cell_type in _first_order_vtk.keys(): 

261 ct = _first_order_vtk[cell_type] 

262 degree = 1 

263 elif cell_type in _arbitrary_lagrange_vtk.keys(): 

264 ct = _arbitrary_lagrange_vtk[cell_type] 

265 degree = _cell_degree(ct, num_nodes=number_of_nodes) 

266 else: 

267 raise ValueError(f"Unknown VTK cell type {cell_type} in {filename}") 

268 perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(ct), number_of_nodes) 

269 topology = topology[:, perm] 

270 # Equispaced is the VTKHDF default; a file io4dolfinx wrote says which 

271 # variant its nodes were actually placed with. 

272 lvar = int(lagrange_variant) 

273 return ReadMeshData( 

274 cells=topology, cell_type=ct, x=points_local.astype(gtype), lvar=lvar, degree=degree 

275 ) 

276 

277 

278def read_point_data( 

279 filename: Path | str, 

280 name: str, 

281 comm: MPI.Comm, 

282 time: float | str | None, 

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

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

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

286 

287 Args: 

288 filename: Path to file 

289 name: Name of point data 

290 comm: Communicator to launch IO on. 

291 time: The time stamp 

292 backend_args: The backend arguments 

293 

294 Returns: 

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

296 """ 

297 backend_args = get_default_backend_args(backend_args) 

298 check_file_exists(filename) 

299 with h5pyfile(filename, "r", comm=comm) as h5file: 

300 hdf = _get_vtk_group(h5file, backend_args["name"]) 

301 point_data = hdf["PointData"] 

302 assert point_data is not None 

303 if name not in point_data.keys(): 

304 raise ValueError(f"No point data named {name} in {filename}.") 

305 func_node = point_data[name] 

306 

307 if time is None: 

308 data_shape = func_node.shape[0] 

309 lr = compute_local_range(comm, data_shape) 

310 data = func_node[slice(*lr)] 

311 return data, lr[0] 

312 else: 

313 time_index = _get_time_index(hdf, time, filename) 

314 stamps = hdf["Steps"]["Values"][:] 

315 step_start = hdf["Steps"]["PointDataOffsets"][name][time_index] 

316 # NOTE: currently, it doesn't seem like we follow: 

317 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata 

318 # As only one num_points is stored irregardless of time steps added. 

319 if hdf["NumberOfPoints"].shape[0] != len(stamps): 

320 num_pts = hdf["NumberOfPoints"][0] 

321 else: 

322 num_pts = hdf["NumberOfPoints"][time_index] 

323 lr = compute_local_range(comm, num_pts) 

324 return func_node[step_start + lr[0] : step_start + lr[1]], lr[0] 

325 

326 

327def read_cell_data( 

328 filename: Path | str, 

329 name: str, 

330 comm: MPI.Comm, 

331 time: str | float | None, 

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

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

334 backend_args = get_default_backend_args(backend_args) 

335 check_file_exists(filename) 

336 with h5pyfile(filename, "r", comm=comm) as h5file: 

337 hdf = _get_vtk_group(h5file, backend_args["name"]) 

338 

339 if "CellData" not in hdf.keys(): 

340 raise RuntimeError(f"No cell data found in {filename}.") 

341 cell_data = hdf["CellData"] 

342 assert cell_data is not None 

343 if name not in cell_data.keys(): 

344 raise ValueError(f"No cell data with name {name} in {filename}") 

345 cell_data_node = cell_data[name] 

346 assert cell_data_node is not None 

347 if time is None: 

348 cell_data_shape = cell_data_node.shape 

349 num_cells_global = hdf["Types"].size 

350 assert num_cells_global == cell_data_shape[0] 

351 local_cell_range = compute_local_range(comm, num_cells_global) 

352 data = cell_data_node[slice(*local_cell_range)] 

353 else: 

354 time_index = _get_time_index(hdf, time, filename) 

355 stamps = hdf["Steps"]["Values"][:] 

356 cd_start = hdf["Steps"]["CellDataOffsets"][name][time_index] 

357 

358 # NOTE: currently, it doesn't seem like we follow: 

359 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata 

360 # As only one num_points is stored irregardless of time steps added. 

361 if hdf["NumberOfCells"].shape[0] != len(stamps): 

362 number_of_cells = hdf["NumberOfCells"][0] 

363 else: 

364 number_of_cells = hdf["NumberOfCells"][time_index] 

365 lr = compute_local_range(comm, number_of_cells) 

366 data = cell_data_node[cd_start + lr[0] : cd_start + lr[1]] 

367 

368 # NOTE: THis could be optimized by hand-coding some communication in 

369 # `read_cell_data` on the frontend side 

370 md = read_mesh_data( 

371 filename, comm, time=time, read_from_partition=False, backend_args=backend_args 

372 ) 

373 if len(data.shape) == 1: 

374 data = data.reshape(-1, 1) 

375 return md.cells, data 

376 

377 

378def write_attributes( 

379 filename: Path | str, 

380 comm: MPI.Comm, 

381 name: str, 

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

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

384): 

385 """Write attributes to file. 

386 

387 Args: 

388 filename: Path to file to write to 

389 comm: MPI communicator used in storage 

390 name: Name of the attribute group 

391 attributes: Dictionary of attributes to write 

392 backend_args: Arguments to backend 

393 """ 

394 raise NotImplementedError("The Pyvista backend cannot write attributes.") 

395 

396 

397def read_attributes( 

398 filename: Path | str, 

399 comm: MPI.Comm, 

400 name: str, 

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

402) -> dict[str, Any]: 

403 """Read attributes from file. 

404 

405 Args: 

406 filename: Path to file to read from 

407 comm: MPI communicator used in storage 

408 name: Name of the attribute group 

409 backend_args: Arguments to backend 

410 

411 Returns: 

412 Dictionary of attributes read from file 

413 """ 

414 raise NotImplementedError("The Pyvista backend cannot read attributes.") 

415 

416 

417def read_timestamps( 

418 filename: Path | str, 

419 comm: MPI.Comm, 

420 function_name: str, 

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

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

423 """Read timestamps from file. 

424 

425 Args: 

426 filename: Path to file to read from 

427 comm: MPI communicator used in storage 

428 function_name: Name of the function to read timestamps for 

429 backend_args: Arguments to backend 

430 

431 Returns: 

432 Numpy array of timestamps read from file 

433 """ 

434 backend_args = get_default_backend_args(backend_args) 

435 check_file_exists(filename) 

436 # Temporal data storage as described in 

437 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-data 

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

439 hdf = h5file["VTKHDF"] 

440 if "Steps" in hdf.keys(): 

441 timestamps = hdf["Steps"]["Values"][:] 

442 # For either point data or cell data, check if function exists, 

443 # and check if offsets in time are changing between steps. 

444 if "CellData" in hdf.keys() and function_name in hdf["CellData"].keys(): 

445 offsets = hdf["Steps"]["CellDataOffsets"][function_name] 

446 step_offsets = offsets[:] 

447 step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1]) 

448 return timestamps[step_diff] 

449 elif "PointData" in hdf.keys() and function_name in hdf["PointData"].keys(): 

450 offsets = hdf["Steps"]["PointDataOffsets"][function_name] 

451 step_offsets = offsets[:] 

452 step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1]) 

453 # This only finds when the offset changes, does not capture first step 

454 return np.hstack([[timestamps[0]], timestamps[step_diff]]) 

455 else: 

456 raise RuntimeError(f"Function {function_name} is not assoicated with a time-stamp.") 

457 else: 

458 raise RuntimeError(f"{filename} does not contain time-stepping information.") 

459 

460 

461def read_function_names( 

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

463) -> list[str]: 

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

465 

466 Args: 

467 filename: Path to file 

468 comm: MPI communicator to launch IO on. 

469 backend_args: Arguments to backend 

470 

471 Returns: 

472 A list of function names. 

473 """ 

474 backend_args = get_default_backend_args(backend_args) 

475 check_file_exists(filename) 

476 with h5pyfile(filename, "r", comm=comm) as h5file: 

477 hdf = h5file["VTKHDF"] 

478 function_names = set() 

479 if "CellData" in hdf.keys(): 

480 for item in hdf["CellData"].keys(): 

481 function_names.add(item) 

482 if "PointData" in hdf.keys(): 

483 for item in hdf["PointData"].keys(): 

484 function_names.add(item) 

485 return list(function_names) 

486 

487 

488def _create_dataset( 

489 root: h5py.File | h5py.Group, 

490 name: str, 

491 shape: tuple[int, ...], 

492 dtype: npt.DTypeLike, 

493 chunks: bool, 

494 maxshape: tuple[int | None, ...], 

495 mode: str, 

496 resize: bool = True, 

497) -> h5py.Dataset: 

498 if name not in root.keys(): 

499 mode = "w" 

500 if mode == "w": 

501 dataset = root.create_dataset( 

502 name, shape=shape, dtype=dtype, chunks=chunks, maxshape=maxshape 

503 ) 

504 elif mode == "a": 

505 dataset = root[name] 

506 old_shape = dataset.shape 

507 # Only resize for dimension 

508 if resize: 

509 if len(old_shape) == 1: 

510 new_shape = (old_shape[0] + shape[0],) 

511 else: 

512 new_shape = (old_shape[0] + shape[0], *old_shape[1:]) 

513 dataset.resize(new_shape) 

514 else: 

515 raise ValueError(f"Unknown file mode '{mode}' when creating dataset {name} in {root}") 

516 return dataset 

517 

518 

519def _create_group(root: h5py.File | h5py.Group, name: str, mode: str) -> h5py.Group: 

520 if name not in root.keys(): 

521 mode = "w" 

522 if mode == "w": 

523 # Track order has to be on to make multiblock work: 

524 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#partitioneddatasetcollection-and-multiblockdataset 

525 group = root.create_group(name, track_order=True) 

526 elif mode == "a": 

527 group = root[name] 

528 else: 

529 raise ValueError("Unknown file mode '{h5_mode}'") 

530 return group 

531 

532 

533def _compute_append_slice( 

534 dataset: h5py.Dataset, input_size: int, original_slice: tuple[int, int] | np.ndarray, mode: str 

535) -> slice: 

536 append_offset = dataset.shape[0] - input_size if mode == "a" else 0 

537 return slice(*(np.asarray(original_slice) + append_offset).astype(np.int64)) 

538 

539 

540def write_mesh( 

541 filename: Path | str, 

542 comm: MPI.Comm, 

543 mesh: MeshData, 

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

545 mode: FileMode, 

546 time: float, 

547): 

548 """ 

549 Write a mesh to file. 

550 

551 Args: 

552 comm: MPI communicator used in storage 

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

554 filename: Path to file to write to 

555 backend_args: Arguments to backend 

556 mode: File-mode to store the mesh 

557 time: Time stamp associated with the mesh 

558 """ 

559 h5_mode = convert_file_mode(mode) 

560 backend_args = get_default_backend_args(backend_args) 

561 name = backend_args["name"] 

562 

563 with h5pyfile(filename, h5_mode, comm=comm) as h5file: 

564 hdf = _create_group(h5file, "/VTKHDF", h5_mode) 

565 hdf.attrs.create("Type", "MultiBlockDataSet") 

566 hdf.attrs["Version"] = _vtk_hdf_version 

567 

568 mesh_group = _create_group(hdf, name, h5_mode) 

569 mesh_group.attrs.create("Type", "UnstructuredGrid") 

570 mesh_group.attrs["Version"] = _vtk_hdf_version 

571 

572 assembly = _create_group(hdf, "Assembly", h5_mode) 

573 

574 mesh_assembly = _create_group(assembly, name, h5_mode) 

575 if name not in mesh_assembly.keys(): 

576 mesh_assembly[name] = h5py.SoftLink(f"/VTKHDF/{name}") 

577 

578 # VTKHDF defines its higher-order cells as equispaced and has nowhere to 

579 # record a Lagrange variant, so it goes in our own top-level group, which 

580 # the specification reserves for exactly this. Without it a degree>1 mesh 

581 # reads back with its coordinates intact but its cells curving 

582 # differently between them. 

583 metadata = _create_group(h5file, _metadata_group, h5_mode) 

584 mesh_metadata = _create_group(metadata, name, h5_mode) 

585 mesh_metadata.attrs[_lagrange_variant_attr] = np.int32(mesh.lagrange_variant) 

586 

587 # Write time dependent points 

588 number_of_points = _create_dataset( 

589 mesh_group, 

590 "NumberOfPoints", 

591 shape=(1,), 

592 dtype=np.int64, 

593 chunks=True, 

594 maxshape=(None,), 

595 mode=h5_mode, 

596 ) 

597 number_of_points[-1] = mesh.num_nodes_global 

598 

599 # Store nodes 

600 points = _create_dataset( 

601 mesh_group, 

602 "Points", 

603 shape=(mesh.num_nodes_global, 3), 

604 dtype=mesh.local_geometry.dtype, 

605 chunks=True, 

606 maxshape=(None, 3), 

607 mode=h5_mode, 

608 ) 

609 insert_slice = _compute_append_slice( 

610 points, mesh.num_nodes_global, mesh.local_geometry_pos, h5_mode 

611 ) 

612 points[insert_slice, : mesh.local_geometry.shape[1]] = mesh.local_geometry 

613 

614 # NOTE: VTKHDF currently does not support reading time dependent topology in 

615 # Paraview: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

616 # Therefore the following data is not resized 

617 time_indep_datasets = {} 

618 for key in ["NumberOfCells", "NumberOfConnectivityIds"]: 

619 time_indep_datasets[key] = _create_dataset( 

620 mesh_group, 

621 key, 

622 shape=(1,), 

623 dtype=np.int64, 

624 chunks=True, 

625 maxshape=(None,), 

626 mode=h5_mode, 

627 resize=False, # Resize should really be True 

628 ) 

629 

630 num_dofs_per_cell = mesh.local_topology.shape[1] 

631 time_indep_datasets["NumberOfCells"][-1] = mesh.num_cells_global 

632 time_indep_datasets["NumberOfConnectivityIds"][-1] = ( 

633 mesh.num_cells_global * num_dofs_per_cell 

634 ) 

635 

636 # NOTE: The following offsets are currently overwriting the existing topology data. 

637 # This is due to the VTKHDF bug. When we switch resize=True this will automatically work 

638 # for time dependent topologies 

639 

640 # Store topology offsets (single celltype assumption) 

641 offsets = _create_dataset( 

642 mesh_group, 

643 "Offsets", 

644 shape=(mesh.num_cells_global + 1,), 

645 dtype=np.int64, 

646 chunks=True, 

647 mode=h5_mode, 

648 maxshape=(None,), 

649 resize=False, # Resize should really be True 

650 ) 

651 offset_data = np.arange(0, mesh.local_topology.size + 1, mesh.local_topology.shape[1]) 

652 offset_data += num_dofs_per_cell * mesh.local_topology_pos[0] 

653 insert_slice = _compute_append_slice( 

654 offsets, 

655 mesh.num_cells_global + 1, 

656 (mesh.local_topology_pos[0], mesh.local_topology_pos[1] + 1), 

657 mode=h5_mode, 

658 ) 

659 offsets[insert_slice] = offset_data 

660 del offset_data 

661 

662 # Permute and store topology data 

663 dx_ct = dolfinx.mesh.to_type(mesh.cell_type) 

664 top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell)) 

665 topology_data = mesh.local_topology[:, top_perm].flatten() 

666 topology = _create_dataset( 

667 mesh_group, 

668 "Connectivity", 

669 shape=(mesh.num_cells_global * num_dofs_per_cell,), 

670 dtype=np.int64, 

671 chunks=True, 

672 maxshape=(None,), 

673 mode=h5_mode, 

674 resize=False, # Resize should really be True, see issue below 

675 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

676 insert_slice = _compute_append_slice( 

677 topology, 

678 mesh.num_cells_global * num_dofs_per_cell, 

679 np.array(mesh.local_topology_pos) * num_dofs_per_cell, 

680 mode=h5_mode, 

681 ) 

682 topology[insert_slice] = topology_data 

683 del topology_data 

684 

685 # Store celltypes 

686 cell_types = np.full( 

687 mesh.local_topology.shape[0], 

688 fenicsx_compat.vtk_cell_type(dx_ct, num_dofs_per_cell), 

689 ) 

690 types = _create_dataset( 

691 mesh_group, 

692 "Types", 

693 shape=(mesh.num_cells_global,), 

694 dtype=np.uint8, 

695 maxshape=(None,), 

696 chunks=True, 

697 mode=h5_mode, 

698 resize=False, # Resize should really be True, see issue below 

699 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

700 insert_slice = _compute_append_slice( 

701 types, mesh.num_cells_global, mesh.local_topology_pos, h5_mode 

702 ) 

703 types[insert_slice] = cell_types 

704 del cell_types 

705 

706 steps = _create_group(mesh_group, "Steps", mode=h5_mode) 

707 # First fetch time-steps to see if we have stored this timestep already 

708 values = _create_dataset( 

709 steps, 

710 "Values", 

711 shape=(1,), 

712 dtype=np.float64, 

713 chunks=True, 

714 maxshape=(None,), 

715 mode=h5_mode, 

716 resize=h5_mode == "a", 

717 ) 

718 

719 if h5_mode == "w": 

720 values[0] = time 

721 else: 

722 existing_steps = values[:-1] 

723 if len(np.flatnonzero(np.isclose(existing_steps, time))) > 0: 

724 raise RuntimeError(f"Mesh already exists at time {time} in {filename}.") 

725 values[-1] = time 

726 steps.attrs.create("NSteps", np.int64(len(values)), dtype=np.int64) 

727 

728 # Write offset data for current time-step 

729 all_parts = {} 

730 for key in [ 

731 "NumberOfParts", 

732 "PartOffsets", 

733 "PointOffsets", 

734 "CellOffsets", 

735 "ConnectivityIdOffsets", 

736 ]: 

737 all_parts[key] = _create_dataset( 

738 steps, 

739 key, 

740 shape=(1,), 

741 dtype=np.int64, 

742 chunks=True, 

743 maxshape=(None,), 

744 mode=h5_mode, 

745 ) 

746 

747 all_parts["NumberOfParts"][-1] = 1 

748 all_parts["PartOffsets"][-1] = 0 

749 all_parts["PointOffsets"][-1] = points.shape[0] - mesh.num_nodes_global 

750 all_parts["CellOffsets"][-1] = types.shape[0] - mesh.num_cells_global 

751 all_parts["ConnectivityIdOffsets"][-1] = ( 

752 topology.shape[0] - mesh.num_cells_global * num_dofs_per_cell 

753 ) 

754 

755 # Update cell-data and point-data offsets by copying over data from previous step 

756 for key in ["CellDataOffsets", "PointDataOffsets"]: 

757 group = _create_group(steps, key, mode=h5_mode) 

758 for cd in group.keys(): 

759 group[cd].resize(steps.attrs["NSteps"], axis=0) 

760 group[cd][-1] = group[cd][-2] 

761 

762 # Update Steps in all other parts of the mesh as well 

763 for key in mesh_assembly.keys(): 

764 if key == name: 

765 continue 

766 # Copy time-dependent geometry info (NumberOfPoints) from mesh to tag 

767 sub_group = mesh_assembly[key] 

768 sub_step = sub_group["Steps"] 

769 sub_step.attrs["NSteps"] = steps.attrs["NSteps"] 

770 

771 # Copy time dependent and partition info from mesh to tag 

772 step_copy_keys = ["Values", "PartOffsets", "NumberOfParts"] 

773 for key in step_copy_keys: 

774 if key in sub_step.keys(): 

775 sub_step[key].resize(steps[key].shape) 

776 sub_step[key][:] = steps[key][:] 

777 else: 

778 raise RuntimeError(f"{sub_step.name} should have {key}/") 

779 

780 # Append value from previous step for meshtags as they are time-independent 

781 append_keys = ["CellOffsets", "ConnectivityIdOffsets"] 

782 for key in append_keys: 

783 if key in sub_step.keys(): 

784 sub_step[key].resize(sub_step[key].size + 1, axis=0) 

785 sub_step[key][-1] = sub_step[key][-2] 

786 else: 

787 raise RuntimeError(f"{sub_step.name} should have {key}/") 

788 # Append value from previous step for meshtags celldata 

789 for key in ["CellDataOffsets", "PointDataOffsets"]: 

790 group = _create_group(sub_step, key, mode=h5_mode) 

791 for cd in group.keys(): 

792 group[cd].resize(sub_step.attrs["NSteps"], axis=0) 

793 group[cd][-1] = group[cd][-2] 

794 

795 

796def write_meshtags( 

797 filename: str | Path, 

798 comm: MPI.Comm, 

799 data: MeshTagsData, 

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

801): 

802 """Write mesh tags to file. 

803 

804 Args: 

805 filename: Path to file to write to 

806 comm: MPI communicator used in storage 

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

808 backend_args: Arguments to backend 

809 """ 

810 backend_args = get_default_backend_args(backend_args) 

811 name = backend_args["name"] 

812 h5_mode = "a" 

813 with h5pyfile(filename, h5_mode, comm=comm) as h5file: 

814 hdf = h5file["VTKHDF"] 

815 file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) 

816 # Check for type of VTKHDF file 

817 if file_type != "MultiBlockDataSet": 

818 raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}") 

819 

820 # We place the mesh-tags in the same subgroup of assembly as the mesh 

821 parent_mesh_group = _get_vtk_group(h5file, name) 

822 block = parent_mesh_group.parent 

823 

824 tag_path = f"/VTKHDF/{name}_{data.name}" 

825 if data.name not in block.keys(): 

826 block[data.name] = h5py.SoftLink(tag_path) 

827 

828 mesh_group = _create_group(hdf, tag_path, mode=h5_mode) 

829 mesh_group.attrs.create("Type", "UnstructuredGrid") 

830 mesh_group.attrs["Version"] = _vtk_hdf_version 

831 

832 cell_data = _create_group(mesh_group, "CellData", mode=h5_mode) 

833 assert data.num_entities_global is not None 

834 dataset = _create_dataset( 

835 cell_data, 

836 data.name, 

837 shape=(data.num_entities_global,), 

838 dtype=data.values.dtype, 

839 chunks=True, 

840 maxshape=(None,), 

841 mode=h5_mode, 

842 ) 

843 assert data.local_start is not None 

844 insert_slice = _compute_append_slice( 

845 dataset, 

846 data.num_entities_global, 

847 np.array([data.local_start, data.local_start + data.indices.shape[0]]), 

848 mode=h5_mode, 

849 ) 

850 dataset[insert_slice] = data.values 

851 

852 # NOTE: The following is more or less a copy from write_mesh, 

853 # except that we pull out the point storage and use a softlink 

854 num_cells = _create_dataset( 

855 mesh_group, 

856 "NumberOfCells", 

857 shape=(1,), 

858 dtype=np.int64, 

859 chunks=True, 

860 maxshape=(None,), 

861 mode=h5_mode, 

862 resize=False, # Resize should really be True, see issue below 

863 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

864 num_cells[-1] = data.num_entities_global 

865 

866 # Hardlink data should also follow hardlink for numbering 

867 for key in ["Points", "NumberOfPoints"]: 

868 if key not in mesh_group.keys(): 

869 mesh_group[key] = parent_mesh_group[key] 

870 

871 # Single celltype assumption 

872 num_dofs_per_cell = data.num_dofs_per_entity 

873 number_of_connectivities = _create_dataset( 

874 mesh_group, 

875 "NumberOfConnectivityIds", 

876 shape=(1,), 

877 dtype=np.int64, 

878 chunks=True, 

879 maxshape=(None,), 

880 mode=h5_mode, 

881 resize=False, # Resize should really be True, see issue below 

882 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

883 assert data.num_entities_global is not None and num_dofs_per_cell is not None 

884 number_of_connectivities[-1] = data.num_entities_global * num_dofs_per_cell 

885 

886 # Store topology offsets (single celltype assumption) 

887 offsets = _create_dataset( 

888 mesh_group, 

889 "Offsets", 

890 shape=(data.num_entities_global + 1,), 

891 dtype=np.int64, 

892 chunks=True, 

893 mode=h5_mode, 

894 maxshape=(None,), 

895 resize=False, # Resize should really be True, see issue below 

896 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

897 offset_data = np.arange(0, data.indices.size + 1, data.indices.shape[1]) 

898 assert data.local_start is not None 

899 offset_data += num_dofs_per_cell * data.local_start 

900 insert_slice = _compute_append_slice( 

901 offsets, 

902 data.num_entities_global + 1, 

903 (data.local_start, data.local_start + data.indices.shape[0] + 1), 

904 mode=h5_mode, 

905 ) 

906 offsets[insert_slice] = offset_data 

907 del offset_data 

908 

909 # Permute and store topology data 

910 if data.cell_type is None: 

911 raise ValueError("Cannot write meshtags without cell-type information") 

912 dx_ct = dolfinx.mesh.to_type(data.cell_type) 

913 top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell)) 

914 topology_data = data.indices[:, top_perm].flatten() 

915 topology = _create_dataset( 

916 mesh_group, 

917 "Connectivity", 

918 shape=(data.num_entities_global * num_dofs_per_cell,), 

919 dtype=np.int64, 

920 chunks=True, 

921 maxshape=(None,), 

922 mode=h5_mode, 

923 resize=False, # Resize should really be True, see issue below 

924 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

925 insert_slice = _compute_append_slice( 

926 topology, 

927 data.num_entities_global * num_dofs_per_cell, 

928 np.array([data.local_start, data.local_start + data.indices.shape[0]]) 

929 * num_dofs_per_cell, 

930 mode=h5_mode, 

931 ) 

932 topology[insert_slice] = topology_data 

933 del topology_data 

934 

935 # Store celltypes 

936 cell_types = np.full( 

937 data.indices.shape[0], 

938 fenicsx_compat.vtk_cell_type(dx_ct, num_dofs_per_cell), 

939 ) 

940 types = _create_dataset( 

941 mesh_group, 

942 "Types", 

943 shape=(data.num_entities_global,), 

944 dtype=np.uint8, 

945 maxshape=(None,), 

946 chunks=True, 

947 mode=h5_mode, 

948 resize=False, # Resize should really be True, see issue below 

949 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 

950 insert_slice = _compute_append_slice( 

951 types, 

952 data.num_entities_global, 

953 (data.local_start, data.local_start + data.indices.shape[0]), 

954 h5_mode, 

955 ) 

956 types[insert_slice] = cell_types 

957 del cell_types 

958 

959 steps = _create_group(mesh_group, "Steps", mode=h5_mode) 

960 

961 # Copy n-step counter 

962 steps.attrs.create("NSteps", parent_mesh_group["Steps"].attrs["NSteps"]) 

963 

964 # Hardlink data that we know is the same across meshes 

965 hardlink_keys = ["NumberOfParts", "PartOffsets", "Values", "PointOffsets"] 

966 for key in hardlink_keys: 

967 steps[key] = parent_mesh_group["Steps"][key] 

968 

969 # Write offset data for current time-step 

970 all_parts = {} 

971 for key in ["CellOffsets", "ConnectivityIdOffsets"]: 

972 all_parts[key] = _create_dataset( 

973 steps, 

974 key, 

975 shape=(1,), 

976 dtype=np.int64, 

977 chunks=True, 

978 maxshape=(None,), 

979 mode=h5_mode, 

980 ) 

981 

982 all_parts["CellOffsets"][-1] = types.shape[0] - data.num_entities_global 

983 all_parts["ConnectivityIdOffsets"][-1] = offsets.shape[0] - (data.num_entities_global + 1) 

984 # CellData requires an offset 

985 cd_off = _create_group(steps, "CellDataOffsets", mode=h5_mode) 

986 cd_data = _create_dataset( 

987 cd_off, 

988 data.name, 

989 shape=parent_mesh_group["Steps"]["CellOffsets"].shape, 

990 dtype=np.int64, 

991 chunks=True, 

992 maxshape=(None,), 

993 mode=h5_mode, 

994 ) 

995 cd_data[-1] = types.shape[0] - data.num_entities_global 

996 

997 

998def read_meshtags_data( 

999 filename: str | Path, 

1000 comm: MPI.Comm, 

1001 name: str, 

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

1003) -> MeshTagsData: 

1004 """Read mesh tags from file. 

1005 

1006 Args: 

1007 filename: Path to file to read from 

1008 comm: MPI communicator used in storage 

1009 name: Name of the mesh tags to read 

1010 backend_args: Arguments to backend 

1011 

1012 Returns: 

1013 Internal data structure for the mesh tags read from file 

1014 """ 

1015 backend_args = get_default_backend_args(backend_args) 

1016 backend_args.update({"name": name}) 

1017 # Reuse reading cell-data 

1018 indices, values = read_cell_data(filename, name, comm, None, backend_args=backend_args) 

1019 # Read cell-type of grid to get topological dimension 

1020 with h5pyfile(filename, "r", comm=comm) as h5file: 

1021 hdf = _get_vtk_group(h5file, backend_args["name"]) 

1022 num_cells_global = hdf["Types"].size 

1023 local_cell_range = compute_local_range(comm, num_cells_global) 

1024 cell_types_local = hdf["Types"][slice(*local_cell_range)] 

1025 unique_cells = find_all_unique_cell_types(comm, cell_types_local, indices.shape[1]) 

1026 if unique_cells.shape[0] > 1: 

1027 raise NotImplementedError("io4dolfinx does not support mixed celltype grids") 

1028 vtk_cell_type = unique_cells[0][0] 

1029 if vtk_cell_type in _first_order_vtk.keys(): 

1030 ct = _first_order_vtk[vtk_cell_type] 

1031 elif vtk_cell_type in _arbitrary_lagrange_vtk.keys(): 

1032 ct = _arbitrary_lagrange_vtk[vtk_cell_type] 

1033 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(ct)) 

1034 return MeshTagsData(name=name, values=values.flatten(), indices=indices, dim=dim) 

1035 

1036 

1037def read_dofmap( 

1038 filename: str | Path, 

1039 comm: MPI.Comm, 

1040 name: str, 

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

1042) -> dolfinx.graph.AdjacencyList: 

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

1044 

1045 Args: 

1046 filename: Path to file to read from 

1047 comm: MPI communicator used in storage 

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

1049 backend_args: Arguments to backend 

1050 

1051 Returns: 

1052 Dofmap as an AdjacencyList 

1053 """ 

1054 raise NotImplementedError("The Pyvista backend cannot make checkpoints.") 

1055 

1056 

1057def read_dofs( 

1058 filename: str | Path, 

1059 comm: MPI.Comm, 

1060 name: str, 

1061 time: float, 

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

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

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

1065 

1066 Args: 

1067 filename: Path to file to read from 

1068 comm: MPI communicator used in storage 

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

1070 time: Time stamp associated with the function to read 

1071 backend_args: Arguments to backend 

1072 

1073 Returns: 

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

1075 and the global starting point on the process. 

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

1077 """ 

1078 raise NotImplementedError("The Pyvista backend cannot make checkpoints.") 

1079 

1080 

1081def read_cell_perms( 

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

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

1084 """ 

1085 Read cell permutation from file with given communicator, 

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

1087 

1088 Args: 

1089 comm: MPI communicator used in storage 

1090 filename: Path to file to read from 

1091 backend_args: Arguments to backend 

1092 

1093 Returns: 

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

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

1096 """ 

1097 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") 

1098 

1099 

1100def write_function( 

1101 filename: Path, 

1102 comm: MPI.Comm, 

1103 u: FunctionData, 

1104 time: float, 

1105 mode: FileMode, 

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

1107): 

1108 """ 

1109 Write a function to file. 

1110 

1111 Args: 

1112 comm: MPI communicator used in storage 

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

1114 filename: Path to file to write to 

1115 time: Time stamp associated with function 

1116 mode: File-mode to store the function 

1117 backend_args: Arguments to backend 

1118 """ 

1119 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") 

1120 

1121 

1122def read_legacy_mesh( 

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

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

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

1126 legacy DOLFIN HDF5-file. 

1127 

1128 Args: 

1129 filename: Path to file to read from 

1130 comm: MPI communicator used in storage 

1131 group: Group in HDF5 file where mesh is stored 

1132 

1133 Returns: 

1134 Tuple containing: 

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

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

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

1138 """ 

1139 raise NotImplementedError("The VTKHDF backend cannot read legacy DOLFIN meshes.") 

1140 

1141 

1142def snapshot_checkpoint( 

1143 filename: Path | str, 

1144 mode: FileMode, 

1145 u: dolfinx.fem.Function, 

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

1147): 

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

1149 

1150 Args: 

1151 filename: Path to file to read from 

1152 mode: File-mode to store the function 

1153 u: dolfinx function to create a snapshot checkpoint for 

1154 backend_args: Arguments to backend 

1155 """ 

1156 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.") 

1157 

1158 

1159def read_hdf5_array( 

1160 comm: MPI.Comm, 

1161 filename: Path | str, 

1162 group: str, 

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

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

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

1166 

1167 Args: 

1168 comm: MPI communicator used in storage 

1169 filename: Path to file to read from 

1170 group: Group in HDF5 file where array is stored 

1171 backend_args: Arguments to backend 

1172 

1173 Returns: 

1174 Tuple containing: 

1175 - Numpy array read from file 

1176 - Global starting point on the process. 

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

1178 """ 

1179 raise NotImplementedError("The VTKHDF backend cannot read HDF5 arrays") 

1180 

1181 

1182def write_data( 

1183 filename: Path | str, 

1184 array_data: ArrayData, 

1185 comm: MPI.Comm, 

1186 time: str | float | None, 

1187 mode: FileMode, 

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

1189): 

1190 """Write function to file by interpolating into geometry nodes. 

1191 

1192 

1193 Args: 

1194 filename: Path to file 

1195 array_data: Data to write to file 

1196 time: Time stamp 

1197 mode: Append or write 

1198 backend_args: The backend arguments 

1199 """ 

1200 h5_mode = convert_file_mode(mode) 

1201 assert h5_mode == "a" 

1202 backend_args = get_default_backend_args(backend_args) 

1203 mesh_name = backend_args["name"] 

1204 extension = array_data.type 

1205 with h5pyfile(filename, h5_mode, comm=comm) as h5file: 

1206 hdf = h5file["VTKHDF"] 

1207 file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) 

1208 # Check for type of VTKHDF file 

1209 if file_type != "MultiBlockDataSet": 

1210 raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}") 

1211 

1212 # Find mesh block to add data to 

1213 block = _get_vtk_group(h5file, mesh_name) 

1214 

1215 data_group = _create_group(block, f"{extension}Data", mode=h5_mode) 

1216 dataset = _create_dataset( 

1217 data_group, 

1218 array_data.name, 

1219 shape=array_data.global_shape, 

1220 dtype=array_data.values.dtype, 

1221 chunks=True, 

1222 maxshape=(None, array_data.values.shape[1]), 

1223 mode=h5_mode, 

1224 ) 

1225 insert_slice = _compute_append_slice( 

1226 dataset, 

1227 array_data.global_shape[0], 

1228 np.array( 

1229 [ 

1230 array_data.local_range[0], 

1231 array_data.local_range[0] + array_data.values.shape[0], 

1232 ] 

1233 ), 

1234 mode=h5_mode, 

1235 ) 

1236 dataset[insert_slice] = array_data.values 

1237 

1238 steps = _create_group(block, "Steps", mode=h5_mode) 

1239 pdo = _create_group(steps, f"{extension}DataOffsets", mode=h5_mode) 

1240 

1241 # Check if time step is already in time-stepping of mesh 

1242 timestamps = steps["Values"][:] 

1243 assert isinstance(timestamps, np.ndarray) 

1244 assert isinstance(time, float) 

1245 time_exists = np.flatnonzero(np.isclose(timestamps, time)) 

1246 if len(time_exists) > 1: 

1247 raise ValueError("Time-step has been written multiple times") 

1248 pdo_u = _create_dataset( 

1249 pdo, 

1250 array_data.name, 

1251 shape=(1,), 

1252 dtype=np.int64, 

1253 chunks=True, 

1254 maxshape=(None,), 

1255 mode=h5_mode, 

1256 resize=False, 

1257 ) 

1258 if len(time_exists) == 1: 

1259 idx = time_exists[0] 

1260 elif len(time_exists) == 0: 

1261 # No mesh written at step, update mesh offsets 

1262 # Geometry has to be time-dependent 

1263 num_points = block["NumberOfPoints"] 

1264 num_points.resize(len(timestamps) + 1, axis=0) 

1265 num_points[-1] = num_points[-2] 

1266 # Even if topology cannot be time dep at the moment, number of cells 

1267 # must be updated 

1268 num_cells = block["NumberOfCells"] 

1269 num_cells.resize(len(timestamps) + 1, axis=0) 

1270 num_cells[-1] = num_cells[-2] 

1271 

1272 steps.attrs.create("NSteps", block["Steps"].attrs["NSteps"] + 1) 

1273 step_vals = _create_dataset( 

1274 steps, 

1275 "Values", 

1276 shape=(1,), 

1277 dtype=np.float64, 

1278 chunks=True, 

1279 maxshape=(None,), 

1280 mode=h5_mode, 

1281 resize=True, 

1282 ) 

1283 step_vals[-1] = time 

1284 idx = step_vals.size - 1 

1285 for key in [ 

1286 "PartOffsets", 

1287 "NumberOfParts", 

1288 "PointOffsets", 

1289 "CellOffsets", 

1290 "ConnectivityIdOffsets", 

1291 ]: 

1292 comp = steps[key] 

1293 comp.resize(steps.attrs["NSteps"], axis=0) 

1294 if comp.size > 1: 

1295 comp[-1] = comp[-2] 

1296 else: 

1297 raise ValueError(f"Time step found multiple times in {filename}") 

1298 

1299 # Update offsets for current variable 

1300 pdo_u.resize(steps.attrs["NSteps"], axis=0) 

1301 pdo_u[idx] = dataset.shape[0] - array_data.global_shape[0] 

1302 

1303 # Point and cell data of all other variables has to be updated to be synced 

1304 for key in [ 

1305 "CellDataOffsets", 

1306 "PointDataOffsets", 

1307 ]: 

1308 comp = steps[key] 

1309 num_steps = steps.attrs["NSteps"] 

1310 for dname, dset in comp.items(): 

1311 # Only resize other data arrays 

1312 if dname != array_data.name: 

1313 # Only resize if it hasn't already been resized 

1314 if dset.size != num_steps: 

1315 dset.resize(num_steps, axis=0) 

1316 # Only update data if we are further than first time step 

1317 if dset.size > 1: 

1318 dset[-1] = dset[-2]