Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/backends/vtkhdf/backend.py: 84%
471 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 18:20 +0000
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 18:20 +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.Comm,
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.Comm,
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.Comm,
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.Comm,
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.Comm,
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.Comm,
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.Comm, 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.Comm,
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.Comm,
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 if data.cell_type is None:
860 raise ValueError("Cannot write meshtags without cell-type information")
861 dx_ct = dolfinx.mesh.to_type(data.cell_type)
862 top_perm = np.argsort(dolfinx.cpp.io.perm_vtk(dx_ct, num_dofs_per_cell))
863 topology_data = data.indices[:, top_perm].flatten()
864 topology = _create_dataset(
865 mesh_group,
866 "Connectivity",
867 shape=(data.num_entities_global * num_dofs_per_cell,),
868 dtype=np.int64,
869 chunks=True,
870 maxshape=(None,),
871 mode=h5_mode,
872 resize=False, # Resize should really be True, see issue below
873 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
874 insert_slice = _compute_append_slice(
875 topology,
876 data.num_entities_global * num_dofs_per_cell,
877 np.array([data.local_start, data.local_start + data.indices.shape[0]])
878 * num_dofs_per_cell,
879 mode=h5_mode,
880 )
881 topology[insert_slice] = topology_data
882 del topology_data
884 # Store celltypes
885 cell_types = np.full(
886 data.indices.shape[0],
887 dolfinx.cpp.io.get_vtk_cell_type(dx_ct, dolfinx.mesh.cell_dim(dx_ct)),
888 )
889 types = _create_dataset(
890 mesh_group,
891 "Types",
892 shape=(data.num_entities_global,),
893 dtype=np.uint8,
894 maxshape=(None,),
895 chunks=True,
896 mode=h5_mode,
897 resize=False, # Resize should really be True, see issue below
898 ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257
899 insert_slice = _compute_append_slice(
900 types,
901 data.num_entities_global,
902 (data.local_start, data.local_start + data.indices.shape[0]),
903 h5_mode,
904 )
905 types[insert_slice] = cell_types
906 del cell_types
908 steps = _create_group(mesh_group, "Steps", mode=h5_mode)
910 # Copy n-step counter
911 steps.attrs.create("NSteps", parent_mesh_group["Steps"].attrs["NSteps"])
913 # Hardlink data that we know is the same across meshes
914 hardlink_keys = ["NumberOfParts", "PartOffsets", "Values", "PointOffsets"]
915 for key in hardlink_keys:
916 steps[key] = parent_mesh_group["Steps"][key]
918 # Write offset data for current time-step
919 all_parts = {}
920 for key in ["CellOffsets", "ConnectivityIdOffsets"]:
921 all_parts[key] = _create_dataset(
922 steps,
923 key,
924 shape=(1,),
925 dtype=np.int64,
926 chunks=True,
927 maxshape=(None,),
928 mode=h5_mode,
929 )
931 all_parts["CellOffsets"][-1] = types.shape[0] - data.num_entities_global
932 all_parts["ConnectivityIdOffsets"][-1] = offsets.shape[0] - (data.num_entities_global + 1)
933 # CellData requires an offset
934 cd_off = _create_group(steps, "CellDataOffsets", mode=h5_mode)
935 cd_data = _create_dataset(
936 cd_off,
937 data.name,
938 shape=parent_mesh_group["Steps"]["CellOffsets"].shape,
939 dtype=np.int64,
940 chunks=True,
941 maxshape=(None,),
942 mode=h5_mode,
943 )
944 cd_data[-1] = types.shape[0] - data.num_entities_global
947def read_meshtags_data(
948 filename: str | Path,
949 comm: MPI.Comm,
950 name: str,
951 backend_args: dict[str, Any] | None,
952) -> MeshTagsData:
953 """Read mesh tags from file.
955 Args:
956 filename: Path to file to read from
957 comm: MPI communicator used in storage
958 name: Name of the mesh tags to read
959 backend_args: Arguments to backend
961 Returns:
962 Internal data structure for the mesh tags read from file
963 """
964 backend_args = get_default_backend_args(backend_args)
965 backend_args.update({"name": name})
966 # Reuse reading cell-data
967 indices, values = read_cell_data(filename, name, comm, None, backend_args=backend_args)
968 # Read cell-type of grid to get topological dimension
969 with h5pyfile(filename, "r", comm=comm) as h5file:
970 hdf = _get_vtk_group(h5file, backend_args["name"])
971 num_cells_global = hdf["Types"].size
972 local_cell_range = compute_local_range(comm, num_cells_global)
973 cell_types_local = hdf["Types"][slice(*local_cell_range)]
974 unique_cells = find_all_unique_cell_types(comm, cell_types_local, indices.shape[1])
975 if unique_cells.shape[0] > 1:
976 raise NotImplementedError("io4dolfinx does not support mixed celltype grids")
977 vtk_cell_type = unique_cells[0][0]
978 if vtk_cell_type in _first_order_vtk.keys():
979 ct = _first_order_vtk[vtk_cell_type]
980 elif vtk_cell_type in _arbitrary_lagrange_vtk.keys():
981 ct = _arbitrary_lagrange_vtk[vtk_cell_type]
982 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(ct))
983 return MeshTagsData(name=name, values=values.flatten(), indices=indices, dim=dim)
986def read_dofmap(
987 filename: str | Path,
988 comm: MPI.Comm,
989 name: str,
990 backend_args: dict[str, Any] | None,
991) -> dolfinx.graph.AdjacencyList:
992 """Read the dofmap of a function with a given name.
994 Args:
995 filename: Path to file to read from
996 comm: MPI communicator used in storage
997 name: Name of the function to read the dofmap for
998 backend_args: Arguments to backend
1000 Returns:
1001 Dofmap as an AdjacencyList
1002 """
1003 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
1006def read_dofs(
1007 filename: str | Path,
1008 comm: MPI.Comm,
1009 name: str,
1010 time: float,
1011 backend_args: dict[str, Any] | None,
1012) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
1013 """Read the dofs (values) of a function with a given name from a given timestep.
1015 Args:
1016 filename: Path to file to read from
1017 comm: MPI communicator used in storage
1018 name: Name of the function to read the dofs for
1019 time: Time stamp associated with the function to read
1020 backend_args: Arguments to backend
1022 Returns:
1023 Contiguous sequence of degrees of freedom (with respect to input data)
1024 and the global starting point on the process.
1025 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1026 """
1027 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
1030def read_cell_perms(
1031 comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None
1032) -> npt.NDArray[np.uint32]:
1033 """
1034 Read cell permutation from file with given communicator,
1035 Split in continuous chunks based on number of cells in the input data.
1037 Args:
1038 comm: MPI communicator used in storage
1039 filename: Path to file to read from
1040 backend_args: Arguments to backend
1042 Returns:
1043 Contiguous sequence of permutations (with respect to input data)
1044 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1045 """
1046 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1049def write_function(
1050 filename: Path,
1051 comm: MPI.Comm,
1052 u: FunctionData,
1053 time: float,
1054 mode: FileMode,
1055 backend_args: dict[str, Any] | None,
1056):
1057 """
1058 Write a function to file.
1060 Args:
1061 comm: MPI communicator used in storage
1062 u: Internal data structure for the function data to save to file
1063 filename: Path to file to write to
1064 time: Time stamp associated with function
1065 mode: File-mode to store the function
1066 backend_args: Arguments to backend
1067 """
1068 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1071def read_legacy_mesh(
1072 filename: Path | str, comm: MPI.Comm, group: str
1073) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
1074 """Read in the mesh topology, geometry and (optionally) cell type from a
1075 legacy DOLFIN HDF5-file.
1077 Args:
1078 filename: Path to file to read from
1079 comm: MPI communicator used in storage
1080 group: Group in HDF5 file where mesh is stored
1082 Returns:
1083 Tuple containing:
1084 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
1085 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
1086 - Cell type as a string (e.g. "tetrahedron") or None if not found
1087 """
1088 raise NotImplementedError("The VTKHDF backend cannot read legacy DOLFIN meshes.")
1091def snapshot_checkpoint(
1092 filename: Path | str,
1093 mode: FileMode,
1094 u: dolfinx.fem.Function,
1095 backend_args: dict[str, Any] | None,
1096):
1097 """Create a snapshot checkpoint of a dolfinx function.
1099 Args:
1100 filename: Path to file to read from
1101 mode: File-mode to store the function
1102 u: dolfinx function to create a snapshot checkpoint for
1103 backend_args: Arguments to backend
1104 """
1105 raise NotImplementedError("The VTKHDF backend cannot make checkpoints.")
1108def read_hdf5_array(
1109 comm: MPI.Comm,
1110 filename: Path | str,
1111 group: str,
1112 backend_args: dict[str, Any] | None,
1113) -> tuple[np.ndarray, int]:
1114 """Read an array from an HDF5 file.
1116 Args:
1117 comm: MPI communicator used in storage
1118 filename: Path to file to read from
1119 group: Group in HDF5 file where array is stored
1120 backend_args: Arguments to backend
1122 Returns:
1123 Tuple containing:
1124 - Numpy array read from file
1125 - Global starting point on the process.
1126 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
1127 """
1128 raise NotImplementedError("The VTKHDF backend cannot read HDF5 arrays")
1131def write_data(
1132 filename: Path | str,
1133 array_data: ArrayData,
1134 comm: MPI.Comm,
1135 time: str | float | None,
1136 mode: FileMode,
1137 backend_args: dict[str, Any] | None,
1138):
1139 """Write function to file by interpolating into geometry nodes.
1142 Args:
1143 filename: Path to file
1144 array_data: Data to write to file
1145 time: Time stamp
1146 mode: Append or write
1147 backend_args: The backend arguments
1148 """
1149 h5_mode = convert_file_mode(mode)
1150 assert h5_mode == "a"
1151 backend_args = get_default_backend_args(backend_args)
1152 mesh_name = backend_args["name"]
1153 extension = array_data.type
1154 with h5pyfile(filename, h5_mode, comm=comm) as h5file:
1155 hdf = h5file["VTKHDF"]
1156 file_type = _decode_bytes_if_needed(hdf.attrs["Type"])
1157 # Check for type of VTKHDF file
1158 if file_type != "MultiBlockDataSet":
1159 raise ValueError(f"Cannot write meshtags to {filename} with VTK type {file_type}")
1161 # Find mesh block to add data to
1162 block = _get_vtk_group(h5file, mesh_name)
1164 data_group = _create_group(block, f"{extension}Data", mode=h5_mode)
1165 dataset = _create_dataset(
1166 data_group,
1167 array_data.name,
1168 shape=array_data.global_shape,
1169 dtype=array_data.values.dtype,
1170 chunks=True,
1171 maxshape=(None, array_data.values.shape[1]),
1172 mode=h5_mode,
1173 )
1174 insert_slice = _compute_append_slice(
1175 dataset,
1176 array_data.global_shape[0],
1177 np.array(
1178 [
1179 array_data.local_range[0],
1180 array_data.local_range[0] + array_data.values.shape[0],
1181 ]
1182 ),
1183 mode=h5_mode,
1184 )
1185 dataset[insert_slice] = array_data.values
1187 steps = _create_group(block, "Steps", mode=h5_mode)
1188 pdo = _create_group(steps, f"{extension}DataOffsets", mode=h5_mode)
1190 # Check if time step is already in time-stepping of mesh
1191 timestamps = steps["Values"][:]
1192 assert isinstance(timestamps, np.ndarray)
1193 assert isinstance(time, float)
1194 time_exists = np.flatnonzero(np.isclose(timestamps, time))
1195 if len(time_exists) > 1:
1196 raise ValueError("Time-step has been written multiple times")
1197 pdo_u = _create_dataset(
1198 pdo,
1199 array_data.name,
1200 shape=(1,),
1201 dtype=np.int64,
1202 chunks=True,
1203 maxshape=(None,),
1204 mode=h5_mode,
1205 resize=False,
1206 )
1207 if len(time_exists) == 1:
1208 idx = time_exists[0]
1209 elif len(time_exists) == 0:
1210 # No mesh written at step, update mesh offsets
1211 # Geometry has to be time-dependent
1212 num_points = block["NumberOfPoints"]
1213 num_points.resize(len(timestamps) + 1, axis=0)
1214 num_points[-1] = num_points[-2]
1215 # Even if topology cannot be time dep at the moment, number of cells
1216 # must be updated
1217 num_cells = block["NumberOfCells"]
1218 num_cells.resize(len(timestamps) + 1, axis=0)
1219 num_cells[-1] = num_cells[-2]
1221 steps.attrs.create("NSteps", block["Steps"].attrs["NSteps"] + 1)
1222 step_vals = _create_dataset(
1223 steps,
1224 "Values",
1225 shape=(1,),
1226 dtype=np.float64,
1227 chunks=True,
1228 maxshape=(None,),
1229 mode=h5_mode,
1230 resize=True,
1231 )
1232 step_vals[-1] = time
1233 idx = step_vals.size - 1
1234 for key in [
1235 "PartOffsets",
1236 "NumberOfParts",
1237 "PointOffsets",
1238 "CellOffsets",
1239 "ConnectivityIdOffsets",
1240 ]:
1241 comp = steps[key]
1242 comp.resize(steps.attrs["NSteps"], axis=0)
1243 if comp.size > 1:
1244 comp[-1] = comp[-2]
1245 else:
1246 raise ValueError(f"Time step found multiple times in {filename}")
1248 # Update offsets for current variable
1249 pdo_u.resize(steps.attrs["NSteps"], axis=0)
1250 pdo_u[idx] = dataset.shape[0] - array_data.global_shape[0]
1252 # Point and cell data of all other variables has to be updated to be synced
1253 for key in [
1254 "CellDataOffsets",
1255 "PointDataOffsets",
1256 ]:
1257 comp = steps[key]
1258 num_steps = steps.attrs["NSteps"]
1259 for dname, dset in comp.items():
1260 # Only resize other data arrays
1261 if dname != array_data.name:
1262 # Only resize if it hasn't already been resized
1263 if dset.size != num_steps:
1264 dset.resize(num_steps, axis=0)
1265 # Only update data if we are further than first time step
1266 if dset.size > 1:
1267 dset[-1] = dset[-2]