Coverage for / dolfinx-env / lib / python3.12 / site-packages / io4dolfinx / backends / vtkhdf / backend.py: 84%
469 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-26 18:16 +0000
1"""
2Module that can read the VTKHDF format using h5py.
3"""
5from pathlib import Path
6from typing import Any
8from mpi4py import MPI
10import basix
11import dolfinx
12import h5py
13import numpy as np
14import numpy.typing as npt
16from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
17from io4dolfinx.utils import check_file_exists, compute_local_range
19from .. import FileMode, ReadMode
20from ..h5py.backend import convert_file_mode, h5pyfile
21from ..pyvista.backend import _arbitrary_lagrange_vtk, _cell_degree, _first_order_vtk
23read_mode = ReadMode.parallel
25_vtk_hdf_version = np.array([2, 1], dtype=np.int32)
28def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]:
29 """Get default backend arguments given a set of input arguments.
31 Args:
32 arguments: Input backend arguments
34 Returns:
35 Updated backend arguments
36 """
37 args = arguments or {"name": "mesh"}
38 return args
41def find_all_unique_cell_types(comm, cell_types, num_nodes):
42 """
43 Given a set of cell types and number of nodes per cell, find all unique cell types
44 across all ranks.
46 Args:
47 comm: MPI communicator
48 cell_types: Local cell types
49 num_nodes: Number of nodes per cell
52 Returns:
53 A 2D array where each row corresponds to a cell type (vtk int)
54 and the number of nodes.
55 """
56 # Combine cell_types, num_nodes as tuple
57 c_hash = np.zeros((2, len(cell_types)), dtype=np.int32)
58 c_hash[0] = cell_types
59 c_hash[1] = num_nodes
60 indexes = np.unique(c_hash.T, axis=0, return_index=True)[1]
61 local_unique_cells = c_hash.T[indexes]
63 all_cell_types = np.vstack(comm.allgather(local_unique_cells))
64 indexes = np.unique(all_cell_types, axis=0, return_index=True)[1]
65 all_unique_cell_types = all_cell_types[indexes]
66 return all_unique_cell_types
69def _decode_bytes_if_needed(value: bytes | str) -> str:
70 """Decode bytes to string if necessary (for h5py compatibility)"""
71 if isinstance(value, bytes):
72 return value.decode("utf-8")
73 return value
76def _get_vtk_group(h5file, name: str) -> h5py.Group:
77 """
78 Navigates the VTKHDF group hierarchy to find the specific first
79 UnstructuredGrid with a specific name.
80 Handles both MultiBlockDataSet and direct UnstructuredGrid types.
81 """
82 hdf = h5file["VTKHDF"]
83 file_type = _decode_bytes_if_needed(hdf.attrs["Type"])
85 if file_type == "MultiBlockDataSet":
86 ass = hdf["Assembly"]
88 def visitor(path):
89 mesh_group = path.rsplit("/", 1)
90 mesh_name = mesh_group[0] if len(mesh_group) == 1 else mesh_group[1]
91 if mesh_name == name:
92 obj = ass.get(path)
93 # Check attributes carefully
94 if "Type" in obj.attrs.keys():
95 attr_type = obj.attrs["Type"]
96 if isinstance(attr_type, bytes):
97 attr_type = attr_type.decode("utf-8")
98 if attr_type == "UnstructuredGrid":
99 return path
100 return None
102 mesh_node = ass.visit_links(visitor)
103 if mesh_node is None:
104 raise RuntimeError(f"Could not find unique mesh named '{name}' in Assembly.")
106 return ass[mesh_node]
108 elif file_type == "UnstructuredGrid":
109 return hdf
110 else:
111 raise RuntimeError(f"Not supported file type {file_type}")
114def _get_time_index(hdf: h5py.Group, time: float | str, filename: str | Path) -> int:
115 """Finds the index of a specific time stamp."""
116 if "Steps" not in hdf.keys():
117 raise RuntimeError(f"No timestepping information found in {filename}.")
118 stamps = hdf["Steps"]["Values"][:]
119 pos = np.flatnonzero(np.isclose(stamps, time))
120 if len(pos) == 0:
121 raise RuntimeError(f"Could not find mesh at t={time} in {filename}.")
122 elif len(pos) > 1:
123 raise RuntimeError(f"Multiple time steps for mesh at t={time} in {filename}")
124 return pos[0]
127def read_mesh_data(
128 filename: Path | str,
129 comm: MPI.Intracomm,
130 time: str | float | None,
131 read_from_partition: bool,
132 backend_args: dict[str, Any] | None,
133) -> ReadMeshData:
134 """Read mesh data from file.
136 Args:
137 filename: Path to file to read from
138 comm: MPI communicator used in storage
139 time: Time stamp associated with the mesh to read
140 read_from_partition: Whether to read partition information
141 backend_args: Arguments to backend
143 Returns:
144 Internal data structure for the mesh data read from file
145 """
146 backend_args = get_default_backend_args(backend_args)
147 check_file_exists(filename)
148 if read_from_partition:
149 raise RuntimeError("Cannot read partition data with VTKHDF")
150 with h5pyfile(filename, "r", comm=comm) as h5file:
151 hdf = _get_vtk_group(h5file, backend_args["name"])
152 if time is None:
153 num_cells_global = hdf["Types"].size
154 local_cell_range = compute_local_range(comm, num_cells_global)
155 cell_types_local = hdf["Types"][slice(*local_cell_range)]
157 num_points_global = hdf["NumberOfPoints"][0]
158 local_point_range = compute_local_range(comm, num_points_global)
159 points_local = hdf["Points"][slice(*local_point_range)]
161 # Connectivity read
162 offsets = hdf["Offsets"]
163 local_connectivity_offset = offsets[local_cell_range[0] : local_cell_range[1] + 1]
164 topology = hdf["Connectivity"][
165 local_connectivity_offset[0] : local_connectivity_offset[-1]
166 ]
167 offset = local_connectivity_offset - local_connectivity_offset[0]
168 else:
169 time_index = _get_time_index(hdf, time, filename)
171 stamps = hdf["Steps"]["Values"][:]
172 # Get number of points
173 point_node = hdf["Points"]
174 step_start = hdf["Steps"]["PointOffsets"][time_index]
176 # NOTE: currently, it doesn't seem like we follow:
177 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata
178 # As only one num_points is stored irregardless of time steps added.
179 if hdf["NumberOfPoints"].shape[0] != len(stamps):
180 num_pts = hdf["NumberOfPoints"][0]
181 else:
182 num_pts = hdf["NumberOfPoints"][time_index]
183 lr = compute_local_range(comm, num_pts)
184 points_local = point_node[step_start + lr[0] : step_start + lr[1]]
186 # Get cell-types in step
187 cell_start = hdf["Steps"]["CellOffsets"][time_index]
188 if hdf["NumberOfCells"].shape[0] != len(stamps):
189 num_cells = hdf["NumberOfCells"][0]
190 else:
191 num_cells = hdf["NumberOfCells"][time_index]
192 local_cell_range = compute_local_range(comm, num_cells)
193 cell_types_local = hdf["Types"][
194 cell_start + local_cell_range[0] : cell_start + local_cell_range[1]
195 ]
197 # Get connectivity in step
198 connectivity_start = hdf["Steps"]["ConnectivityIdOffsets"][time_index]
199 # Connectivity read
200 offsets = hdf["Offsets"]
201 local_connectivity_offset = offsets[
202 connectivity_start + local_cell_range[0] : connectivity_start
203 + local_cell_range[1]
204 + 1
205 ]
206 topology = hdf["Connectivity"][
207 local_connectivity_offset[0] : local_connectivity_offset[-1]
208 ]
209 offset = local_connectivity_offset - local_connectivity_offset[0]
211 # NOTE: Currently we limit ourselfs to a single celltype, as it makes life easier,
212 # other things have to change in `MeshReadData` to support this.
213 num_nodes_per_cell = offset[1:] - offset[:-1]
214 unique_cells = find_all_unique_cell_types(MPI.COMM_WORLD, cell_types_local, num_nodes_per_cell)
215 if unique_cells.shape[0] > 1:
216 raise NotImplementedError("io4dolfinx does not support mixed celltype grids")
217 topology = topology.reshape(-1, num_nodes_per_cell[0])
218 cell_type, number_of_nodes = unique_cells[0]
219 gtype = backend_args.get("dtype", points_local.dtype)
220 if cell_type in _first_order_vtk.keys():
221 ct = _first_order_vtk[cell_type]
222 degree = 1
223 elif cell_type in _arbitrary_lagrange_vtk.keys():
224 ct = _arbitrary_lagrange_vtk[cell_type]
225 degree = _cell_degree(ct, num_nodes=number_of_nodes)
226 else:
227 raise ValueError(f"Unknown VTK cell type {cell_type} in {filename}")
228 perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(ct), number_of_nodes)
229 topology = topology[:, perm]
230 lvar = int(basix.LagrangeVariant.equispaced)
231 return ReadMeshData(
232 cells=topology, cell_type=ct, x=points_local.astype(gtype), lvar=lvar, degree=degree
233 )
236def read_point_data(
237 filename: Path | str,
238 name: str,
239 comm: MPI.Intracomm,
240 time: float | str | None,
241 backend_args: dict[str, Any] | None,
242) -> tuple[np.ndarray, int]:
243 """Read data from the nodes of a mesh.
245 Args:
246 filename: Path to file
247 name: Name of point data
248 comm: Communicator to launch IO on.
249 time: The time stamp
250 backend_args: The backend arguments
252 Returns:
253 Data local to process (contiguous, no mpi comm) and local start range
254 """
255 backend_args = get_default_backend_args(backend_args)
256 check_file_exists(filename)
257 with h5pyfile(filename, "r", comm=comm) as h5file:
258 hdf = _get_vtk_group(h5file, backend_args["name"])
259 point_data = hdf["PointData"]
260 assert point_data is not None
261 if name not in point_data.keys():
262 raise ValueError(f"No point data named {name} in {filename}.")
263 func_node = point_data[name]
265 if time is None:
266 data_shape = func_node.shape[0]
267 lr = compute_local_range(comm, data_shape)
268 data = func_node[slice(*lr)]
269 return data, lr[0]
270 else:
271 time_index = _get_time_index(hdf, time, filename)
272 stamps = hdf["Steps"]["Values"][:]
273 step_start = hdf["Steps"]["PointDataOffsets"][name][time_index]
274 # NOTE: currently, it doesn't seem like we follow:
275 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata
276 # As only one num_points is stored irregardless of time steps added.
277 if hdf["NumberOfPoints"].shape[0] != len(stamps):
278 num_pts = hdf["NumberOfPoints"][0]
279 else:
280 num_pts = hdf["NumberOfPoints"][time_index]
281 lr = compute_local_range(comm, num_pts)
282 return func_node[step_start + lr[0] : step_start + lr[1]], lr[0]
285def read_cell_data(
286 filename: Path | str,
287 name: str,
288 comm: MPI.Intracomm,
289 time: str | float | None,
290 backend_args: dict[str, Any] | None,
291) -> tuple[npt.NDArray[np.int64], np.ndarray]:
292 backend_args = get_default_backend_args(backend_args)
293 check_file_exists(filename)
294 with h5pyfile(filename, "r", comm=comm) as h5file:
295 hdf = _get_vtk_group(h5file, backend_args["name"])
297 if "CellData" not in hdf.keys():
298 raise RuntimeError(f"No cell data found in {filename}.")
299 cell_data = hdf["CellData"]
300 assert cell_data is not None
301 if name not in cell_data.keys():
302 raise ValueError(f"No cell data with name {name} in {filename}")
303 cell_data_node = cell_data[name]
304 assert cell_data_node is not None
305 if time is None:
306 cell_data_shape = cell_data_node.shape
307 num_cells_global = hdf["Types"].size
308 assert num_cells_global == cell_data_shape[0]
309 local_cell_range = compute_local_range(comm, num_cells_global)
310 data = cell_data_node[slice(*local_cell_range)]
311 else:
312 time_index = _get_time_index(hdf, time, filename)
313 stamps = hdf["Steps"]["Values"][:]
314 cd_start = hdf["Steps"]["CellDataOffsets"][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["NumberOfCells"].shape[0] != len(stamps):
320 number_of_cells = hdf["NumberOfCells"][0]
321 else:
322 number_of_cells = hdf["NumberOfCells"][time_index]
323 lr = compute_local_range(comm, number_of_cells)
324 data = cell_data_node[cd_start + lr[0] : cd_start + lr[1]]
326 # NOTE: THis could be optimized by hand-coding some communication in
327 # `read_cell_data` on the frontend side
328 md = read_mesh_data(
329 filename, comm, time=time, read_from_partition=False, backend_args=backend_args
330 )
331 if len(data.shape) == 1:
332 data = data.reshape(-1, 1)
333 return md.cells, data
336def write_attributes(
337 filename: Path | str,
338 comm: MPI.Intracomm,
339 name: str,
340 attributes: dict[str, np.ndarray],
341 backend_args: dict[str, Any] | None,
342):
343 """Write attributes to file.
345 Args:
346 filename: Path to file to write to
347 comm: MPI communicator used in storage
348 name: Name of the attribute group
349 attributes: Dictionary of attributes to write
350 backend_args: Arguments to backend
351 """
352 raise NotImplementedError("The Pyvista backend cannot write attributes.")
355def read_attributes(
356 filename: Path | str,
357 comm: MPI.Intracomm,
358 name: str,
359 backend_args: dict[str, Any] | None,
360) -> dict[str, Any]:
361 """Read attributes from file.
363 Args:
364 filename: Path to file to read from
365 comm: MPI communicator used in storage
366 name: Name of the attribute group
367 backend_args: Arguments to backend
369 Returns:
370 Dictionary of attributes read from file
371 """
372 raise NotImplementedError("The Pyvista backend cannot read attributes.")
375def read_timestamps(
376 filename: Path | str,
377 comm: MPI.Intracomm,
378 function_name: str,
379 backend_args: dict[str, Any] | None,
380) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
381 """Read timestamps from file.
383 Args:
384 filename: Path to file to read from
385 comm: MPI communicator used in storage
386 function_name: Name of the function to read timestamps for
387 backend_args: Arguments to backend
389 Returns:
390 Numpy array of timestamps read from file
391 """
392 backend_args = get_default_backend_args(backend_args)
393 check_file_exists(filename)
394 # Temporal data storage as described in
395 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-data
396 with h5pyfile(filename, "r", comm=comm) as h5file:
397 hdf = h5file["VTKHDF"]
398 if "Steps" in hdf.keys():
399 timestamps = hdf["Steps"]["Values"][:]
400 # For either point data or cell data, check if function exists,
401 # and check if offsets in time are changing between steps.
402 if "CellData" in hdf.keys() and function_name in hdf["CellData"].keys():
403 offsets = hdf["Steps"]["CellDataOffsets"][function_name]
404 step_offsets = offsets[:]
405 step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1])
406 return timestamps[step_diff]
407 elif "PointData" in hdf.keys() and function_name in hdf["PointData"].keys():
408 offsets = hdf["Steps"]["PointDataOffsets"][function_name]
409 step_offsets = offsets[:]
410 step_diff = np.flatnonzero(step_offsets[1:] - step_offsets[:-1])
411 # This only finds when the offset changes, does not capture first step
412 return np.hstack([[timestamps[0]], timestamps[step_diff]])
413 else:
414 raise RuntimeError(f"Function {function_name} is not assoicated with a time-stamp.")
415 else:
416 raise RuntimeError(f"{filename} does not contain time-stepping information.")
419def read_function_names(
420 filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None
421) -> list[str]:
422 """Read all function names from a file.
424 Args:
425 filename: Path to file
426 comm: MPI communicator to launch IO on.
427 backend_args: Arguments to backend
429 Returns:
430 A list of function names.
431 """
432 backend_args = get_default_backend_args(backend_args)
433 check_file_exists(filename)
434 with h5pyfile(filename, "r", comm=comm) as h5file:
435 hdf = h5file["VTKHDF"]
436 function_names = set()
437 if "CellData" in hdf.keys():
438 for item in hdf["CellData"].keys():
439 function_names.add(item)
440 if "PointData" in hdf.keys():
441 for item in hdf["PointData"].keys():
442 function_names.add(item)
443 return list(function_names)
446def _create_dataset(
447 root: h5py.File | h5py.Group,
448 name: str,
449 shape: tuple[int, ...],
450 dtype: npt.DTypeLike,
451 chunks: bool,
452 maxshape: tuple[int | None, ...],
453 mode: str,
454 resize: bool = True,
455) -> h5py.Dataset:
456 if name not in root.keys():
457 mode = "w"
458 if mode == "w":
459 dataset = root.create_dataset(
460 name, shape=shape, dtype=dtype, chunks=chunks, maxshape=maxshape
461 )
462 elif mode == "a":
463 dataset = root[name]
464 old_shape = dataset.shape
465 # Only resize for dimension
466 if resize:
467 if len(old_shape) == 1:
468 new_shape = (old_shape[0] + shape[0],)
469 else:
470 new_shape = (old_shape[0] + shape[0], *old_shape[1:])
471 dataset.resize(new_shape)
472 else:
473 raise ValueError(f"Unknown file mode '{mode}' when creating dataset {name} in {root}")
474 return dataset
477def _create_group(root: h5py.File | h5py.Group, name: str, mode: str) -> h5py.Group:
478 if name not in root.keys():
479 mode = "w"
480 if mode == "w":
481 # Track order has to be on to make multiblock work:
482 # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#partitioneddatasetcollection-and-multiblockdataset
483 group = root.create_group(name, track_order=True)
484 elif mode == "a":
485 group = root[name]
486 else:
487 raise ValueError("Unknown file mode '{h5_mode}'")
488 return group
491def _compute_append_slice(
492 dataset: h5py.Dataset, input_size: int, original_slice: tuple[int, int] | np.ndarray, mode: str
493) -> slice:
494 append_offset = dataset.shape[0] - input_size if mode == "a" else 0
495 return slice(*(np.asarray(original_slice) + append_offset).astype(np.int64))
498def write_mesh(
499 filename: Path | str,
500 comm: MPI.Intracomm,
501 mesh: MeshData,
502 backend_args: dict[str, Any] | None,
503 mode: FileMode,
504 time: float,
505):
506 """
507 Write a mesh to file.
509 Args:
510 comm: MPI communicator used in storage
511 mesh: Internal data structure for the mesh data to save to file
512 filename: Path to file to write to
513 backend_args: Arguments to backend
514 mode: File-mode to store the mesh
515 time: Time stamp associated with the mesh
516 """
517 h5_mode = convert_file_mode(mode)
518 backend_args = get_default_backend_args(backend_args)
519 name = backend_args["name"]
521 with h5pyfile(filename, h5_mode, comm=comm) as h5file:
522 hdf = _create_group(h5file, "/VTKHDF", h5_mode)
523 hdf.attrs.create("Type", "MultiBlockDataSet")
524 hdf.attrs["Version"] = _vtk_hdf_version
526 mesh_group = _create_group(hdf, name, h5_mode)
527 mesh_group.attrs.create("Type", "UnstructuredGrid")
528 mesh_group.attrs["Version"] = _vtk_hdf_version
530 assembly = _create_group(hdf, "Assembly", h5_mode)
532 mesh_assembly = _create_group(assembly, name, h5_mode)
533 if name not in mesh_assembly.keys():
534 mesh_assembly[name] = h5py.SoftLink(f"/VTKHDF/{name}")
536 # Write time dependent points
537 number_of_points = _create_dataset(
538 mesh_group,
539 "NumberOfPoints",
540 shape=(1,),
541 dtype=np.int64,
542 chunks=True,
543 maxshape=(None,),
544 mode=h5_mode,
545 )
546 number_of_points[-1] = mesh.num_nodes_global
548 # Store nodes
549 points = _create_dataset(
550 mesh_group,
551 "Points",
552 shape=(mesh.num_nodes_global, 3),
553 dtype=mesh.local_geometry.dtype,
554 chunks=True,
555 maxshape=(None, 3),
556 mode=h5_mode,
557 )
558 insert_slice = _compute_append_slice(
559 points, mesh.num_nodes_global, mesh.local_geometry_pos, h5_mode
560 )
561 points[insert_slice, : mesh.local_geometry.shape[1]] = mesh.local_geometry
563 # NOTE: VTKHDF currently does not support reading time dependent topology in
564 # Paraview: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
565 # Therefore the following data is not resized
566 time_indep_datasets = {}
567 for key in ["NumberOfCells", "NumberOfConnectivityIds"]:
568 time_indep_datasets[key] = _create_dataset(
569 mesh_group,
570 key,
571 shape=(1,),
572 dtype=np.int64,
573 chunks=True,
574 maxshape=(None,),
575 mode=h5_mode,
576 resize=False, # Resize should really be True
577 )
579 num_dofs_per_cell = mesh.local_topology.shape[1]
580 time_indep_datasets["NumberOfCells"][-1] = mesh.num_cells_global
581 time_indep_datasets["NumberOfConnectivityIds"][-1] = (
582 mesh.num_cells_global * num_dofs_per_cell
583 )
585 # NOTE: The following offsets are currently overwriting the existing topology data.
586 # This is due to the VTKHDF bug. When we switch resize=True this will automatically work
587 # for time dependent topologies
589 # Store topology offsets (single celltype assumption)
590 offsets = _create_dataset(
591 mesh_group,
592 "Offsets",
593 shape=(mesh.num_cells_global + 1,),
594 dtype=np.int64,
595 chunks=True,
596 mode=h5_mode,
597 maxshape=(None,),
598 resize=False, # Resize should really be True
599 )
600 offset_data = np.arange(0, mesh.local_topology.size + 1, mesh.local_topology.shape[1])
601 offset_data += num_dofs_per_cell * mesh.local_topology_pos[0]
602 insert_slice = _compute_append_slice(
603 offsets,
604 mesh.num_cells_global + 1,
605 (mesh.local_topology_pos[0], mesh.local_topology_pos[1] + 1),
606 mode=h5_mode,
607 )
608 offsets[insert_slice] = offset_data
609 del offset_data
611 # Permute and store topology data
612 dx_ct = dolfinx.mesh.to_type(mesh.cell_type)
613 top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell))
614 topology_data = mesh.local_topology[:, top_perm].flatten()
615 topology = _create_dataset(
616 mesh_group,
617 "Connectivity",
618 shape=(mesh.num_cells_global * num_dofs_per_cell,),
619 dtype=np.int64,
620 chunks=True,
621 maxshape=(None,),
622 mode=h5_mode,
623 resize=False, # Resize should really be True, see issue below
624 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
625 insert_slice = _compute_append_slice(
626 topology,
627 mesh.num_cells_global * num_dofs_per_cell,
628 np.array(mesh.local_topology_pos) * num_dofs_per_cell,
629 mode=h5_mode,
630 )
631 topology[insert_slice] = topology_data
632 del topology_data
634 # Store celltypes
635 cell_types = np.full(
636 mesh.local_topology.shape[0],
637 dolfinx.cpp.io.get_vtk_cell_type(dx_ct, dolfinx.mesh.cell_dim(dx_ct)),
638 )
639 types = _create_dataset(
640 mesh_group,
641 "Types",
642 shape=(mesh.num_cells_global,),
643 dtype=np.uint8,
644 maxshape=(None,),
645 chunks=True,
646 mode=h5_mode,
647 resize=False, # Resize should really be True, see issue below
648 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
649 insert_slice = _compute_append_slice(
650 types, mesh.num_cells_global, mesh.local_topology_pos, h5_mode
651 )
652 types[insert_slice] = cell_types
653 del cell_types
655 steps = _create_group(mesh_group, "Steps", mode=h5_mode)
656 # First fetch time-steps to see if we have stored this timestep already
657 values = _create_dataset(
658 steps,
659 "Values",
660 shape=(1,),
661 dtype=np.float64,
662 chunks=True,
663 maxshape=(None,),
664 mode=h5_mode,
665 resize=h5_mode == "a",
666 )
668 if h5_mode == "w":
669 values[0] = time
670 else:
671 existing_steps = values[:-1]
672 if len(np.flatnonzero(np.isclose(existing_steps, time))) > 0:
673 raise RuntimeError(f"Mesh already exists at time {time} in {filename}.")
674 values[-1] = time
675 steps.attrs.create("NSteps", np.int64(len(values)), dtype=np.int64)
677 # Write offset data for current time-step
678 all_parts = {}
679 for key in [
680 "NumberOfParts",
681 "PartOffsets",
682 "PointOffsets",
683 "CellOffsets",
684 "ConnectivityIdOffsets",
685 ]:
686 all_parts[key] = _create_dataset(
687 steps,
688 key,
689 shape=(1,),
690 dtype=np.int64,
691 chunks=True,
692 maxshape=(None,),
693 mode=h5_mode,
694 )
696 all_parts["NumberOfParts"][-1] = 1
697 all_parts["PartOffsets"][-1] = 0
698 all_parts["PointOffsets"][-1] = points.shape[0] - mesh.num_nodes_global
699 all_parts["CellOffsets"][-1] = types.shape[0] - mesh.num_cells_global
700 all_parts["ConnectivityIdOffsets"][-1] = (
701 topology.shape[0] - mesh.num_cells_global * num_dofs_per_cell
702 )
704 # Update cell-data and point-data offsets by copying over data from previous step
705 for key in ["CellDataOffsets", "PointDataOffsets"]:
706 group = _create_group(steps, key, mode=h5_mode)
707 for cd in group.keys():
708 group[cd].resize(steps.attrs["NSteps"], axis=0)
709 group[cd][-1] = group[cd][-2]
711 # Update Steps in all other parts of the mesh as well
712 for key in mesh_assembly.keys():
713 if key == name:
714 continue
715 # Copy time-dependent geometry info (NumberOfPoints) from mesh to tag
716 sub_group = mesh_assembly[key]
717 sub_step = sub_group["Steps"]
718 sub_step.attrs["NSteps"] = steps.attrs["NSteps"]
720 # Copy time dependent and partition info from mesh to tag
721 step_copy_keys = ["Values", "PartOffsets", "NumberOfParts"]
722 for key in step_copy_keys:
723 if key in sub_step.keys():
724 sub_step[key].resize(steps[key].shape)
725 sub_step[key][:] = steps[key][:]
726 else:
727 raise RuntimeError(f"{sub_step.name} should have {key}/")
729 # Append value from previous step for meshtags as they are time-independent
730 append_keys = ["CellOffsets", "ConnectivityIdOffsets"]
731 for key in append_keys:
732 if key in sub_step.keys():
733 sub_step[key].resize(sub_step[key].size + 1, axis=0)
734 sub_step[key][-1] = sub_step[key][-2]
735 else:
736 raise RuntimeError(f"{sub_step.name} should have {key}/")
737 # Append value from previous step for meshtags celldata
738 for key in ["CellDataOffsets", "PointDataOffsets"]:
739 group = _create_group(sub_step, key, mode=h5_mode)
740 for cd in group.keys():
741 group[cd].resize(sub_step.attrs["NSteps"], axis=0)
742 group[cd][-1] = group[cd][-2]
745def write_meshtags(
746 filename: str | Path,
747 comm: MPI.Intracomm,
748 data: MeshTagsData,
749 backend_args: dict[str, Any] | None,
750):
751 """Write mesh tags to file.
753 Args:
754 filename: Path to file to write to
755 comm: MPI communicator used in storage
756 data: Internal data structure for the mesh tags to save to file
757 backend_args: Arguments to backend
758 """
759 backend_args = get_default_backend_args(backend_args)
760 name = backend_args["name"]
761 h5_mode = "a"
762 with h5pyfile(filename, h5_mode, comm=comm) as h5file:
763 hdf = h5file["VTKHDF"]
764 file_type = _decode_bytes_if_needed(hdf.attrs["Type"])
765 # Check for type of VTKHDF file
766 if file_type != "MultiBlockDataSet":
767 raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}")
769 # We place the mesh-tags in the same subgroup of assembly as the mesh
770 parent_mesh_group = _get_vtk_group(h5file, name)
771 block = parent_mesh_group.parent
773 tag_path = f"/VTKHDF/{name}_{data.name}"
774 if data.name not in block.keys():
775 block[data.name] = h5py.SoftLink(tag_path)
777 mesh_group = _create_group(hdf, tag_path, mode=h5_mode)
778 mesh_group.attrs.create("Type", "UnstructuredGrid")
779 mesh_group.attrs["Version"] = _vtk_hdf_version
781 cell_data = _create_group(mesh_group, "CellData", mode=h5_mode)
782 assert data.num_entities_global is not None
783 dataset = _create_dataset(
784 cell_data,
785 data.name,
786 shape=(data.num_entities_global,),
787 dtype=data.values.dtype,
788 chunks=True,
789 maxshape=(None,),
790 mode=h5_mode,
791 )
792 assert data.local_start is not None
793 insert_slice = _compute_append_slice(
794 dataset,
795 data.num_entities_global,
796 np.array([data.local_start, data.local_start + data.indices.shape[0]]),
797 mode=h5_mode,
798 )
799 dataset[insert_slice] = data.values
801 # NOTE: The following is more or less a copy from write_mesh,
802 # except that we pull out the point storage and use a softlink
803 num_cells = _create_dataset(
804 mesh_group,
805 "NumberOfCells",
806 shape=(1,),
807 dtype=np.int64,
808 chunks=True,
809 maxshape=(None,),
810 mode=h5_mode,
811 resize=False, # Resize should really be True, see issue below
812 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
813 num_cells[-1] = data.num_entities_global
815 # Hardlink data should also follow hardlink for numbering
816 for key in ["Points", "NumberOfPoints"]:
817 if key not in mesh_group.keys():
818 mesh_group[key] = parent_mesh_group[key]
820 # Single celltype assumption
821 num_dofs_per_cell = data.num_dofs_per_entity
822 number_of_connectivities = _create_dataset(
823 mesh_group,
824 "NumberOfConnectivityIds",
825 shape=(1,),
826 dtype=np.int64,
827 chunks=True,
828 maxshape=(None,),
829 mode=h5_mode,
830 resize=False, # Resize should really be True, see issue below
831 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
832 assert data.num_entities_global is not None and num_dofs_per_cell is not None
833 number_of_connectivities[-1] = data.num_entities_global * num_dofs_per_cell
835 # Store topology offsets (single celltype assumption)
836 offsets = _create_dataset(
837 mesh_group,
838 "Offsets",
839 shape=(data.num_entities_global + 1,),
840 dtype=np.int64,
841 chunks=True,
842 mode=h5_mode,
843 maxshape=(None,),
844 resize=False, # Resize should really be True, see issue below
845 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
846 offset_data = np.arange(0, data.indices.size + 1, data.indices.shape[1])
847 assert data.local_start is not None
848 offset_data += num_dofs_per_cell * data.local_start
849 insert_slice = _compute_append_slice(
850 offsets,
851 data.num_entities_global + 1,
852 (data.local_start, data.local_start + data.indices.shape[0] + 1),
853 mode=h5_mode,
854 )
855 offsets[insert_slice] = offset_data
856 del offset_data
858 # Permute and store topology data
859 dx_ct = dolfinx.mesh.to_type(data.cell_type)
860 top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell))
861 topology_data = data.indices[:, top_perm].flatten()
862 topology = _create_dataset(
863 mesh_group,
864 "Connectivity",
865 shape=(data.num_entities_global * num_dofs_per_cell,),
866 dtype=np.int64,
867 chunks=True,
868 maxshape=(None,),
869 mode=h5_mode,
870 resize=False, # Resize should really be True, see issue below
871 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
872 insert_slice = _compute_append_slice(
873 topology,
874 data.num_entities_global * num_dofs_per_cell,
875 np.array([data.local_start, data.local_start + data.indices.shape[0]])
876 * num_dofs_per_cell,
877 mode=h5_mode,
878 )
879 topology[insert_slice] = topology_data
880 del topology_data
882 # Store celltypes
883 cell_types = np.full(
884 data.indices.shape[0],
885 dolfinx.cpp.io.get_vtk_cell_type(dx_ct, dolfinx.mesh.cell_dim(dx_ct)),
886 )
887 types = _create_dataset(
888 mesh_group,
889 "Types",
890 shape=(data.num_entities_global,),
891 dtype=np.uint8,
892 maxshape=(None,),
893 chunks=True,
894 mode=h5_mode,
895 resize=False, # Resize should really be True, see issue below
896 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
897 insert_slice = _compute_append_slice(
898 types,
899 data.num_entities_global,
900 (data.local_start, data.local_start + data.indices.shape[0]),
901 h5_mode,
902 )
903 types[insert_slice] = cell_types
904 del cell_types
906 steps = _create_group(mesh_group, "Steps", mode=h5_mode)
908 # Copy n-step counter
909 steps.attrs.create("NSteps", parent_mesh_group["Steps"].attrs["NSteps"])
911 # Hardlink data that we know is the same across meshes
912 hardlink_keys = ["NumberOfParts", "PartOffsets", "Values", "PointOffsets"]
913 for key in hardlink_keys:
914 steps[key] = parent_mesh_group["Steps"][key]
916 # Write offset data for current time-step
917 all_parts = {}
918 for key in ["CellOffsets", "ConnectivityIdOffsets"]:
919 all_parts[key] = _create_dataset(
920 steps,
921 key,
922 shape=(1,),
923 dtype=np.int64,
924 chunks=True,
925 maxshape=(None,),
926 mode=h5_mode,
927 )
929 all_parts["CellOffsets"][-1] = types.shape[0] - data.num_entities_global
930 all_parts["ConnectivityIdOffsets"][-1] = offsets.shape[0] - (data.num_entities_global + 1)
931 # CellData requires an offset
932 cd_off = _create_group(steps, "CellDataOffsets", mode=h5_mode)
933 cd_data = _create_dataset(
934 cd_off,
935 data.name,
936 shape=parent_mesh_group["Steps"]["CellOffsets"].shape,
937 dtype=np.int64,
938 chunks=True,
939 maxshape=(None,),
940 mode=h5_mode,
941 )
942 cd_data[-1] = types.shape[0] - data.num_entities_global
945def read_meshtags_data(
946 filename: str | Path,
947 comm: MPI.Intracomm,
948 name: str,
949 backend_args: dict[str, Any] | None,
950) -> MeshTagsData:
951 """Read mesh tags from file.
953 Args:
954 filename: Path to file to read from
955 comm: MPI communicator used in storage
956 name: Name of the mesh tags to read
957 backend_args: Arguments to backend
959 Returns:
960 Internal data structure for the mesh tags read from file
961 """
962 backend_args = get_default_backend_args(backend_args)
963 backend_args.update({"name": name})
964 # Reuse reading cell-data
965 indices, values = read_cell_data(filename, name, comm, None, backend_args=backend_args)
966 # Read cell-type of grid to get topological dimension
967 with h5pyfile(filename, "r", comm=comm) as h5file:
968 hdf = _get_vtk_group(h5file, backend_args["name"])
969 num_cells_global = hdf["Types"].size
970 local_cell_range = compute_local_range(comm, num_cells_global)
971 cell_types_local = hdf["Types"][slice(*local_cell_range)]
972 unique_cells = find_all_unique_cell_types(comm, cell_types_local, indices.shape[1])
973 if unique_cells.shape[0] > 1:
974 raise NotImplementedError("io4dolfinx does not support mixed celltype grids")
975 vtk_cell_type = unique_cells[0][0]
976 if vtk_cell_type in _first_order_vtk.keys():
977 ct = _first_order_vtk[vtk_cell_type]
978 elif vtk_cell_type in _arbitrary_lagrange_vtk.keys():
979 ct = _arbitrary_lagrange_vtk[vtk_cell_type]
980 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(ct))
981 return MeshTagsData(name=name, values=values.flatten(), indices=indices, dim=dim)
984def read_dofmap(
985 filename: str | Path,
986 comm: MPI.Intracomm,
987 name: str,
988 backend_args: dict[str, Any] | None,
989) -> dolfinx.graph.AdjacencyList:
990 """Read the dofmap of a function with a given name.
992 Args:
993 filename: Path to file to read from
994 comm: MPI communicator used in storage
995 name: Name of the function to read the dofmap for
996 backend_args: Arguments to backend
998 Returns:
999 Dofmap as an AdjacencyList
1000 """
1001 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
1004def read_dofs(
1005 filename: str | Path,
1006 comm: MPI.Intracomm,
1007 name: str,
1008 time: float,
1009 backend_args: dict[str, Any] | None,
1010) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
1011 """Read the dofs (values) of a function with a given name from a given timestep.
1013 Args:
1014 filename: Path to file to read from
1015 comm: MPI communicator used in storage
1016 name: Name of the function to read the dofs for
1017 time: Time stamp associated with the function to read
1018 backend_args: Arguments to backend
1020 Returns:
1021 Contiguous sequence of degrees of freedom (with respect to input data)
1022 and the global starting point on the process.
1023 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1024 """
1025 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
1028def read_cell_perms(
1029 comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None
1030) -> npt.NDArray[np.uint32]:
1031 """
1032 Read cell permutation from file with given communicator,
1033 Split in continuous chunks based on number of cells in the input data.
1035 Args:
1036 comm: MPI communicator used in storage
1037 filename: Path to file to read from
1038 backend_args: Arguments to backend
1040 Returns:
1041 Contiguous sequence of permutations (with respect to input data)
1042 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1043 """
1044 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1047def write_function(
1048 filename: Path,
1049 comm: MPI.Intracomm,
1050 u: FunctionData,
1051 time: float,
1052 mode: FileMode,
1053 backend_args: dict[str, Any] | None,
1054):
1055 """
1056 Write a function to file.
1058 Args:
1059 comm: MPI communicator used in storage
1060 u: Internal data structure for the function data to save to file
1061 filename: Path to file to write to
1062 time: Time stamp associated with function
1063 mode: File-mode to store the function
1064 backend_args: Arguments to backend
1065 """
1066 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1069def read_legacy_mesh(
1070 filename: Path | str, comm: MPI.Intracomm, group: str
1071) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
1072 """Read in the mesh topology, geometry and (optionally) cell type from a
1073 legacy DOLFIN HDF5-file.
1075 Args:
1076 filename: Path to file to read from
1077 comm: MPI communicator used in storage
1078 group: Group in HDF5 file where mesh is stored
1080 Returns:
1081 Tuple containing:
1082 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
1083 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
1084 - Cell type as a string (e.g. "tetrahedron") or None if not found
1085 """
1086 raise NotImplementedError("The VTKHDF backend cannot read legacy DOLFIN meshes.")
1089def snapshot_checkpoint(
1090 filename: Path | str,
1091 mode: FileMode,
1092 u: dolfinx.fem.Function,
1093 backend_args: dict[str, Any] | None,
1094):
1095 """Create a snapshot checkpoint of a dolfinx function.
1097 Args:
1098 filename: Path to file to read from
1099 mode: File-mode to store the function
1100 u: dolfinx function to create a snapshot checkpoint for
1101 backend_args: Arguments to backend
1102 """
1103 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1106def read_hdf5_array(
1107 comm: MPI.Intracomm,
1108 filename: Path | str,
1109 group: str,
1110 backend_args: dict[str, Any] | None,
1111) -> tuple[np.ndarray, int]:
1112 """Read an array from an HDF5 file.
1114 Args:
1115 comm: MPI communicator used in storage
1116 filename: Path to file to read from
1117 group: Group in HDF5 file where array is stored
1118 backend_args: Arguments to backend
1120 Returns:
1121 Tuple containing:
1122 - Numpy array read from file
1123 - Global starting point on the process.
1124 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1125 """
1126 raise NotImplementedError("The VTKHDF backend cannot read HDF5 arrays")
1129def write_data(
1130 filename: Path | str,
1131 array_data: ArrayData,
1132 comm: MPI.Intracomm,
1133 time: str | float | None,
1134 mode: FileMode,
1135 backend_args: dict[str, Any] | None,
1136):
1137 """Write function to file by interpolating into geometry nodes.
1140 Args:
1141 filename: Path to file
1142 array_data: Data to write to file
1143 time: Time stamp
1144 mode: Append or write
1145 backend_args: The backend arguments
1146 """
1147 h5_mode = convert_file_mode(mode)
1148 assert h5_mode == "a"
1149 backend_args = get_default_backend_args(backend_args)
1150 mesh_name = backend_args["name"]
1151 extension = array_data.type
1152 with h5pyfile(filename, h5_mode, comm=comm) as h5file:
1153 hdf = h5file["VTKHDF"]
1154 file_type = _decode_bytes_if_needed(hdf.attrs["Type"])
1155 # Check for type of VTKHDF file
1156 if file_type != "MultiBlockDataSet":
1157 raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}")
1159 # Find mesh block to add data to
1160 block = _get_vtk_group(h5file, mesh_name)
1162 data_group = _create_group(block, f"{extension}Data", mode=h5_mode)
1163 dataset = _create_dataset(
1164 data_group,
1165 array_data.name,
1166 shape=array_data.global_shape,
1167 dtype=array_data.values.dtype,
1168 chunks=True,
1169 maxshape=(None, array_data.values.shape[1]),
1170 mode=h5_mode,
1171 )
1172 insert_slice = _compute_append_slice(
1173 dataset,
1174 array_data.global_shape[0],
1175 np.array(
1176 [
1177 array_data.local_range[0],
1178 array_data.local_range[0] + array_data.values.shape[0],
1179 ]
1180 ),
1181 mode=h5_mode,
1182 )
1183 dataset[insert_slice] = array_data.values
1185 steps = _create_group(block, "Steps", mode=h5_mode)
1186 pdo = _create_group(steps, f"{extension}DataOffsets", mode=h5_mode)
1188 # Check if time step is already in time-stepping of mesh
1189 timestamps = steps["Values"][:]
1190 assert isinstance(timestamps, np.ndarray)
1191 assert isinstance(time, float)
1192 time_exists = np.flatnonzero(np.isclose(timestamps, time))
1193 if len(time_exists) > 1:
1194 raise ValueError("Time-step has been written multiple times")
1195 pdo_u = _create_dataset(
1196 pdo,
1197 array_data.name,
1198 shape=(1,),
1199 dtype=np.int64,
1200 chunks=True,
1201 maxshape=(None,),
1202 mode=h5_mode,
1203 resize=False,
1204 )
1205 if len(time_exists) == 1:
1206 idx = time_exists[0]
1207 elif len(time_exists) == 0:
1208 # No mesh written at step, update mesh offsets
1209 # Geometry has to be time-dependent
1210 num_points = block["NumberOfPoints"]
1211 num_points.resize(len(timestamps) + 1, axis=0)
1212 num_points[-1] = num_points[-2]
1213 # Even if topology cannot be time dep at the moment, number of cells
1214 # must be updated
1215 num_cells = block["NumberOfCells"]
1216 num_cells.resize(len(timestamps) + 1, axis=0)
1217 num_cells[-1] = num_cells[-2]
1219 steps.attrs.create("NSteps", block["Steps"].attrs["NSteps"] + 1)
1220 step_vals = _create_dataset(
1221 steps,
1222 "Values",
1223 shape=(1,),
1224 dtype=np.float64,
1225 chunks=True,
1226 maxshape=(None,),
1227 mode=h5_mode,
1228 resize=True,
1229 )
1230 step_vals[-1] = time
1231 idx = step_vals.size - 1
1232 for key in [
1233 "PartOffsets",
1234 "NumberOfParts",
1235 "PointOffsets",
1236 "CellOffsets",
1237 "ConnectivityIdOffsets",
1238 ]:
1239 comp = steps[key]
1240 comp.resize(steps.attrs["NSteps"], axis=0)
1241 if comp.size > 1:
1242 comp[-1] = comp[-2]
1243 else:
1244 raise ValueError(f"Time step found multiple times in {filename}")
1246 # Update offsets for current variable
1247 pdo_u.resize(steps.attrs["NSteps"], axis=0)
1248 pdo_u[idx] = dataset.shape[0] - array_data.global_shape[0]
1250 # Point and cell data of all other variables has to be updated to be synced
1251 for key in [
1252 "CellDataOffsets",
1253 "PointDataOffsets",
1254 ]:
1255 comp = steps[key]
1256 num_steps = steps.attrs["NSteps"]
1257 for dname, dset in comp.items():
1258 # Only resize other data arrays
1259 if dname != array_data.name:
1260 # Only resize if it hasn't already been resized
1261 if dset.size != num_steps:
1262 dset.resize(num_steps, axis=0)
1263 # Only update data if we are further than first time step
1264 if dset.size > 1:
1265 dset[-1] = dset[-2]