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