Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/readers.py: 97%
202 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1# Copyright (C) 2023 Jørgen Schartum Dokken
2#
3# This file is part of io4dolfinx
4#
5# SPDX-License-Identifier: MIT
7from __future__ import annotations
9import logging
10import pathlib
11import typing
12from pathlib import Path
13from typing import Any
15from mpi4py import MPI
17import basix
18import dolfinx
19import fenicsx_compat
20import numpy as np
21import numpy.typing as npt
22import ufl
24from .backends import ReadMode, get_backend
25from .comm_helpers import all_to_all, send_dofs_and_recv_values
26from .utils import (
27 check_file_exists,
28 compute_dofmap_pos,
29 compute_insert_position,
30 compute_local_range,
31 index_owner,
32)
34__all__ = ["read_mesh_from_legacy_h5", "read_function_from_legacy_h5", "read_point_data"]
35logger = logging.getLogger(__name__)
38def _with_mesh_name(
39 backend_args: dict[str, Any] | None, mesh_name: str | None
40) -> dict[str, Any] | None:
41 """Select which mesh in the file a backend call refers to.
43 See :func:`io4dolfinx.checkpointing._with_mesh_name`; duplicated here to keep
44 :mod:`io4dolfinx.readers` free of a circular import.
45 """
46 if mesh_name is None:
47 return backend_args
48 args = dict(backend_args) if backend_args else {}
49 args["name"] = mesh_name
50 return args
53def map_dofmap(dofmap: dolfinx.graph.AdjacencyList, bs: int | np.int64) -> npt.NDArray[np.int64]:
54 """
55 Map xxxyyyzzz to xyzxyz
56 """
58 in_dofmap = dofmap.array
59 in_offsets = dofmap.offsets
61 mapped_dofmap = np.empty_like(in_dofmap)
62 for i in range(len(in_offsets) - 1):
63 pos_begin, pos_end = (
64 in_offsets[i] - in_offsets[0],
65 in_offsets[i + 1] - in_offsets[0],
66 )
67 dofs_i = in_dofmap[pos_begin:pos_end]
68 assert (pos_end - pos_begin) % bs == 0
69 num_dofs_local = int((pos_end - pos_begin) // bs)
70 for k in range(bs):
71 for j in range(num_dofs_local):
72 mapped_dofmap[int(pos_begin + j * bs + k)] = dofs_i[int(num_dofs_local * k + j)]
73 return mapped_dofmap.astype(np.int64)
76def send_cells_and_receive_dofmap_index(
77 filename: pathlib.Path,
78 comm: MPI.Comm,
79 source_ranks: npt.NDArray[np.int32],
80 dest_ranks: npt.NDArray[np.int32],
81 dest_size: npt.NDArray[np.int32],
82 output_owners: npt.NDArray[np.int32],
83 input_cells: npt.NDArray[np.int64],
84 dofmap_pos: npt.NDArray[np.int32],
85 num_cells_global: int | np.int64,
86 dofmap_path: str,
87 xdofmap_path: str,
88 bs: int | np.int64,
89 backend: str | None,
90) -> npt.NDArray[np.int64]:
91 """
92 Given a set of positions in input dofmap, give the global input index of this dofmap entry
93 in input file.
94 """
95 check_file_exists(filename)
97 recv_size = np.zeros(len(source_ranks), dtype=np.int32)
98 assert isinstance(comm, MPI.Intracomm)
99 mesh_to_data_comm = comm.Create_dist_graph_adjacent(
100 source_ranks.tolist(), dest_ranks.tolist(), reorder=False
101 )
102 # Send sizes to create data structures for receiving from NeighAlltoAllv
103 all_to_all(mesh_to_data_comm, dest_size, recv_size)
105 # Sort output for sending and fill send data
106 out_cells = np.zeros(len(output_owners), dtype=np.int64)
107 out_pos = np.zeros(len(output_owners), dtype=np.int32)
108 proc_to_dof = np.zeros_like(input_cells, dtype=np.int32)
109 insertion_array = compute_insert_position(output_owners, dest_ranks, dest_size)
110 out_cells[insertion_array] = input_cells
111 out_pos[insertion_array] = dofmap_pos
112 proc_to_dof[insertion_array] = np.arange(len(input_cells), dtype=np.int32)
113 del insertion_array
115 # Prepare data-structures for receiving
116 total_incoming = sum(recv_size)
117 inc_cells = np.zeros(total_incoming, dtype=np.int64)
118 inc_pos = np.zeros(total_incoming, dtype=np.intc)
120 # Send data
121 s_msg = [out_cells, dest_size, MPI.INT64_T]
122 r_msg = [inc_cells, recv_size, MPI.INT64_T]
123 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
125 s_msg = [out_pos, dest_size, MPI.INT32_T]
126 r_msg = [inc_pos, recv_size, MPI.INT32_T]
127 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
128 mesh_to_data_comm.Free()
130 backend_cls = get_backend(backend)
131 # Read dofmap from file
132 backend_args = {"dofmap": dofmap_path, "offsets": xdofmap_path}
133 if backend == "adios2":
134 backend_args.update({"engine": "HDF5"})
135 input_dofs = backend_cls.read_dofmap(filename, comm, name="", backend_args=backend_args)
136 # Map to xyz
137 mapped_dofmap = map_dofmap(input_dofs, bs).astype(np.int64)
139 # Extract dofmap data
140 local_cell_range = compute_local_range(comm, num_cells_global)
141 input_cell_positions = inc_cells - local_cell_range[0]
142 in_offsets = input_dofs.offsets
143 read_pos = (in_offsets[input_cell_positions] + inc_pos - in_offsets[0]).astype(np.int32)
144 input_dofs = mapped_dofmap[read_pos]
145 del input_cell_positions, read_pos
147 # Send input dofs back to owning process
148 data_to_mesh_comm = comm.Create_dist_graph_adjacent(
149 dest_ranks.tolist(), source_ranks.tolist(), reorder=False
150 )
152 incoming_global_dofs = np.zeros(sum(dest_size), dtype=np.int64)
153 s_msg = [input_dofs, recv_size, MPI.INT64_T]
154 r_msg = [incoming_global_dofs, dest_size, MPI.INT64_T]
155 data_to_mesh_comm.Neighbor_alltoallv(s_msg, r_msg)
157 # Sort incoming global dofs as they were inputted
158 sorted_global_dofs = np.zeros_like(incoming_global_dofs, dtype=np.int64)
159 assert len(incoming_global_dofs) == len(input_cells)
160 sorted_global_dofs[proc_to_dof] = incoming_global_dofs
161 data_to_mesh_comm.Free()
162 return sorted_global_dofs
165def read_mesh_from_legacy_h5(
166 filename: pathlib.Path,
167 comm: MPI.Comm,
168 group: str,
169 cell_type: str = "tetrahedron",
170 backend: str | None = None,
171 max_facet_to_cell_links: int = 2,
172) -> dolfinx.mesh.Mesh:
173 """
174 Read mesh from `h5`-file generated by legacy DOLFIN `HDF5File.write` or `XDMF.write_checkpoint`.
176 Args:
177 comm: MPI communicator to distribute mesh over
178 filename: Path to `h5` or `xdmf` file
179 group: Name of mesh in `h5`-file
180 cell_type: What type of cell type, by default tetrahedron.
181 backend: The IO backend to use when reading the mesh (must
182 support legacy mesh reading, e.g., "adios2").
183 max_facet_to_cell_links: Maximum number of cells a facet
184 can be connected to.
185 """
186 logger.debug(f"Reading mesh from {filename} at group {group}")
187 logger.debug(f"Using backend {backend} with max_facet_to_cell_links {max_facet_to_cell_links}")
188 # Make sure we use the HDF5File and check that the file is present
189 check_file_exists(filename)
191 backend_cls = get_backend(backend)
192 mesh_topology, mesh_geometry, ct = backend_cls.read_legacy_mesh(filename, comm, group)
193 if ct is not None:
194 cell_type = ct
195 # Create DOLFINx mesh
196 element = basix.ufl.element(
197 basix.ElementFamily.P,
198 cell_type,
199 1,
200 basix.LagrangeVariant.equispaced,
201 shape=(mesh_geometry.shape[1],),
202 )
203 domain = ufl.Mesh(element)
205 return dolfinx.mesh.create_mesh(
206 comm=MPI.COMM_WORLD,
207 cells=mesh_topology,
208 x=mesh_geometry,
209 e=domain,
210 partitioner=None,
211 max_facet_to_cell_links=max_facet_to_cell_links,
212 )
215def read_function_from_legacy_h5(
216 filename: pathlib.Path,
217 comm: MPI.Comm,
218 u: dolfinx.fem.Function,
219 group: str = "mesh",
220 step: typing.Optional[int] = None,
221 vector_group: str | None = None,
222 backend: str | None = None,
223):
224 """
225 Read function from a `h5`-file generated by legacy DOLFIN `HDF5File.write`
226 or `XDMF.write_checkpoint`.
229 Args:
230 comm : MPI communicator to distribute mesh over
231 filename : Path to `h5` or `xdmf` file
232 u : The function used to stored the read values
233 group : Group within the `h5` file where the function is stored, by default "mesh"
234 step : The time step used when saving the checkpoint. If not provided it will assume that
235 the function is saved as a regular function (i.e with `HDF5File.write`)
236 backend: The IO backend
237 """
238 logger.debug(f"Reading function from {filename} at group {group}")
239 logger.debug(f"Using backend {backend} with group {group} and step {step}")
240 # Make sure we use the HDF5File and check that the file is present
241 filename = pathlib.Path(filename)
242 if filename.suffix == ".xdmf":
243 filename = filename.with_suffix(".h5")
244 if not filename.is_file():
245 raise FileNotFoundError(f"File {filename} does not exist")
247 V = u.function_space
248 mesh = u.function_space.mesh
249 if u.function_space.element.needs_dof_transformations:
250 raise RuntimeError(
251 "Function-spaces requiring dof permutations are not compatible with legacy data"
252 )
253 # ----------------------Step 1---------------------------------
254 # Compute index of input cells, and position in input dofmap
255 local_cells, dof_pos = compute_dofmap_pos(u.function_space)
256 input_cells = mesh.topology.original_cell_index[local_cells]
258 # Compute mesh->input communicator
259 # 1.1 Compute mesh->input communicator
260 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
261 backend_cls = get_backend(backend)
262 owners: npt.NDArray[np.int32]
263 if backend_cls.read_mode == ReadMode.serial:
264 owners = np.zeros(len(input_cells), dtype=np.int32)
265 elif backend_cls.read_mode == ReadMode.parallel:
266 owners = index_owner(V.mesh.comm, input_cells, num_cells_global)
267 else:
268 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
270 unique_owners, owner_count = np.unique(owners, return_counts=True)
271 # FIXME: In C++ use NBX to find neighbourhood
272 assert isinstance(mesh.comm, MPI.Intracomm)
273 _tmp_comm = mesh.comm.Create_dist_graph(
274 [mesh.comm.rank], [len(unique_owners)], unique_owners.tolist(), reorder=False
275 )
276 source, dest, _ = _tmp_comm.Get_dist_neighbors()
277 _tmp_comm.Free()
278 # Strip out any /
279 group = group.strip("/")
280 if step is not None:
281 group = f"{group}/{group}_{step}"
282 vector_group = vector_group or "vector"
283 else:
284 vector_group = vector_group or "vector_0"
286 # ----------------------Step 2--------------------------------
287 # Get global dofmap indices from input process
288 bs = V.dofmap.bs
289 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
290 dofmap_indices = send_cells_and_receive_dofmap_index(
291 filename,
292 comm,
293 np.asarray(source, dtype=np.int32),
294 np.asarray(dest, dtype=np.int32),
295 owner_count.astype(np.int32),
296 owners,
297 input_cells,
298 dof_pos,
299 num_cells_global,
300 f"/{group}/cell_dofs",
301 f"/{group}/x_cell_dofs",
302 bs,
303 backend=backend,
304 )
306 # ----------------------Step 3---------------------------------
307 dof_owner: npt.NDArray[np.int32]
308 if backend_cls.read_mode == ReadMode.serial:
309 dof_owner = np.zeros(len(dofmap_indices), dtype=np.int32)
310 elif backend_cls.read_mode == ReadMode.parallel:
311 # Compute owner of global dof on distributed input data
312 num_dof_global = V.dofmap.index_map_bs * V.dofmap.index_map.size_global
313 dof_owner = index_owner(comm=mesh.comm, indices=dofmap_indices, N=num_dof_global)
314 else:
315 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
317 # Create MPI neigh comm to owner.
318 # NOTE: USE NBX in C++
320 # Read input data
321 local_array, starting_pos = backend_cls.read_hdf5_array(
322 comm, filename, f"/{group}/{vector_group}", backend_args=None
323 )
325 # Send global dof indices to correct input process, and receive value of given dof
326 local_values = send_dofs_and_recv_values(
327 dofmap_indices, dof_owner, comm, local_array, starting_pos
328 )
330 # ----------------------Step 4---------------------------------
331 # Populate local part of array and scatter forward
332 u.x.array[: len(local_values)] = local_values
333 u.x.scatter_forward()
336def create_geometry_function_space(mesh: dolfinx.mesh.Mesh, N: int) -> dolfinx.fem.FunctionSpace:
337 """Reconstruct a vector space with the N components using the geometry dofmap to ensure
338 a 1-1 mapping between mesh nodes and DOFs."""
339 geom_imap = mesh.geometry.index_map()
340 geom_dofmap = fenicsx_compat.dofmap(mesh)
341 ufl_domain = mesh.ufl_domain()
342 assert ufl_domain is not None
343 sub_el = ufl_domain.ufl_coordinate_element().sub_elements[0]
344 adj_list = dolfinx.cpp.graph.AdjacencyList_int32(geom_dofmap)
346 value_shape: tuple[int, ...]
347 if N == 1:
348 ufl_el = sub_el
349 value_shape = ()
350 else:
351 ufl_el = basix.ufl.blocked_element(sub_el, shape=(N,))
352 value_shape = (N,)
354 _fe_constructor: (
355 type[dolfinx.cpp.fem.FiniteElement_float32] | type[dolfinx.cpp.fem.FiniteElement_float64]
356 )
357 _fem_constructor: (
358 type[dolfinx.cpp.fem.FunctionSpace_float32] | type[dolfinx.cpp.fem.FunctionSpace_float64]
359 )
360 if ufl_el.dtype == np.float32:
361 _fe_constructor = dolfinx.cpp.fem.FiniteElement_float32
362 _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float32
363 elif ufl_el.dtype == np.float64:
364 _fe_constructor = dolfinx.cpp.fem.FiniteElement_float64
365 _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float64
366 else:
367 raise RuntimeError(f"Unsupported type {ufl_el.dtype}")
368 cpp_el = fenicsx_compat.finite_element_ctor_kwargs(
369 _fe_constructor,
370 ufl_el.basix_element._e,
371 gdim=mesh.geometry.dim,
372 value_shape=value_shape,
373 )
374 dof_layout = dolfinx.cpp.fem.create_element_dof_layout(cpp_el, [])
375 cpp_dofmap = dolfinx.cpp.fem.DofMap(
376 dof_layout, fenicsx_compat.unwrap_index_map(geom_imap), N, adj_list, N
377 )
379 # Create function space
380 cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap) # type: ignore[arg-type]
382 return dolfinx.fem.FunctionSpace(mesh, ufl_el, cpp_space)
385def read_point_data(
386 filename: Path | str,
387 name: str,
388 mesh: dolfinx.mesh.Mesh,
389 time: float | None = None,
390 backend_args: dict[str, Any] | None = None,
391 backend: str = "xdmf",
392 mesh_name: str | None = None,
393) -> dolfinx.fem.Function:
394 """Read data from the nodes of a mesh.
396 Note:
397 Backend has to implement :meth:`io4dolfinx.backends.IOBackend.read_point_data`.
399 Args:
400 filename: Path to file
401 name: Name of point data
402 mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`.
403 time: Time-step to read from.
404 mesh_name: Name of the mesh in the file the data belongs to. Defaults
405 to the mesh written without an explicit name.
407 Returns:
408 A function in the space equivalent to the mesh
409 coordinate element (up to shape).
410 """
412 backend_args = _with_mesh_name(backend_args, mesh_name)
413 logger.debug(f"Reading point data from {filename} with name {name} at time {time}")
414 logger.debug(f"Using backend {backend} with arguments {backend_args}")
415 backend_cls = get_backend(backend)
416 dataset, local_range_start = backend_cls.read_point_data(
417 filename=filename, name=name, comm=mesh.comm, time=time, backend_args=backend_args
418 )
420 num_components = dataset.shape[1]
422 # Create appropriate function space (based on coordinate map)
423 V = create_geometry_function_space(mesh, num_components)
424 uh = dolfinx.fem.Function(V, name=name, dtype=dataset.dtype)
425 # Assume that mesh is first order for now
426 x_dofmap = fenicsx_compat.dofmap(mesh)
427 igi = np.array(mesh.geometry.input_global_indices, dtype=np.int64)
429 # This is dependent on how the data is read in. If distributed equally this is correct
430 global_geom_input = igi[x_dofmap]
432 if backend_cls.read_mode == ReadMode.parallel:
433 num_nodes_global = mesh.geometry.index_map().size_global
434 global_geom_owner = index_owner(mesh.comm, global_geom_input.reshape(-1), num_nodes_global)
435 elif backend_cls.read_mode == ReadMode.serial:
436 # This is correct if everything is read in on rank 0
437 global_geom_owner = np.zeros(len(global_geom_input.flatten()), dtype=np.int32)
438 else:
439 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
441 for i in range(num_components):
442 arr_i = send_dofs_and_recv_values(
443 global_geom_input.reshape(-1),
444 global_geom_owner,
445 mesh.comm,
446 dataset[:, i],
447 local_range_start,
448 )
449 dof_pos = x_dofmap.reshape(-1) * num_components + i
450 uh.x.array[dof_pos] = arr_i
451 uh.x.scatter_forward()
452 return uh
455def read_cell_data(
456 filename: Path | str,
457 name: str,
458 mesh: dolfinx.mesh.Mesh,
459 time: float | None = None,
460 backend_args: dict[str, Any] | None = None,
461 backend: str = "xdmf",
462 mesh_name: str | None = None,
463) -> dolfinx.fem.Function:
464 """Read data from the cells of a mesh.
466 Note:
467 Backend has to implement :meth:`io4dolfinx.backends.IOBackend.read_cell_data`.
469 Args:
470 filename: Path to file
471 name: Name of point data
472 mesh: The corresponding :py:class:`dolfinx.mesh.Mesh`.
473 time: Time-step to read from.
474 mesh_name: Name of the mesh in the file the data belongs to. Defaults
475 to the mesh written without an explicit name.
477 Returns:
478 A function in a DG-0 space on the mesh. The cells not found in input is set to zero.
479 """
480 backend_args = _with_mesh_name(backend_args, mesh_name)
482 backend_cls = get_backend(backend)
484 topology, dofs = backend_cls.read_cell_data(
485 filename=filename, name=name, comm=mesh.comm, time=time, backend_args=backend_args
486 )
487 num_components = dofs.shape[1]
488 shape: tuple[int, ...]
489 if num_components == 1:
490 shape = ()
491 else:
492 shape = (num_components,)
493 V = dolfinx.fem.functionspace(mesh, ("DG", 0, shape))
494 u = dolfinx.fem.Function(V, dtype=dofs.dtype)
495 data_array = u.x.array.reshape(-1, num_components)
496 for i in range(num_components):
497 local_entities, local_values = dolfinx.io.distribute_entity_data(
498 mesh, mesh.topology.dim, topology, dofs[:, i].copy()
499 )
500 adj = dolfinx.graph.adjacencylist(local_entities)
501 order = np.arange(len(local_values), dtype=np.int32)
502 mt = dolfinx.mesh.meshtags_from_entities(mesh, mesh.topology.dim, adj, order)
503 data_array[mt.indices, i] = local_values[mt.values]
504 u.x.scatter_forward()
505 return u