Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/checkpointing.py: 97%
263 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-23 13:26 +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 indices = data.indices
279 if indices.shape[0] == 0:
280 # A backend without entities may not know the number of geometry nodes per entity
281 layout = fenicsx_compat.cmap(mesh).create_dof_layout()
282 num_nodes = len(layout.entity_closure_dofs(int(data.dim), 0))
283 indices = np.zeros((0, num_nodes), dtype=np.int64)
284 local_entities, local_values = dolfinx.io.distribute_entity_data(
285 mesh, int(data.dim), indices, data.values
286 )
287 mesh.topology.create_connectivity(data.dim, 0)
288 mesh.topology.create_connectivity(data.dim, mesh.topology.dim)
290 adj = dolfinx.graph.adjacencylist(local_entities)
292 local_values = np.array(local_values, dtype=np.int32)
294 mt = dolfinx.mesh.meshtags_from_entities(mesh, int(data.dim), adj, local_values)
295 mt.name = meshtag_name
296 return mt
299def read_function(
300 filename: Path | str,
301 u: dolfinx.fem.Function,
302 time: float = 0.0,
303 name: str | None = None,
304 backend_args: dict[str, Any] | None = None,
305 backend: str | None = None,
306 mesh_name: str | None = None,
307):
308 """
309 Read checkpoint from file and fill it into `u`.
311 Args:
312 filename: Path to checkpoint
313 u: Function to fill
314 time: Time-stamp associated with checkpoint
315 name: If not provided, `u.name` is used to search through the input file for the function
316 mesh_name: Name of the mesh in the file to act on. Defaults to the
317 mesh written without an explicit name, so existing checkpoints and
318 call sites are unaffected.
319 """
320 backend_args = _with_mesh_name(backend_args, mesh_name)
321 logger.debug(
322 f"Reading function checkpoint from {filename} for function {name or u.name} at time {time}"
323 )
324 logger.debug(
325 f"Using {backend} backend with arguments {backend_args} to read function checkpoint"
326 )
327 check_file_exists(filename)
329 mesh = u.function_space.mesh
330 comm = mesh.comm
331 if name is None:
332 name = u.name
334 check_file_exists(filename)
335 backend_cls = get_backend(backend)
336 backend_args = backend_cls.get_default_backend_args(backend_args)
338 # Compute index of input cells and get cell permutation
339 num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
340 input_cells = mesh.topology.original_cell_index[:num_owned_cells]
341 cell_perm = fenicsx_compat.cell_permutation_info(mesh.topology)[:num_owned_cells]
343 # Compute mesh->input communicator
344 # 1.1 Compute mesh->input communicator
345 owners: npt.NDArray[np.int32]
346 if backend_cls.read_mode == ReadMode.serial:
347 owners = np.zeros(input_cells, dtype=np.int32)
348 elif backend_cls.read_mode == ReadMode.parallel:
349 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
350 owners = index_owner(mesh.comm, input_cells, num_cells_global)
351 else:
352 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
353 # -------------------Step 2------------------------------------
354 # Send and receive global cell index and cell perm
355 inc_cells, inc_perms = send_and_recv_cell_perm(input_cells, cell_perm, owners, mesh.comm)
357 input_dofmap = backend_cls.read_dofmap(filename, comm, name, backend_args)
359 # Compute owner of dofs in dofmap
360 dof_owner: npt.NDArray[np.int32]
361 if backend_cls.read_mode == ReadMode.serial:
362 dof_owner = np.zeros(len(input_dofmap.array), dtype=np.int32)
363 elif backend_cls.read_mode == ReadMode.parallel:
364 num_dofs_global = (
365 u.function_space.dofmap.index_map.size_global * u.function_space.dofmap.index_map_bs
366 )
367 dof_owner = index_owner(comm, input_dofmap.array.astype(np.int64), num_dofs_global)
368 else:
369 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
371 # --------------------Step 4-----------------------------------
372 # Read array from file and communicate them to input dofmap process
373 input_array, starting_pos = backend_cls.read_dofs(filename, comm, name, time, backend_args)
375 recv_array = send_dofs_and_recv_values(
376 input_dofmap.array.astype(np.int64), dof_owner, comm, input_array, starting_pos
377 )
379 # -------------------Step 5--------------------------------------
380 # Invert permutation of input data based on input perm
381 # Then apply current permutation to the local data
382 element = u.function_space.element
383 if element.needs_dof_transformations:
384 bs = u.function_space.dofmap.bs
386 # Read input cell permutations on dofmap process
387 local_input_range = compute_local_range(comm, num_cells_global)
388 input_local_cell_index = inc_cells - local_input_range[0]
389 input_perms = backend_cls.read_cell_perms(comm, filename, backend_args)
391 # Start by sorting data array by cell permutation
392 num_dofs_per_cell = input_dofmap.offsets[1:] - input_dofmap.offsets[:-1]
393 assert np.allclose(num_dofs_per_cell, num_dofs_per_cell[0])
395 # Sort dofmap by input local cell index
396 input_perms_sorted = input_perms[input_local_cell_index]
397 unrolled_dofmap_position = unroll_insert_position(
398 input_local_cell_index, num_dofs_per_cell[0]
399 )
400 dofmap_sorted_by_input = recv_array[unrolled_dofmap_position]
402 # First invert input data to reference element then transform to current mesh
403 element.Tt_apply(dofmap_sorted_by_input, input_perms_sorted, bs)
404 element.Tt_inv_apply(dofmap_sorted_by_input, inc_perms, bs)
405 # Compute invert permutation
406 inverted_perm = np.empty_like(unrolled_dofmap_position)
407 inverted_perm[unrolled_dofmap_position] = np.arange(
408 len(unrolled_dofmap_position), dtype=inverted_perm.dtype
409 )
410 recv_array = dofmap_sorted_by_input[inverted_perm]
412 # ------------------Step 6----------------------------------------
413 # For each dof owned by a process, find the local position in the dofmap.
414 V = u.function_space
415 local_cells, dof_pos = compute_dofmap_pos(V)
416 input_cells = V.mesh.topology.original_cell_index[local_cells]
417 num_cells_global = V.mesh.topology.index_map(V.mesh.topology.dim).size_global
419 if backend_cls.read_mode == ReadMode.serial:
420 owners = np.zeros(len(input_cells), dtype=np.int32)
421 elif backend_cls.read_mode == ReadMode.parallel:
422 owners = index_owner(V.mesh.comm, input_cells, num_cells_global)
423 else:
424 raise NotImplementedError(f"{backend_cls.read_mode} not implemented")
426 unique_owners, owner_count = np.unique(owners, return_counts=True)
427 # FIXME: In C++ use NBX to find neighbourhood
428 assert isinstance(V.mesh.comm, MPI.Intracomm)
429 sub_comm = V.mesh.comm.Create_dist_graph(
430 [V.mesh.comm.rank], [len(unique_owners)], unique_owners.tolist(), reorder=False
431 )
432 source, dest, _ = sub_comm.Get_dist_neighbors()
433 sub_comm.Free()
435 owned_values = send_dofmap_and_recv_values(
436 comm,
437 np.asarray(source, dtype=np.int32),
438 np.asarray(dest, dtype=np.int32),
439 owners,
440 owner_count.astype(np.int32),
441 input_cells,
442 dof_pos,
443 num_cells_global,
444 recv_array,
445 input_dofmap.offsets,
446 )
447 u.x.array[: len(owned_values)] = owned_values
448 u.x.scatter_forward()
451def read_mesh(
452 filename: Path | str,
453 comm: MPI.Comm,
454 ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
455 time: float | str | None = 0.0,
456 read_from_partition: bool = False,
457 backend_args: dict[str, Any] | None = None,
458 backend: str | None = None,
459 max_facet_to_cell_links: int = 2,
460 mesh_name: str | None = None,
461) -> dolfinx.mesh.Mesh:
462 """
463 Read an ADIOS2 mesh into DOLFINx.
465 Args:
466 filename: Path to input file
467 comm: The MPI communciator to distribute the mesh over
468 ghost_mode: Ghost mode to use for mesh. If `read_from_partition`
469 is set to `True` this option is ignored.
470 time: Time stamp associated with mesh
471 read_from_partition: Read mesh with partition from file
472 backend_args: List of arguments to reader backend
473 max_facet_to_cell_links: Maximum number of cells a facet
474 can be connected to.
475 mesh_name: Name of the mesh in the file to act on. Defaults to the
476 mesh written without an explicit name, so existing checkpoints and
477 call sites are unaffected.
478 Returns:
479 The distributed mesh
480 """
481 backend_args = _with_mesh_name(backend_args, mesh_name)
482 logger.debug(f"Reading mesh from {filename}")
483 logger.debug(f"Using {backend} backend with arguments {backend_args}")
484 logger.debug(f"Time {time} and read_from_partition {read_from_partition}")
485 # Read in data in a distributed fashin
486 check_file_exists(filename)
487 backend_cls = get_backend(backend)
488 backend_args = backend_cls.get_default_backend_args(backend_args)
490 # Let each backend handle what should be default behavior when reading mesh
491 # with or without time stamp.
492 dist_in_data = backend_cls.read_mesh_data(
493 filename,
494 comm,
495 time=time,
496 read_from_partition=read_from_partition,
497 backend_args=backend_args,
498 )
500 # Create DOLFINx mesh
501 element: basix.ufl._ElementBase
502 if dist_in_data.serendipity:
503 try:
504 from dolfinx.io.utils import _serendipity_coordinate_element
505 except ImportError as e:
506 raise NotImplementedError(
507 "Meshes with serendipity geometry (8-node quadrilaterals, 20-node "
508 "hexahedra) require the main branch of DOLFINx."
509 ) from e
510 element = basix.ufl.blocked_element(
511 _serendipity_coordinate_element(
512 dolfinx.mesh.to_type(dist_in_data.cell_type), dist_in_data.x.dtype
513 ),
514 shape=(dist_in_data.x.shape[1],),
515 )
516 else:
517 element = basix.ufl.element(
518 basix.ElementFamily.P,
519 dist_in_data.cell_type,
520 dist_in_data.degree,
521 basix.LagrangeVariant(int(dist_in_data.lvar)),
522 shape=(dist_in_data.x.shape[1],),
523 dtype=dist_in_data.x.dtype,
524 )
525 domain = ufl.Mesh(element)
526 # `dolfinx.mesh.PartitioningFunc` is not available in all supported
527 # versions of DOLFINx, and the accepted callback signature varies (see
528 # below), so type the callback permissively.
529 partitioner: Callable[..., dolfinx.cpp.graph.AdjacencyList_int32] | None = None
530 if (partition_graph := dist_in_data.partition_graph) is not None:
531 # The arguments DOLFINx passes to a partitioner callback have changed
532 # over time: 0.11 calls it with
533 # ``(comm, nparts, cell_types, local_graph)``, while newer versions call
534 # it with ``(comm, nparts, dual_graph, cell_weights, edge_weights,
535 # ghosting)`` (see https://github.com/FEniCS/dolfinx/pull/4403).
536 # We read the partitioning from file, so the returned graph does not
537 # depend on any of these arguments. Accept them all and stay agnostic to
538 # the calling convention rather than branching on `dolfinx.__version__`,
539 # which does not distinguish pre-releases and post-releases reliably.
540 def _custom_partitioner(*args: Any, **kwargs: Any) -> dolfinx.cpp.graph.AdjacencyList_int32:
541 if hasattr(partition_graph, "_cpp_object"):
542 cpp_obj = partition_graph._cpp_object
543 assert isinstance(cpp_obj, dolfinx.cpp.graph.AdjacencyList_int32)
544 return cpp_obj
545 else:
546 assert isinstance(partition_graph, dolfinx.cpp.graph.AdjacencyList_int32)
547 return partition_graph
549 partitioner = _custom_partitioner
551 return fenicsx_compat.create_mesh(
552 comm,
553 dist_in_data.cells,
554 domain,
555 dist_in_data.x,
556 partitioner=partitioner,
557 ghost_mode=ghost_mode,
558 max_facet_to_cell_links=max_facet_to_cell_links,
559 )
562def write_mesh(
563 filename: Path,
564 mesh: dolfinx.mesh.Mesh,
565 mode: FileMode = FileMode.write,
566 time: float = 0.0,
567 store_partition_info: bool = False,
568 backend_args: dict[str, Any] | None = None,
569 backend: str | None = None,
570 mesh_name: str | None = None,
571):
572 """
573 Write a mesh to file.
575 Args:
576 filename: Path to save mesh (without file-extension)
577 mesh: The mesh to write to file
579 store_partition_info: Store mesh partitioning (including ghosting) to file
580 mesh_name: Name of the mesh in the file to act on. Defaults to the
581 mesh written without an explicit name, so existing checkpoints and
582 call sites are unaffected.
583 """
584 backend_args = _with_mesh_name(backend_args, mesh_name)
585 logger.debug(f"Writing mesh to {filename}")
586 logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}")
587 mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info)
588 logger.debug(f"Write mesh using {backend} backend, with arguments {backend_args}")
589 logger.debug(f"Mode {mode} and time {time}")
590 _internal_mesh_writer(
591 filename,
592 mesh.comm,
593 mesh_data=mesh_data,
594 time=time,
595 backend_args=backend_args,
596 backend=backend,
597 mode=mode,
598 )
601def write_function(
602 filename: Path | str,
603 u: dolfinx.fem.Function,
604 time: float = 0.0,
605 mode: FileMode = FileMode.append,
606 name: str | None = None,
607 backend_args: dict[str, Any] | None = None,
608 backend: str | None = None,
609 mesh_name: str | None = None,
610):
611 """
612 Write function checkpoint to file.
614 Args:
615 u: Function to write to file
616 time: Time-stamp for simulation
617 filename: Path to write to
618 mode: Write or append.
619 name: Name of function to write. If None, the name of the function is used.
620 backend_args: Arguments to the IO backend.
621 backend: The backend to use
622 mesh_name: Name of the mesh in the file to act on. Defaults to the
623 mesh written without an explicit name, so existing checkpoints and
624 call sites are unaffected.
625 """
626 backend_args = _with_mesh_name(backend_args, mesh_name)
627 n = u.name if name is None else name
628 logger.debug(f"Writing function checkpoint to {filename} for function {n} at time {time}")
629 logger.debug(f"Using {backend} backend with arguments {backend_args}")
630 dofmap = u.function_space.dofmap
631 values = u.x.array
632 mesh = u.function_space.mesh
633 comm = mesh.comm
634 cell_perm = fenicsx_compat.cell_permutation_info(mesh.topology)
635 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
636 local_cell_range = mesh.topology.index_map(mesh.topology.dim).local_range
637 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
639 # Convert local dofmap into global_dofmap
640 dmap = dofmap.list
641 num_dofs_per_cell = dmap.shape[1]
642 dofmap_bs = dofmap.bs
643 num_dofs_local_dmap = num_cells_local * num_dofs_per_cell * dofmap_bs
644 index_map_bs = dofmap.index_map_bs
646 # Unroll dofmap for block size
647 unrolled_dofmap = unroll_dofmap(dofmap.list[:num_cells_local, :], dofmap_bs)
648 dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1)
649 dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1)
651 # Convert imap index to global index
652 imap_global = dofmap.index_map.local_to_global(dmap_loc)
653 dofmap_global = imap_global * index_map_bs + dmap_rem
654 dofmap_imap = fenicsx_compat.create_index_map(
655 mesh.comm,
656 num_dofs_local_dmap,
657 np.empty(0, dtype=np.int64),
658 np.empty(0, dtype=np.int32),
659 )
661 # Compute dofmap offsets
662 local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64)
663 local_dofmap_offsets[:] *= num_dofs_per_cell * dofmap_bs
664 local_dofmap_offsets += dofmap_imap.local_range[0]
666 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs
667 local_dof_range = np.asarray(dofmap.index_map.local_range) * dofmap.index_map_bs
668 num_dofs_local = local_dof_range[1] - local_dof_range[0]
670 # Create internal data structure for function data to write to file
671 function_data = FunctionData(
672 cell_permutations=cell_perm[:num_cells_local].copy(),
673 local_cell_range=local_cell_range,
674 num_cells_global=num_cells_global,
675 dofmap_array=dofmap_global,
676 dofmap_offsets=local_dofmap_offsets,
677 dofmap_range=dofmap_imap.local_range,
678 global_dofs_in_dofmap=dofmap_imap.size_global,
679 values=values[:num_dofs_local].copy(),
680 dof_range=(local_dof_range[0], local_dof_range[1]),
681 num_dofs_global=num_dofs_global,
682 name=name or u.name,
683 )
684 # Write to file
685 fname = Path(filename)
686 _internal_function_writer(
687 fname, comm, function_data, time, backend_args=backend_args, backend=backend, mode=mode
688 )
691def read_function_names(
692 filename: Path | str,
693 comm: MPI.Comm,
694 backend_args: dict[str, Any] | None = None,
695 backend: str = "h5py",
696 mesh_name: str | None = None,
697) -> list[str]:
698 """Read all function names from a file.
700 Args:
701 filename: Path to file
702 comm: MPI communicator to launch IO on.
703 backend_args: Arguments to backend
704 mesh_name: Name of the mesh in the file to act on. Defaults to the
705 mesh written without an explicit name, so existing checkpoints and
706 call sites are unaffected.
708 Returns:
709 A list of function names.
710 """
711 backend_args = _with_mesh_name(backend_args, mesh_name)
712 logger.debug(f"Reading function names from {filename}")
713 logger.debug(f"Using {backend} backend with arguments {backend_args} to read function names")
714 check_file_exists(filename)
715 backend_cls = get_backend(backend)
716 return backend_cls.read_function_names(filename, comm, backend_args=backend_args)
719def write_point_data(
720 filename: Path | str,
721 u: dolfinx.fem.Function,
722 time: str | float | None,
723 mode: FileMode,
724 backend_args: dict[str, Any] | None,
725 backend: str = "vtkhdf",
726 mesh_name: str | None = None,
727):
728 """Write function to file by interpolating into geometry nodes.
731 Args:
732 filename: Path to file
733 u: The function to store
734 time: Time stamp
735 mode: Append or write
736 backend_args: The backend arguments
737 backend: Which backend to use.
738 mesh_name: Name of the mesh in the file to write the data against.
739 Defaults to the mesh written without an explicit name.
740 """
741 backend_args = _with_mesh_name(backend_args, mesh_name)
742 logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}")
743 V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape)))
744 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
745 v_out.interpolate(u)
746 comm = v_out.function_space.mesh.comm
747 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
748 local_range = V.dofmap.index_map.local_range
749 num_dofs_local = V.dofmap.index_map.size_local
750 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
751 ad = ArrayData(
752 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Point"
753 )
754 logger.debug(
755 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write point data"
756 )
757 backend_cls = get_backend(backend)
758 return backend_cls.write_data(
759 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
760 )
763def write_cell_data(
764 filename: Path | str,
765 u: dolfinx.fem.Function,
766 time: str | float | None,
767 mode: FileMode,
768 backend_args: dict[str, Any] | None,
769 backend: str = "vtkhdf",
770 mesh_name: str | None = None,
771):
772 """Write function to file by interpolating into cell midpoints.
775 Args:
776 filename: Path to file
777 point_data: Data to write to file
778 time: Time stamp
779 mode: Append or write
780 backend_args: The backend arguments
781 mesh_name: Name of the mesh in the file to write the data against.
782 Defaults to the mesh written without an explicit name.
783 """
784 backend_args = _with_mesh_name(backend_args, mesh_name)
785 logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}")
786 V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape))
787 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
788 v_out.interpolate(u)
789 comm = v_out.function_space.mesh.comm
790 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
791 local_range = V.dofmap.index_map.local_range
792 num_dofs_local = V.dofmap.index_map.size_local
793 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
795 ad = ArrayData(
796 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Cell"
797 )
798 logger.debug(
799 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write cell data"
800 )
801 backend_cls = get_backend(backend)
803 return backend_cls.write_data(
804 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
805 )