Coverage for / dolfinx-env / lib / python3.12 / site-packages / io4dolfinx / backends / adios2 / backend.py: 93%
382 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-12 11:21 +0000
1import warnings
2from pathlib import Path
3from typing import Any
5from mpi4py import MPI
7import adios2
8import dolfinx
9import numpy as np
10import numpy.typing as npt
12from ...structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
13from ...utils import check_file_exists, compute_local_range
14from .. import FileMode, ReadMode
15from .helpers import (
16 ADIOSFile,
17 adios_to_numpy_dtype,
18 check_variable_exists,
19 read_adjacency_list,
20 read_array,
21 resolve_adios_scope,
22)
24adios2 = 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 = arguments or {}
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 return args
39def convert_file_mode(mode: FileMode) -> adios2.Mode: # type: ignore[override]
40 match mode:
41 case FileMode.append:
42 return adios2.Mode.Append
43 case FileMode.write:
44 return adios2.Mode.Write
45 case FileMode.read:
46 return adios2.Mode.Read
47 case _:
48 raise NotImplementedError(f"FileMode {mode} not implemented.")
51def write_attributes(
52 filename: Path | str,
53 comm: MPI.Intracomm,
54 name: str,
55 attributes: dict[str, np.ndarray],
56 backend_args: dict[str, Any] | None = None,
57):
58 """Write attributes to file using ADIOS2.
60 Args:
61 filename: Path to file to write to
62 comm: MPI communicator used in storage
63 name: Name of the attributes
64 attributes: Dictionary of attributes to write to file
65 engine: ADIOS2 engine to use
66 """
68 adios = adios2.ADIOS(comm)
69 backend_args = get_default_backend_args(backend_args)
71 with ADIOSFile(
72 adios=adios,
73 filename=filename,
74 mode=adios2.Mode.Append,
75 io_name="AttributeWriter",
76 engine=backend_args["engine"],
77 ) as adios_file:
78 adios_file.file.BeginStep()
80 for k, v in attributes.items():
81 adios_file.io.DefineAttribute(f"{name}_{k}", v)
83 adios_file.file.PerformPuts()
84 adios_file.file.EndStep()
87def read_attributes(
88 filename: Path | str,
89 comm: MPI.Intracomm,
90 name: str,
91 backend_args: dict[str, Any] | None = None,
92) -> dict[str, np.ndarray]:
93 """Read attributes from file using ADIOS2.
95 Args:
96 filename: Path to file to read from
97 comm: MPI communicator used in storage
98 name: Name of the attributes
99 engine: ADIOS2 engine to use
100 Returns:
101 The attributes
102 """
103 check_file_exists(filename)
104 adios = adios2.ADIOS(comm)
105 backend_args = get_default_backend_args(backend_args)
106 with ADIOSFile(
107 adios=adios,
108 filename=filename,
109 mode=adios2.Mode.Read,
110 engine=backend_args["engine"],
111 io_name="AttributesReader",
112 ) as adios_file:
113 adios_file.file.BeginStep()
114 attributes = {}
115 for k in adios_file.io.AvailableAttributes().keys():
116 if k.startswith(f"{name}_"):
117 a = adios_file.io.InquireAttribute(k)
118 attributes[k[len(name) + 1 :]] = a.Data()
119 adios_file.file.EndStep()
120 return attributes
123def read_timestamps(
124 filename: Path | str,
125 comm: MPI.Intracomm,
126 function_name: str,
127 backend_args: dict[str, Any] | None = None,
128) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
129 """Read time-stamps from a checkpoint file.
131 Args:
132 comm: MPI communicator
133 filename: Path to file
134 function_name: Name of the function to read time-stamps for
135 backend_args: Arguments for backend, for instance file type.
136 backend: What backend to use for writing.
137 Returns:
138 The time-stamps
139 """
140 check_file_exists(filename)
142 adios = adios2.ADIOS(comm)
143 backend_args = get_default_backend_args(backend_args)
144 with ADIOSFile(
145 adios=adios,
146 filename=filename,
147 mode=adios2.Mode.Read,
148 engine=backend_args["engine"],
149 io_name="TimestepReader",
150 ) as adios_file:
151 time_name = f"{function_name}_time"
152 time_stamps = []
153 for _ in range(adios_file.file.Steps()):
154 adios_file.file.BeginStep()
155 if time_name in adios_file.io.AvailableVariables().keys():
156 arr = adios_file.io.InquireVariable(time_name)
157 time_shape = arr.Shape()
158 arr.SetSelection([[0], [time_shape[0]]])
159 times = np.empty(
160 time_shape[0],
161 dtype=adios_to_numpy_dtype[arr.Type()],
162 )
163 adios_file.file.Get(arr, times, adios2.Mode.Sync)
164 time_stamps.append(times[0])
165 adios_file.file.EndStep()
167 return np.array(time_stamps)
170def write_mesh(
171 filename: Path | str,
172 comm: MPI.Intracomm,
173 mesh: MeshData,
174 backend_args: dict[str, Any] | None = None,
175 mode: FileMode = FileMode.write,
176 time: float = 0.0,
177):
178 """Write a mesh to file using ADIOS2.
180 Args:
181 comm: MPI communicator used in storage
182 mesh: Internal data structure for the mesh data to save to file
183 filename: Path to file to write to
184 backend_args: File mode and potentially the io-name.
185 mode: Mode to use (write or append)
186 time: Time stamp
187 """
188 backend_args = get_default_backend_args(backend_args)
189 if "io_name" not in backend_args.keys():
190 backend_args["io_name"] = "MeshWriter"
192 mode = convert_file_mode(mode)
193 gdim = mesh.local_geometry.shape[1]
194 adios = adios2.ADIOS(comm)
195 with ADIOSFile(
196 adios=adios,
197 filename=filename,
198 mode=mode,
199 comm=comm,
200 engine=backend_args["engine"],
201 io_name=backend_args["io_name"],
202 ) as adios_file:
203 adios_file.file.BeginStep()
204 # Write geometry
205 pointvar = adios_file.io.DefineVariable(
206 "Points",
207 mesh.local_geometry,
208 shape=[mesh.num_nodes_global, gdim],
209 start=[mesh.local_geometry_pos[0], 0],
210 count=[mesh.local_geometry_pos[1] - mesh.local_geometry_pos[0], gdim],
211 )
212 adios_file.file.Put(pointvar, mesh.local_geometry, adios2.Mode.Sync)
214 if mode == adios2.Mode.Write:
215 adios_file.io.DefineAttribute("CellType", mesh.cell_type)
216 adios_file.io.DefineAttribute("Degree", np.array([mesh.degree], dtype=np.int32))
217 adios_file.io.DefineAttribute(
218 "LagrangeVariant", np.array([mesh.lagrange_variant], dtype=np.int32)
219 )
220 # Write topology (on;y on first write as topology is constant)
221 num_dofs_per_cell = mesh.local_topology.shape[1]
222 dvar = adios_file.io.DefineVariable(
223 "Topology",
224 mesh.local_topology,
225 shape=[mesh.num_cells_global, num_dofs_per_cell],
226 start=[mesh.local_topology_pos[0], 0],
227 count=[
228 mesh.local_topology_pos[1] - mesh.local_topology_pos[0],
229 num_dofs_per_cell,
230 ],
231 )
232 adios_file.file.Put(dvar, mesh.local_topology)
234 # Add partitioning data
235 if mesh.store_partition:
236 assert mesh.partition_range is not None
237 par_data = adios_file.io.DefineVariable(
238 "PartitioningData",
239 mesh.ownership_array,
240 shape=[mesh.partition_global],
241 start=[mesh.partition_range[0]],
242 count=[
243 mesh.partition_range[1] - mesh.partition_range[0],
244 ],
245 )
246 adios_file.file.Put(par_data, mesh.ownership_array)
247 assert mesh.ownership_offset is not None
248 par_offset = adios_file.io.DefineVariable(
249 "PartitioningOffset",
250 mesh.ownership_offset,
251 shape=[mesh.num_cells_global + 1],
252 start=[mesh.local_topology_pos[0]],
253 count=[mesh.local_topology_pos[1] - mesh.local_topology_pos[0] + 1],
254 )
255 adios_file.file.Put(par_offset, mesh.ownership_offset)
256 assert mesh.partition_processes is not None
257 adios_file.io.DefineAttribute(
258 "PartitionProcesses", np.array([mesh.partition_processes], dtype=np.int32)
259 )
260 if mode == adios2.Mode.Append and mesh.store_partition:
261 warnings.warn("Partitioning data is not written in append mode")
263 # Add time step to file
264 t_arr = np.array([time], dtype=np.float64)
265 time_var = adios_file.io.DefineVariable(
266 "MeshTime",
267 t_arr,
268 shape=[1],
269 start=[0],
270 count=[1 if comm.rank == 0 else 0],
271 )
272 adios_file.file.Put(time_var, t_arr)
274 adios_file.file.PerformPuts()
275 adios_file.file.EndStep()
278def read_mesh_data(
279 filename: Path | str,
280 comm: MPI.Intracomm,
281 time: str | float | None = 0.0,
282 read_from_partition: bool = False,
283 backend_args: dict[str, Any] | None = None,
284) -> ReadMeshData:
285 """Read an ADIOS2 mesh data for use with DOLFINx.
287 Args:
288 filename: Path to input file
289 comm: The MPI communciator to distribute the mesh over
290 engine: ADIOS engine to use for reading (BP4, BP5 or HDF5)
291 time: Time stamp associated with mesh
292 legacy: If checkpoint was made prior to time-dependent mesh-writer set to True
293 read_from_partition: Read mesh with partition from file
294 Returns:
295 The mesh topology, geometry, UFL domain and partition function
296 """
298 adios = adios2.ADIOS(comm)
299 backend_args = get_default_backend_args(backend_args)
300 legacy = backend_args.get("legacy", False)
301 io_name = backend_args.get("io_name", "MeshReader")
302 engine = backend_args["engine"]
303 with ADIOSFile(
304 adios=adios,
305 filename=filename,
306 mode=adios2.Mode.Read,
307 engine=engine,
308 io_name=io_name,
309 ) as adios_file:
310 # Get time independent mesh variables (mesh topology and cell type info) first
311 adios_file.file.BeginStep()
312 # Get mesh topology (distributed)
313 if "Topology" not in adios_file.io.AvailableVariables().keys():
314 raise KeyError(f"Mesh topology not found at Topology in {filename}")
315 topology = adios_file.io.InquireVariable("Topology")
316 shape = topology.Shape()
317 local_range = compute_local_range(comm, shape[0])
318 topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
319 mesh_topology = np.empty((local_range[1] - local_range[0], shape[1]), dtype=np.int64)
320 adios_file.file.Get(topology, mesh_topology, adios2.Mode.Deferred)
322 # Check validity of partitioning information
323 if read_from_partition:
324 if "PartitionProcesses" not in adios_file.io.AvailableAttributes().keys():
325 raise KeyError(f"Partitioning information not found in {filename}")
326 par_num_procs = adios_file.io.InquireAttribute("PartitionProcesses")
327 num_procs = par_num_procs.Data()[0]
328 if num_procs != comm.size:
329 raise ValueError(f"Number of processes in file ({num_procs})!=({comm.size=})")
331 # Get mesh cell type
332 if "CellType" not in adios_file.io.AvailableAttributes().keys():
333 raise KeyError(f"Mesh cell type not found at CellType in {filename}")
334 celltype = adios_file.io.InquireAttribute("CellType")
335 cell_type = celltype.DataString()[0]
337 # Get basix info
338 if "LagrangeVariant" not in adios_file.io.AvailableAttributes().keys():
339 raise KeyError(f"Mesh LagrangeVariant not found in {filename}")
340 lvar = adios_file.io.InquireAttribute("LagrangeVariant").Data()[0]
341 if "Degree" not in adios_file.io.AvailableAttributes().keys():
342 raise KeyError(f"Mesh degree not found in {filename}")
343 degree = adios_file.io.InquireAttribute("Degree").Data()[0]
345 if not legacy:
346 time_name = "MeshTime"
347 for i in range(adios_file.file.Steps()):
348 if i > 0:
349 adios_file.file.BeginStep()
350 if time_name in adios_file.io.AvailableVariables().keys():
351 arr = adios_file.io.InquireVariable(time_name)
352 time_shape = arr.Shape()
353 arr.SetSelection([[0], [time_shape[0]]])
354 times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()])
355 adios_file.file.Get(arr, times, adios2.Mode.Sync)
356 if times[0] == time:
357 break
358 if i == adios_file.file.Steps() - 1:
359 raise KeyError(
360 f"No data associated with {time_name}={time} found in {filename}"
361 )
363 adios_file.file.EndStep()
365 if time_name not in adios_file.io.AvailableVariables().keys():
366 raise KeyError(f"No data associated with {time_name}={time} found in {filename}")
368 # Get mesh geometry
369 if "Points" not in adios_file.io.AvailableVariables().keys():
370 raise KeyError(f"Mesh coordinates not found at Points in {filename}")
371 geometry = adios_file.io.InquireVariable("Points")
372 x_shape = geometry.Shape()
373 geometry_range = compute_local_range(comm, x_shape[0])
374 geometry.SetSelection(
375 [
376 [geometry_range[0], 0],
377 [geometry_range[1] - geometry_range[0], x_shape[1]],
378 ]
379 )
380 mesh_geometry = np.empty(
381 (geometry_range[1] - geometry_range[0], x_shape[1]),
382 dtype=adios_to_numpy_dtype[geometry.Type()],
383 )
384 adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Deferred)
385 adios_file.file.PerformGets()
386 adios_file.file.EndStep()
388 if read_from_partition:
389 partition_graph = read_adjacency_list(
390 adios,
391 comm,
392 filename,
393 "PartitioningData",
394 "PartitioningOffset",
395 backend_args["engine"],
396 )
397 else:
398 partition_graph = None
400 return ReadMeshData(
401 cells=mesh_topology,
402 cell_type=cell_type,
403 x=mesh_geometry,
404 degree=degree,
405 lvar=lvar,
406 partition_graph=partition_graph,
407 )
410def write_meshtags(
411 filename: str | Path,
412 comm: MPI.Intracomm,
413 data: MeshTagsData,
414 backend_args: dict[str, Any] | None = None,
415):
416 """Write mesh tags to file.
418 Args:
419 filename: Path to file to write to
420 comm: MPI communicator used in storage
421 data: Internal data structure for the mesh tags to save to file
422 backend_args: Arguments to backend
423 """
424 backend_args = {} if backend_args is None else backend_args
425 io_name = backend_args.get("io_name", "MeshTagWriter")
426 engine = backend_args.get("engine", "BP4")
427 adios = adios2.ADIOS(comm)
428 with ADIOSFile(
429 adios=adios,
430 filename=filename,
431 mode=adios2.Mode.Append,
432 engine=engine,
433 io_name=io_name,
434 ) as adios_file:
435 adios_file.file.BeginStep()
437 # Write meshtag topology
438 topology_var = adios_file.io.DefineVariable(
439 data.name + "_topology",
440 data.indices,
441 shape=[data.num_entities_global, data.num_dofs_per_entity],
442 start=[data.local_start, 0],
443 count=[len(data.indices), data.num_dofs_per_entity],
444 )
445 adios_file.file.Put(topology_var, data.indices, adios2.Mode.Sync)
447 # Write meshtag values
448 vals = np.array(data.values)
449 values_var = adios_file.io.DefineVariable(
450 data.name + "_values",
451 vals,
452 shape=[data.num_entities_global],
453 start=[data.local_start],
454 count=[len(data.indices)],
455 )
456 adios_file.file.Put(values_var, vals, adios2.Mode.Sync)
458 # Write meshtag dim
459 adios_file.io.DefineAttribute(data.name + "_dim", np.array([data.dim], dtype=np.uint8))
460 adios_file.file.PerformPuts()
461 adios_file.file.EndStep()
464def read_meshtags_data(
465 filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None
466) -> MeshTagsData:
467 """Read mesh tags from file.
469 Args:
470 filename: Path to file to read from
471 comm: MPI communicator used in storage
472 name: Name of the mesh tags to read
473 backend_args: Arguments to backend
475 Returns:
476 Internal data structure for the mesh tags read from file
477 """
479 adios = adios2.ADIOS(comm)
480 backend_args = get_default_backend_args(backend_args)
481 io_name = backend_args.get("io_name", "MeshTagsReader")
482 engine = backend_args["engine"]
483 legacy = backend_args["legacy"]
484 with ADIOSFile(
485 adios=adios,
486 filename=filename,
487 mode=adios2.Mode.Read,
488 engine=engine,
489 io_name=io_name,
490 ) as adios_file:
491 if not legacy:
492 # Get mesh cell type
493 dim_attr_name = f"{name}_dim"
494 step = 0
495 for i in range(adios_file.file.Steps()):
496 adios_file.file.BeginStep()
497 if dim_attr_name in adios_file.io.AvailableAttributes().keys():
498 step = i
499 break
500 adios_file.file.EndStep()
501 if dim_attr_name not in adios_file.io.AvailableAttributes().keys():
502 raise KeyError(f"{dim_attr_name} not found in {filename}")
504 m_dim = adios_file.io.InquireAttribute(dim_attr_name)
505 dim = int(m_dim.Data()[0])
507 # Get mesh tags entites
508 topology_name = f"{name}_topology"
509 for i in range(step, adios_file.file.Steps()):
510 if i > step:
511 adios_file.file.BeginStep()
512 if topology_name in adios_file.io.AvailableVariables().keys():
513 break
514 adios_file.file.EndStep()
515 if topology_name not in adios_file.io.AvailableVariables().keys():
516 raise KeyError(f"{topology_name} not found in {filename}")
518 topology = adios_file.io.InquireVariable(topology_name)
519 top_shape = topology.Shape()
520 topology_range = compute_local_range(comm, top_shape[0])
522 topology.SetSelection(
523 [
524 [topology_range[0], 0],
525 [topology_range[1] - topology_range[0], top_shape[1]],
526 ]
527 )
528 mesh_entities = np.empty(
529 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
530 )
531 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
533 # Get mesh tags values
534 values_name = f"{name}_values"
535 if values_name not in adios_file.io.AvailableVariables().keys():
536 raise KeyError(f"{values_name} not found")
538 values = adios_file.io.InquireVariable(values_name)
539 val_shape = values.Shape()
540 assert val_shape[0] == top_shape[0]
541 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
542 tag_values = np.empty(
543 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
544 )
545 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
547 adios_file.file.PerformGets()
548 adios_file.file.EndStep()
549 else:
550 # Get mesh cell type
551 dim_attr_name = f"{name}_dim"
552 assert adios_file.file.Steps() == 0
553 if (ct_key := f"/{name}/topology/celltype") in adios_file.io.AvailableAttributes():
554 cell_type = adios_file.io.InquireAttribute(ct_key)
555 else:
556 raise ValueError(f"Celltype not found for meshtags {name} in {filename}.")
557 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type.DataString()[0]))
559 # Get mesh tags entites
560 if (top_key := f"/{name}/topology") in adios_file.io.AvailableVariables():
561 topology = adios_file.io.InquireVariable(top_key)
562 else:
563 raise ValueError(f"Topology not found for meshtags {name} in {filename}.")
565 top_shape = topology.Shape()
566 topology_range = compute_local_range(comm, top_shape[0])
568 topology.SetSelection(
569 [
570 [topology_range[0], 0],
571 [topology_range[1] - topology_range[0], top_shape[1]],
572 ]
573 )
574 mesh_entities = np.empty(
575 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
576 )
577 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
579 # Get mesh tags values
580 if (val_key := f"/{name}/values") in adios_file.io.AvailableVariables():
581 values = adios_file.io.InquireVariable(val_key)
582 else:
583 raise ValueError(f"Values not found for meshtags {name} in {filename}.")
585 val_shape = values.Shape()
586 assert val_shape[0] == top_shape[0]
588 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
589 tag_values = np.empty(
590 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
591 )
592 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
594 adios_file.file.PerformGets()
595 adios_file.file.EndStep()
597 return MeshTagsData(
598 name=name, values=tag_values.astype(np.int32), indices=mesh_entities, dim=dim
599 )
602def read_dofmap(
603 filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None
604) -> dolfinx.graph.AdjacencyList:
605 """Read the dofmap of a function with a given name.
607 Args:
608 filename: Path to file to read from
609 comm: MPI communicator used in storage
610 name: Name of the function to read the dofmap for
611 backend_args: Arguments to backend
613 Returns:
614 Dofmap as an AdjacencyList
615 """
616 backend_args = {} if backend_args is None else backend_args
618 # Handles legacy io4dolfinx files, modern files, and custom location of dofmap.
619 legacy = backend_args.get("legacy", False)
620 xdofmap_path: str | None
621 dofmap_path: str | None
622 if (dofmap_path := backend_args.get("dofmap", None)) is None:
623 if legacy:
624 dofmap_path = "Dofmap"
625 else:
626 dofmap_path = f"{name}_dofmap"
628 if (xdofmap_path := backend_args.get("offsets", None)) is None:
629 if legacy:
630 xdofmap_path = "XDofmap"
631 else:
632 xdofmap_path = f"{name}_XDofmap"
634 engine = backend_args.get("engine", "BP4")
636 adios = adios2.ADIOS(comm)
637 check_file_exists(filename)
638 assert isinstance(xdofmap_path, str)
639 return read_adjacency_list(adios, comm, filename, dofmap_path, xdofmap_path, engine=engine)
642def read_dofs(
643 filename: str | Path,
644 comm: MPI.Intracomm,
645 name: str,
646 time: float,
647 backend_args: dict[str, Any] | None = None,
648) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
649 """Read the dofs (values) of a function with a given name from a given timestep.
651 Args:
652 filename: Path to file to read from
653 comm: MPI communicator used in storage
654 name: Name of the function to read the dofs for
655 time: Time stamp associated with the function to read
656 backend_args: Arguments to backend
658 Returns:
659 Contiguous sequence of degrees of freedom (with respect to input data)
660 and the global starting point on the process.
661 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
662 """
664 backend_args = {} if backend_args is None else backend_args
665 legacy = backend_args.get("legacy", False)
666 engine = backend_args.get("engine", "BP4")
667 io_name = backend_args.get("io_name", f"{name}_FunctionReader")
668 # Check that file contains the function to read
669 adios = adios2.ADIOS(comm)
670 check_file_exists(filename)
672 if not legacy:
673 with ADIOSFile(
674 adios=adios,
675 filename=filename,
676 mode=adios2.Mode.Read,
677 engine=engine,
678 io_name=io_name,
679 ) as adios_file:
680 variables = set(
681 sorted(
682 map(
683 lambda x: x.split("_time")[0],
684 filter(lambda x: x.endswith("_time"), adios_file.io.AvailableVariables()),
685 )
686 )
687 )
688 if name not in variables:
689 raise KeyError(f"{name} not found in {filename}. Did you mean one of {variables}?")
691 if legacy:
692 array_path = "Values"
693 else:
694 array_path = f"{name}_values"
696 time_name = f"{name}_time"
697 return read_array(adios, filename, array_path, engine, comm, time, time_name, legacy=legacy)
700def read_cell_perms(
701 comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None = None
702) -> npt.NDArray[np.uint32]:
703 """
704 Read cell permutation from file with given communicator,
705 Split in continuous chunks based on number of cells in the mesh (global).
707 Args:
708 adios: The ADIOS instance
709 comm: The MPI communicator used to read the data
710 filename: Path to input file
711 variable: Name of cell-permutation variable
712 num_cells_global: Number of cells in the mesh (global)
713 engine: Type of ADIOS engine to use for reading data
715 Returns:
716 Cell-permutations local to the process
718 .. note::
719 No MPI communication is done during this call
720 """
721 adios = adios2.ADIOS(comm)
722 check_file_exists(filename)
724 # Open ADIOS engine
725 backend_args = {} if backend_args is None else backend_args
726 engine = backend_args.get("engine", "BP4")
728 cell_perms, _ = read_array(
729 adios, filename, "CellPermutations", engine=engine, comm=comm, legacy=True
730 )
732 return cell_perms.astype(np.uint32)
735def read_hdf5_array(
736 comm: MPI.Intracomm,
737 filename: Path | str,
738 group: str,
739 backend_args: dict[str, Any] | None = None,
740) -> tuple[np.ndarray, int]:
741 adios = adios2.ADIOS(comm)
742 """Read an array from an HDF5 file.
744 Args:
745 comm: MPI communicator used in storage
746 filename: Path to file to read from
747 group: Group in HDF5 file where array is stored
748 backend_args: Arguments to backend
750 Returns:
751 Tuple containing:
752 - Numpy array read from file
753 - Global starting point on the process.
754 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
755 """
756 return read_array(adios, filename, group, engine="HDF5", comm=comm, legacy=True)
759def write_function(
760 filename: Path,
761 comm: MPI.Intracomm,
762 u: FunctionData,
763 time: float = 0.0,
764 mode: FileMode = FileMode.append,
765 backend_args: dict[str, Any] | None = None,
766):
767 """
768 Write a function to file using ADIOS2
770 Args:
771 comm: MPI communicator used in storage
772 u: Internal data structure for the function data to save to file
773 filename: Path to file to write to
774 engine: ADIOS2 engine to use
775 mode: ADIOS2 mode to use (write or append)
776 time: Time stamp associated with function
777 io_name: Internal name used for the ADIOS IO object
778 """
779 adios_mode = convert_file_mode(mode)
780 backend_args = get_default_backend_args(backend_args)
781 engine = backend_args["engine"]
782 io_name = backend_args.get("io_name", "{name}_writer")
784 adios = adios2.ADIOS(comm)
785 cell_permutations_exists = False
786 dofmap_exists = False
787 XDofmap_exists = False
788 if mode == adios2.Mode.Append:
789 cell_permutations_exists = check_variable_exists(
790 adios, filename, "CellPermutations", engine=engine
791 )
792 dofmap_exists = check_variable_exists(adios, filename, f"{u.name}_dofmap", engine=engine)
793 XDofmap_exists = check_variable_exists(adios, filename, f"{u.name}_XDofmap", engine=engine)
795 with ADIOSFile(
796 adios=adios, filename=filename, mode=adios_mode, engine=engine, io_name=io_name, comm=comm
797 ) as adios_file:
798 adios_file.file.BeginStep()
800 if not cell_permutations_exists:
801 # Add mesh permutations
802 pvar = adios_file.io.DefineVariable(
803 "CellPermutations",
804 u.cell_permutations,
805 shape=[u.num_cells_global],
806 start=[u.local_cell_range[0]],
807 count=[u.local_cell_range[1] - u.local_cell_range[0]],
808 )
809 adios_file.file.Put(pvar, u.cell_permutations)
811 if not dofmap_exists:
812 # Add dofmap
813 dofmap_var = adios_file.io.DefineVariable(
814 f"{u.name}_dofmap",
815 u.dofmap_array,
816 shape=[u.global_dofs_in_dofmap],
817 start=[u.dofmap_range[0]],
818 count=[u.dofmap_range[1] - u.dofmap_range[0]],
819 )
820 adios_file.file.Put(dofmap_var, u.dofmap_array)
822 if not XDofmap_exists:
823 # Add XDofmap
824 xdofmap_var = adios_file.io.DefineVariable(
825 f"{u.name}_XDofmap",
826 u.dofmap_offsets,
827 shape=[u.num_cells_global + 1],
828 start=[u.local_cell_range[0]],
829 count=[u.local_cell_range[1] - u.local_cell_range[0] + 1],
830 )
831 adios_file.file.Put(xdofmap_var, u.dofmap_offsets)
833 val_var = adios_file.io.DefineVariable(
834 f"{u.name}_values",
835 u.values,
836 shape=[u.num_dofs_global],
837 start=[u.dof_range[0]],
838 count=[u.dof_range[1] - u.dof_range[0]],
839 )
840 adios_file.file.Put(val_var, u.values)
842 # Add time step to file
843 t_arr = np.array([time], dtype=np.float64)
844 time_var = adios_file.io.DefineVariable(
845 f"{u.name}_time",
846 t_arr,
847 shape=[1],
848 start=[0],
849 count=[1 if comm.rank == 0 else 0],
850 )
851 adios_file.file.Put(time_var, t_arr)
852 adios_file.file.PerformPuts()
853 adios_file.file.EndStep()
856def read_legacy_mesh(
857 filename: Path | str, comm: MPI.Intracomm, group: str
858) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
859 """Read in the mesh topology, geometry and (optionally) cell type from a
860 legacy DOLFIN HDF5-file.
862 Args:
863 filename: Path to file to read from
864 comm: MPI communicator used in storage
865 group: Group in HDF5 file where mesh is stored
867 Returns:
868 Tuple containing:
869 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
870 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
871 - Cell type as a string (e.g. "tetrahedron") or None if not found
872 """
873 # Create ADIOS2 reader
874 adios = adios2.ADIOS(comm)
875 with ADIOSFile(
876 adios=adios,
877 filename=filename,
878 mode=adios2.Mode.Read,
879 io_name="Mesh reader",
880 engine="HDF5",
881 ) as adios_file:
882 # Get mesh topology (distributed)
883 if f"{group}/topology" not in adios_file.io.AvailableVariables().keys():
884 raise KeyError(f"Mesh topology not found at '{group}/topology'")
885 topology = adios_file.io.InquireVariable(f"{group}/topology")
886 shape = topology.Shape()
887 local_range = compute_local_range(comm, shape[0])
888 topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
890 mesh_topology = np.empty(
891 (local_range[1] - local_range[0], shape[1]),
892 dtype=topology.Type().strip("_t"),
893 )
894 adios_file.file.Get(topology, mesh_topology, adios2.Mode.Sync)
896 # Get mesh cell type
897 cell_type = None
898 if f"{group}/topology/celltype" in adios_file.io.AvailableAttributes().keys():
899 celltype = adios_file.io.InquireAttribute(f"{group}/topology/celltype")
900 cell_type = celltype.DataString()[0]
902 # Get mesh geometry
904 for geometry_key in [f"{group}/geometry", f"{group}/coordinates"]:
905 if geometry_key in adios_file.io.AvailableVariables().keys():
906 break
907 else:
908 raise KeyError(
909 f"Mesh coordinates not found at '{group}/coordinates' or '{group}/geometry'"
910 )
912 geometry = adios_file.io.InquireVariable(geometry_key)
913 shape = geometry.Shape()
914 local_range = compute_local_range(comm, shape[0])
915 geometry.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
916 mesh_geometry = np.empty(
917 (local_range[1] - local_range[0], shape[1]),
918 dtype=adios_to_numpy_dtype[geometry.Type()],
919 )
920 adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Sync)
922 return mesh_topology, mesh_geometry, cell_type
925def snapshot_checkpoint(
926 filename: Path | str,
927 mode: FileMode,
928 u: dolfinx.fem.Function,
929 backend_args: dict[str, Any] | None,
930):
931 """Create a snapshot checkpoint of a dolfinx function.
933 Args:
934 filename: Path to file to read from
935 mode: File-mode to store the function
936 u: dolfinx function to create a snapshot checkpoint for
937 backend_args: Arguments to backend
938 """
939 adios_mode = convert_file_mode(mode)
940 adios = adios2.ADIOS(u.function_space.mesh.comm)
941 backend_args = {} if backend_args is None else backend_args
942 io_name = backend_args.get("io_name", "SnapshotCheckPoint")
943 engine = backend_args.get("engine", "BP4")
944 with ADIOSFile(
945 adios=adios,
946 filename=filename,
947 mode=adios_mode,
948 io_name=io_name,
949 engine=engine,
950 ) as adios_file:
951 if adios_mode == adios2.Mode.Write:
952 dofmap = u.function_space.dofmap
953 num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs
954 local_dofs = u.x.array[:num_dofs_local].copy()
956 # Write to file
957 adios_file.file.BeginStep()
958 dofs = adios_file.io.DefineVariable("dofs", local_dofs, count=[num_dofs_local])
959 adios_file.file.Put(dofs, local_dofs, adios2.Mode.Sync)
960 adios_file.file.EndStep()
961 elif adios_mode == adios2.Mode.Read:
962 adios_file.file.BeginStep()
963 in_variable = adios_file.io.InquireVariable("dofs")
964 in_variable.SetBlockSelection(u.function_space.mesh.comm.rank)
965 adios_file.file.Get(in_variable, u.x.array, adios2.Mode.Sync)
966 adios_file.file.EndStep()
967 u.x.scatter_forward()
968 else:
969 raise NotImplementedError(f"Mode {mode} is not implemented for snapshot checkpoint")
972def read_function_names(
973 filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None
974) -> list[str]:
975 """Read all function names from a file.
977 Args:
978 filename: Path to file
979 comm: MPI communicator to launch IO on.
980 backend_args: Arguments to backend
982 Returns:
983 A list of function names.
984 """
985 raise NotImplementedError("Reading function names are not implemented with ADIOS2")
988def read_point_data(
989 filename: Path | str,
990 name: str,
991 comm: MPI.Intracomm,
992 time: float | str | None,
993 backend_args: dict[str, Any] | None,
994) -> tuple[np.ndarray, int]:
995 """Read data from the nodes of a mesh.
997 Args:
998 filename: Path to file
999 name: Name of point data
1000 comm: Communicator to launch IO on.
1001 time: The time stamp
1002 backend_args: The backend arguments
1004 Returns:
1005 Data local to process (contiguous, no mpi comm) and local start range
1006 """
1007 raise NotImplementedError("The ADIOS2 backend cannot read point data.")
1010def read_cell_data(
1011 filename: Path | str,
1012 name: str,
1013 comm: MPI.Intracomm,
1014 time: str | float | None,
1015 backend_args: dict[str, Any] | None,
1016) -> tuple[npt.NDArray[np.int64], np.ndarray]:
1017 """Read data from the cells of a mesh.
1019 Args:
1020 filename: Path to file
1021 name: Name of point data
1022 comm: Communicator to launch IO on.
1023 time: The time stamp
1024 backend_args: The backend arguments
1025 Returns:
1026 A tuple (topology, dofs) where topology contains the
1027 vertex indices of the cells, dofs the degrees of
1028 freedom within that cell.
1029 """
1030 raise NotImplementedError("The ADIOS2 backend does not support reading cell data.")
1033def write_data(
1034 filename: Path | str,
1035 array_data: ArrayData,
1036 comm: MPI.Intracomm,
1037 time: str | float | None,
1038 mode: FileMode,
1039 backend_args: dict[str, Any] | None,
1040):
1041 """Write a 2D-array to file (distributed across proceses with MPI).
1044 Args:
1045 filename: Path to file
1046 array_data: Data to write to file
1047 comm: MPI communicator to open the file with.
1048 time: Time stamp
1049 mode: Append or write
1050 backend_args: The backend arguments
1051 """
1052 raise NotImplementedError("ADIOS2 has not implemented this yet")