Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/adios2/backend.py: 94%
408 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1import warnings
2from pathlib import Path
3from typing import Any
5from mpi4py import MPI
7import adios2
8import dolfinx
9import fenicsx_compat
10import numpy as np
11import numpy.typing as npt
13from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
14from ...utils import check_file_exists, compute_local_range
15from .. import DEFAULT_MESH_NAME, FileMode, ReadMode, get_mesh_name
16from .helpers import (
17 ADIOSFile,
18 adios_to_numpy_dtype,
19 check_variable_exists,
20 read_adjacency_list,
21 read_array,
22)
24adios2 = fenicsx_compat.resolve_adios_scope(adios2)
26read_mode = ReadMode.parallel
29def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]:
30 """Get default arguements (sets engine to BP4)."""
31 args = dict(arguments) if arguments else {} # Copy: do not mutate the caller's dict
32 if "engine" not in args.keys():
33 args["engine"] = "BP4"
34 if "legacy" not in args.keys():
35 args["legacy"] = False # Only used for legacy HDF5 meshtags
36 if "name" not in args.keys():
37 args["name"] = DEFAULT_MESH_NAME
38 return args
41def mesh_prefix(backend_args: dict[str, Any] | None) -> str:
42 """ADIOS2 variable-name prefix for the mesh named in ``backend_args``.
44 The default mesh gets an *empty* prefix, so that it is stored under exactly
45 the names used before checkpoints could hold more than one mesh. Files
46 written by earlier versions therefore stay readable, and a file holding only
47 the default mesh stays readable by them. Any other mesh is namespaced under
48 ``"<name>/"``.
50 .. note::
51 This is io4dolfinx's own checkpoint schema, not ADIOS2's VTX
52 visualization schema -- no ``vtk.xml`` attribute is written, and these
53 files are not readable by ParaView's VTX reader whether they hold one
54 mesh or several. VTX supports only a single mesh per file, so holding
55 several here costs nothing that was ever available. For a
56 visualizable file with several meshes use the ``vtkhdf`` backend, whose
57 ``MultiBlockDataSet`` is designed for it.
59 Args:
60 backend_args: Arguments to backend
62 Returns:
63 The prefix to prepend to every variable and attribute name of that mesh.
64 """
65 name = get_mesh_name(backend_args)
66 return "" if name == DEFAULT_MESH_NAME else f"{name}/"
69def _seek_step(adios_file, variable: str, start: int = 0) -> int:
70 """Advance the reader to the first step at or after ``start`` holding ``variable``.
72 With several meshes in one file a step no longer holds every variable, so a
73 reader cannot assume the one it wants is in step 0. Leaves the file *inside*
74 the matching step (the caller is responsible for ``EndStep``).
76 Args:
77 adios_file: The open :class:`ADIOSFile`
78 variable: Name of the variable to seek
79 start: First step to consider
81 Returns:
82 Index of the step holding ``variable``.
84 Raises:
85 KeyError: If no step holds ``variable``.
86 """
87 for i in range(start, adios_file.file.Steps()):
88 adios_file.file.BeginStep()
89 if variable in adios_file.io.AvailableVariables().keys():
90 return i
91 adios_file.file.EndStep()
92 raise KeyError(f"'{variable}' not found in file")
95def convert_file_mode(mode: FileMode) -> adios2.Mode: # type: ignore[override]
96 match mode:
97 case FileMode.append:
98 return adios2.Mode.Append
99 case FileMode.write:
100 return adios2.Mode.Write
101 case FileMode.read:
102 return adios2.Mode.Read
103 case _:
104 raise NotImplementedError(f"FileMode {mode} not implemented.")
107def write_attributes(
108 filename: Path | str,
109 comm: MPI.Comm,
110 name: str,
111 attributes: dict[str, np.ndarray],
112 backend_args: dict[str, Any] | None = None,
113):
114 """Write attributes to file using ADIOS2.
116 Args:
117 filename: Path to file to write to
118 comm: MPI communicator used in storage
119 name: Name of the attributes
120 attributes: Dictionary of attributes to write to file
121 engine: ADIOS2 engine to use
122 """
124 adios = adios2.ADIOS(comm)
125 backend_args = get_default_backend_args(backend_args)
127 with ADIOSFile(
128 adios=adios,
129 filename=filename,
130 mode=adios2.Mode.Append,
131 io_name="AttributeWriter",
132 engine=backend_args["engine"],
133 ) as adios_file:
134 adios_file.file.BeginStep()
136 for k, v in attributes.items():
137 adios_file.io.DefineAttribute(f"{name}_{k}", v)
139 adios_file.file.PerformPuts()
140 adios_file.file.EndStep()
143def read_attributes(
144 filename: Path | str,
145 comm: MPI.Comm,
146 name: str,
147 backend_args: dict[str, Any] | None = None,
148) -> dict[str, np.ndarray]:
149 """Read attributes from file using ADIOS2.
151 Args:
152 filename: Path to file to read from
153 comm: MPI communicator used in storage
154 name: Name of the attributes
155 engine: ADIOS2 engine to use
156 Returns:
157 The attributes
158 """
159 check_file_exists(filename)
160 adios = adios2.ADIOS(comm)
161 backend_args = get_default_backend_args(backend_args)
162 with ADIOSFile(
163 adios=adios,
164 filename=filename,
165 mode=adios2.Mode.Read,
166 engine=backend_args["engine"],
167 io_name="AttributesReader",
168 ) as adios_file:
169 adios_file.file.BeginStep()
170 attributes = {}
171 for k in adios_file.io.AvailableAttributes().keys():
172 if k.startswith(f"{name}_"):
173 a = adios_file.io.InquireAttribute(k)
174 attributes[k[len(name) + 1 :]] = a.Data()
175 adios_file.file.EndStep()
176 return attributes
179def read_timestamps(
180 filename: Path | str,
181 comm: MPI.Comm,
182 function_name: str,
183 backend_args: dict[str, Any] | None = None,
184) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
185 """Read time-stamps from a checkpoint file.
187 Args:
188 comm: MPI communicator
189 filename: Path to file
190 function_name: Name of the function to read time-stamps for
191 backend_args: Arguments for backend, for instance file type.
192 backend: What backend to use for writing.
193 Returns:
194 The time-stamps
195 """
196 check_file_exists(filename)
198 adios = adios2.ADIOS(comm)
199 backend_args = get_default_backend_args(backend_args)
200 prefix = mesh_prefix(backend_args)
201 with ADIOSFile(
202 adios=adios,
203 filename=filename,
204 mode=adios2.Mode.Read,
205 engine=backend_args["engine"],
206 io_name="TimestepReader",
207 ) as adios_file:
208 time_name = f"{prefix}{function_name}_time"
209 time_stamps = []
210 for _ in range(adios_file.file.Steps()):
211 adios_file.file.BeginStep()
212 if time_name in adios_file.io.AvailableVariables().keys():
213 arr = adios_file.io.InquireVariable(time_name)
214 time_shape = arr.Shape()
215 arr.SetSelection([[0], [time_shape[0]]])
216 times = np.empty(
217 time_shape[0],
218 dtype=adios_to_numpy_dtype[arr.Type()],
219 )
220 adios_file.file.Get(arr, times, adios2.Mode.Sync)
221 time_stamps.append(times[0])
222 adios_file.file.EndStep()
224 return np.array(time_stamps)
227def write_mesh(
228 filename: Path | str,
229 comm: MPI.Comm,
230 mesh: MeshData,
231 backend_args: dict[str, Any] | None = None,
232 mode: FileMode = FileMode.write,
233 time: float = 0.0,
234):
235 """Write a mesh to file using ADIOS2.
237 Args:
238 comm: MPI communicator used in storage
239 mesh: Internal data structure for the mesh data to save to file
240 filename: Path to file to write to
241 backend_args: File mode and potentially the io-name.
242 mode: Mode to use (write or append)
243 time: Time stamp
244 """
245 backend_args = get_default_backend_args(backend_args)
246 if "io_name" not in backend_args.keys():
247 backend_args["io_name"] = "MeshWriter"
248 prefix = mesh_prefix(backend_args)
250 mode = convert_file_mode(mode)
251 gdim = mesh.local_geometry.shape[1]
252 adios = adios2.ADIOS(comm)
254 # Topology is constant in time and written once. In append mode it may
255 # already be present either because this mesh was written at an earlier time
256 # step, or because a *different* mesh was written first -- so the test is
257 # whether this mesh's topology exists, not whether the file is new.
258 topology_exists = False
259 if mode == adios2.Mode.Append:
260 topology_exists = check_variable_exists(
261 adios, filename, f"{prefix}Topology", engine=backend_args["engine"]
262 )
264 with ADIOSFile(
265 adios=adios,
266 filename=filename,
267 mode=mode,
268 comm=comm,
269 engine=backend_args["engine"],
270 io_name=backend_args["io_name"],
271 ) as adios_file:
272 adios_file.file.BeginStep()
273 # Write geometry
274 pointvar = adios_file.io.DefineVariable(
275 f"{prefix}Points",
276 mesh.local_geometry,
277 shape=[mesh.num_nodes_global, gdim],
278 start=[mesh.local_geometry_pos[0], 0],
279 count=[mesh.local_geometry_pos[1] - mesh.local_geometry_pos[0], gdim],
280 )
281 adios_file.file.Put(pointvar, mesh.local_geometry, adios2.Mode.Sync)
283 if not topology_exists:
284 adios_file.io.DefineAttribute(f"{prefix}CellType", mesh.cell_type)
285 adios_file.io.DefineAttribute(
286 f"{prefix}Degree", np.array([mesh.degree], dtype=np.int32)
287 )
288 adios_file.io.DefineAttribute(
289 f"{prefix}LagrangeVariant", np.array([mesh.lagrange_variant], dtype=np.int32)
290 )
291 # Write topology (on;y on first write as topology is constant)
292 num_dofs_per_cell = mesh.local_topology.shape[1]
293 dvar = adios_file.io.DefineVariable(
294 f"{prefix}Topology",
295 mesh.local_topology,
296 shape=[mesh.num_cells_global, num_dofs_per_cell],
297 start=[mesh.local_topology_pos[0], 0],
298 count=[
299 mesh.local_topology_pos[1] - mesh.local_topology_pos[0],
300 num_dofs_per_cell,
301 ],
302 )
303 adios_file.file.Put(dvar, mesh.local_topology)
305 # Add partitioning data
306 if mesh.store_partition:
307 assert mesh.partition_range is not None
308 par_data = adios_file.io.DefineVariable(
309 f"{prefix}PartitioningData",
310 mesh.ownership_array,
311 shape=[mesh.partition_global],
312 start=[mesh.partition_range[0]],
313 count=[
314 mesh.partition_range[1] - mesh.partition_range[0],
315 ],
316 )
317 adios_file.file.Put(par_data, mesh.ownership_array)
318 assert mesh.ownership_offset is not None
319 par_offset = adios_file.io.DefineVariable(
320 f"{prefix}PartitioningOffset",
321 mesh.ownership_offset,
322 shape=[mesh.num_cells_global + 1],
323 start=[mesh.local_topology_pos[0]],
324 count=[mesh.local_topology_pos[1] - mesh.local_topology_pos[0] + 1],
325 )
326 adios_file.file.Put(par_offset, mesh.ownership_offset)
327 assert mesh.partition_processes is not None
328 adios_file.io.DefineAttribute(
329 f"{prefix}PartitionProcesses",
330 np.array([mesh.partition_processes], dtype=np.int32),
331 )
332 if topology_exists and mesh.store_partition:
333 warnings.warn("Partitioning data is not written in append mode")
335 # Add time step to file
336 t_arr = np.array([time], dtype=np.float64)
337 time_var = adios_file.io.DefineVariable(
338 f"{prefix}MeshTime",
339 t_arr,
340 shape=[1],
341 start=[0],
342 count=[1 if comm.rank == 0 else 0],
343 )
344 adios_file.file.Put(time_var, t_arr)
346 adios_file.file.PerformPuts()
347 adios_file.file.EndStep()
350def read_mesh_data(
351 filename: Path | str,
352 comm: MPI.Comm,
353 time: str | float | None = 0.0,
354 read_from_partition: bool = False,
355 backend_args: dict[str, Any] | None = None,
356) -> ReadMeshData:
357 """Read an ADIOS2 mesh data for use with DOLFINx.
359 Args:
360 filename: Path to input file
361 comm: The MPI communciator to distribute the mesh over
362 engine: ADIOS engine to use for reading (BP4, BP5 or HDF5)
363 time: Time stamp associated with mesh
364 legacy: If checkpoint was made prior to time-dependent mesh-writer set to True
365 read_from_partition: Read mesh with partition from file
366 Returns:
367 The mesh topology, geometry, UFL domain and partition function
368 """
370 adios = adios2.ADIOS(comm)
371 backend_args = get_default_backend_args(backend_args)
372 legacy = backend_args.get("legacy", False)
373 io_name = backend_args.get("io_name", "MeshReader")
374 engine = backend_args["engine"]
375 prefix = mesh_prefix(backend_args)
376 with ADIOSFile(
377 adios=adios,
378 filename=filename,
379 mode=adios2.Mode.Read,
380 engine=engine,
381 io_name=io_name,
382 ) as adios_file:
383 # Get time independent mesh variables (mesh topology and cell type info)
384 # first. With several meshes in the file this mesh's topology need not be
385 # in step 0, so seek the step holding it.
386 try:
387 step = _seek_step(adios_file, f"{prefix}Topology")
388 except KeyError as e:
389 raise KeyError(f"Mesh topology not found at {prefix}Topology in {filename}") from e
390 topology = adios_file.io.InquireVariable(f"{prefix}Topology")
391 shape = topology.Shape()
392 local_range = compute_local_range(comm, shape[0])
393 topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
394 mesh_topology = np.empty((local_range[1] - local_range[0], shape[1]), dtype=np.int64)
395 adios_file.file.Get(topology, mesh_topology, adios2.Mode.Deferred)
397 # Check validity of partitioning information
398 if read_from_partition:
399 if f"{prefix}PartitionProcesses" not in adios_file.io.AvailableAttributes().keys():
400 raise KeyError(f"Partitioning information not found in {filename}")
401 par_num_procs = adios_file.io.InquireAttribute(f"{prefix}PartitionProcesses")
402 num_procs = par_num_procs.Data()[0]
403 if num_procs != comm.size:
404 raise ValueError(f"Number of processes in file ({num_procs})!=({comm.size=})")
406 # Get mesh cell type
407 if f"{prefix}CellType" not in adios_file.io.AvailableAttributes().keys():
408 raise KeyError(f"Mesh cell type not found at {prefix}CellType in {filename}")
409 celltype = adios_file.io.InquireAttribute(f"{prefix}CellType")
410 cell_type = celltype.DataString()[0]
412 # Get basix info
413 if f"{prefix}LagrangeVariant" not in adios_file.io.AvailableAttributes().keys():
414 raise KeyError(f"Mesh LagrangeVariant not found in {filename}")
415 lvar = adios_file.io.InquireAttribute(f"{prefix}LagrangeVariant").Data()[0]
416 if f"{prefix}Degree" not in adios_file.io.AvailableAttributes().keys():
417 raise KeyError(f"Mesh degree not found in {filename}")
418 degree = adios_file.io.InquireAttribute(f"{prefix}Degree").Data()[0]
420 if not legacy:
421 # Geometry for this time stamp is written in the same step as the
422 # topology or a later one, so resume the scan from there.
423 time_name = f"{prefix}MeshTime"
424 for i in range(step, adios_file.file.Steps()):
425 if i > step:
426 adios_file.file.BeginStep()
427 if time_name in adios_file.io.AvailableVariables().keys():
428 arr = adios_file.io.InquireVariable(time_name)
429 time_shape = arr.Shape()
430 arr.SetSelection([[0], [time_shape[0]]])
431 times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()])
432 adios_file.file.Get(arr, times, adios2.Mode.Sync)
433 if times[0] == time:
434 break
435 if i == adios_file.file.Steps() - 1:
436 raise KeyError(
437 f"No data associated with {time_name}={time} found in {filename}"
438 )
440 adios_file.file.EndStep()
442 if time_name not in adios_file.io.AvailableVariables().keys():
443 raise KeyError(f"No data associated with {time_name}={time} found in {filename}")
445 # Get mesh geometry
446 if f"{prefix}Points" not in adios_file.io.AvailableVariables().keys():
447 raise KeyError(f"Mesh coordinates not found at {prefix}Points in {filename}")
448 geometry = adios_file.io.InquireVariable(f"{prefix}Points")
449 x_shape = geometry.Shape()
450 geometry_range = compute_local_range(comm, x_shape[0])
451 geometry.SetSelection(
452 [
453 [geometry_range[0], 0],
454 [geometry_range[1] - geometry_range[0], x_shape[1]],
455 ]
456 )
457 mesh_geometry = np.empty(
458 (geometry_range[1] - geometry_range[0], x_shape[1]),
459 dtype=adios_to_numpy_dtype[geometry.Type()],
460 )
461 adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Deferred)
462 adios_file.file.PerformGets()
463 adios_file.file.EndStep()
465 if read_from_partition:
466 partition_graph = read_adjacency_list(
467 adios,
468 comm,
469 filename,
470 f"{prefix}PartitioningData",
471 f"{prefix}PartitioningOffset",
472 backend_args["engine"],
473 )
474 else:
475 partition_graph = None
477 return ReadMeshData(
478 cells=mesh_topology,
479 cell_type=cell_type,
480 x=mesh_geometry,
481 degree=degree,
482 lvar=lvar,
483 partition_graph=partition_graph,
484 )
487def write_meshtags(
488 filename: str | Path,
489 comm: MPI.Comm,
490 data: MeshTagsData,
491 backend_args: dict[str, Any] | None = None,
492):
493 """Write mesh tags to file.
495 Args:
496 filename: Path to file to write to
497 comm: MPI communicator used in storage
498 data: Internal data structure for the mesh tags to save to file
499 backend_args: Arguments to backend
500 """
501 backend_args = {} if backend_args is None else backend_args
502 io_name = backend_args.get("io_name", "MeshTagWriter")
503 engine = backend_args.get("engine", "BP4")
504 prefix = mesh_prefix(backend_args)
505 adios = adios2.ADIOS(comm)
506 with ADIOSFile(
507 adios=adios,
508 filename=filename,
509 mode=adios2.Mode.Append,
510 engine=engine,
511 io_name=io_name,
512 ) as adios_file:
513 adios_file.file.BeginStep()
515 # Write meshtag topology
516 topology_var = adios_file.io.DefineVariable(
517 f"{prefix}{data.name}_topology",
518 data.indices,
519 shape=[data.num_entities_global, data.num_dofs_per_entity],
520 start=[data.local_start, 0],
521 count=[len(data.indices), data.num_dofs_per_entity],
522 )
523 adios_file.file.Put(topology_var, data.indices, adios2.Mode.Sync)
525 # Write meshtag values
526 vals = np.array(data.values)
527 values_var = adios_file.io.DefineVariable(
528 f"{prefix}{data.name}_values",
529 vals,
530 shape=[data.num_entities_global],
531 start=[data.local_start],
532 count=[len(data.indices)],
533 )
534 adios_file.file.Put(values_var, vals, adios2.Mode.Sync)
536 # Write meshtag dim
537 adios_file.io.DefineAttribute(
538 f"{prefix}{data.name}_dim", np.array([data.dim], dtype=np.uint8)
539 )
540 adios_file.file.PerformPuts()
541 adios_file.file.EndStep()
544def read_meshtags_data(
545 filename: str | Path, comm: MPI.Comm, name: str, backend_args: dict[str, Any] | None = None
546) -> MeshTagsData:
547 """Read mesh tags from file.
549 Args:
550 filename: Path to file to read from
551 comm: MPI communicator used in storage
552 name: Name of the mesh tags to read
553 backend_args: Arguments to backend
555 Returns:
556 Internal data structure for the mesh tags read from file
557 """
559 adios = adios2.ADIOS(comm)
560 backend_args = get_default_backend_args(backend_args)
561 io_name = backend_args.get("io_name", "MeshTagsReader")
562 engine = backend_args["engine"]
563 legacy = backend_args["legacy"]
564 prefix = mesh_prefix(backend_args)
565 with ADIOSFile(
566 adios=adios,
567 filename=filename,
568 mode=adios2.Mode.Read,
569 engine=engine,
570 io_name=io_name,
571 ) as adios_file:
572 if not legacy:
573 # Get mesh cell type
574 dim_attr_name = f"{prefix}{name}_dim"
575 step = 0
576 for i in range(adios_file.file.Steps()):
577 adios_file.file.BeginStep()
578 if dim_attr_name in adios_file.io.AvailableAttributes().keys():
579 step = i
580 break
581 adios_file.file.EndStep()
582 if dim_attr_name not in adios_file.io.AvailableAttributes().keys():
583 raise KeyError(f"{dim_attr_name} not found in {filename}")
585 m_dim = adios_file.io.InquireAttribute(dim_attr_name)
586 dim = int(m_dim.Data()[0])
588 # Get mesh tags entites
589 topology_name = f"{prefix}{name}_topology"
590 for i in range(step, adios_file.file.Steps()):
591 if i > step:
592 adios_file.file.BeginStep()
593 if topology_name in adios_file.io.AvailableVariables().keys():
594 break
595 adios_file.file.EndStep()
596 if topology_name not in adios_file.io.AvailableVariables().keys():
597 raise KeyError(f"{topology_name} not found in {filename}")
599 topology = adios_file.io.InquireVariable(topology_name)
600 top_shape = topology.Shape()
601 topology_range = compute_local_range(comm, top_shape[0])
603 topology.SetSelection(
604 [
605 [topology_range[0], 0],
606 [topology_range[1] - topology_range[0], top_shape[1]],
607 ]
608 )
609 mesh_entities = np.empty(
610 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
611 )
612 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
614 # Get mesh tags values
615 values_name = f"{prefix}{name}_values"
616 if values_name not in adios_file.io.AvailableVariables().keys():
617 raise KeyError(f"{values_name} not found")
619 values = adios_file.io.InquireVariable(values_name)
620 val_shape = values.Shape()
621 assert val_shape[0] == top_shape[0]
622 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
623 tag_values = np.empty(
624 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
625 )
626 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
628 adios_file.file.PerformGets()
629 adios_file.file.EndStep()
630 else:
631 # Get mesh cell type
632 dim_attr_name = f"{name}_dim"
633 assert adios_file.file.Steps() == 0
634 if (ct_key := f"/{name}/topology/celltype") in adios_file.io.AvailableAttributes():
635 cell_type = adios_file.io.InquireAttribute(ct_key)
636 else:
637 raise ValueError(f"Celltype not found for meshtags {name} in {filename}.")
638 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type.DataString()[0]))
640 # Get mesh tags entites
641 if (top_key := f"/{name}/topology") in adios_file.io.AvailableVariables():
642 topology = adios_file.io.InquireVariable(top_key)
643 else:
644 raise ValueError(f"Topology not found for meshtags {name} in {filename}.")
646 top_shape = topology.Shape()
647 topology_range = compute_local_range(comm, top_shape[0])
649 topology.SetSelection(
650 [
651 [topology_range[0], 0],
652 [topology_range[1] - topology_range[0], top_shape[1]],
653 ]
654 )
655 mesh_entities = np.empty(
656 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
657 )
658 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
660 # Get mesh tags values
661 if (val_key := f"/{name}/values") in adios_file.io.AvailableVariables():
662 values = adios_file.io.InquireVariable(val_key)
663 else:
664 raise ValueError(f"Values not found for meshtags {name} in {filename}.")
666 val_shape = values.Shape()
667 assert val_shape[0] == top_shape[0]
669 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
670 tag_values = np.empty(
671 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
672 )
673 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
675 adios_file.file.PerformGets()
676 adios_file.file.EndStep()
678 # Tag values are int32 unless the caller asks for something wider. The
679 # submesh post codes are global cell indices, which need int64.
680 values_dtype = np.dtype(backend_args.get("values_dtype", np.int32))
681 return MeshTagsData(
682 name=name, values=tag_values.astype(values_dtype), indices=mesh_entities, dim=dim
683 )
686def read_dofmap(
687 filename: str | Path, comm: MPI.Comm, name: str, backend_args: dict[str, Any] | None = None
688) -> dolfinx.graph.AdjacencyList:
689 """Read the dofmap of a function with a given name.
691 Args:
692 filename: Path to file to read from
693 comm: MPI communicator used in storage
694 name: Name of the function to read the dofmap for
695 backend_args: Arguments to backend
697 Returns:
698 Dofmap as an AdjacencyList
699 """
700 backend_args = {} if backend_args is None else backend_args
701 prefix = mesh_prefix(backend_args)
703 # Handles legacy io4dolfinx files, modern files, and custom location of dofmap.
704 legacy = backend_args.get("legacy", False)
705 xdofmap_path: str | None
706 dofmap_path: str | None
707 if (dofmap_path := backend_args.get("dofmap", None)) is None:
708 if legacy:
709 dofmap_path = "Dofmap"
710 else:
711 dofmap_path = f"{prefix}{name}_dofmap"
713 if (xdofmap_path := backend_args.get("offsets", None)) is None:
714 if legacy:
715 xdofmap_path = "XDofmap"
716 else:
717 xdofmap_path = f"{prefix}{name}_XDofmap"
719 engine = backend_args.get("engine", "BP4")
721 adios = adios2.ADIOS(comm)
722 check_file_exists(filename)
723 assert isinstance(xdofmap_path, str)
724 return read_adjacency_list(adios, comm, filename, dofmap_path, xdofmap_path, engine=engine)
727def read_dofs(
728 filename: str | Path,
729 comm: MPI.Comm,
730 name: str,
731 time: float,
732 backend_args: dict[str, Any] | None = None,
733) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
734 """Read the dofs (values) of a function with a given name from a given timestep.
736 Args:
737 filename: Path to file to read from
738 comm: MPI communicator used in storage
739 name: Name of the function to read the dofs for
740 time: Time stamp associated with the function to read
741 backend_args: Arguments to backend
743 Returns:
744 Contiguous sequence of degrees of freedom (with respect to input data)
745 and the global starting point on the process.
746 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
747 """
749 backend_args = {} if backend_args is None else backend_args
750 legacy = backend_args.get("legacy", False)
751 engine = backend_args.get("engine", "BP4")
752 io_name = backend_args.get("io_name", f"{name}_FunctionReader")
753 prefix = mesh_prefix(backend_args)
754 # Check that file contains the function to read
755 adios = adios2.ADIOS(comm)
756 check_file_exists(filename)
758 if not legacy:
759 with ADIOSFile(
760 adios=adios,
761 filename=filename,
762 mode=adios2.Mode.Read,
763 engine=engine,
764 io_name=io_name,
765 ) as adios_file:
766 variables = set(
767 sorted(
768 map(
769 lambda x: x[len(prefix) :].split("_time")[0],
770 filter(
771 lambda x: x.startswith(prefix) and x.endswith("_time"),
772 adios_file.io.AvailableVariables(),
773 ),
774 )
775 )
776 )
777 if name not in variables:
778 raise KeyError(f"{name} not found in {filename}. Did you mean one of {variables}?")
780 if legacy:
781 array_path = "Values"
782 else:
783 array_path = f"{prefix}{name}_values"
785 time_name = f"{prefix}{name}_time"
786 return read_array(adios, filename, array_path, engine, comm, time, time_name, legacy=legacy)
789def read_cell_perms(
790 comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None = None
791) -> npt.NDArray[np.uint32]:
792 """
793 Read cell permutation from file with given communicator,
794 Split in continuous chunks based on number of cells in the mesh (global).
796 Args:
797 adios: The ADIOS instance
798 comm: The MPI communicator used to read the data
799 filename: Path to input file
800 variable: Name of cell-permutation variable
801 num_cells_global: Number of cells in the mesh (global)
802 engine: Type of ADIOS engine to use for reading data
804 Returns:
805 Cell-permutations local to the process
807 .. note::
808 No MPI communication is done during this call
809 """
810 adios = adios2.ADIOS(comm)
811 check_file_exists(filename)
813 # Open ADIOS engine
814 backend_args = {} if backend_args is None else backend_args
815 engine = backend_args.get("engine", "BP4")
816 prefix = mesh_prefix(backend_args)
818 cell_perms, _ = read_array(
819 adios, filename, f"{prefix}CellPermutations", engine=engine, comm=comm, legacy=True
820 )
822 return cell_perms.astype(np.uint32)
825def read_hdf5_array(
826 comm: MPI.Comm,
827 filename: Path | str,
828 group: str,
829 backend_args: dict[str, Any] | None = None,
830) -> tuple[np.ndarray, int]:
831 adios = adios2.ADIOS(comm)
832 """Read an array from an HDF5 file.
834 Args:
835 comm: MPI communicator used in storage
836 filename: Path to file to read from
837 group: Group in HDF5 file where array is stored
838 backend_args: Arguments to backend
840 Returns:
841 Tuple containing:
842 - Numpy array read from file
843 - Global starting point on the process.
844 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
845 """
846 return read_array(adios, filename, group, engine="HDF5", comm=comm, legacy=True)
849def write_function(
850 filename: Path,
851 comm: MPI.Comm,
852 u: FunctionData,
853 time: float = 0.0,
854 mode: FileMode = FileMode.append,
855 backend_args: dict[str, Any] | None = None,
856):
857 """
858 Write a function to file using ADIOS2
860 Args:
861 comm: MPI communicator used in storage
862 u: Internal data structure for the function data to save to file
863 filename: Path to file to write to
864 engine: ADIOS2 engine to use
865 mode: ADIOS2 mode to use (write or append)
866 time: Time stamp associated with function
867 io_name: Internal name used for the ADIOS IO object
868 """
869 adios_mode = convert_file_mode(mode)
870 backend_args = get_default_backend_args(backend_args)
871 engine = backend_args["engine"]
872 io_name = backend_args.get("io_name", "{name}_writer")
873 prefix = mesh_prefix(backend_args)
875 adios = adios2.ADIOS(comm)
876 cell_permutations_exists = False
877 dofmap_exists = False
878 XDofmap_exists = False
879 if adios_mode == adios2.Mode.Append:
880 cell_permutations_exists = check_variable_exists(
881 adios, filename, f"{prefix}CellPermutations", engine=engine
882 )
883 dofmap_exists = check_variable_exists(
884 adios, filename, f"{prefix}{u.name}_dofmap", engine=engine
885 )
886 XDofmap_exists = check_variable_exists(
887 adios, filename, f"{prefix}{u.name}_XDofmap", engine=engine
888 )
890 with ADIOSFile(
891 adios=adios, filename=filename, mode=adios_mode, engine=engine, io_name=io_name, comm=comm
892 ) as adios_file:
893 adios_file.file.BeginStep()
895 if not cell_permutations_exists:
896 # Add mesh permutations. These belong to the mesh, not the function,
897 # so they are shared by every function on that mesh.
898 pvar = adios_file.io.DefineVariable(
899 f"{prefix}CellPermutations",
900 u.cell_permutations,
901 shape=[u.num_cells_global],
902 start=[u.local_cell_range[0]],
903 count=[u.local_cell_range[1] - u.local_cell_range[0]],
904 )
905 adios_file.file.Put(pvar, u.cell_permutations)
907 if not dofmap_exists:
908 # Add dofmap
909 dofmap_var = adios_file.io.DefineVariable(
910 f"{prefix}{u.name}_dofmap",
911 u.dofmap_array,
912 shape=[u.global_dofs_in_dofmap],
913 start=[u.dofmap_range[0]],
914 count=[u.dofmap_range[1] - u.dofmap_range[0]],
915 )
916 adios_file.file.Put(dofmap_var, u.dofmap_array)
918 if not XDofmap_exists:
919 # Add XDofmap
920 xdofmap_var = adios_file.io.DefineVariable(
921 f"{prefix}{u.name}_XDofmap",
922 u.dofmap_offsets,
923 shape=[u.num_cells_global + 1],
924 start=[u.local_cell_range[0]],
925 count=[u.local_cell_range[1] - u.local_cell_range[0] + 1],
926 )
927 adios_file.file.Put(xdofmap_var, u.dofmap_offsets)
929 val_var = adios_file.io.DefineVariable(
930 f"{prefix}{u.name}_values",
931 u.values,
932 shape=[u.num_dofs_global],
933 start=[u.dof_range[0]],
934 count=[u.dof_range[1] - u.dof_range[0]],
935 )
936 adios_file.file.Put(val_var, u.values)
938 # Add time step to file
939 t_arr = np.array([time], dtype=np.float64)
940 time_var = adios_file.io.DefineVariable(
941 f"{prefix}{u.name}_time",
942 t_arr,
943 shape=[1],
944 start=[0],
945 count=[1 if comm.rank == 0 else 0],
946 )
947 adios_file.file.Put(time_var, t_arr)
948 adios_file.file.PerformPuts()
949 adios_file.file.EndStep()
952def read_legacy_mesh(
953 filename: Path | str, comm: MPI.Comm, group: str
954) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
955 """Read in the mesh topology, geometry and (optionally) cell type from a
956 legacy DOLFIN HDF5-file.
958 Args:
959 filename: Path to file to read from
960 comm: MPI communicator used in storage
961 group: Group in HDF5 file where mesh is stored
963 Returns:
964 Tuple containing:
965 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
966 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
967 - Cell type as a string (e.g. "tetrahedron") or None if not found
968 """
969 # Create ADIOS2 reader
970 adios = adios2.ADIOS(comm)
971 with ADIOSFile(
972 adios=adios,
973 filename=filename,
974 mode=adios2.Mode.Read,
975 io_name="Mesh reader",
976 engine="HDF5",
977 ) as adios_file:
978 # Get mesh topology (distributed)
979 if f"{group}/topology" not in adios_file.io.AvailableVariables().keys():
980 raise KeyError(f"Mesh topology not found at '{group}/topology'")
981 topology = adios_file.io.InquireVariable(f"{group}/topology")
982 shape = topology.Shape()
983 local_range = compute_local_range(comm, shape[0])
984 topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
986 mesh_topology = np.empty(
987 (local_range[1] - local_range[0], shape[1]),
988 dtype=topology.Type().strip("_t"),
989 )
990 adios_file.file.Get(topology, mesh_topology, adios2.Mode.Sync)
992 # Get mesh cell type
993 cell_type = None
994 if f"{group}/topology/celltype" in adios_file.io.AvailableAttributes().keys():
995 celltype = adios_file.io.InquireAttribute(f"{group}/topology/celltype")
996 cell_type = celltype.DataString()[0]
998 # Get mesh geometry
1000 for geometry_key in [f"{group}/geometry", f"{group}/coordinates"]:
1001 if geometry_key in adios_file.io.AvailableVariables().keys():
1002 break
1003 else:
1004 raise KeyError(
1005 f"Mesh coordinates not found at '{group}/coordinates' or '{group}/geometry'"
1006 )
1008 geometry = adios_file.io.InquireVariable(geometry_key)
1009 shape = geometry.Shape()
1010 local_range = compute_local_range(comm, shape[0])
1011 geometry.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
1012 mesh_geometry = np.empty(
1013 (local_range[1] - local_range[0], shape[1]),
1014 dtype=adios_to_numpy_dtype[geometry.Type()],
1015 )
1016 adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Sync)
1018 return mesh_topology, mesh_geometry, cell_type
1021def snapshot_checkpoint(
1022 filename: Path | str,
1023 mode: FileMode,
1024 u: dolfinx.fem.Function,
1025 backend_args: dict[str, Any] | None,
1026):
1027 """Create a snapshot checkpoint of a dolfinx function.
1029 Args:
1030 filename: Path to file to read from
1031 mode: File-mode to store the function
1032 u: dolfinx function to create a snapshot checkpoint for
1033 backend_args: Arguments to backend
1034 """
1035 adios_mode = convert_file_mode(mode)
1036 adios = adios2.ADIOS(u.function_space.mesh.comm)
1037 backend_args = {} if backend_args is None else backend_args
1038 io_name = backend_args.get("io_name", "SnapshotCheckPoint")
1039 engine = backend_args.get("engine", "BP4")
1040 with ADIOSFile(
1041 adios=adios,
1042 filename=filename,
1043 mode=adios_mode,
1044 io_name=io_name,
1045 engine=engine,
1046 ) as adios_file:
1047 if adios_mode == adios2.Mode.Write:
1048 dofmap = u.function_space.dofmap
1049 num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs
1050 local_dofs = u.x.array[:num_dofs_local].copy()
1052 # Write to file
1053 adios_file.file.BeginStep()
1054 dofs = adios_file.io.DefineVariable("dofs", local_dofs, count=[num_dofs_local])
1055 adios_file.file.Put(dofs, local_dofs, adios2.Mode.Sync)
1056 adios_file.file.EndStep()
1057 elif adios_mode == adios2.Mode.Read:
1058 adios_file.file.BeginStep()
1059 in_variable = adios_file.io.InquireVariable("dofs")
1060 in_variable.SetBlockSelection(u.function_space.mesh.comm.rank)
1061 adios_file.file.Get(in_variable, u.x.array, adios2.Mode.Sync)
1062 adios_file.file.EndStep()
1063 u.x.scatter_forward()
1064 else:
1065 raise NotImplementedError(f"Mode {mode} is not implemented for snapshot checkpoint")
1068def read_function_names(
1069 filename: Path | str, comm: MPI.Comm, backend_args: dict[str, Any] | None
1070) -> list[str]:
1071 """Read all function names from a file.
1073 Args:
1074 filename: Path to file
1075 comm: MPI communicator to launch IO on.
1076 backend_args: Arguments to backend
1078 Returns:
1079 A list of function names.
1080 """
1081 raise NotImplementedError("Reading function names are not implemented with ADIOS2")
1084def read_point_data(
1085 filename: Path | str,
1086 name: str,
1087 comm: MPI.Comm,
1088 time: float | str | None,
1089 backend_args: dict[str, Any] | None,
1090) -> tuple[np.ndarray, int]:
1091 """Read data from the nodes of a mesh.
1093 Args:
1094 filename: Path to file
1095 name: Name of point data
1096 comm: Communicator to launch IO on.
1097 time: The time stamp
1098 backend_args: The backend arguments
1100 Returns:
1101 Data local to process (contiguous, no mpi comm) and local start range
1102 """
1103 raise NotImplementedError("The ADIOS2 backend cannot read point data.")
1106def read_cell_data(
1107 filename: Path | str,
1108 name: str,
1109 comm: MPI.Comm,
1110 time: str | float | None,
1111 backend_args: dict[str, Any] | None,
1112) -> tuple[npt.NDArray[np.int64], np.ndarray]:
1113 """Read data from the cells of a mesh.
1115 Args:
1116 filename: Path to file
1117 name: Name of point data
1118 comm: Communicator to launch IO on.
1119 time: The time stamp
1120 backend_args: The backend arguments
1121 Returns:
1122 A tuple (topology, dofs) where topology contains the
1123 vertex indices of the cells, dofs the degrees of
1124 freedom within that cell.
1125 """
1126 raise NotImplementedError("The ADIOS2 backend does not support reading cell data.")
1129def write_data(
1130 filename: Path | str,
1131 array_data: ArrayData,
1132 comm: MPI.Comm,
1133 time: str | float | None,
1134 mode: FileMode,
1135 backend_args: dict[str, Any] | None,
1136):
1137 """Write a 2D-array to file (distributed across proceses with MPI).
1140 Args:
1141 filename: Path to file
1142 array_data: Data to write to file
1143 comm: MPI communicator to open the file with.
1144 time: Time stamp
1145 mode: Append or write
1146 backend_args: The backend arguments
1147 """
1148 raise NotImplementedError("ADIOS2 has not implemented this yet")