Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/checkpointing.py: 98%
253 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +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 fenicsx_compat
19import numpy as np
20import numpy.typing as npt
21import ufl
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 _with_mesh_name(
58 backend_args: dict[str, Any] | None, mesh_name: str | None
59) -> dict[str, Any] | None:
60 """Select which mesh in the file a backend call refers to.
62 A checkpoint may hold several meshes -- a mesh and a submesh of it, say --
63 and backends take the one to act on from ``backend_args["name"]``. This
64 folds an explicit ``mesh_name`` in without mutating the caller's dictionary.
65 ``mesh_name=None`` leaves ``backend_args`` untouched, so an existing
66 ``backend_args["name"]`` still wins and files written before named meshes
67 existed stay readable.
69 Args:
70 backend_args: Arguments for the backend, or None
71 mesh_name: Name of the mesh to act on, or None for the default
73 Returns:
74 Backend arguments naming the requested mesh.
75 """
76 if mesh_name is None:
77 return backend_args
78 args = dict(backend_args) if backend_args else {}
79 args["name"] = mesh_name
80 return args
83def write_attributes(
84 filename: Path | str,
85 comm: MPI.Comm,
86 name: str,
87 attributes: dict[str, np.ndarray],
88 backend_args: dict[str, typing.Any] | None = None,
89 backend: str | None = None,
90):
91 """Write attributes to file.
93 Args:
94 filename: Path to file to write to
95 comm: MPI communicator used in storage
96 name: Name of the attributes
97 attributes: Dictionary of attributes to write to file
98 backend_args: Arguments for backend, for instance file type.
99 backend: What backend to use for writing.
100 """
101 logger.debug(f"Writing attributes to {filename} for attribute {name}")
102 logger.debug(f"Using {backend} backend with arguments {backend_args} to write attributes")
103 backend_cls = get_backend(backend)
104 backend_args = backend_cls.get_default_backend_args(backend_args)
105 backend_cls.write_attributes(filename, comm, name, attributes, backend_args)
108def read_attributes(
109 filename: Path | str,
110 comm: MPI.Comm,
111 name: str,
112 backend_args: dict[str, typing.Any] | None = None,
113 backend: str | None = None,
114) -> dict[str, typing.Any]:
115 """Read attributes from file.
117 Args:
118 filename: Path to file to read from
119 comm: MPI communicator used in storage
120 name: Name of the attributes
121 backend_args: Arguments for backend, for instance file type.
122 backend: What backend to use for writing.
123 Returns:
124 The attributes
125 """
126 logger.debug(f"Reading attributes from {filename} for attribute {name}")
127 logger.debug(f"Using {backend} backend with arguments {backend_args} to read attributes")
128 backend_cls = get_backend(backend)
129 backend_args = backend_cls.get_default_backend_args(backend_args)
130 return backend_cls.read_attributes(filename, comm, name, backend_args)
133def read_timestamps(
134 filename: Path | str,
135 comm: MPI.Comm,
136 function_name: str,
137 backend_args: dict[str, typing.Any] | None = None,
138 backend: str | None = None,
139 mesh_name: str | None = None,
140) -> npt.NDArray[np.float64 | str]: # type: ignore[type-var]
141 """
142 Read time-stamps from a checkpoint file.
144 Args:
145 comm: MPI communicator
146 filename: Path to file
147 function_name: Name of the function to read time-stamps for
148 backend_args: Arguments for backend, for instance file type.
149 backend: What backend to use for writing.
150 mesh_name: Name of the mesh in the file to act on. Defaults to the
151 mesh written without an explicit name, so existing checkpoints and
152 call sites are unaffected.
153 Returns:
154 The time-stamps
155 """
156 backend_args = _with_mesh_name(backend_args, mesh_name)
157 logger.debug(f"Reading time-stamps from {filename} for function {function_name}")
158 logger.debug(f"Using {backend} backend with arguments {backend_args} to read time-stamps")
159 check_file_exists(filename)
160 backend_cls = get_backend(backend)
161 backend_args = backend_cls.get_default_backend_args(backend_args)
162 return backend_cls.read_timestamps(filename, comm, function_name, backend_args)
165def write_meshtags(
166 filename: Path | str,
167 mesh: dolfinx.mesh.Mesh,
168 meshtags: dolfinx.mesh.MeshTags,
169 meshtag_name: typing.Optional[str] = None,
170 backend_args: dict[str, Any] | None = None,
171 backend: str | None = None,
172 on_input_mesh: bool = False,
173 mesh_name: str | None = None,
174):
175 """
176 Write meshtags associated with input mesh to file.
178 .. note::
179 For this checkpoint to work, the mesh must be written to file
180 using :func:`write_mesh` before calling this function.
182 Args:
183 filename: Path to save meshtags (with file-extension)
184 mesh: The mesh associated with the meshtags
185 meshtags: The meshtags to write to file
186 meshtag_name: Name of the meshtag. If None, the meshtag name is used.
187 backend_args: Option to IO backend.
188 backend: IO backend
189 on_input_mesh: If True, the meshtags are written with the node ordering
190 of the input mesh.
191 mesh_name: Name of the mesh in the file to act on. Defaults to the
192 mesh written without an explicit name, so existing checkpoints and
193 call sites are unaffected.
194 """
195 backend_args = _with_mesh_name(backend_args, mesh_name)
196 logger.debug(f"Writing meshtags to {filename} for meshtag {meshtag_name or meshtags.name}")
197 logger.debug(f"Using {backend} backend with arguments {backend_args} to write meshtags")
199 # Extract data from meshtags (convert to global geometry node indices for each entity)
200 tag_entities = meshtags.indices
201 dim = meshtags.dim
202 num_tag_entities_local = mesh.topology.index_map(dim).size_local
203 local_tag_entities = tag_entities[tag_entities < num_tag_entities_local]
204 local_values = meshtags.values[: len(local_tag_entities)]
206 num_saved_tag_entities = len(local_tag_entities)
207 assert isinstance(mesh.comm, MPI.Intracomm)
208 local_start = mesh.comm.exscan(num_saved_tag_entities, op=MPI.SUM)
209 local_start = local_start if mesh.comm.rank != 0 else 0
210 global_num_tag_entities = mesh.comm.allreduce(num_saved_tag_entities, op=MPI.SUM)
212 dof_layout = fenicsx_compat.cmap(mesh).create_dof_layout()
213 num_dofs_per_entity = len(dof_layout.entity_closure_dofs(dim, 0))
214 mesh.topology.create_connectivity(dim, mesh.topology.dim)
215 mesh.topology.create_connectivity(0, mesh.topology.dim)
216 entities_to_geometry = dolfinx.cpp.mesh.entities_to_geometry(
217 mesh._cpp_object, dim, local_tag_entities, False
218 )
220 if on_input_mesh:
221 indices = mesh.geometry.input_global_indices[entities_to_geometry]
222 else:
223 indices = (
224 mesh.geometry.index_map()
225 .local_to_global(entities_to_geometry.reshape(-1))
226 .reshape(entities_to_geometry.shape)
227 )
228 name = meshtag_name or meshtags.name
230 tag_ct = dolfinx.cpp.mesh.cell_entity_type(mesh.topology.cell_type, dim, 0).name
231 tag_data = MeshTagsData(
232 values=local_values,
233 num_entities_global=global_num_tag_entities,
234 num_dofs_per_entity=num_dofs_per_entity,
235 indices=indices,
236 name=name,
237 local_start=local_start,
238 dim=meshtags.dim,
239 cell_type=tag_ct,
240 )
242 # Get backend and default arguments
243 backend_cls = get_backend(backend)
244 backend_args = backend_cls.get_default_backend_args(backend_args)
245 return backend_cls.write_meshtags(filename, mesh.comm, tag_data, backend_args=backend_args)
248def read_meshtags(
249 filename: Path | str,
250 mesh: dolfinx.mesh.Mesh,
251 meshtag_name: str,
252 backend_args: dict[str, Any] | None = None,
253 backend: str | None = None,
254 mesh_name: str | None = None,
255) -> dolfinx.mesh.MeshTags:
256 """
257 Read meshtags from file and return a :class:`dolfinx.mesh.MeshTags` object.
259 Args:
260 filename: Path to meshtags file (with file-extension)
261 mesh: The mesh associated with the meshtags
262 meshtag_name: The name of the meshtag to read
263 engine: Adios2 Engine
264 mesh_name: Name of the mesh in the file to act on. Defaults to the
265 mesh written without an explicit name, so existing checkpoints and
266 call sites are unaffected.
267 Returns:
268 The meshtags
269 """
270 backend_args = _with_mesh_name(backend_args, mesh_name)
271 logger.debug(f"Reading meshtags from {filename} for meshtag {meshtag_name}")
272 logger.debug(f"Using {backend} backend with arguments {backend_args} to read meshtags")
273 check_file_exists(filename)
274 backend_cls = get_backend(backend)
275 backend_args = backend_cls.get_default_backend_args(backend_args)
276 data = backend_cls.read_meshtags_data(filename, mesh.comm, meshtag_name, backend_args)
278 local_entities, local_values = dolfinx.io.distribute_entity_data(
279 mesh, int(data.dim), data.indices, data.values
280 )
281 mesh.topology.create_connectivity(data.dim, 0)
282 mesh.topology.create_connectivity(data.dim, mesh.topology.dim)
284 adj = dolfinx.graph.adjacencylist(local_entities)
286 local_values = np.array(local_values, dtype=np.int32)
288 mt = dolfinx.mesh.meshtags_from_entities(mesh, int(data.dim), adj, local_values)
289 mt.name = meshtag_name
290 return mt
293def read_function(
294 filename: Path | str,
295 u: dolfinx.fem.Function,
296 time: float = 0.0,
297 name: str | None = None,
298 backend_args: dict[str, Any] | None = None,
299 backend: str | None = None,
300 mesh_name: str | None = None,
301):
302 """
303 Read checkpoint from file and fill it into `u`.
305 Args:
306 filename: Path to checkpoint
307 u: Function to fill
308 time: Time-stamp associated with checkpoint
309 name: If not provided, `u.name` is used to search through the input file for the function
310 mesh_name: Name of the mesh in the file to act on. Defaults to the
311 mesh written without an explicit name, so existing checkpoints and
312 call sites are unaffected.
313 """
314 backend_args = _with_mesh_name(backend_args, mesh_name)
315 logger.debug(
316 f"Reading function checkpoint from {filename} for function {name or u.name} at time {time}"
317 )
318 logger.debug(
319 f"Using {backend} backend with arguments {backend_args} to read function checkpoint"
320 )
321 check_file_exists(filename)
323 mesh = u.function_space.mesh
324 comm = mesh.comm
325 if name is None:
326 name = u.name
328 check_file_exists(filename)
329 backend_cls = get_backend(backend)
330 backend_args = backend_cls.get_default_backend_args(backend_args)
332 # Compute index of input cells and get cell permutation
333 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
334 input_cells = mesh.topology.original_cell_index[:num_owned_cells]
335 cell_perm = fenicsx_compat.cell_permutation_info(mesh.topology)[:num_owned_cells]
337 # Compute mesh->input communicator
338 # 1.1 Compute mesh->input communicator
339 owners: npt.NDArray[np.int32]
340 if backend_cls.read_mode == ReadMode.serial:
341 owners = np.zeros(input_cells, dtype=np.int32)
342 elif backend_cls.read_mode == ReadMode.parallel:
343 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
344 owners = index_owner(mesh.comm, input_cells, num_cells_global)
345 else:
346 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
347 # -------------------Step 2------------------------------------
348 # Send and receive global cell index and cell perm
349 inc_cells, inc_perms = send_and_recv_cell_perm(input_cells, cell_perm, owners, mesh.comm)
351 input_dofmap = backend_cls.read_dofmap(filename, comm, name, backend_args)
353 # Compute owner of dofs in dofmap
354 dof_owner: npt.NDArray[np.int32]
355 if backend_cls.read_mode == ReadMode.serial:
356 dof_owner = np.zeros(len(input_dofmap.array), dtype=np.int32)
357 elif backend_cls.read_mode == ReadMode.parallel:
358 num_dofs_global = (
359 u.function_space.dofmap.index_map.size_global * u.function_space.dofmap.index_map_bs
360 )
361 dof_owner = index_owner(comm, input_dofmap.array.astype(np.int64), num_dofs_global)
362 else:
363 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
365 # --------------------Step 4-----------------------------------
366 # Read array from file and communicate them to input dofmap process
367 input_array, starting_pos = backend_cls.read_dofs(filename, comm, name, time, backend_args)
369 recv_array = send_dofs_and_recv_values(
370 input_dofmap.array.astype(np.int64), dof_owner, comm, input_array, starting_pos
371 )
373 # -------------------Step 5--------------------------------------
374 # Invert permutation of input data based on input perm
375 # Then apply current permutation to the local data
376 element = u.function_space.element
377 if element.needs_dof_transformations:
378 bs = u.function_space.dofmap.bs
380 # Read input cell permutations on dofmap process
381 local_input_range = compute_local_range(comm, num_cells_global)
382 input_local_cell_index = inc_cells - local_input_range[0]
383 input_perms = backend_cls.read_cell_perms(comm, filename, backend_args)
385 # Start by sorting data array by cell permutation
386 num_dofs_per_cell = input_dofmap.offsets[1:] - input_dofmap.offsets[:-1]
387 assert np.allclose(num_dofs_per_cell, num_dofs_per_cell[0])
389 # Sort dofmap by input local cell index
390 input_perms_sorted = input_perms[input_local_cell_index]
391 unrolled_dofmap_position = unroll_insert_position(
392 input_local_cell_index, num_dofs_per_cell[0]
393 )
394 dofmap_sorted_by_input = recv_array[unrolled_dofmap_position]
396 # First invert input data to reference element then transform to current mesh
397 element.Tt_apply(dofmap_sorted_by_input, input_perms_sorted, bs)
398 element.Tt_inv_apply(dofmap_sorted_by_input, inc_perms, bs)
399 # Compute invert permutation
400 inverted_perm = np.empty_like(unrolled_dofmap_position)
401 inverted_perm[unrolled_dofmap_position] = np.arange(
402 len(unrolled_dofmap_position), dtype=inverted_perm.dtype
403 )
404 recv_array = dofmap_sorted_by_input[inverted_perm]
406 # ------------------Step 6----------------------------------------
407 # For each dof owned by a process, find the local position in the dofmap.
408 V = u.function_space
409 local_cells, dof_pos = compute_dofmap_pos(V)
410 input_cells = V.mesh.topology.original_cell_index[local_cells]
411 num_cells_global = V.mesh.topology.index_map(V.mesh.topology.dim).size_global
413 if backend_cls.read_mode == ReadMode.serial:
414 owners = np.zeros(len(input_cells), dtype=np.int32)
415 elif backend_cls.read_mode == ReadMode.parallel:
416 owners = index_owner(V.mesh.comm, input_cells, num_cells_global)
417 else:
418 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
420 unique_owners, owner_count = np.unique(owners, return_counts=True)
421 # FIXME: In C++ use NBX to find neighbourhood
422 assert isinstance(V.mesh.comm, MPI.Intracomm)
423 sub_comm = V.mesh.comm.Create_dist_graph(
424 [V.mesh.comm.rank], [len(unique_owners)], unique_owners.tolist(), reorder=False
425 )
426 source, dest, _ = sub_comm.Get_dist_neighbors()
427 sub_comm.Free()
429 owned_values = send_dofmap_and_recv_values(
430 comm,
431 np.asarray(source, dtype=np.int32),
432 np.asarray(dest, dtype=np.int32),
433 owners,
434 owner_count.astype(np.int32),
435 input_cells,
436 dof_pos,
437 num_cells_global,
438 recv_array,
439 input_dofmap.offsets,
440 )
441 u.x.array[: len(owned_values)] = owned_values
442 u.x.scatter_forward()
445def read_mesh(
446 filename: Path | str,
447 comm: MPI.Comm,
448 ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
449 time: float | str | None = 0.0,
450 read_from_partition: bool = False,
451 backend_args: dict[str, Any] | None = None,
452 backend: str | None = None,
453 max_facet_to_cell_links: int = 2,
454 mesh_name: str | None = None,
455) -> dolfinx.mesh.Mesh:
456 """
457 Read an ADIOS2 mesh into DOLFINx.
459 Args:
460 filename: Path to input file
461 comm: The MPI communciator to distribute the mesh over
462 ghost_mode: Ghost mode to use for mesh. If `read_from_partition`
463 is set to `True` this option is ignored.
464 time: Time stamp associated with mesh
465 read_from_partition: Read mesh with partition from file
466 backend_args: List of arguments to reader backend
467 max_facet_to_cell_links: Maximum number of cells a facet
468 can be connected to.
469 mesh_name: Name of the mesh in the file to act on. Defaults to the
470 mesh written without an explicit name, so existing checkpoints and
471 call sites are unaffected.
472 Returns:
473 The distributed mesh
474 """
475 backend_args = _with_mesh_name(backend_args, mesh_name)
476 logger.debug(f"Reading mesh from {filename}")
477 logger.debug(f"Using {backend} backend with arguments {backend_args}")
478 logger.debug(f"Time {time} and read_from_partition {read_from_partition}")
479 # Read in data in a distributed fashin
480 check_file_exists(filename)
481 backend_cls = get_backend(backend)
482 backend_args = backend_cls.get_default_backend_args(backend_args)
484 # Let each backend handle what should be default behavior when reading mesh
485 # with or without time stamp.
486 dist_in_data = backend_cls.read_mesh_data(
487 filename,
488 comm,
489 time=time,
490 read_from_partition=read_from_partition,
491 backend_args=backend_args,
492 )
494 # Create DOLFINx mesh
495 element = basix.ufl.element(
496 basix.ElementFamily.P,
497 dist_in_data.cell_type,
498 dist_in_data.degree,
499 basix.LagrangeVariant(int(dist_in_data.lvar)),
500 shape=(dist_in_data.x.shape[1],),
501 dtype=dist_in_data.x.dtype,
502 )
503 domain = ufl.Mesh(element)
504 # `dolfinx.mesh.PartitioningFunc` is not available in all supported
505 # versions of DOLFINx, and the accepted callback signature varies (see
506 # below), so type the callback permissively.
507 partitioner: Callable[..., dolfinx.cpp.graph.AdjacencyList_int32] | None = None
508 if (partition_graph := dist_in_data.partition_graph) is not None:
509 # The arguments DOLFINx passes to a partitioner callback have changed
510 # over time: 0.11 calls it with
511 # ``(comm, nparts, cell_types, local_graph)``, while newer versions call
512 # it with ``(comm, nparts, dual_graph, cell_weights, edge_weights,
513 # ghosting)`` (see https://github.com/FEniCS/dolfinx/pull/4403).
514 # We read the partitioning from file, so the returned graph does not
515 # depend on any of these arguments. Accept them all and stay agnostic to
516 # the calling convention rather than branching on `dolfinx.__version__`,
517 # which does not distinguish pre-releases and post-releases reliably.
518 def _custom_partitioner(*args: Any, **kwargs: Any) -> dolfinx.cpp.graph.AdjacencyList_int32:
519 if hasattr(partition_graph, "_cpp_object"):
520 cpp_obj = partition_graph._cpp_object
521 assert isinstance(cpp_obj, dolfinx.cpp.graph.AdjacencyList_int32)
522 return cpp_obj
523 else:
524 assert isinstance(partition_graph, dolfinx.cpp.graph.AdjacencyList_int32)
525 return partition_graph
527 partitioner = _custom_partitioner
529 return fenicsx_compat.create_mesh(
530 comm,
531 dist_in_data.cells,
532 domain,
533 dist_in_data.x,
534 partitioner=partitioner,
535 ghost_mode=ghost_mode,
536 max_facet_to_cell_links=max_facet_to_cell_links,
537 )
540def write_mesh(
541 filename: Path,
542 mesh: dolfinx.mesh.Mesh,
543 mode: FileMode = FileMode.write,
544 time: float = 0.0,
545 store_partition_info: bool = False,
546 backend_args: dict[str, Any] | None = None,
547 backend: str | None = None,
548 mesh_name: str | None = None,
549):
550 """
551 Write a mesh to file.
553 Args:
554 filename: Path to save mesh (without file-extension)
555 mesh: The mesh to write to file
557 store_partition_info: Store mesh partitioning (including ghosting) to file
558 mesh_name: Name of the mesh in the file to act on. Defaults to the
559 mesh written without an explicit name, so existing checkpoints and
560 call sites are unaffected.
561 """
562 backend_args = _with_mesh_name(backend_args, mesh_name)
563 logger.debug(f"Writing mesh to {filename}")
564 logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}")
565 mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info)
566 logger.debug(f"Write mesh using {backend} backend, with arguments {backend_args}")
567 logger.debug(f"Mode {mode} and time {time}")
568 _internal_mesh_writer(
569 filename,
570 mesh.comm,
571 mesh_data=mesh_data,
572 time=time,
573 backend_args=backend_args,
574 backend=backend,
575 mode=mode,
576 )
579def write_function(
580 filename: Path | str,
581 u: dolfinx.fem.Function,
582 time: float = 0.0,
583 mode: FileMode = FileMode.append,
584 name: str | None = None,
585 backend_args: dict[str, Any] | None = None,
586 backend: str | None = None,
587 mesh_name: str | None = None,
588):
589 """
590 Write function checkpoint to file.
592 Args:
593 u: Function to write to file
594 time: Time-stamp for simulation
595 filename: Path to write to
596 mode: Write or append.
597 name: Name of function to write. If None, the name of the function is used.
598 backend_args: Arguments to the IO backend.
599 backend: The backend to use
600 mesh_name: Name of the mesh in the file to act on. Defaults to the
601 mesh written without an explicit name, so existing checkpoints and
602 call sites are unaffected.
603 """
604 backend_args = _with_mesh_name(backend_args, mesh_name)
605 n = u.name if name is None else name
606 logger.debug(f"Writing function checkpoint to {filename} for function {n} at time {time}")
607 logger.debug(f"Using {backend} backend with arguments {backend_args}")
608 dofmap = u.function_space.dofmap
609 values = u.x.array
610 mesh = u.function_space.mesh
611 comm = mesh.comm
612 cell_perm = fenicsx_compat.cell_permutation_info(mesh.topology)
613 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
614 local_cell_range = mesh.topology.index_map(mesh.topology.dim).local_range
615 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
617 # Convert local dofmap into global_dofmap
618 dmap = dofmap.list
619 num_dofs_per_cell = dmap.shape[1]
620 dofmap_bs = dofmap.bs
621 num_dofs_local_dmap = num_cells_local * num_dofs_per_cell * dofmap_bs
622 index_map_bs = dofmap.index_map_bs
624 # Unroll dofmap for block size
625 unrolled_dofmap = unroll_dofmap(dofmap.list[:num_cells_local, :], dofmap_bs)
626 dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1)
627 dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1)
629 # Convert imap index to global index
630 imap_global = dofmap.index_map.local_to_global(dmap_loc)
631 dofmap_global = imap_global * index_map_bs + dmap_rem
632 dofmap_imap = fenicsx_compat.create_index_map(
633 mesh.comm,
634 num_dofs_local_dmap,
635 np.empty(0, dtype=np.int64),
636 np.empty(0, dtype=np.int32),
637 )
639 # Compute dofmap offsets
640 local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64)
641 local_dofmap_offsets[:] *= num_dofs_per_cell * dofmap_bs
642 local_dofmap_offsets += dofmap_imap.local_range[0]
644 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs
645 local_dof_range = np.asarray(dofmap.index_map.local_range) * dofmap.index_map_bs
646 num_dofs_local = local_dof_range[1] - local_dof_range[0]
648 # Create internal data structure for function data to write to file
649 function_data = FunctionData(
650 cell_permutations=cell_perm[:num_cells_local].copy(),
651 local_cell_range=local_cell_range,
652 num_cells_global=num_cells_global,
653 dofmap_array=dofmap_global,
654 dofmap_offsets=local_dofmap_offsets,
655 dofmap_range=dofmap_imap.local_range,
656 global_dofs_in_dofmap=dofmap_imap.size_global,
657 values=values[:num_dofs_local].copy(),
658 dof_range=(local_dof_range[0], local_dof_range[1]),
659 num_dofs_global=num_dofs_global,
660 name=name or u.name,
661 )
662 # Write to file
663 fname = Path(filename)
664 _internal_function_writer(
665 fname, comm, function_data, time, backend_args=backend_args, backend=backend, mode=mode
666 )
669def read_function_names(
670 filename: Path | str,
671 comm: MPI.Comm,
672 backend_args: dict[str, Any] | None = None,
673 backend: str = "h5py",
674 mesh_name: str | None = None,
675) -> list[str]:
676 """Read all function names from a file.
678 Args:
679 filename: Path to file
680 comm: MPI communicator to launch IO on.
681 backend_args: Arguments to backend
682 mesh_name: Name of the mesh in the file to act on. Defaults to the
683 mesh written without an explicit name, so existing checkpoints and
684 call sites are unaffected.
686 Returns:
687 A list of function names.
688 """
689 backend_args = _with_mesh_name(backend_args, mesh_name)
690 logger.debug(f"Reading function names from {filename}")
691 logger.debug(f"Using {backend} backend with arguments {backend_args} to read function names")
692 check_file_exists(filename)
693 backend_cls = get_backend(backend)
694 return backend_cls.read_function_names(filename, comm, backend_args=backend_args)
697def write_point_data(
698 filename: Path | str,
699 u: dolfinx.fem.Function,
700 time: str | float | None,
701 mode: FileMode,
702 backend_args: dict[str, Any] | None,
703 backend: str = "vtkhdf",
704 mesh_name: str | None = None,
705):
706 """Write function to file by interpolating into geometry nodes.
709 Args:
710 filename: Path to file
711 u: The function to store
712 time: Time stamp
713 mode: Append or write
714 backend_args: The backend arguments
715 backend: Which backend to use.
716 mesh_name: Name of the mesh in the file to write the data against.
717 Defaults to the mesh written without an explicit name.
718 """
719 backend_args = _with_mesh_name(backend_args, mesh_name)
720 logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}")
721 V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape)))
722 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
723 v_out.interpolate(u)
724 comm = v_out.function_space.mesh.comm
725 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
726 local_range = V.dofmap.index_map.local_range
727 num_dofs_local = V.dofmap.index_map.size_local
728 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
729 ad = ArrayData(
730 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Point"
731 )
732 logger.debug(
733 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write point data"
734 )
735 backend_cls = get_backend(backend)
736 return backend_cls.write_data(
737 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
738 )
741def write_cell_data(
742 filename: Path | str,
743 u: dolfinx.fem.Function,
744 time: str | float | None,
745 mode: FileMode,
746 backend_args: dict[str, Any] | None,
747 backend: str = "vtkhdf",
748 mesh_name: str | None = None,
749):
750 """Write function to file by interpolating into cell midpoints.
753 Args:
754 filename: Path to file
755 point_data: Data to write to file
756 time: Time stamp
757 mode: Append or write
758 backend_args: The backend arguments
759 mesh_name: Name of the mesh in the file to write the data against.
760 Defaults to the mesh written without an explicit name.
761 """
762 backend_args = _with_mesh_name(backend_args, mesh_name)
763 logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}")
764 V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape))
765 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
766 v_out.interpolate(u)
767 comm = v_out.function_space.mesh.comm
768 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
769 local_range = V.dofmap.index_map.local_range
770 num_dofs_local = V.dofmap.index_map.size_local
771 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
773 ad = ArrayData(
774 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Cell"
775 )
776 logger.debug(
777 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write cell data"
778 )
779 backend_cls = get_backend(backend)
781 return backend_cls.write_data(
782 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
783 )