Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/__init__.py: 88%
65 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1from enum import Enum
2from importlib import import_module
3from pathlib import Path
4from typing import Any, Protocol
6from mpi4py import MPI
8import dolfinx
9import numpy as np
10import numpy.typing as npt
12from ..structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
14__all__ = [
15 "DEFAULT_MESH_NAME",
16 "FileMode",
17 "IOBackend",
18 "get_backend",
19 "get_mesh_name",
20 "set_default_backend",
21]
23_DEFAULT_BACKEND = "adios2"
25#: Name used for a mesh when the caller does not supply one. A file written
26#: before named meshes existed holds exactly this mesh, so backends must store
27#: it exactly as they did before (see :func:`get_mesh_name`).
28DEFAULT_MESH_NAME = "mesh"
31def get_mesh_name(backend_args: dict[str, Any] | None) -> str:
32 """Get the name of the mesh a backend call refers to.
34 A checkpoint may hold several meshes -- for instance a mesh and a submesh of
35 it -- distinguished by name. Backends take the name from
36 ``backend_args["name"]``.
38 Args:
39 backend_args: Arguments to backend, or None
41 Returns:
42 The mesh name, :data:`DEFAULT_MESH_NAME` if none was given.
43 """
44 if backend_args is None:
45 return DEFAULT_MESH_NAME
46 name = backend_args.get("name", DEFAULT_MESH_NAME)
47 return DEFAULT_MESH_NAME if name is None else str(name)
50def set_default_backend(backend: str):
51 """Set the global default backend for io4dolfinx."""
52 global _DEFAULT_BACKEND
53 if backend not in BUILTIN_BAKENDS:
54 try:
55 get_backend(backend)
56 except ImportError as e:
57 raise ImportError(f"Backend {backend} not found.") from e
58 _DEFAULT_BACKEND = backend
61def get_default_backend():
62 """Get the global default backend for io4dolfinx."""
63 return _DEFAULT_BACKEND
66class ReadMode(Enum):
67 serial = 10 # This means that all data is read in on root rank
69 # Total number of data P, num processes = i + 1.
70 # All processes reads at least `P // (i+1)` items
71 # The first j=P%(i+1) processes reads `P // (i+1) + 1` items
72 # ```python
73 # def compute_partitioning(P, J):
74 # min_num = P // J
75 # num_per_proc = np.full(J, min_num)
76 # rem = P % J
77 # num_per_proc[:int(rem)] += 1
78 # assert(sum(num_per_proc)) == P
79 # return num_per_proc
80 # ```
81 parallel = 20
84class FileMode(Enum):
85 """Filen mode used for opening files."""
87 append = 10 #: Append data to file
88 write = 20 #: Write data to file
89 read = 30 #: Read data from file
92# See https://peps.python.org/pep-0544/#modules-as-implementations-of-protocols
93class IOBackend(Protocol):
94 """Interface a storage backend must provide.
96 .. note::
97 A checkpoint may hold more than one mesh, each with its own topology,
98 geometry, meshtags, cell permutations and functions. Every method that
99 touches mesh-associated data takes the mesh it refers to from
100 ``backend_args["name"]``; use :func:`get_mesh_name` to read it, so that
101 omitting it selects :data:`DEFAULT_MESH_NAME` consistently. A backend
102 must store that default mesh exactly as it did before named meshes
103 existed, so that older checkpoints stay readable.
104 """
106 read_mode: ReadMode
108 def get_default_backend_args(self, arguments: dict[str, Any] | None) -> dict[str, Any]:
109 """Get default backend arguments given a set of input arguments.
111 Args:
112 arguments: Input backend arguments
114 Returns:
115 Updated backend arguments
116 """
118 def write_attributes(
119 self,
120 filename: Path | str,
121 comm: MPI.Comm,
122 name: str,
123 attributes: dict[str, np.ndarray],
124 backend_args: dict[str, Any] | None,
125 ):
126 """Write attributes to file.
128 Args:
129 filename: Path to file to write to
130 comm: MPI communicator used in storage
131 name: Name of the attribute group
132 attributes: Dictionary of attributes to write
133 backend_args: Arguments to backend
134 """
136 def read_attributes(
137 self,
138 filename: Path | str,
139 comm: MPI.Comm,
140 name: str,
141 backend_args: dict[str, Any] | None,
142 ) -> dict[str, Any]:
143 """Read attributes from file.
145 Args:
146 filename: Path to file to read from
147 comm: MPI communicator used in storage
148 name: Name of the attribute group
149 backend_args: Arguments to backend
151 Returns:
152 Dictionary of attributes read from file
153 """
155 def read_timestamps(
156 self,
157 filename: Path | str,
158 comm: MPI.Comm,
159 function_name: str,
160 backend_args: dict[str, Any] | None,
161 ) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
162 """Read timestamps from file.
164 Args:
165 filename: Path to file to read from
166 comm: MPI communicator used in storage
167 function_name: Name of the function to read timestamps for
168 backend_args: Arguments to backend
170 Returns:
171 Numpy array of timestamps read from file
172 """
174 def write_mesh(
175 self,
176 filename: Path | str,
177 comm: MPI.Comm,
178 mesh: MeshData,
179 backend_args: dict[str, Any] | None,
180 mode: FileMode,
181 time: float,
182 ):
183 """
184 Write a mesh to file.
186 Args:
187 comm: MPI communicator used in storage
188 mesh: Internal data structure for the mesh data to save to file
189 filename: Path to file to write to
190 backend_args: Arguments to backend
191 mode: File-mode to store the mesh
192 time: Time stamp associated with the mesh
193 """
195 def write_meshtags(
196 self,
197 filename: str | Path,
198 comm: MPI.Comm,
199 data: MeshTagsData,
200 backend_args: dict[str, Any] | None,
201 ):
202 """Write mesh tags to file.
204 Args:
205 filename: Path to file to write to
206 comm: MPI communicator used in storage
207 data: Internal data structure for the mesh tags to save to file
208 backend_args: Arguments to backend
209 """
211 def read_mesh_data(
212 self,
213 filename: Path | str,
214 comm: MPI.Comm,
215 time: str | float | None,
216 read_from_partition: bool,
217 backend_args: dict[str, Any] | None,
218 ) -> ReadMeshData:
219 """Read mesh data from file.
221 Args:
222 filename: Path to file to read from
223 comm: MPI communicator used in storage
224 time: Time stamp associated with the mesh to read
225 read_from_partition: Whether to read partition information
226 backend_args: Arguments to backend
228 Returns:
229 Internal data structure for the mesh data read from file
230 """
232 def read_meshtags_data(
233 self,
234 filename: str | Path,
235 comm: MPI.Comm,
236 name: str,
237 backend_args: dict[str, Any] | None,
238 ) -> MeshTagsData:
239 """Read mesh tags from file.
241 Args:
242 filename: Path to file to read from
243 comm: MPI communicator used in storage
244 name: Name of the mesh tags to read
245 backend_args: Arguments to backend
247 Returns:
248 Internal data structure for the mesh tags read from file
249 """
251 def read_dofmap(
252 self,
253 filename: str | Path,
254 comm: MPI.Comm,
255 name: str,
256 backend_args: dict[str, Any] | None,
257 ) -> dolfinx.graph.AdjacencyList:
258 """Read the dofmap of a function with a given name.
260 Args:
261 filename: Path to file to read from
262 comm: MPI communicator used in storage
263 name: Name of the function to read the dofmap for
264 backend_args: Arguments to backend
266 Returns:
267 Dofmap as an :class:`dolfinx.graph.AdjacencyList`
268 """
270 def read_dofs(
271 self,
272 filename: str | Path,
273 comm: MPI.Comm,
274 name: str,
275 time: float,
276 backend_args: dict[str, Any] | None,
277 ) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
278 """Read the dofs (values) of a function with a given name from a given timestep.
280 Args:
281 filename: Path to file to read from
282 comm: MPI communicator used in storage
283 name: Name of the function to read the dofs for
284 time: Time stamp associated with the function to read
285 backend_args: Arguments to backend
287 Returns:
288 Contiguous sequence of degrees of freedom (with respect to input data)
289 and the global starting point on the process.
290 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
291 """
293 def read_cell_perms(
294 self, comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None
295 ) -> npt.NDArray[np.uint32]:
296 """
297 Read cell permutation from file with given communicator,
298 Split in continuous chunks based on number of cells in the input data.
300 Cell permutations belong to a mesh, not to a function, so which mesh is
301 meant is taken from ``backend_args["name"]``.
303 Args:
304 comm: MPI communicator used in storage
305 filename: Path to file to read from
306 backend_args: Arguments to backend
308 Returns:
309 Contiguous sequence of permutations (with respect to input data)
310 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
311 """
313 def write_function(
314 self,
315 filename: Path,
316 comm: MPI.Comm,
317 u: FunctionData,
318 time: float,
319 mode: FileMode,
320 backend_args: dict[str, Any] | None,
321 ):
322 """Write a function to file.
324 Args:
325 comm: MPI communicator used in storage
326 u: Internal data structure for the function data to save to file
327 filename: Path to file to write to
328 time: Time stamp associated with function
329 mode: File-mode to store the function
330 backend_args: Arguments to backend
331 """
333 def read_legacy_mesh(
334 self, filename: Path | str, comm: MPI.Comm, group: str
335 ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
336 """Read in the mesh topology, geometry and (optionally) cell type from a
337 legacy DOLFIN HDF5-file.
339 Args:
340 filename: Path to file to read from
341 comm: MPI communicator used in storage
342 group: Group in HDF5 file where mesh is stored
344 Returns:
345 Tuple containing:
346 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
347 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
348 - Cell type as a string (e.g. "tetrahedron") or None if not found
349 """
351 def snapshot_checkpoint(
352 self,
353 filename: Path | str,
354 mode: FileMode,
355 u: dolfinx.fem.Function,
356 backend_args: dict[str, Any] | None,
357 ):
358 """Create a snapshot checkpoint of a dolfinx function.
360 Args:
361 filename: Path to file to read from
362 mode: File-mode to store the function
363 u: dolfinx function to create a snapshot checkpoint for
364 backend_args: Arguments to backend
365 """
367 def read_hdf5_array(
368 self,
369 comm: MPI.Comm,
370 filename: Path | str,
371 group: str,
372 backend_args: dict[str, Any] | None,
373 ) -> tuple[np.ndarray, int]:
374 """Read an array from an HDF5 file.
376 Args:
377 comm: MPI communicator used in storage
378 filename: Path to file to read from
379 group: Group in HDF5 file where array is stored
380 backend_args: Arguments to backend
382 Returns:
383 Tuple containing:
384 - Numpy array read from file
385 - Global starting point on the process.
386 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
387 """
389 def read_point_data(
390 self,
391 filename: Path | str,
392 name: str,
393 comm: MPI.Comm,
394 time: str | float | None,
395 backend_args: dict[str, Any] | None,
396 ) -> tuple[np.ndarray, int]:
397 """Read data from the nodes of a mesh.
399 Args:
400 filename: Path to file
401 name: Name of point data
402 comm: Communicator to launch IO on.
403 time: The time stamp
404 backend_args: The backend arguments
405 Returns:
406 Data local to process (contiguous, no mpi comm) and local start range
407 """
408 ...
410 def read_function_names(
411 self, filename: Path | str, comm: MPI.Comm, backend_args: dict[str, Any] | None
412 ) -> list[str]:
413 """Read all function names from a file.
415 Args:
416 filename: Path to file
417 comm: MPI communicator to launch IO on.
418 backend_args: Arguments to backend
420 Returns:
421 A list of function names.
422 """
423 ...
425 def read_cell_data(
426 self,
427 filename: Path | str,
428 name: str,
429 comm: MPI.Comm,
430 time: str | float | None,
431 backend_args: dict[str, Any] | None,
432 ) -> tuple[npt.NDArray[np.int64], np.ndarray]:
433 """Read data from the cells of a mesh.
435 Args:
436 filename: Path to file
437 name: Name of point data
438 comm: Communicator to launch IO on.
439 time: The time stamp
440 backend_args: The backend arguments
441 Returns:
442 A tuple (topology, dofs) where topology contains the
443 vertex indices of the cells, dofs the degrees of
444 freedom within that cell.
445 """
446 ...
448 def write_data(
449 self,
450 filename: Path | str,
451 array_data: ArrayData,
452 comm: MPI.Comm,
453 time: str | float | None,
454 mode: FileMode,
455 backend_args: dict[str, Any] | None,
456 ):
457 """Write a 2D-array to file.
460 Args:
461 filename: Path to file
462 array_data: Data to write to file.
463 comm: The MPI communicator to open the writer with.
464 time: The time stamp
465 mode: Append or write
466 backend_args: The backend arguments
467 """
468 ...
471def get_backend(backend: str | None = None) -> IOBackend:
472 """Get backend class from backend name.
474 Args:
475 backend: Name of the backend to get
477 Returns:
478 Backend class
479 """
480 if backend is None:
481 backend = _DEFAULT_BACKEND
483 if backend == "h5py":
484 from .h5py import backend as H5PYInterface
486 return H5PYInterface
487 elif backend == "adios2":
488 from .adios2 import backend as ADIOS2Interface
490 return ADIOS2Interface
491 elif backend == "pyvista":
492 from .pyvista import backend as PYVISTAInterface
494 return PYVISTAInterface
495 elif backend == "xdmf":
496 from .xdmf import backend as XDMFInterface
498 return XDMFInterface
499 elif backend == "vtkhdf":
500 from .vtkhdf import backend as VTKDHFInterface
502 return VTKDHFInterface
503 elif backend == "exodus":
504 from .exodus import backend as EXODUSInterface
506 return EXODUSInterface
507 else:
508 return import_module(backend)
511BUILTIN_BAKENDS = ("h5py", "adios2", "pyvista", "xdmf", "vtkhdf", "exodus")
514def list_builtin_backends() -> list[str]:
515 """List available builtin backends.
517 Returns:
518 List of available backends
519 """
520 lst = []
521 for backend in BUILTIN_BAKENDS:
522 try:
523 get_backend(backend)
524 except ImportError:
525 continue
526 lst.append(backend)
527 return lst