Coverage for / dolfinx-env / lib / python3.12 / site-packages / io4dolfinx / backends / adios2 / backend.py: 93%
381 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-26 18:16 +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 values_var = adios_file.io.DefineVariable(
449 data.name + "_values",
450 data.values,
451 shape=[data.num_entities_global],
452 start=[data.local_start],
453 count=[len(data.indices)],
454 )
455 adios_file.file.Put(values_var, data.values, adios2.Mode.Sync)
457 # Write meshtag dim
458 adios_file.io.DefineAttribute(data.name + "_dim", np.array([data.dim], dtype=np.uint8))
459 adios_file.file.PerformPuts()
460 adios_file.file.EndStep()
463def read_meshtags_data(
464 filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None
465) -> MeshTagsData:
466 """Read mesh tags from file.
468 Args:
469 filename: Path to file to read from
470 comm: MPI communicator used in storage
471 name: Name of the mesh tags to read
472 backend_args: Arguments to backend
474 Returns:
475 Internal data structure for the mesh tags read from file
476 """
478 adios = adios2.ADIOS(comm)
479 backend_args = get_default_backend_args(backend_args)
480 io_name = backend_args.get("io_name", "MeshTagsReader")
481 engine = backend_args["engine"]
482 legacy = backend_args["legacy"]
483 with ADIOSFile(
484 adios=adios,
485 filename=filename,
486 mode=adios2.Mode.Read,
487 engine=engine,
488 io_name=io_name,
489 ) as adios_file:
490 if not legacy:
491 # Get mesh cell type
492 dim_attr_name = f"{name}_dim"
493 step = 0
494 for i in range(adios_file.file.Steps()):
495 adios_file.file.BeginStep()
496 if dim_attr_name in adios_file.io.AvailableAttributes().keys():
497 step = i
498 break
499 adios_file.file.EndStep()
500 if dim_attr_name not in adios_file.io.AvailableAttributes().keys():
501 raise KeyError(f"{dim_attr_name} not found in {filename}")
503 m_dim = adios_file.io.InquireAttribute(dim_attr_name)
504 dim = int(m_dim.Data()[0])
506 # Get mesh tags entites
507 topology_name = f"{name}_topology"
508 for i in range(step, adios_file.file.Steps()):
509 if i > step:
510 adios_file.file.BeginStep()
511 if topology_name in adios_file.io.AvailableVariables().keys():
512 break
513 adios_file.file.EndStep()
514 if topology_name not in adios_file.io.AvailableVariables().keys():
515 raise KeyError(f"{topology_name} not found in {filename}")
517 topology = adios_file.io.InquireVariable(topology_name)
518 top_shape = topology.Shape()
519 topology_range = compute_local_range(comm, top_shape[0])
521 topology.SetSelection(
522 [
523 [topology_range[0], 0],
524 [topology_range[1] - topology_range[0], top_shape[1]],
525 ]
526 )
527 mesh_entities = np.empty(
528 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
529 )
530 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
532 # Get mesh tags values
533 values_name = f"{name}_values"
534 if values_name not in adios_file.io.AvailableVariables().keys():
535 raise KeyError(f"{values_name} not found")
537 values = adios_file.io.InquireVariable(values_name)
538 val_shape = values.Shape()
539 assert val_shape[0] == top_shape[0]
540 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
541 tag_values = np.empty(
542 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
543 )
544 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
546 adios_file.file.PerformGets()
547 adios_file.file.EndStep()
548 else:
549 # Get mesh cell type
550 dim_attr_name = f"{name}_dim"
551 assert adios_file.file.Steps() == 0
552 if (ct_key := f"/{name}/topology/celltype") in adios_file.io.AvailableAttributes():
553 cell_type = adios_file.io.InquireAttribute(ct_key)
554 else:
555 raise ValueError(f"Celltype not found for meshtags {name} in {filename}.")
556 dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type.DataString()[0]))
558 # Get mesh tags entites
559 if (top_key := f"/{name}/topology") in adios_file.io.AvailableVariables():
560 topology = adios_file.io.InquireVariable(top_key)
561 else:
562 raise ValueError(f"Topology not found for meshtags {name} in {filename}.")
564 top_shape = topology.Shape()
565 topology_range = compute_local_range(comm, top_shape[0])
567 topology.SetSelection(
568 [
569 [topology_range[0], 0],
570 [topology_range[1] - topology_range[0], top_shape[1]],
571 ]
572 )
573 mesh_entities = np.empty(
574 (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
575 )
576 adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
578 # Get mesh tags values
579 if (val_key := f"/{name}/values") in adios_file.io.AvailableVariables():
580 values = adios_file.io.InquireVariable(val_key)
581 else:
582 raise ValueError(f"Values not found for meshtags {name} in {filename}.")
584 val_shape = values.Shape()
585 assert val_shape[0] == top_shape[0]
587 values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
588 tag_values = np.empty(
589 (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
590 )
591 adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
593 adios_file.file.PerformGets()
594 adios_file.file.EndStep()
596 return MeshTagsData(
597 name=name, values=tag_values.astype(np.int32), indices=mesh_entities, dim=dim
598 )
601def read_dofmap(
602 filename: str | Path, comm: MPI.Intracomm, name: str, backend_args: dict[str, Any] | None = None
603) -> dolfinx.graph.AdjacencyList:
604 """Read the dofmap of a function with a given name.
606 Args:
607 filename: Path to file to read from
608 comm: MPI communicator used in storage
609 name: Name of the function to read the dofmap for
610 backend_args: Arguments to backend
612 Returns:
613 Dofmap as an AdjacencyList
614 """
615 backend_args = {} if backend_args is None else backend_args
617 # Handles legacy io4dolfinx files, modern files, and custom location of dofmap.
618 legacy = backend_args.get("legacy", False)
619 xdofmap_path: str | None
620 dofmap_path: str | None
621 if (dofmap_path := backend_args.get("dofmap", None)) is None:
622 if legacy:
623 dofmap_path = "Dofmap"
624 else:
625 dofmap_path = f"{name}_dofmap"
627 if (xdofmap_path := backend_args.get("offsets", None)) is None:
628 if legacy:
629 xdofmap_path = "XDofmap"
630 else:
631 xdofmap_path = f"{name}_XDofmap"
633 engine = backend_args.get("engine", "BP4")
635 adios = adios2.ADIOS(comm)
636 check_file_exists(filename)
637 assert isinstance(xdofmap_path, str)
638 return read_adjacency_list(adios, comm, filename, dofmap_path, xdofmap_path, engine=engine)
641def read_dofs(
642 filename: str | Path,
643 comm: MPI.Intracomm,
644 name: str,
645 time: float,
646 backend_args: dict[str, Any] | None = None,
647) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
648 """Read the dofs (values) of a function with a given name from a given timestep.
650 Args:
651 filename: Path to file to read from
652 comm: MPI communicator used in storage
653 name: Name of the function to read the dofs for
654 time: Time stamp associated with the function to read
655 backend_args: Arguments to backend
657 Returns:
658 Contiguous sequence of degrees of freedom (with respect to input data)
659 and the global starting point on the process.
660 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
661 """
663 backend_args = {} if backend_args is None else backend_args
664 legacy = backend_args.get("legacy", False)
665 engine = backend_args.get("engine", "BP4")
666 io_name = backend_args.get("io_name", f"{name}_FunctionReader")
667 # Check that file contains the function to read
668 adios = adios2.ADIOS(comm)
669 check_file_exists(filename)
671 if not legacy:
672 with ADIOSFile(
673 adios=adios,
674 filename=filename,
675 mode=adios2.Mode.Read,
676 engine=engine,
677 io_name=io_name,
678 ) as adios_file:
679 variables = set(
680 sorted(
681 map(
682 lambda x: x.split("_time")[0],
683 filter(lambda x: x.endswith("_time"), adios_file.io.AvailableVariables()),
684 )
685 )
686 )
687 if name not in variables:
688 raise KeyError(f"{name} not found in {filename}. Did you mean one of {variables}?")
690 if legacy:
691 array_path = "Values"
692 else:
693 array_path = f"{name}_values"
695 time_name = f"{name}_time"
696 return read_array(adios, filename, array_path, engine, comm, time, time_name, legacy=legacy)
699def read_cell_perms(
700 comm: MPI.Intracomm, filename: Path | str, backend_args: dict[str, Any] | None = None
701) -> npt.NDArray[np.uint32]:
702 """
703 Read cell permutation from file with given communicator,
704 Split in continuous chunks based on number of cells in the mesh (global).
706 Args:
707 adios: The ADIOS instance
708 comm: The MPI communicator used to read the data
709 filename: Path to input file
710 variable: Name of cell-permutation variable
711 num_cells_global: Number of cells in the mesh (global)
712 engine: Type of ADIOS engine to use for reading data
714 Returns:
715 Cell-permutations local to the process
717 .. note::
718 No MPI communication is done during this call
719 """
720 adios = adios2.ADIOS(comm)
721 check_file_exists(filename)
723 # Open ADIOS engine
724 backend_args = {} if backend_args is None else backend_args
725 engine = backend_args.get("engine", "BP4")
727 cell_perms, _ = read_array(
728 adios, filename, "CellPermutations", engine=engine, comm=comm, legacy=True
729 )
731 return cell_perms.astype(np.uint32)
734def read_hdf5_array(
735 comm: MPI.Intracomm,
736 filename: Path | str,
737 group: str,
738 backend_args: dict[str, Any] | None = None,
739) -> tuple[np.ndarray, int]:
740 adios = adios2.ADIOS(comm)
741 """Read an array from an HDF5 file.
743 Args:
744 comm: MPI communicator used in storage
745 filename: Path to file to read from
746 group: Group in HDF5 file where array is stored
747 backend_args: Arguments to backend
749 Returns:
750 Tuple containing:
751 - Numpy array read from file
752 - Global starting point on the process.
753 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
754 """
755 return read_array(adios, filename, group, engine="HDF5", comm=comm, legacy=True)
758def write_function(
759 filename: Path,
760 comm: MPI.Intracomm,
761 u: FunctionData,
762 time: float = 0.0,
763 mode: FileMode = FileMode.append,
764 backend_args: dict[str, Any] | None = None,
765):
766 """
767 Write a function to file using ADIOS2
769 Args:
770 comm: MPI communicator used in storage
771 u: Internal data structure for the function data to save to file
772 filename: Path to file to write to
773 engine: ADIOS2 engine to use
774 mode: ADIOS2 mode to use (write or append)
775 time: Time stamp associated with function
776 io_name: Internal name used for the ADIOS IO object
777 """
778 adios_mode = convert_file_mode(mode)
779 backend_args = get_default_backend_args(backend_args)
780 engine = backend_args["engine"]
781 io_name = backend_args.get("io_name", "{name}_writer")
783 adios = adios2.ADIOS(comm)
784 cell_permutations_exists = False
785 dofmap_exists = False
786 XDofmap_exists = False
787 if mode == adios2.Mode.Append:
788 cell_permutations_exists = check_variable_exists(
789 adios, filename, "CellPermutations", engine=engine
790 )
791 dofmap_exists = check_variable_exists(adios, filename, f"{u.name}_dofmap", engine=engine)
792 XDofmap_exists = check_variable_exists(adios, filename, f"{u.name}_XDofmap", engine=engine)
794 with ADIOSFile(
795 adios=adios, filename=filename, mode=adios_mode, engine=engine, io_name=io_name, comm=comm
796 ) as adios_file:
797 adios_file.file.BeginStep()
799 if not cell_permutations_exists:
800 # Add mesh permutations
801 pvar = adios_file.io.DefineVariable(
802 "CellPermutations",
803 u.cell_permutations,
804 shape=[u.num_cells_global],
805 start=[u.local_cell_range[0]],
806 count=[u.local_cell_range[1] - u.local_cell_range[0]],
807 )
808 adios_file.file.Put(pvar, u.cell_permutations)
810 if not dofmap_exists:
811 # Add dofmap
812 dofmap_var = adios_file.io.DefineVariable(
813 f"{u.name}_dofmap",
814 u.dofmap_array,
815 shape=[u.global_dofs_in_dofmap],
816 start=[u.dofmap_range[0]],
817 count=[u.dofmap_range[1] - u.dofmap_range[0]],
818 )
819 adios_file.file.Put(dofmap_var, u.dofmap_array)
821 if not XDofmap_exists:
822 # Add XDofmap
823 xdofmap_var = adios_file.io.DefineVariable(
824 f"{u.name}_XDofmap",
825 u.dofmap_offsets,
826 shape=[u.num_cells_global + 1],
827 start=[u.local_cell_range[0]],
828 count=[u.local_cell_range[1] - u.local_cell_range[0] + 1],
829 )
830 adios_file.file.Put(xdofmap_var, u.dofmap_offsets)
832 val_var = adios_file.io.DefineVariable(
833 f"{u.name}_values",
834 u.values,
835 shape=[u.num_dofs_global],
836 start=[u.dof_range[0]],
837 count=[u.dof_range[1] - u.dof_range[0]],
838 )
839 adios_file.file.Put(val_var, u.values)
841 # Add time step to file
842 t_arr = np.array([time], dtype=np.float64)
843 time_var = adios_file.io.DefineVariable(
844 f"{u.name}_time",
845 t_arr,
846 shape=[1],
847 start=[0],
848 count=[1 if comm.rank == 0 else 0],
849 )
850 adios_file.file.Put(time_var, t_arr)
851 adios_file.file.PerformPuts()
852 adios_file.file.EndStep()
855def read_legacy_mesh(
856 filename: Path | str, comm: MPI.Intracomm, group: str
857) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
858 """Read in the mesh topology, geometry and (optionally) cell type from a
859 legacy DOLFIN HDF5-file.
861 Args:
862 filename: Path to file to read from
863 comm: MPI communicator used in storage
864 group: Group in HDF5 file where mesh is stored
866 Returns:
867 Tuple containing:
868 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
869 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
870 - Cell type as a string (e.g. "tetrahedron") or None if not found
871 """
872 # Create ADIOS2 reader
873 adios = adios2.ADIOS(comm)
874 with ADIOSFile(
875 adios=adios,
876 filename=filename,
877 mode=adios2.Mode.Read,
878 io_name="Mesh reader",
879 engine="HDF5",
880 ) as adios_file:
881 # Get mesh topology (distributed)
882 if f"{group}/topology" not in adios_file.io.AvailableVariables().keys():
883 raise KeyError(f"Mesh topology not found at '{group}/topology'")
884 topology = adios_file.io.InquireVariable(f"{group}/topology")
885 shape = topology.Shape()
886 local_range = compute_local_range(comm, shape[0])
887 topology.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
889 mesh_topology = np.empty(
890 (local_range[1] - local_range[0], shape[1]),
891 dtype=topology.Type().strip("_t"),
892 )
893 adios_file.file.Get(topology, mesh_topology, adios2.Mode.Sync)
895 # Get mesh cell type
896 cell_type = None
897 if f"{group}/topology/celltype" in adios_file.io.AvailableAttributes().keys():
898 celltype = adios_file.io.InquireAttribute(f"{group}/topology/celltype")
899 cell_type = celltype.DataString()[0]
901 # Get mesh geometry
903 for geometry_key in [f"{group}/geometry", f"{group}/coordinates"]:
904 if geometry_key in adios_file.io.AvailableVariables().keys():
905 break
906 else:
907 raise KeyError(
908 f"Mesh coordinates not found at '{group}/coordinates' or '{group}/geometry'"
909 )
911 geometry = adios_file.io.InquireVariable(geometry_key)
912 shape = geometry.Shape()
913 local_range = compute_local_range(comm, shape[0])
914 geometry.SetSelection([[local_range[0], 0], [local_range[1] - local_range[0], shape[1]]])
915 mesh_geometry = np.empty(
916 (local_range[1] - local_range[0], shape[1]),
917 dtype=adios_to_numpy_dtype[geometry.Type()],
918 )
919 adios_file.file.Get(geometry, mesh_geometry, adios2.Mode.Sync)
921 return mesh_topology, mesh_geometry, cell_type
924def snapshot_checkpoint(
925 filename: Path | str,
926 mode: FileMode,
927 u: dolfinx.fem.Function,
928 backend_args: dict[str, Any] | None,
929):
930 """Create a snapshot checkpoint of a dolfinx function.
932 Args:
933 filename: Path to file to read from
934 mode: File-mode to store the function
935 u: dolfinx function to create a snapshot checkpoint for
936 backend_args: Arguments to backend
937 """
938 adios_mode = convert_file_mode(mode)
939 adios = adios2.ADIOS(u.function_space.mesh.comm)
940 backend_args = {} if backend_args is None else backend_args
941 io_name = backend_args.get("io_name", "SnapshotCheckPoint")
942 engine = backend_args.get("engine", "BP4")
943 with ADIOSFile(
944 adios=adios,
945 filename=filename,
946 mode=adios_mode,
947 io_name=io_name,
948 engine=engine,
949 ) as adios_file:
950 if adios_mode == adios2.Mode.Write:
951 dofmap = u.function_space.dofmap
952 num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs
953 local_dofs = u.x.array[:num_dofs_local].copy()
955 # Write to file
956 adios_file.file.BeginStep()
957 dofs = adios_file.io.DefineVariable("dofs", local_dofs, count=[num_dofs_local])
958 adios_file.file.Put(dofs, local_dofs, adios2.Mode.Sync)
959 adios_file.file.EndStep()
960 elif adios_mode == adios2.Mode.Read:
961 adios_file.file.BeginStep()
962 in_variable = adios_file.io.InquireVariable("dofs")
963 in_variable.SetBlockSelection(u.function_space.mesh.comm.rank)
964 adios_file.file.Get(in_variable, u.x.array, adios2.Mode.Sync)
965 adios_file.file.EndStep()
966 u.x.scatter_forward()
967 else:
968 raise NotImplementedError(f"Mode {mode} is not implemented for snapshot checkpoint")
971def read_function_names(
972 filename: Path | str, comm: MPI.Intracomm, backend_args: dict[str, Any] | None
973) -> list[str]:
974 """Read all function names from a file.
976 Args:
977 filename: Path to file
978 comm: MPI communicator to launch IO on.
979 backend_args: Arguments to backend
981 Returns:
982 A list of function names.
983 """
984 raise NotImplementedError("Reading function names are not implemented with ADIOS2")
987def read_point_data(
988 filename: Path | str,
989 name: str,
990 comm: MPI.Intracomm,
991 time: float | str | None,
992 backend_args: dict[str, Any] | None,
993) -> tuple[np.ndarray, int]:
994 """Read data from the nodes of a mesh.
996 Args:
997 filename: Path to file
998 name: Name of point data
999 comm: Communicator to launch IO on.
1000 time: The time stamp
1001 backend_args: The backend arguments
1003 Returns:
1004 Data local to process (contiguous, no mpi comm) and local start range
1005 """
1006 raise NotImplementedError("The ADIOS2 backend cannot read point data.")
1009def read_cell_data(
1010 filename: Path | str,
1011 name: str,
1012 comm: MPI.Intracomm,
1013 time: str | float | None,
1014 backend_args: dict[str, Any] | None,
1015) -> tuple[npt.NDArray[np.int64], np.ndarray]:
1016 """Read data from the cells of a mesh.
1018 Args:
1019 filename: Path to file
1020 name: Name of point data
1021 comm: Communicator to launch IO on.
1022 time: The time stamp
1023 backend_args: The backend arguments
1024 Returns:
1025 A tuple (topology, dofs) where topology contains the
1026 vertex indices of the cells, dofs the degrees of
1027 freedom within that cell.
1028 """
1029 raise NotImplementedError("The ADIOS2 backend does not support reading cell data.")
1032def write_data(
1033 filename: Path | str,
1034 array_data: ArrayData,
1035 comm: MPI.Intracomm,
1036 time: str | float | None,
1037 mode: FileMode,
1038 backend_args: dict[str, Any] | None,
1039):
1040 """Write a 2D-array to file (distributed across proceses with MPI).
1043 Args:
1044 filename: Path to file
1045 array_data: Data to write to file
1046 comm: MPI communicator to open the file with.
1047 time: Time stamp
1048 mode: Append or write
1049 backend_args: The backend arguments
1050 """
1051 raise NotImplementedError("ADIOS2 has not implemented this yet")