Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/adios2/helpers.py: 94%
125 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1"""
2Helpers reading/writing data with ADIOS2
3"""
5from __future__ import annotations
7import shutil
8from contextlib import contextmanager
9from pathlib import Path
10from typing import NamedTuple
12from mpi4py import MPI
14import adios2
15import dolfinx.graph
16import fenicsx_compat
17import numpy as np
18import numpy.typing as npt
20from ...utils import compute_local_range, valid_function_types
22adios2 = fenicsx_compat.resolve_adios_scope(adios2)
25__all__ = [
26 "AdiosFile",
27 "ADIOSFile",
28 "check_variable_exists",
29 "read_array",
30 "read_adjacency_list",
31 "adios_to_numpy_dtype",
32]
34adios_to_numpy_dtype = {
35 "float": np.float32,
36 "double": np.float64,
37 "float complex": np.complex64,
38 "double complex": np.complex128,
39 "uint32_t": np.uint32,
40}
43class AdiosFile(NamedTuple):
44 io: adios2.IO
45 file: adios2.Engine
48@contextmanager
49def ADIOSFile(
50 adios: adios2.ADIOS,
51 filename: Path | str,
52 engine: str,
53 mode: adios2.Mode,
54 io_name: str,
55 comm: MPI.Comm | None = None,
56):
57 io = adios.DeclareIO(io_name)
58 io.SetEngine(engine)
59 # ADIOS2 sometimes struggles with existing files/folders it should overwrite
60 if mode == adios2.Mode.Write:
61 filename = Path(filename)
62 if filename.exists() and comm is not None and comm.rank == 0:
63 if filename.is_dir():
64 shutil.rmtree(filename)
65 else:
66 filename.unlink()
67 if comm is not None:
68 comm.Barrier()
70 file = io.Open(str(filename), mode)
71 try:
72 yield AdiosFile(io=io, file=file)
73 finally:
74 file.Close()
75 adios.RemoveIO(io_name)
78def check_variable_exists(
79 adios: adios2.ADIOS,
80 filename: Path | str,
81 variable: str,
82 engine: str,
83) -> bool:
84 io_name = f"{variable}_reader"
86 if not Path(filename).exists():
87 return False
89 variable_found = False
90 with ADIOSFile(
91 adios=adios,
92 engine=engine,
93 filename=filename,
94 mode=adios2.Mode.Read,
95 io_name=io_name,
96 ) as adios_file:
97 # Find step that has cell permutation
98 for _ in range(adios_file.file.Steps()):
99 adios_file.file.BeginStep()
100 if variable in adios_file.io.AvailableVariables().keys():
101 variable_found = True
102 break
103 adios_file.file.EndStep()
105 # Not sure if this is needed, but just in case
106 if variable in adios_file.io.AvailableVariables().keys():
107 variable_found = True
108 return variable_found
111def read_adjacency_list(
112 adios: adios2.ADIOS,
113 comm: MPI.Comm,
114 filename: Path | str,
115 data_name: str,
116 offsets_name: str,
117 engine: str,
118) -> dolfinx.graph.AdjacencyList:
119 """
120 Read an adjacency-list from an ADIOS file with given communicator.
121 The adjancency list is split in to a flat array (data) and its corresponding offset.
123 Args:
124 adios: The ADIOS instance
125 comm: The MPI communicator used to read the data
126 filename: Path to input file
127 data_name: Name of variable containing the indices of the adjacencylist
128 dofmap_offsets: Name of variable containing offsets of the adjacencylist
129 engine: Type of ADIOS engine to use for reading data
131 Returns:
132 The local part of dofmap from input dofs
134 .. note::
135 No MPI communication is done during this call
136 """
138 # Open ADIOS engine
139 io_name = f"{data_name=}_reader"
141 with ADIOSFile(
142 adios=adios,
143 engine=engine,
144 filename=filename,
145 mode=adios2.Mode.Read,
146 io_name=io_name,
147 ) as adios_file:
148 # First find step with dofmap offsets, to be able to read
149 # in a full row of the dofmap
150 for _ in range(adios_file.file.Steps()):
151 adios_file.file.BeginStep()
152 if offsets_name in adios_file.io.AvailableVariables().keys():
153 break
154 adios_file.file.EndStep()
155 if offsets_name not in adios_file.io.AvailableVariables().keys():
156 raise KeyError(f"Dof offsets not found at '{offsets_name}' in {filename}")
158 # Get global shape of dofmap-offset, and read in data with an overlap
159 d_offsets = adios_file.io.InquireVariable(offsets_name)
160 shape = d_offsets.Shape()
161 num_nodes = shape[0] - 1
162 local_range = compute_local_range(comm, num_nodes)
164 # As the offsets are one longer than the number of cells, we need to read in with an overlap
165 if len(shape) == 1:
166 d_offsets.SetSelection([[local_range[0]], [local_range[1] + 1 - local_range[0]]])
167 in_offsets = np.empty(
168 local_range[1] + 1 - local_range[0],
169 dtype=d_offsets.Type().strip("_t"),
170 )
171 else:
172 d_offsets.SetSelection(
173 [
174 [local_range[0], 0],
175 [local_range[1] + 1 - local_range[0], shape[1]],
176 ]
177 )
178 in_offsets = np.empty(
179 (local_range[1] + 1 - local_range[0], shape[1]),
180 dtype=d_offsets.Type().strip("_t"),
181 )
183 adios_file.file.Get(d_offsets, in_offsets, adios2.Mode.Sync)
184 in_offsets = in_offsets.squeeze()
186 # Assuming dofmap is saved in stame step
187 # Get the relevant part of the dofmap
188 if data_name not in adios_file.io.AvailableVariables().keys():
189 raise KeyError(f"Dofs not found at {data_name} in {filename}")
190 cell_dofs = adios_file.io.InquireVariable(data_name)
191 if len(shape) == 1:
192 cell_dofs.SetSelection([[in_offsets[0]], [in_offsets[-1] - in_offsets[0]]])
193 in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t"))
194 else:
195 cell_dofs.SetSelection([[in_offsets[0], 0], [in_offsets[-1] - in_offsets[0], shape[1]]])
196 in_dofmap = np.empty(
197 (in_offsets[-1] - in_offsets[0], shape[1]),
198 dtype=cell_dofs.Type().strip("_t"),
199 )
200 assert shape[1] == 1
202 in_dofmap = np.empty(in_offsets[-1] - in_offsets[0], dtype=cell_dofs.Type().strip("_t"))
203 adios_file.file.Get(cell_dofs, in_dofmap, adios2.Mode.Sync)
204 in_offsets -= in_offsets[0]
205 adios_file.file.EndStep()
207 # Return local dofmap
208 return dolfinx.graph.adjacencylist(in_dofmap, in_offsets.astype(np.int32))
211def read_array(
212 adios: adios2.ADIOS,
213 filename: Path | str,
214 array_name: str,
215 engine: str,
216 comm: MPI.Comm,
217 time: float = 0.0,
218 time_name: str = "",
219 legacy: bool = False,
220) -> tuple[npt.NDArray[valid_function_types], int]:
221 """
222 Read an array from file, return the global starting position of the local array
224 Args:
225 adios: The ADIOS instance
226 filename: Path to file to read array from
227 array_name: Name of array in file
228 engine: Name of engine to use to read file
229 comm: MPI communicator used for reading the data
230 time_name: Name of time variable for modern checkpoints
231 legacy: If True ignore time_name and read the first available step
232 Returns:
233 Local part of array and its global starting position
234 """
236 with ADIOSFile(
237 adios=adios,
238 engine=engine,
239 filename=filename,
240 mode=adios2.Mode.Read,
241 io_name="ArrayReader",
242 ) as adios_file:
243 # Get time-stamp from first available step
244 if legacy:
245 for i in range(adios_file.file.Steps()):
246 adios_file.file.BeginStep()
247 if array_name in adios_file.io.AvailableVariables().keys():
248 break
249 adios_file.file.EndStep()
250 if array_name not in adios_file.io.AvailableVariables().keys():
251 raise KeyError(f"No array found at {array_name}")
252 else:
253 for i in range(adios_file.file.Steps()):
254 adios_file.file.BeginStep()
255 if time_name in adios_file.io.AvailableVariables().keys():
256 arr = adios_file.io.InquireVariable(time_name)
257 time_shape = arr.Shape()
258 arr.SetSelection([[0], [time_shape[0]]])
259 times = np.empty(time_shape[0], dtype=adios_to_numpy_dtype[arr.Type()])
260 adios_file.file.Get(arr, times, adios2.Mode.Sync)
261 if times[0] == time:
262 break
263 if i == adios_file.file.Steps() - 1:
264 raise KeyError(
265 f"No data associated with {time_name}={time} found in {filename}"
266 )
268 adios_file.file.EndStep()
270 if time_name not in adios_file.io.AvailableVariables().keys():
271 raise KeyError(f"No data associated with {time_name}={time} found in {filename}")
273 if array_name not in adios_file.io.AvailableVariables().keys():
274 raise KeyError(f"No array found at {time=} for {array_name}")
276 arr = adios_file.io.InquireVariable(array_name)
277 arr_shape = arr.Shape()
278 # TODO: Should we always pick the first element?
279 assert len(arr_shape) >= 1
280 arr_range = compute_local_range(comm, arr_shape[0])
282 if len(arr_shape) == 1:
283 arr.SetSelection([[arr_range[0]], [arr_range[1] - arr_range[0]]])
284 vals = np.empty(arr_range[1] - arr_range[0], dtype=adios_to_numpy_dtype[arr.Type()])
285 else:
286 arr.SetSelection([[arr_range[0], 0], [arr_range[1] - arr_range[0], arr_shape[1]]])
287 vals = np.empty(
288 (arr_range[1] - arr_range[0], arr_shape[1]),
289 dtype=adios_to_numpy_dtype[arr.Type()],
290 )
291 assert arr_shape[1] == 1
293 adios_file.file.Get(arr, vals, adios2.Mode.Sync)
294 adios_file.file.EndStep()
296 return vals.reshape(-1), arr_range[0]