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