Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/backends/pyvista/backend.py: 79%
187 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1"""
2Module that uses pyvista to import grids.
3"""
5from typing import Any
7import numpy as np
8import numpy.typing as npt
10try:
11 import pyvista
12except ImportError:
13 raise ModuleNotFoundError("The PyVista-backend requires pyvista")
14from pathlib import Path
16from mpi4py import MPI
18import basix
19import dolfinx
20import fenicsx_compat
22from io4dolfinx.structures import ArrayData, FunctionData, MeshData, MeshTagsData, ReadMeshData
23from io4dolfinx.utils import check_file_exists
25from .. import FileMode, ReadMode
27# Cell types can be found at
28# https://vtk.org/doc/nightly/html/vtkCellType_8h_source.html
29_first_order_vtk = {
30 1: "point",
31 3: "interval",
32 5: "triangle",
33 9: "quadrilateral",
34 10: "tetrahedron",
35 12: "hexahedron",
36}
38_arbitrary_lagrange_vtk = {
39 68: "interval",
40 69: "triangle",
41 70: "quadrilateral",
42 71: "tetrahedron",
43 72: "hexahedron",
44 73: "prism",
45 74: "pyramid",
46}
48_quadratric_vtk = {
49 21: "interval",
50 22: "triangle",
51 23: "quadrilateral",
52 24: "tetrahedron",
53 25: "hexahedron",
54}
57read_mode = ReadMode.serial
60def _cell_degree(ct: str, num_nodes: int) -> int:
61 if ct == "point":
62 return 1
63 elif ct == "interval":
64 return int(num_nodes - 1)
65 elif ct == "triangle":
66 n = (np.sqrt(1 + 8 * num_nodes) - 1) / 2
67 if 2 * num_nodes != n * (n + 1):
68 raise ValueError(f"Unknown triangle layout. Number of nodes: {num_nodes}")
69 return int(n - 1)
70 elif ct == "tetrahedron":
71 n = 0
72 while n * (n + 1) * (n + 2) < 6 * num_nodes:
73 n += 1
74 if n * (n + 1) * (n + 2) != 6 * num_nodes:
75 raise ValueError(f"Unknown tetrahedron layout. Number of nodes: {num_nodes}")
76 return int(n - 1)
78 elif ct == "quadrilateral":
79 n = np.sqrt(num_nodes)
80 if num_nodes != n * n:
81 raise ValueError(f"Unknown quadrilateral layout. Number of nodes: {num_nodes}")
82 return int(n - 1)
83 elif ct == "hexahedron":
84 n = np.cbrt(num_nodes)
85 if num_nodes != n * n * n:
86 raise ValueError(f"Unknown hexahedron layout. Number of nodes: {num_nodes}")
87 return int(n - 1)
88 elif ct == "prism":
89 if num_nodes == 6:
90 return 1
91 elif num_nodes == 15:
92 return 2
93 else:
94 raise ValueError(f"Unknown prism layout. Number of nodes: {num_nodes}")
95 elif ct == "pyramid":
96 if num_nodes == 5:
97 return 1
98 elif num_nodes == 13:
99 return 2
100 else:
101 raise ValueError(f"Unknown pyramid layout. Number of nodes: {num_nodes}")
102 else:
103 raise ValueError(f"Unknown cell type {ct} with {num_nodes=}.")
106def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]:
107 """Get default backend arguments given a set of input arguments.
109 Args:
110 arguments: Input backend arguments
112 Returns:
113 Updated backend arguments
114 """
115 args = arguments or {}
116 return args
119def read_mesh_data(
120 filename: Path | str,
121 comm: MPI.Comm,
122 time: str | float | None = None,
123 read_from_partition: bool = False,
124 backend_args: dict[str, Any] | None = None,
125) -> ReadMeshData:
126 """Read mesh data from file.
128 Args:
129 filename: Path to file to read from
130 comm: MPI communicator used in storage
131 time: Time stamp associated with the mesh to read
132 read_from_partition: Whether to read partition information
133 backend_args: Arguments to backend
135 Returns:
136 Internal data structure for the mesh data read from file
137 """
138 backend_args = get_default_backend_args(backend_args)
139 check_file_exists(filename)
140 if read_from_partition:
141 raise RuntimeError("Cannot read partition data with Pyvista")
142 cells: npt.NDArray[np.int64]
143 geom: npt.NDArray[np.float64 | np.float32]
144 if comm.rank == 0:
145 in_data = pyvista.read(filename)
146 if isinstance(in_data, pyvista.UnstructuredGrid):
147 grid = in_data
148 elif isinstance(in_data, pyvista.core.composite.MultiBlock):
149 # To handle multiblock like pvd
150 fenicsx_compat.pyvista_allow_snake_case(pyvista)
151 number_of_blocks = in_data.number_of_blocks
152 assert number_of_blocks == 1
153 b0 = in_data.get_block(0)
154 assert isinstance(b0, pyvista.UnstructuredGrid)
155 grid = b0
156 else:
157 raise RuntimeError(f"Unknown data type {type(in_data)}")
158 geom = grid.points
159 num_cells_global = grid.number_of_cells
160 cells = grid.cells.reshape(num_cells_global, -1).astype(np.int64)
161 nodes_per_cell_type = cells[:, 0]
162 assert np.allclose(nodes_per_cell_type, nodes_per_cell_type[0]), "Single celltype support"
163 cells = cells[:, 1:].astype(np.int64)
164 cell_types = grid.celltypes
165 assert len(np.unique(cell_types)) == 1
166 if (cell_type := cell_types[0]) in _first_order_vtk.keys():
167 cell_type = _first_order_vtk[cell_type]
168 order = 1
169 elif cell_type in _quadratric_vtk.keys():
170 cell_type = _quadratric_vtk[cell_type]
171 order = 2
172 elif cell_type in _arbitrary_lagrange_vtk.keys():
173 cell_type = _arbitrary_lagrange_vtk[cell_type]
174 order = _cell_degree(cell_type, cells.shape[1])
175 else:
176 raise NotImplementedError(f"Cell type {cell_type} not supported in Pyvista backend.")
177 perm = dolfinx.cpp.io.perm_vtk(dolfinx.mesh.to_type(cell_type), cells.shape[1])
178 cells = cells[:, perm]
179 lvar = int(basix.LagrangeVariant.equispaced)
180 gtype = backend_args.get("dtype", geom.dtype)
181 order, lvar, nodes_per_cell, cell_type, gtype, gdim = comm.bcast(
182 (order, lvar, cells.shape[1], cell_type, gtype, geom.shape[1]), root=0
183 )
184 else:
185 order, lvar, nodes_per_cell, cell_type, gtype, gdim = comm.bcast(None, root=0)
186 geom = np.zeros((0, gdim), dtype=gtype)
187 cells = np.zeros((0, nodes_per_cell), dtype=np.int64)
189 return ReadMeshData(
190 cells=cells, cell_type=cell_type, x=geom.astype(gtype), lvar=lvar, degree=order
191 )
194def read_point_data(
195 filename: Path | str,
196 name: str,
197 comm: MPI.Comm,
198 time: float | str | None,
199 backend_args: dict[str, Any] | None,
200) -> tuple[np.ndarray, int]:
201 """Read data from the nodes of a mesh.
203 Args:
204 filename: Path to file
205 name: Name of point data
206 comm: Communicator to launch IO on.
207 time: The time stamp
208 backend_args: The backend arguments
210 Returns:
211 Data local to process (contiguous, no mpi comm) and local start range
212 """
213 dataset: np.ndarray
214 if MPI.COMM_WORLD.rank == 0:
215 in_data = pyvista.read(filename)
216 if isinstance(in_data, pyvista.UnstructuredGrid):
217 grid = in_data
218 elif isinstance(in_data, pyvista.core.composite.MultiBlock):
219 # To handle multiblock like pvd
220 fenicsx_compat.pyvista_allow_snake_case(pyvista)
221 number_of_blocks = in_data.number_of_blocks
222 assert number_of_blocks == 1
223 b0 = in_data.get_block(0)
224 assert isinstance(b0, pyvista.UnstructuredGrid)
225 grid = b0
227 dataset = grid.point_data[name]
228 if len(dataset.shape) == 1:
229 num_components = 1
230 dataset = dataset.reshape(-1, num_components)
231 else:
232 num_components = dataset.shape[1]
233 if np.issubdtype(dataset.dtype, np.integer):
234 gtype = grid.points.dtype
235 dataset = dataset.astype(gtype)
236 else:
237 gtype = dataset.dtype
238 num_components, gtype = comm.bcast((num_components, gtype), root=0)
239 local_range_start = 0
240 else:
241 num_components, gtype = comm.bcast(None, root=0)
242 dataset = np.zeros((0, num_components), dtype=gtype)
243 local_range_start = 0
245 return dataset, int(local_range_start)
248def read_cell_data(
249 filename: Path | str,
250 name: str,
251 comm: MPI.Comm,
252 time: str | float | None,
253 backend_args: dict[str, Any] | None,
254) -> tuple[npt.NDArray[np.int64], np.ndarray]:
255 dataset: np.ndarray
256 topology: np.ndarray
257 if MPI.COMM_WORLD.rank == 0:
258 in_data = pyvista.read(filename)
259 if isinstance(in_data, pyvista.UnstructuredGrid):
260 grid = in_data
261 elif isinstance(in_data, pyvista.core.composite.MultiBlock):
262 # To handle multiblock like pvd
263 fenicsx_compat.pyvista_allow_snake_case(pyvista)
264 number_of_blocks = in_data.number_of_blocks
265 assert number_of_blocks == 1
266 b0 = in_data.get_block(0)
267 assert isinstance(b0, pyvista.UnstructuredGrid)
268 grid = b0
270 dataset = grid.cell_data[name]
271 if len(dataset.shape) == 1:
272 num_components = 1
273 dataset = dataset.reshape(-1, num_components)
274 else:
275 num_components = dataset.shape[1]
277 if np.issubdtype(dataset.dtype, np.integer):
278 gtype = in_data.points.dtype
279 dataset = dataset.astype(gtype)
280 else:
281 gtype = dataset.dtype
282 num_components, gtype = comm.bcast((num_components, gtype), root=0)
283 else:
284 num_components, gtype = comm.bcast(None, root=0)
285 dataset = np.zeros((0, num_components), dtype=gtype)
286 _time = float(time) if time is not None else None
287 topology = read_mesh_data(filename, comm, _time, False, backend_args=None).cells
288 return topology, dataset
291def write_attributes(
292 filename: Path | str,
293 comm: MPI.Comm,
294 name: str,
295 attributes: dict[str, np.ndarray],
296 backend_args: dict[str, Any] | None,
297):
298 """Write attributes to file.
300 Args:
301 filename: Path to file to write to
302 comm: MPI communicator used in storage
303 name: Name of the attribute group
304 attributes: Dictionary of attributes to write
305 backend_args: Arguments to backend
306 """
307 raise NotImplementedError("The Pyvista backend cannot write attributes.")
310def read_attributes(
311 filename: Path | str,
312 comm: MPI.Comm,
313 name: str,
314 backend_args: dict[str, Any] | None,
315) -> dict[str, Any]:
316 """Read attributes from file.
318 Args:
319 filename: Path to file to read from
320 comm: MPI communicator used in storage
321 name: Name of the attribute group
322 backend_args: Arguments to backend
324 Returns:
325 Dictionary of attributes read from file
326 """
327 raise NotImplementedError("The Pyvista backend cannot read attributes.")
330def read_timestamps(
331 filename: Path | str,
332 comm: MPI.Comm,
333 function_name: str,
334 backend_args: dict[str, Any] | None,
335) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
336 """Read timestamps from file.
338 Args:
339 filename: Path to file to read from
340 comm: MPI communicator used in storage
341 function_name: Name of the function to read timestamps for
342 backend_args: Arguments to backend
344 Returns:
345 Numpy array of timestamps read from file
346 """
347 raise NotImplementedError("The Pyvista backend cannot read timestamps.")
350def read_function_names(
351 filename: Path | str, comm: MPI.Comm, backend_args: dict[str, Any] | None
352) -> list[str]:
353 """Read all function names from a file.
355 Args:
356 filename: Path to file
357 comm: MPI communicator to launch IO on.
358 backend_args: Arguments to backend
360 Returns:
361 A list of function names.
362 """
363 in_data = pyvista.read(filename)
364 if isinstance(in_data, pyvista.UnstructuredGrid):
365 grid = in_data
366 elif isinstance(in_data, pyvista.core.composite.MultiBlock):
367 # To handle multiblock like pvd
368 fenicsx_compat.pyvista_allow_snake_case(pyvista)
369 number_of_blocks = in_data.number_of_blocks
370 assert number_of_blocks == 1
371 b0 = in_data.get_block(0)
372 assert isinstance(b0, pyvista.UnstructuredGrid)
373 grid = b0
375 point_data = list(grid.point_data.keys())
376 cell_data = list(grid.cell_data.keys())
377 return point_data + cell_data
380def write_mesh(
381 filename: Path | str,
382 comm: MPI.Comm,
383 mesh: MeshData,
384 backend_args: dict[str, Any] | None,
385 mode: FileMode,
386 time: float,
387):
388 """
389 Write a mesh to file.
391 Args:
392 comm: MPI communicator used in storage
393 mesh: Internal data structure for the mesh data to save to file
394 filename: Path to file to write to
395 backend_args: Arguments to backend
396 mode: File-mode to store the mesh
397 time: Time stamp associated with the mesh
398 """
399 raise NotImplementedError("The Pyvista backend cannot write meshes.")
402def write_meshtags(
403 filename: str | Path,
404 comm: MPI.Comm,
405 data: MeshTagsData,
406 backend_args: dict[str, Any] | None,
407):
408 """Write mesh tags to file.
410 Args:
411 filename: Path to file to write to
412 comm: MPI communicator used in storage
413 data: Internal data structure for the mesh tags to save to file
414 backend_args: Arguments to backend
415 """
416 raise NotImplementedError("The Pyvista backend cannot write meshtags.")
419def read_meshtags_data(
420 filename: str | Path,
421 comm: MPI.Comm,
422 name: str,
423 backend_args: dict[str, Any] | None,
424) -> MeshTagsData:
425 """Read mesh tags from file.
427 Args:
428 filename: Path to file to read from
429 comm: MPI communicator used in storage
430 name: Name of the mesh tags to read
431 backend_args: Arguments to backend
433 Returns:
434 Internal data structure for the mesh tags read from file
435 """
436 raise NotImplementedError("The Pyvista backend cannot read meshtags.")
439def read_dofmap(
440 filename: str | Path,
441 comm: MPI.Comm,
442 name: str,
443 backend_args: dict[str, Any] | None,
444) -> dolfinx.graph.AdjacencyList:
445 """Read the dofmap of a function with a given name.
447 Args:
448 filename: Path to file to read from
449 comm: MPI communicator used in storage
450 name: Name of the function to read the dofmap for
451 backend_args: Arguments to backend
453 Returns:
454 Dofmap as an AdjacencyList
455 """
456 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
459def read_dofs(
460 filename: str | Path,
461 comm: MPI.Comm,
462 name: str,
463 time: float,
464 backend_args: dict[str, Any] | None,
465) -> tuple[npt.NDArray[np.float32 | np.float64 | np.complex64 | np.complex128], int]:
466 """Read the dofs (values) of a function with a given name from a given timestep.
468 Args:
469 filename: Path to file to read from
470 comm: MPI communicator used in storage
471 name: Name of the function to read the dofs for
472 time: Time stamp associated with the function to read
473 backend_args: Arguments to backend
475 Returns:
476 Contiguous sequence of degrees of freedom (with respect to input data)
477 and the global starting point on the process.
478 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
479 """
480 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
483def read_cell_perms(
484 comm: MPI.Comm, filename: Path | str, backend_args: dict[str, Any] | None
485) -> npt.NDArray[np.uint32]:
486 """
487 Read cell permutation from file with given communicator,
488 Split in continuous chunks based on number of cells in the input data.
490 Args:
491 comm: MPI communicator used in storage
492 filename: Path to file to read from
493 backend_args: Arguments to backend
495 Returns:
496 Contiguous sequence of permutations (with respect to input data)
497 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
498 """
499 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
502def write_function(
503 filename: Path,
504 comm: MPI.Comm,
505 u: FunctionData,
506 time: float,
507 mode: FileMode,
508 backend_args: dict[str, Any] | None,
509):
510 """
511 Write a function to file.
513 Args:
514 comm: MPI communicator used in storage
515 u: Internal data structure for the function data to save to file
516 filename: Path to file to write to
517 time: Time stamp associated with function
518 mode: File-mode to store the function
519 backend_args: Arguments to backend
520 """
521 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
524def read_legacy_mesh(
525 filename: Path | str, comm: MPI.Comm, group: str
526) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating], str | None]:
527 """Read in the mesh topology, geometry and (optionally) cell type from a
528 legacy DOLFIN HDF5-file.
530 Args:
531 filename: Path to file to read from
532 comm: MPI communicator used in storage
533 group: Group in HDF5 file where mesh is stored
535 Returns:
536 Tuple containing:
537 - Topology as a (num_cells, num_vertices_per_cell) array of global vertex indices
538 - Geometry as a (num_vertices, geometric_dimension) array of vertex coordinates
539 - Cell type as a string (e.g. "tetrahedron") or None if not found
540 """
541 raise NotImplementedError("The Pyvista backend cannot read legacy DOLFIN meshes.")
544def snapshot_checkpoint(
545 filename: Path | str,
546 mode: FileMode,
547 u: dolfinx.fem.Function,
548 backend_args: dict[str, Any] | None,
549):
550 """Create a snapshot checkpoint of a dolfinx function.
552 Args:
553 filename: Path to file to read from
554 mode: File-mode to store the function
555 u: dolfinx function to create a snapshot checkpoint for
556 backend_args: Arguments to backend
557 """
558 raise NotImplementedError("The Pyvista backend cannot make checkpoints.")
561def read_hdf5_array(
562 comm: MPI.Comm,
563 filename: Path | str,
564 group: str,
565 backend_args: dict[str, Any] | None,
566) -> tuple[np.ndarray, int]:
567 """Read an array from an HDF5 file.
569 Args:
570 comm: MPI communicator used in storage
571 filename: Path to file to read from
572 group: Group in HDF5 file where array is stored
573 backend_args: Arguments to backend
575 Returns:
576 Tuple containing:
577 - Numpy array read from file
578 - Global starting point on the process.
579 Process 0 has [0, M), process 1 [M, N), process 2 [N, O) etc.
580 """
581 raise NotImplementedError("The Pyvista backend cannot read HDF5 arrays")
584def write_data(
585 filename: Path | str,
586 array_data: ArrayData,
587 comm: MPI.Comm,
588 time: str | float | None,
589 mode: FileMode,
590 backend_args: dict[str, Any] | None,
591):
592 """Write a 2D-array to file (distributed across proceses with MPI).
594 Args:
595 filename: Path to file
596 array_data: Data to write to file
597 comm: MPI communicator to open the file with
598 time: Time stamp
599 mode: Append or write
600 backend_args: The backend arguments
601 """
602 raise NotImplementedError("The pyvista backend does not support writing point data")