Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/xdmf/backend.py: 79%
175 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:46 +0000
1"""
2Module that uses DOLFINx/H5py to import XDMF files.
3"""
5import re
6from pathlib import Path
7from typing import Any
8from xml.etree import ElementTree
10from mpi4py import MPI
12import basix
13import dolfinx
14import numpy as np
15import numpy.typing as npt
17from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
18from io4dolfinx.utils import check_file_exists, compute_local_range
20from .. import FileMode, ReadMode
21from ..h5py.backend import h5pyfile
23read_mode = ReadMode.parallel
26def extract_function_names_and_timesteps(filename: Path | str) -> dict[str, list[str]]:
27 tree = ElementTree.parse(filename)
28 root = tree.getroot()
29 mesh_nodes = root.findall(".//Grid[@CollectionType='Temporal']")
30 function_names = []
31 for mesh in mesh_nodes:
32 function_names.append(mesh.attrib["Name"])
34 time_stamps: dict[str, list[str]] = {name: [] for name in function_names}
35 for name in function_names:
36 time_steps = root.findall(f".//Grid[@Name='{name}']")
37 for time in time_steps:
38 step = time.find(".//Time")
39 if step is not None:
40 val = step.attrib["Value"]
41 time_stamps[name].append(val)
42 for name in function_names:
43 float_steps = np.argsort(np.array(list(set(time_stamps[name])), dtype=np.float64))
44 time_stamps[name] = np.array(list(set(time_stamps[name])), dtype=str)[float_steps].tolist()
45 return time_stamps
48def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]:
49 """Get default backend arguments given a set of input arguments.
51 Args:
52 arguments: Input backend arguments
54 Returns:
55 Updated backend arguments
56 """
57 args = arguments or {}
58 return args
61def read_mesh_data(
62 filename: Path | str,
63 comm: MPI.Comm,
64 time: str | float | None,
65 read_from_partition: bool,
66 backend_args: dict[str, Any] | None,
67) -> ReadMeshData:
68 """Read mesh data from file.
70 Args:
71 filename: Path to file to read from
72 comm: MPI communicator used in storage
73 time: Time stamp associated with the mesh to read
74 read_from_partition: Whether to read partition information
75 backend_args: Arguments to backend
77 Returns:
78 Internal data structure for the mesh data read from file
79 """
80 assert not read_from_partition
81 check_file_exists(filename)
82 with dolfinx.io.XDMFFile(comm, filename, "r") as file:
83 cell_shape, cell_degree = file.read_cell_type()
84 cells = file.read_topology_data()
85 x = file.read_geometry_data()
86 return ReadMeshData(
87 cells=cells,
88 cell_type=cell_shape.name,
89 x=x,
90 lvar=int(basix.LagrangeVariant.equispaced),
91 degree=cell_degree,
92 )
95def read_point_data(
96 filename: Path | str,
97 name: str,
98 comm: MPI.Comm,
99 time: float | str | None,
100 backend_args: dict[str, Any] | None,
101) -> tuple[np.ndarray, int]:
102 """Read data from the nodes of a mesh.
104 Args:
105 filename: Path to file
106 name: Name of point data
107 comm: Communicator to launch IO on.
108 time: The time stamp
109 backend_args: The backend arguments
111 Returns:
112 Data local to process (contiguous, no mpi comm) and local start range
113 """
114 check_file_exists(filename)
115 filename = Path(filename)
117 root = ElementTree.parse(filename).getroot()
118 backend_args = get_default_backend_args(backend_args)
120 grid = _time_step_grid(root, name, time, filename)
121 attributes = _function_attributes(grid, name)
122 local_range_start, dataset = _read_data_item(filename, _data_item(attributes[0]), comm)
123 if len(attributes) == 2:
124 _, imag = _read_data_item(filename, _data_item(attributes[1]), comm)
125 # complex64 from float32 parts, complex128 from float64
126 dataset = dataset + 1j * imag
127 return dataset, local_range_start
130def _time_step_grid(
131 root: ElementTree.Element, name: str, time: str | float | None, filename: Path
132) -> ElementTree.Element:
133 """Grid holding function ``name`` at ``time``, or at the first time step if ``None``.
135 Both DOLFINx and legacy DOLFIN write one grid per time step, with a `Time`
136 child, into a temporal collection named after the function. DOLFINx names
137 the time-step grids ``name``, legacy DOLFIN ``{name}_{counter}``.
139 Raises:
140 RuntimeError: If there is no such grid.
141 """
142 collection = root.find(f".//Grid[@CollectionType='Temporal'][@Name='{name}']")
143 if collection is not None:
144 for grid in collection.findall("./Grid"):
145 step_node = grid.find("./Time")
146 if step_node is None:
147 continue
148 if time is None or np.isclose(float(step_node.attrib["Value"]), float(time)):
149 return grid
150 at_time = "" if time is None else f"(t={time})"
151 raise RuntimeError(f"Function {name}{at_time} not found in {filename}")
154def _function_attributes(grid: ElementTree.Element, name: str) -> list[ElementTree.Element]:
155 """Attributes of ``grid`` holding function ``name``.
157 DOLFINx writes a real function as attribute ``name``, and a complex one as
158 two attributes ``real_{name}`` and ``imag_{name}``.
160 Returns:
161 ``[data]`` for a real function, ``[real, imag]`` for a complex one.
163 Raises:
164 RuntimeError: If ``grid`` holds neither form, or only one complex part.
165 """
166 if (node := grid.find(f"./Attribute[@Name='{name}']")) is not None:
167 return [node]
168 real = grid.find(f"./Attribute[@Name='real_{name}']")
169 imag = grid.find(f"./Attribute[@Name='imag_{name}']")
170 if real is not None and imag is not None:
171 return [real, imag]
172 if real is not None or imag is not None:
173 present, missing = ("real", "imag") if real is not None else ("imag", "real")
174 raise RuntimeError(
175 f"Complex function {name} is incomplete: found '{present}_{name}'"
176 f" but not '{missing}_{name}'."
177 )
178 found = [attribute.attrib.get("Name") for attribute in grid.findall("./Attribute")]
179 raise RuntimeError(
180 f"Function {name} not found: expected attribute '{name}', or 'real_{name}'"
181 f" and 'imag_{name}'. Attributes present: {found}."
182 )
185def _data_item(node: ElementTree.Element) -> ElementTree.Element:
186 """The `DataItem` directly under ``node``.
188 Raises:
189 RuntimeError: If ``node`` has none.
190 """
191 if (data_item := node.find("./DataItem")) is None:
192 raise RuntimeError(f"<{node.tag} Name='{node.attrib.get('Name')}'> has no DataItem.")
193 return data_item
196def _read_data_item(
197 filename: Path, data_item: ElementTree.Element, comm: MPI.Comm
198) -> tuple[int, np.ndarray]:
199 """Read this process's contiguous block of rows of an HDF5-backed `DataItem`.
201 Args:
202 filename: The XDMF file, which HDF5 paths are relative to
203 data_item: `DataItem` whose text is ``"file.h5:/path/to/dataset"``
204 comm: Communicator the rows are split over
206 Returns:
207 The first global row read and the rows themselves.
209 Raises:
210 RuntimeError: If the dataset's shape does not match the `Dimensions`.
211 """
212 location = data_item.text.strip() if data_item.text is not None else ""
213 data_file, sep, data_loc = location.partition(":")
214 if data_item.attrib.get("Format") != "HDF" or not sep:
215 raise RuntimeError(f"Expected an HDF5 DataItem 'file.h5:/path', got '{location}'.")
216 dimensions = [int(d) for d in data_item.attrib["Dimensions"].split()]
217 with h5pyfile(filename.parent / data_file, "r", comm=comm) as h5file:
218 data = h5file[data_loc]
219 # Shapes may differ by a trailing unit dimension, so compare rows and size.
220 if data.shape[0] != dimensions[0] or data.size != int(np.prod(dimensions)):
221 raise RuntimeError(
222 f"Dataset {location} has shape {data.shape}, but its DataItem"
223 f" has Dimensions {dimensions}."
224 )
225 lr = compute_local_range(comm, data.shape[0])
226 return lr[0], data[slice(*lr)]
229def read_attributes(
230 filename: Path | str,
231 comm: MPI.Comm,
232 name: str,
233 backend_args: dict[str, Any] | None,
234) -> dict[str, Any]:
235 """Read attributes from file.
237 Args:
238 filename: Path to file to read from
239 comm: MPI communicator used in storage
240 name: Name of the attribute group
241 backend_args: Arguments to backend
243 Returns:
244 Dictionary of attributes read from file
245 """
246 raise NotImplementedError("The XDMF backend cannot read attributes.")
249def read_timestamps(
250 filename: Path | str,
251 comm: MPI.Comm,
252 function_name: str,
253 backend_args: dict[str, Any] | None,
254) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
255 """Read timestamps from file.
257 Args:
258 filename: Path to file to read from
259 comm: MPI communicator used in storage
260 function_name: Name of the function to read timestamps for
261 backend_args: Arguments to backend
263 Returns:
264 Numpy array of timestamps read from file
265 """
266 tree = ElementTree.parse(filename)
267 root = tree.getroot()
268 time_stamps = []
269 time_steps = root.findall(f".//Grid[@Name='{function_name}']")
270 for time in time_steps:
271 step = time.find(".//Time")
272 if step is not None:
273 val = step.attrib["Value"]
274 time_stamps.append(val)
275 float_steps = np.argsort(np.array(list(set(time_stamps)), dtype=np.float64))
276 return np.array(list(set(time_stamps)), dtype=str)[float_steps].tolist()
279def write_attributes(
280 filename: Path | str,
281 comm: MPI.Comm,
282 name: str,
283 attributes: dict[str, np.ndarray],
284 backend_args: dict[str, Any] | None,
285):
286 """Write attributes to file.
288 Args:
289 filename: Path to file to write to
290 comm: MPI communicator used in storage
291 name: Name of the attribute group
292 attributes: Dictionary of attributes to write
293 backend_args: Arguments to backend
294 """
295 raise NotImplementedError("The XDMF backend cannot write attributes.")
298def write_mesh(
299 filename: Path | str,
300 comm: MPI.Comm,
301 mesh: MeshData,
302 backend_args: dict[str, Any] | None,
303 mode: FileMode,
304 time: float,
305):
306 """
307 Write a mesh to file.
309 Args:
310 comm: MPI communicator used in storage
311 mesh: Internal data structure for the mesh data to save to file
312 filename: Path to file to write to
313 backend_args: Arguments to backend
314 mode: File-mode to store the mesh
315 time: Time stamp associated with the mesh
316 """
317 raise NotImplementedError("The XDMF backend cannot write meshes.")
320def write_meshtags(
321 filename: str | Path,
322 comm: MPI.Comm,
323 data: MeshTagsData,
324 backend_args: dict[str, Any] | None,
325):
326 """Write mesh tags to file.
328 Args:
329 filename: Path to file to write to
330 comm: MPI communicator used in storage
331 data: Internal data structure for the mesh tags to save to file
332 backend_args: Arguments to backend
333 """
334 raise NotImplementedError("The XDMF backend cannot write meshtags.")
337def read_meshtags_data(
338 filename: str | Path,
339 comm: MPI.Comm,
340 name: str,
341 backend_args: dict[str, Any] | None,
342) -> MeshTagsData:
343 """Read mesh tags from file.
345 Args:
346 filename: Path to file to read from
347 comm: MPI communicator used in storage
348 name: Name of the mesh tags to read
349 backend_args: Arguments to backend
351 Returns:
352 Internal data structure for the mesh tags read from file
353 """
354 raise NotImplementedError("The XDMF backend cannot read meshtags.")
357def read_dofmap(
358 filename: str | Path,
359 comm: MPI.Comm,
360 name: str,
361 backend_args: dict[str, Any] | None,
362) -> dolfinx.graph.AdjacencyList:
363 """Read the dofmap of a function with a given name.
365 Args:
366 filename: Path to file to read from
367 comm: MPI communicator used in storage
368 name: Name of the function to read the dofmap for
369 backend_args: Arguments to backend
371 Returns:
372 Dofmap as an AdjacencyList
373 """
374 raise NotImplementedError("The XDMF backend cannot make checkpoints.")
377def read_dofs(
378 filename: str | Path,
379 comm: MPI.Comm,
380 name: str,
381 time: float,
382 backend_args: dict[str, Any] | None,
383) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
384 """Read the dofs (values) of a function with a given name from a given timestep.
386 Args:
387 filename: Path to file to read from
388 comm: MPI communicator used in storage
389 name: Name of the function to read the dofs for
390 time: Time stamp associated with the function to read
391 backend_args: Arguments to backend
393 Returns:
394 Contiguous sequence of degrees of freedom (with respect to input data)
395 and the global starting point on the process.
396 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
397 """
398 raise NotImplementedError("The XDMF backend cannot make checkpoints.")
401def read_cell_perms(
402 comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None
403) -> npt.NDArray[np.uint32]:
404 """
405 Read cell permutation from file with given communicator,
406 Split in continuous chunks based on number of cells in the input data.
408 Args:
409 comm: MPI communicator used in storage
410 filename: Path to file to read from
411 backend_args: Arguments to backend
413 Returns:
414 Contiguous sequence of permutations (with respect to input data)
415 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
416 """
417 raise NotImplementedError("The XDMF backend cannot make checkpoints.")
420def write_function(
421 filename: Path,
422 comm: MPI.Comm,
423 u: FunctionData,
424 time: float,
425 mode: FileMode,
426 backend_args: dict[str, Any] | None,
427):
428 """
429 Write a function to file.
431 Args:
432 comm: MPI communicator used in storage
433 u: Internal data structure for the function data to save to file
434 filename: Path to file to write to
435 time: Time stamp associated with function
436 mode: File-mode to store the function
437 backend_args: Arguments to backend
438 """
439 raise NotImplementedError("The XDMF backend cannot make checkpoints.")
442def read_legacy_mesh(
443 filename: Path | str, comm: MPI.Comm, group: str
444) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
445 """Read in the mesh topology, geometry and (optionally) cell type from a
446 legacy DOLFIN HDF5-file.
448 Args:
449 filename: Path to file to read from
450 comm: MPI communicator used in storage
451 group: Group in HDF5 file where mesh is stored
453 Returns:
454 Tuple containing:
455 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
456 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
457 - Cell type as a string (e.g. "tetrahedron") or None if not found
458 """
459 raise NotImplementedError("The XDMF backend cannot read legacy DOLFIN meshes.")
462def snapshot_checkpoint(
463 filename: Path | str,
464 mode: FileMode,
465 u: dolfinx.fem.Function,
466 backend_args: dict[str, Any] | None,
467):
468 """Create a snapshot checkpoint of a dolfinx function.
470 Args:
471 filename: Path to file to read from
472 mode: File-mode to store the function
473 u: dolfinx function to create a snapshot checkpoint for
474 backend_args: Arguments to backend
475 """
476 raise NotImplementedError("The XDMF backend cannot make checkpoints.")
479def read_hdf5_array(
480 comm: MPI.Comm,
481 filename: Path | str,
482 group: str,
483 backend_args: dict[str, Any] | None,
484) -> tuple[np.ndarray, int]:
485 """Read an array from an HDF5 file.
487 Args:
488 comm: MPI communicator used in storage
489 filename: Path to file to read from
490 group: Group in HDF5 file where array is stored
491 backend_args: Arguments to backend
493 Returns:
494 Tuple containing:
495 - Numpy array read from file
496 - Global starting point on the process.
497 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
498 """
499 raise NotImplementedError("The XDMF backend cannot read HDF5 arrays")
502def read_function_names(
503 filename: Path | str, comm: MPI.Comm, backend_args: dict[str, Any] | None
504) -> list[str]:
505 """Read all function names from a file.
507 Args:
508 filename: Path to file
509 comm: MPI communicator to launch IO on.
510 backend_args: Arguments to backend
512 Returns:
513 A list of function names.
514 """
515 # Find all functions in xml tree
516 check_file_exists(filename)
517 filename = Path(filename)
519 tree = ElementTree.parse(filename)
520 root = tree.getroot()
521 backend_args = get_default_backend_args(backend_args)
522 # Functions in checkpoint format
523 checkpoint_funcs = root.findall(".//Attribute[@ItemType='FiniteElementFunction']")
524 names = [func.attrib["Name"] for func in checkpoint_funcs]
525 # Temporal funcs
526 temporal_funcs = root.findall(".//Grid[@GridType='Collection']")
527 for func in temporal_funcs:
528 names.append(func.attrib["Name"])
529 return list(set(names))
532def read_cell_data(
533 filename: Path | str,
534 name: str,
535 comm: MPI.Comm,
536 time: str | float | None,
537 backend_args: dict[str, Any] | None,
538) -> tuple[npt.NDArray[np.int64], np.ndarray]:
539 """Read data from the cells of a mesh.
541 Reads cell-centred functions written by DOLFINx (DG-0), including complex
542 ones, and DG-0 checkpoints written by legacy DOLFIN.
544 Args:
545 filename: Path to file
546 name: Name of cell data
547 comm: Communicator to launch IO on.
548 time: The time stamp
549 backend_args: The backend arguments
550 Returns:
551 A tuple (topology, dofs) where topology contains the
552 vertex indices of the cells, dofs the degrees of
553 freedom within that cell.
554 """
555 check_file_exists(filename)
556 filename = Path(filename)
558 root = ElementTree.parse(filename).getroot()
559 backend_args = get_default_backend_args(backend_args)
561 grid = _time_step_grid(root, name, time, filename)
562 attributes = _function_attributes(grid, name)
563 values = [_cell_values(attribute) for attribute in attributes]
565 _, dofs = _read_data_item(filename, values[0], comm)
566 if len(values) == 2:
567 _, imag = _read_data_item(filename, values[1], comm)
568 dofs = dofs + 1j * imag
569 _, topology = _read_data_item(filename, _grid_topology(root, grid), comm)
570 # One row per cell in both, so both processes' blocks cover the same cells.
571 if topology.shape[0] != dofs.shape[0]:
572 raise RuntimeError(
573 f"Cell data {name} has {dofs.shape[0]} rows on this process, but the"
574 f" topology has {topology.shape[0]}."
575 )
576 return topology.astype(np.int64), dofs.reshape(dofs.shape[0], -1)
579def _cell_values(attribute: ElementTree.Element) -> ElementTree.Element:
580 """`DataItem` with one value row per cell, from a cell-data attribute.
582 Raises:
583 ValueError: If the attribute is not DG-0 data.
584 """
585 if attribute.attrib.get("ItemType") == "FiniteElementFunction":
586 # Legacy DOLFIN checkpoint: the dof values are the `.../vector` item.
587 family = attribute.attrib["ElementFamily"]
588 degree = int(attribute.attrib["ElementDegree"])
589 if family != "DG" or degree != 0:
590 raise ValueError(
591 f"Cannot read in finite element function ({family}, {degree}) as cell data."
592 )
593 for data_item in attribute.findall("./DataItem"):
594 if (data_item.text or "").strip().endswith("/vector"):
595 return data_item
596 raise RuntimeError(f"Checkpoint {attribute.attrib['Name']} has no vector DataItem.")
597 if (center := attribute.attrib.get("Center")) != "Cell":
598 raise ValueError(
599 f"Attribute {attribute.attrib['Name']} has Center='{center}'; cell data"
600 " needs Center='Cell'."
601 )
602 return _data_item(attribute)
605def _grid_topology(root: ElementTree.Element, grid: ElementTree.Element) -> ElementTree.Element:
606 """Topology `DataItem` of ``grid``, following DOLFINx's ``xi:include`` of the mesh.
608 Legacy DOLFIN writes the topology into each grid. DOLFINx refers to the mesh
609 grid with ``xpointer(<grid path>/*[self::Topology or self::Geometry])``.
611 Raises:
612 RuntimeError: If the topology cannot be found.
613 """
614 if (topology := grid.find("./Topology")) is not None:
615 return _data_item(topology)
616 # Any namespace: DOLFINx declares XInclude's as https, the standard says http.
617 if (include := grid.find("./{*}include")) is not None:
618 pointer = include.attrib["xpointer"]
619 match = re.fullmatch(r"xpointer\(/Xdmf(/.+?)(?:\[(\d+)\])?/\*\[.*\]\)", pointer)
620 if match is not None:
621 # ElementTree applies a position predicate to all same-tag siblings,
622 # XPath to the elements left by the preceding predicates.
623 path, position = match.groups()
624 mesh_grids = root.findall(f".{path}")
625 index = 0 if position is None else int(position) - 1
626 if index < len(mesh_grids):
627 if (topology := mesh_grids[index].find("./Topology")) is not None:
628 return _data_item(topology)
629 raise RuntimeError(f"No topology found for grid {grid.attrib.get('Name')}.")
632def write_data(
633 filename: Path | str,
634 point_data: ArrayData,
635 comm: MPI.Comm,
636 time: str | float | None,
637 mode: FileMode,
638 backend_args: dict[str, Any] | None,
639):
640 """Write a 2D-array to file (distributed across proceses with MPI).
643 Args:
644 filename: Path to file
645 point_data: Data to write to file
646 time: Time stamp
647 mode: Append or write
648 backend_args: The backend arguments
649 """
650 raise NotImplementedError("XDMF has not implemented this yet")