Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/checkpointing.py: 97%
246 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 18:21 +0000
1# Copyright (C) 2023-2026 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 typing
11from pathlib import Path
12from typing import Any, Callable
14from mpi4py import MPI
16import basix
17import dolfinx
18import numpy as np
19import numpy.typing as npt
20import ufl
22from . import compat
23from .backends import FileMode, ReadMode, get_backend
24from .comm_helpers import (
25 send_and_recv_cell_perm,
26 send_dofmap_and_recv_values,
27 send_dofs_and_recv_values,
28)
29from .readers import create_geometry_function_space
30from .structures import ArrayData, FunctionData, MeshTagsData
31from .utils import (
32 check_file_exists,
33 compute_dofmap_pos,
34 compute_local_range,
35 index_owner,
36 unroll_dofmap,
37 unroll_insert_position,
38)
39from .writers import prepare_meshdata_for_storage
40from .writers import write_function as _internal_function_writer
41from .writers import write_mesh as _internal_mesh_writer
43__all__ = [
44 "read_mesh",
45 "write_function",
46 "read_function",
47 "write_mesh",
48 "read_meshtags",
49 "write_meshtags",
50 "read_attributes",
51 "write_attributes",
52]
54logger = logging.getLogger(__name__)
57def write_attributes(
58 filename: Path | str,
59 comm: MPI.Intracomm,
60 name: str,
61 attributes: dict[str, np.ndarray],
62 backend_args: dict[str, typing.Any] | None = None,
63 backend: str | None = None,
64):
65 """Write attributes to file.
67 Args:
68 filename: Path to file to write to
69 comm: MPI communicator used in storage
70 name: Name of the attributes
71 attributes: Dictionary of attributes to write to file
72 backend_args: Arguments for backend, for instance file type.
73 backend: What backend to use for writing.
74 """
75 logger.debug(f"Writing attributes to {filename} for attribute {name}")
76 logger.debug(f"Using {backend} backend with arguments {backend_args} to write attributes")
77 backend_cls = get_backend(backend)
78 backend_args = backend_cls.get_default_backend_args(backend_args)
79 backend_cls.write_attributes(filename, comm, name, attributes, backend_args)
82def read_attributes(
83 filename: Path | str,
84 comm: MPI.Intracomm,
85 name: str,
86 backend_args: dict[str, typing.Any] | None = None,
87 backend: str | None = None,
88) -> dict[str, typing.Any]:
89 """Read attributes from file.
91 Args:
92 filename: Path to file to read from
93 comm: MPI communicator used in storage
94 name: Name of the attributes
95 backend_args: Arguments for backend, for instance file type.
96 backend: What backend to use for writing.
97 Returns:
98 The attributes
99 """
100 logger.debug(f"Reading attributes from {filename} for attribute {name}")
101 logger.debug(f"Using {backend} backend with arguments {backend_args} to read attributes")
102 backend_cls = get_backend(backend)
103 backend_args = backend_cls.get_default_backend_args(backend_args)
104 return backend_cls.read_attributes(filename, comm, name, backend_args)
107def read_timestamps(
108 filename: Path | str,
109 comm: MPI.Intracomm,
110 function_name: str,
111 backend_args: dict[str, typing.Any] | None = None,
112 backend: str | None = None,
113) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
114 """
115 Read time-stamps from a checkpoint file.
117 Args:
118 comm: MPI communicator
119 filename: Path to file
120 function_name: Name of the function to read time-stamps for
121 backend_args: Arguments for backend, for instance file type.
122 backend: What backend to use for writing.
123 Returns:
124 The time-stamps
125 """
126 logger.debug(f"Reading time-stamps from {filename} for function {function_name}")
127 logger.debug(f"Using {backend} backend with arguments {backend_args} to read time-stamps")
128 check_file_exists(filename)
129 backend_cls = get_backend(backend)
130 backend_args = backend_cls.get_default_backend_args(backend_args)
131 return backend_cls.read_timestamps(filename, comm, function_name, backend_args)
134def write_meshtags(
135 filename: Path | str,
136 mesh: dolfinx.mesh.Mesh,
137 meshtags: dolfinx.mesh.MeshTags,
138 meshtag_name: typing.Optional[str] = None,
139 backend_args: dict[str, Any] | None = None,
140 backend: str | None = None,
141 on_input_mesh: bool = False,
142):
143 """
144 Write meshtags associated with input mesh to file.
146 .. note::
147 For this checkpoint to work, the mesh must be written to file
148 using :func:`write_mesh` before calling this function.
150 Args:
151 filename: Path to save meshtags (with file-extension)
152 mesh: The mesh associated with the meshtags
153 meshtags: The meshtags to write to file
154 meshtag_name: Name of the meshtag. If None, the meshtag name is used.
155 backend_args: Option to IO backend.
156 backend: IO backend
157 on_input_mesh: If True, the meshtags are written with the node ordering
158 of the input mesh.
159 """
160 logger.debug(f"Writing meshtags to {filename} for meshtag {meshtag_name or meshtags.name}")
161 logger.debug(f"Using {backend} backend with arguments {backend_args} to write meshtags")
163 # Extract data from meshtags (convert to global geometry node indices for each entity)
164 tag_entities = meshtags.indices
165 dim = meshtags.dim
166 num_tag_entities_local = mesh.topology.index_map(dim).size_local
167 local_tag_entities = tag_entities[tag_entities < num_tag_entities_local]
168 local_values = meshtags.values[: len(local_tag_entities)]
170 num_saved_tag_entities = len(local_tag_entities)
171 assert isinstance(mesh.comm, MPI.Intracomm)
172 local_start = mesh.comm.exscan(num_saved_tag_entities, op=MPI.SUM)
173 local_start = local_start if mesh.comm.rank != 0 else 0
174 global_num_tag_entities = mesh.comm.allreduce(num_saved_tag_entities, op=MPI.SUM)
176 dof_layout = compat.cmap(mesh).create_dof_layout()
177 if hasattr(dof_layout, "num_entity_closure_dofs"):
178 num_dofs_per_entity = dof_layout.num_entity_closure_dofs(dim)
179 else:
180 num_dofs_per_entity = len(dof_layout.entity_closure_dofs(dim, 0))
181 mesh.topology.create_connectivity(dim, mesh.topology.dim)
182 mesh.topology.create_connectivity(0, mesh.topology.dim)
183 entities_to_geometry = dolfinx.cpp.mesh.entities_to_geometry(
184 mesh._cpp_object, dim, local_tag_entities, False
185 )
187 if on_input_mesh:
188 indices = mesh.geometry.input_global_indices[entities_to_geometry]
189 else:
190 indices = (
191 mesh.geometry.index_map()
192 .local_to_global(entities_to_geometry.reshape(-1))
193 .reshape(entities_to_geometry.shape)
194 )
195 name = meshtag_name or meshtags.name
197 tag_ct = dolfinx.cpp.mesh.cell_entity_type(mesh.topology.cell_type, dim, 0).name
198 tag_data = MeshTagsData(
199 values=local_values,
200 num_entities_global=global_num_tag_entities,
201 num_dofs_per_entity=num_dofs_per_entity,
202 indices=indices,
203 name=name,
204 local_start=local_start,
205 dim=meshtags.dim,
206 cell_type=tag_ct,
207 )
209 # Get backend and default arguments
210 backend_cls = get_backend(backend)
211 backend_args = backend_cls.get_default_backend_args(backend_args)
212 return backend_cls.write_meshtags(filename, mesh.comm, tag_data, backend_args=backend_args)
215def read_meshtags(
216 filename: Path | str,
217 mesh: dolfinx.mesh.Mesh,
218 meshtag_name: str,
219 backend_args: dict[str, Any] | None = None,
220 backend: str | None = None,
221) -> dolfinx.mesh.MeshTags:
222 """
223 Read meshtags from file and return a :class:`dolfinx.mesh.MeshTags` object.
225 Args:
226 filename: Path to meshtags file (with file-extension)
227 mesh: The mesh associated with the meshtags
228 meshtag_name: The name of the meshtag to read
229 engine: Adios2 Engine
230 Returns:
231 The meshtags
232 """
233 logger.debug(f"Reading meshtags from {filename} for meshtag {meshtag_name}")
234 logger.debug(f"Using {backend} backend with arguments {backend_args} to read meshtags")
235 check_file_exists(filename)
236 backend_cls = get_backend(backend)
237 backend_args = backend_cls.get_default_backend_args(backend_args)
238 data = backend_cls.read_meshtags_data(filename, mesh.comm, meshtag_name, backend_args)
240 local_entities, local_values = dolfinx.io.distribute_entity_data(
241 mesh, int(data.dim), data.indices, data.values
242 )
243 mesh.topology.create_connectivity(data.dim, 0)
244 mesh.topology.create_connectivity(data.dim, mesh.topology.dim)
246 adj = dolfinx.graph.adjacencylist(local_entities)
248 local_values = np.array(local_values, dtype=np.int32)
250 mt = dolfinx.mesh.meshtags_from_entities(mesh, int(data.dim), adj, local_values)
251 mt.name = meshtag_name
252 return mt
255def read_function(
256 filename: Path | str,
257 u: dolfinx.fem.Function,
258 time: float = 0.0,
259 name: str | None = None,
260 backend_args: dict[str, Any] | None = None,
261 backend: str | None = None,
262):
263 """
264 Read checkpoint from file and fill it into `u`.
266 Args:
267 filename: Path to checkpoint
268 u: Function to fill
269 time: Time-stamp associated with checkpoint
270 name: If not provided, `u.name` is used to search through the input file for the function
271 """
272 logger.debug(
273 f"Reading function checkpoint from {filename} for function {name or u.name} at time {time}"
274 )
275 logger.debug(
276 f"Using {backend} backend with arguments {backend_args} to read function checkpoint"
277 )
278 check_file_exists(filename)
280 mesh = u.function_space.mesh
281 comm = mesh.comm
282 if name is None:
283 name = u.name
285 check_file_exists(filename)
286 backend_cls = get_backend(backend)
287 backend_args = backend_cls.get_default_backend_args(backend_args)
289 # Compute index of input cells and get cell permutation
290 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
291 input_cells = mesh.topology.original_cell_index[:num_owned_cells]
292 mesh.topology.create_entity_permutations()
293 cell_perm = mesh.topology.get_cell_permutation_info()[:num_owned_cells]
295 # Compute mesh->input communicator
296 # 1.1 Compute mesh->input communicator
297 owners: npt.NDArray[np.int32]
298 if backend_cls.read_mode == ReadMode.serial:
299 owners = np.zeros(input_cells, dtype=np.int32)
300 elif backend_cls.read_mode == ReadMode.parallel:
301 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
302 owners = index_owner(mesh.comm, input_cells, num_cells_global)
303 else:
304 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
305 # -------------------Step 2------------------------------------
306 # Send and receive global cell index and cell perm
307 inc_cells, inc_perms = send_and_recv_cell_perm(input_cells, cell_perm, owners, mesh.comm)
309 input_dofmap = backend_cls.read_dofmap(filename, comm, name, backend_args)
311 # Compute owner of dofs in dofmap
312 dof_owner: npt.NDArray[np.int32]
313 if backend_cls.read_mode == ReadMode.serial:
314 dof_owner = np.zeros(len(input_dofmap.array), dtype=np.int32)
315 elif backend_cls.read_mode == ReadMode.parallel:
316 num_dofs_global = (
317 u.function_space.dofmap.index_map.size_global * u.function_space.dofmap.index_map_bs
318 )
319 dof_owner = index_owner(comm, input_dofmap.array.astype(np.int64), num_dofs_global)
320 else:
321 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
323 # --------------------Step 4-----------------------------------
324 # Read array from file and communicate them to input dofmap process
325 input_array, starting_pos = backend_cls.read_dofs(filename, comm, name, time, backend_args)
327 recv_array = send_dofs_and_recv_values(
328 input_dofmap.array.astype(np.int64), dof_owner, comm, input_array, starting_pos
329 )
331 # -------------------Step 5--------------------------------------
332 # Invert permutation of input data based on input perm
333 # Then apply current permutation to the local data
334 element = u.function_space.element
335 if element.needs_dof_transformations:
336 bs = u.function_space.dofmap.bs
338 # Read input cell permutations on dofmap process
339 local_input_range = compute_local_range(comm, num_cells_global)
340 input_local_cell_index = inc_cells - local_input_range[0]
341 input_perms = backend_cls.read_cell_perms(comm, filename, backend_args)
343 # Start by sorting data array by cell permutation
344 num_dofs_per_cell = input_dofmap.offsets[1:] - input_dofmap.offsets[:-1]
345 assert np.allclose(num_dofs_per_cell, num_dofs_per_cell[0])
347 # Sort dofmap by input local cell index
348 input_perms_sorted = input_perms[input_local_cell_index]
349 unrolled_dofmap_position = unroll_insert_position(
350 input_local_cell_index, num_dofs_per_cell[0]
351 )
352 dofmap_sorted_by_input = recv_array[unrolled_dofmap_position]
354 # First invert input data to reference element then transform to current mesh
355 element.Tt_apply(dofmap_sorted_by_input, input_perms_sorted, bs)
356 element.Tt_inv_apply(dofmap_sorted_by_input, inc_perms, bs)
357 # Compute invert permutation
358 inverted_perm = np.empty_like(unrolled_dofmap_position)
359 inverted_perm[unrolled_dofmap_position] = np.arange(
360 len(unrolled_dofmap_position), dtype=inverted_perm.dtype
361 )
362 recv_array = dofmap_sorted_by_input[inverted_perm]
364 # ------------------Step 6----------------------------------------
365 # For each dof owned by a process, find the local position in the dofmap.
366 V = u.function_space
367 local_cells, dof_pos = compute_dofmap_pos(V)
368 input_cells = V.mesh.topology.original_cell_index[local_cells]
369 num_cells_global = V.mesh.topology.index_map(V.mesh.topology.dim).size_global
371 if backend_cls.read_mode == ReadMode.serial:
372 owners = np.zeros(len(input_cells), dtype=np.int32)
373 elif backend_cls.read_mode == ReadMode.parallel:
374 owners = index_owner(V.mesh.comm, input_cells, num_cells_global)
375 else:
376 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
378 unique_owners, owner_count = np.unique(owners, return_counts=True)
379 # FIXME: In C++ use NBX to find neighbourhood
380 assert isinstance(V.mesh.comm, MPI.Intracomm)
381 sub_comm = V.mesh.comm.Create_dist_graph(
382 [V.mesh.comm.rank], [len(unique_owners)], unique_owners.tolist(), reorder=False
383 )
384 source, dest, _ = sub_comm.Get_dist_neighbors()
385 sub_comm.Free()
387 owned_values = send_dofmap_and_recv_values(
388 comm,
389 np.asarray(source, dtype=np.int32),
390 np.asarray(dest, dtype=np.int32),
391 owners,
392 owner_count.astype(np.int32),
393 input_cells,
394 dof_pos,
395 num_cells_global,
396 recv_array,
397 input_dofmap.offsets,
398 )
399 u.x.array[: len(owned_values)] = owned_values
400 u.x.scatter_forward()
403def read_mesh(
404 filename: Path | str,
405 comm: MPI.Intracomm,
406 ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
407 time: float | str | None = 0.0,
408 read_from_partition: bool = False,
409 backend_args: dict[str, Any] | None = None,
410 backend: str | None = None,
411 max_facet_to_cell_links: int = 2,
412) -> dolfinx.mesh.Mesh:
413 """
414 Read an ADIOS2 mesh into DOLFINx.
416 Args:
417 filename: Path to input file
418 comm: The MPI communciator to distribute the mesh over
419 ghost_mode: Ghost mode to use for mesh. If `read_from_partition`
420 is set to `True` this option is ignored.
421 time: Time stamp associated with mesh
422 read_from_partition: Read mesh with partition from file
423 backend_args: List of arguments to reader backend
424 max_facet_to_cell_links: Maximum number of cells a facet
425 can be connected to.
426 Returns:
427 The distributed mesh
428 """
429 logger.debug(f"Reading mesh from {filename}")
430 logger.debug(f"Using {backend} backend with arguments {backend_args}")
431 logger.debug(f"Time {time} and read_from_partition {read_from_partition}")
432 # Read in data in a distributed fashin
433 check_file_exists(filename)
434 backend_cls = get_backend(backend)
435 backend_args = backend_cls.get_default_backend_args(backend_args)
437 # Let each backend handle what should be default behavior when reading mesh
438 # with or without time stamp.
439 dist_in_data = backend_cls.read_mesh_data(
440 filename,
441 comm,
442 time=time,
443 read_from_partition=read_from_partition,
444 backend_args=backend_args,
445 )
447 # Create DOLFINx mesh
448 element = basix.ufl.element(
449 basix.ElementFamily.P,
450 dist_in_data.cell_type,
451 dist_in_data.degree,
452 basix.LagrangeVariant(int(dist_in_data.lvar)),
453 shape=(dist_in_data.x.shape[1],),
454 dtype=dist_in_data.x.dtype,
455 )
456 domain = ufl.Mesh(element)
457 PartitionerType = Callable[
458 [MPI.Comm, int, list[dolfinx.mesh.CellType], list[npt.NDArray[np.int64]]],
459 dolfinx.cpp.graph.AdjacencyList_int32,
460 ]
461 partitioner: PartitionerType
462 if (partition_graph := dist_in_data.partition_graph) is not None:
464 def _custom_partitioner(
465 comm: MPI.Comm,
466 nparts: int,
467 cell_types: list[dolfinx.mesh.CellType],
468 local_graph: list[npt.NDArray[np.int64]],
469 ) -> dolfinx.cpp.graph.AdjacencyList_int32:
470 assert len(local_graph[0]) % (len(partition_graph.offsets) - 1) == 0
471 if hasattr(partition_graph, "_cpp_object"):
472 cpp_obj = partition_graph._cpp_object
473 assert isinstance(cpp_obj, dolfinx.cpp.graph.AdjacencyList_int32)
474 return cpp_obj
475 else:
476 assert isinstance(partition_graph, dolfinx.cpp.graph.AdjacencyList_int32)
477 return partition_graph
479 partitioner = _custom_partitioner
480 else:
481 try:
482 partitioner = dolfinx.cpp.mesh.create_cell_partitioner(
483 ghost_mode, max_facet_to_cell_links=max_facet_to_cell_links
484 )
485 except TypeError:
486 partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode) # type: ignore[call-overload]
488 # Should change to the commented code below when we require python
489 # minimum version to be >=3.12 see https://github.com/python/cpython/pull/116198
490 # import inspect
491 # sig = inspect.signature(dolfinx.mesh.create_cell_partitioner)
492 # part_kwargs = {}
493 # if "max_facet_to_cell_links" in list(sig.parameters.keys()):
494 # part_kwargs["max_facet_to_cell_links"] = max_facet_to_cell_links
495 # partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode, **part_kwargs)
497 return dolfinx.mesh.create_mesh(
498 comm,
499 cells=dist_in_data.cells,
500 x=dist_in_data.x,
501 e=domain,
502 partitioner=partitioner,
503 )
506def write_mesh(
507 filename: Path,
508 mesh: dolfinx.mesh.Mesh,
509 mode: FileMode = FileMode.write,
510 time: float = 0.0,
511 store_partition_info: bool = False,
512 backend_args: dict[str, Any] | None = None,
513 backend: str | None = None,
514):
515 """
516 Write a mesh to file.
518 Args:
519 filename: Path to save mesh (without file-extension)
520 mesh: The mesh to write to file
522 store_partition_info: Store mesh partitioning (including ghosting) to file
523 """
524 logger.debug(f"Writing mesh to {filename}")
525 logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}")
526 mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info)
527 logger.debug(f"Write mesh using {backend} backend, with arguments {backend_args}")
528 logger.debug(f"Mode {mode} and time {time}")
529 _internal_mesh_writer(
530 filename,
531 mesh.comm,
532 mesh_data=mesh_data,
533 time=time,
534 backend_args=backend_args,
535 backend=backend,
536 mode=mode,
537 )
540def write_function(
541 filename: Path | str,
542 u: dolfinx.fem.Function,
543 time: float = 0.0,
544 mode: FileMode = FileMode.append,
545 name: str | None = None,
546 backend_args: dict[str, Any] | None = None,
547 backend: str | None = None,
548):
549 """
550 Write function checkpoint to file.
552 Args:
553 u: Function to write to file
554 time: Time-stamp for simulation
555 filename: Path to write to
556 mode: Write or append.
557 name: Name of function to write. If None, the name of the function is used.
558 backend_args: Arguments to the IO backend.
559 backend: The backend to use
560 """
561 n = u.name if name is None else name
562 logger.debug(f"Writing function checkpoint to {filename} for function {n} at time {time}")
563 logger.debug(f"Using {backend} backend with arguments {backend_args}")
564 dofmap = u.function_space.dofmap
565 values = u.x.array
566 mesh = u.function_space.mesh
567 comm = mesh.comm
568 mesh.topology.create_entity_permutations()
569 cell_perm = mesh.topology.get_cell_permutation_info()
570 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
571 local_cell_range = mesh.topology.index_map(mesh.topology.dim).local_range
572 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
574 # Convert local dofmap into global_dofmap
575 dmap = dofmap.list
576 num_dofs_per_cell = dmap.shape[1]
577 dofmap_bs = dofmap.bs
578 num_dofs_local_dmap = num_cells_local * num_dofs_per_cell * dofmap_bs
579 index_map_bs = dofmap.index_map_bs
581 # Unroll dofmap for block size
582 unrolled_dofmap = unroll_dofmap(dofmap.list[:num_cells_local, :], dofmap_bs)
583 dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1)
584 dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1)
586 # Convert imap index to global index
587 imap_global = dofmap.index_map.local_to_global(dmap_loc)
588 dofmap_global = imap_global * index_map_bs + dmap_rem
589 dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap)
591 # Compute dofmap offsets
592 local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64)
593 local_dofmap_offsets[:] *= num_dofs_per_cell * dofmap_bs
594 local_dofmap_offsets += dofmap_imap.local_range[0]
596 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs
597 local_dof_range = np.asarray(dofmap.index_map.local_range) * dofmap.index_map_bs
598 num_dofs_local = local_dof_range[1] - local_dof_range[0]
600 # Create internal data structure for function data to write to file
601 function_data = FunctionData(
602 cell_permutations=cell_perm[:num_cells_local].copy(),
603 local_cell_range=local_cell_range,
604 num_cells_global=num_cells_global,
605 dofmap_array=dofmap_global,
606 dofmap_offsets=local_dofmap_offsets,
607 dofmap_range=dofmap_imap.local_range,
608 global_dofs_in_dofmap=dofmap_imap.size_global,
609 values=values[:num_dofs_local].copy(),
610 dof_range=(local_dof_range[0], local_dof_range[1]),
611 num_dofs_global=num_dofs_global,
612 name=name or u.name,
613 )
614 # Write to file
615 fname = Path(filename)
616 _internal_function_writer(
617 fname, comm, function_data, time, backend_args=backend_args, backend=backend, mode=mode
618 )
621def read_function_names(
622 filename: Path | str,
623 comm: MPI.Intracomm,
624 backend_args: dict[str, Any] | None = None,
625 backend: str = "h5py",
626) -> list[str]:
627 """Read all function names from a file.
629 Args:
630 filename: Path to file
631 comm: MPI communicator to launch IO on.
632 backend_args: Arguments to backend
634 Returns:
635 A list of function names.
636 """
637 logger.debug(f"Reading function names from {filename}")
638 logger.debug(f"Using {backend} backend with arguments {backend_args} to read function names")
639 check_file_exists(filename)
640 backend_cls = get_backend(backend)
641 return backend_cls.read_function_names(filename, comm, backend_args=backend_args)
644def write_point_data(
645 filename: Path | str,
646 u: dolfinx.fem.Function,
647 time: str | float | None,
648 mode: FileMode,
649 backend_args: dict[str, Any] | None,
650 backend: str = "vtkhdf",
651):
652 """Write function to file by interpolating into geometry nodes.
655 Args:
656 filename: Path to file
657 u: The function to store
658 time: Time stamp
659 mode: Append or write
660 backend_args: The backend arguments
661 backend: Which backend to use.
662 """
663 logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}")
664 V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape)))
665 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
666 v_out.interpolate(u)
667 comm = v_out.function_space.mesh.comm
668 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
669 local_range = V.dofmap.index_map.local_range
670 num_dofs_local = V.dofmap.index_map.size_local
671 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
672 ad = ArrayData(
673 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Point"
674 )
675 logger.debug(
676 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write point data"
677 )
678 backend_cls = get_backend(backend)
679 return backend_cls.write_data(
680 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
681 )
684def write_cell_data(
685 filename: Path | str,
686 u: dolfinx.fem.Function,
687 time: str | float | None,
688 mode: FileMode,
689 backend_args: dict[str, Any] | None,
690 backend: str = "vtkhdf",
691):
692 """Write function to file by interpolating into cell midpoints.
695 Args:
696 filename: Path to file
697 point_data: Data to write to file
698 time: Time stamp
699 mode: Append or write
700 backend_args: The backend arguments
701 """
702 logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}")
703 V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape))
704 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
705 v_out.interpolate(u)
706 comm = v_out.function_space.mesh.comm
707 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
708 local_range = V.dofmap.index_map.local_range
709 num_dofs_local = V.dofmap.index_map.size_local
710 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
712 ad = ArrayData(
713 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Cell"
714 )
715 logger.debug(
716 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write cell data"
717 )
718 backend_cls = get_backend(backend)
720 return backend_cls.write_data(
721 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
722 )