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