Coverage for /dolfinx-env/lib/python3.12/site-packages/io4dolfinx/checkpointing.py: 97%
237 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-08 09:09 +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
14from mpi4py import MPI
16import basix
17import dolfinx
18import numpy as np
19import numpy.typing as npt
20import ufl
21from packaging.version import Version
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.Intracomm,
61 name: str,
62 attributes: dict[str, np.ndarray],
63 backend_args: dict[str, typing.Any] | None = None,
64 backend: str = "adios2",
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.Intracomm,
86 name: str,
87 backend_args: dict[str, typing.Any] | None = None,
88 backend: str = "adios2",
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.Intracomm,
111 function_name: str,
112 backend_args: dict[str, typing.Any] | None = None,
113 backend: str = "adios2",
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 = "adios2",
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 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 = "adios2",
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 = "adios2",
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 sub_comm = V.mesh.comm.Create_dist_graph(
381 [V.mesh.comm.rank], [len(unique_owners)], unique_owners, reorder=False
382 )
383 source, dest, _ = sub_comm.Get_dist_neighbors()
384 sub_comm.Free()
386 owned_values = send_dofmap_and_recv_values(
387 comm,
388 np.asarray(source, dtype=np.int32),
389 np.asarray(dest, dtype=np.int32),
390 owners,
391 owner_count.astype(np.int32),
392 input_cells,
393 dof_pos,
394 num_cells_global,
395 recv_array,
396 input_dofmap.offsets,
397 )
398 u.x.array[: len(owned_values)] = owned_values
399 u.x.scatter_forward()
402def read_mesh(
403 filename: Path | str,
404 comm: MPI.Intracomm,
405 ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
406 time: float | str | None = 0.0,
407 read_from_partition: bool = False,
408 backend_args: dict[str, Any] | None = None,
409 backend: str = "adios2",
410 max_facet_to_cell_links: int = 2,
411) -> dolfinx.mesh.Mesh:
412 """
413 Read an ADIOS2 mesh into DOLFINx.
415 Args:
416 filename: Path to input file
417 comm: The MPI communciator to distribute the mesh over
418 ghost_mode: Ghost mode to use for mesh. If `read_from_partition`
419 is set to `True` this option is ignored.
420 time: Time stamp associated with mesh
421 read_from_partition: Read mesh with partition from file
422 backend_args: List of arguments to reader backend
423 max_facet_to_cell_links: Maximum number of cells a facet
424 can be connected to.
425 Returns:
426 The distributed mesh
427 """
428 logger.debug(f"Reading mesh from {filename}")
429 logger.debug(
430 f"Using {backend} backend with arguments {backend_args}, "
431 f"time {time} and read_from_partition {read_from_partition}"
432 )
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 if (partition_graph := dist_in_data.partition_graph) is not None:
460 def partitioner(comm: MPI.Intracomm, n, m, topo):
461 assert len(topo[0]) % (len(partition_graph.offsets) - 1) == 0
462 if Version(dolfinx.__version__) > Version("0.9.0"):
463 return partition_graph._cpp_object
464 else:
465 return partition_graph
466 else:
467 try:
468 partitioner = dolfinx.cpp.mesh.create_cell_partitioner(
469 ghost_mode, max_facet_to_cell_links=max_facet_to_cell_links
470 )
471 except TypeError:
472 partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode)
474 # Should change to the commented code below when we require python
475 # minimum version to be >=3.12 see https://github.com/python/cpython/pull/116198
476 # import inspect
477 # sig = inspect.signature(dolfinx.mesh.create_cell_partitioner)
478 # part_kwargs = {}
479 # if "max_facet_to_cell_links" in list(sig.parameters.keys()):
480 # part_kwargs["max_facet_to_cell_links"] = max_facet_to_cell_links
481 # partitioner = dolfinx.cpp.mesh.create_cell_partitioner(ghost_mode, **part_kwargs)
483 return dolfinx.mesh.create_mesh(
484 comm,
485 cells=dist_in_data.cells,
486 x=dist_in_data.x,
487 e=domain,
488 partitioner=partitioner,
489 )
492def write_mesh(
493 filename: Path,
494 mesh: dolfinx.mesh.Mesh,
495 mode: FileMode = FileMode.write,
496 time: float = 0.0,
497 store_partition_info: bool = False,
498 backend_args: dict[str, Any] | None = None,
499 backend: str = "adios2",
500):
501 """
502 Write a mesh to file.
504 Args:
505 filename: Path to save mesh (without file-extension)
506 mesh: The mesh to write to file
508 store_partition_info: Store mesh partitioning (including ghosting) to file
509 """
510 logger.debug(f"Writing mesh to {filename}")
511 logger.debug(f"Preparing mesh data for storage storing partition info: {store_partition_info}")
512 mesh_data = prepare_meshdata_for_storage(mesh=mesh, store_partition_info=store_partition_info)
513 logger.debug(
514 f"Write mesh using {backend} backend, with arguments {backend_args}, "
515 f"mode {mode} and time {time}"
516 )
517 _internal_mesh_writer(
518 filename,
519 mesh.comm,
520 mesh_data=mesh_data,
521 time=time,
522 backend_args=backend_args,
523 backend=backend,
524 mode=mode,
525 )
528def write_function(
529 filename: Path | str,
530 u: dolfinx.fem.Function,
531 time: float = 0.0,
532 mode: FileMode = FileMode.append,
533 name: str | None = None,
534 backend_args: dict[str, Any] | None = None,
535 backend: str = "adios2",
536):
537 """
538 Write function checkpoint to file.
540 Args:
541 u: Function to write to file
542 time: Time-stamp for simulation
543 filename: Path to write to
544 mode: Write or append.
545 name: Name of function to write. If None, the name of the function is used.
546 backend_args: Arguments to the IO backend.
547 backend: The backend to use
548 """
549 logger.debug(
550 f"Writing function checkpoint to {filename} for function {name or u.name} at time {time}"
551 )
552 logger.debug(
553 f"Extracting data from function and dofmap for storage using {backend} "
554 f"backend with arguments {backend_args}"
555 )
556 dofmap = u.function_space.dofmap
557 values = u.x.array
558 mesh = u.function_space.mesh
559 comm = mesh.comm
560 mesh.topology.create_entity_permutations()
561 cell_perm = mesh.topology.get_cell_permutation_info()
562 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
563 local_cell_range = mesh.topology.index_map(mesh.topology.dim).local_range
564 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
566 # Convert local dofmap into global_dofmap
567 dmap = dofmap.list
568 num_dofs_per_cell = dmap.shape[1]
569 dofmap_bs = dofmap.bs
570 num_dofs_local_dmap = num_cells_local * num_dofs_per_cell * dofmap_bs
571 index_map_bs = dofmap.index_map_bs
573 # Unroll dofmap for block size
574 unrolled_dofmap = unroll_dofmap(dofmap.list[:num_cells_local, :], dofmap_bs)
575 dmap_loc = (unrolled_dofmap // index_map_bs).reshape(-1)
576 dmap_rem = (unrolled_dofmap % index_map_bs).reshape(-1)
578 # Convert imap index to global index
579 imap_global = dofmap.index_map.local_to_global(dmap_loc)
580 dofmap_global = imap_global * index_map_bs + dmap_rem
581 dofmap_imap = dolfinx.common.IndexMap(mesh.comm, num_dofs_local_dmap)
583 # Compute dofmap offsets
584 local_dofmap_offsets = np.arange(num_cells_local + 1, dtype=np.int64)
585 local_dofmap_offsets[:] *= num_dofs_per_cell * dofmap_bs
586 local_dofmap_offsets += dofmap_imap.local_range[0]
588 num_dofs_global = dofmap.index_map.size_global * dofmap.index_map_bs
589 local_dof_range = np.asarray(dofmap.index_map.local_range) * dofmap.index_map_bs
590 num_dofs_local = local_dof_range[1] - local_dof_range[0]
592 # Create internal data structure for function data to write to file
593 function_data = FunctionData(
594 cell_permutations=cell_perm[:num_cells_local].copy(),
595 local_cell_range=local_cell_range,
596 num_cells_global=num_cells_global,
597 dofmap_array=dofmap_global,
598 dofmap_offsets=local_dofmap_offsets,
599 dofmap_range=dofmap_imap.local_range,
600 global_dofs_in_dofmap=dofmap_imap.size_global,
601 values=values[:num_dofs_local].copy(),
602 dof_range=local_dof_range,
603 num_dofs_global=num_dofs_global,
604 name=name or u.name,
605 )
606 # Write to file
607 fname = Path(filename)
608 _internal_function_writer(
609 fname, comm, function_data, time, backend_args=backend_args, backend=backend, mode=mode
610 )
613def read_function_names(
614 filename: Path | str,
615 comm: MPI.Intracomm,
616 backend_args: dict[str, Any] | None = None,
617 backend: str = "h5py",
618) -> list[str]:
619 """Read all function names from a file.
621 Args:
622 filename: Path to file
623 comm: MPI communicator to launch IO on.
624 backend_args: Arguments to backend
626 Returns:
627 A list of function names.
628 """
629 logger.debug(f"Reading function names from {filename}")
630 logger.debug(f"Using {backend} backend with arguments {backend_args} to read function names")
631 check_file_exists(filename)
632 backend_cls = get_backend(backend)
633 return backend_cls.read_function_names(filename, comm, backend_args=backend_args)
636def write_point_data(
637 filename: Path | str,
638 u: dolfinx.fem.Function,
639 time: str | float | None,
640 mode: FileMode,
641 backend_args: dict[str, Any] | None,
642 backend: str = "vtkhdf",
643):
644 """Write function to file by interpolating into geometry nodes.
647 Args:
648 filename: Path to file
649 u: The function to store
650 time: Time stamp
651 mode: Append or write
652 backend_args: The backend arguments
653 backend: Which backend to use.
654 """
655 logger.debug(f"Writing point data to {filename} for function {u.name} at time {time}")
656 V = create_geometry_function_space(u.function_space.mesh, int(np.prod(u.ufl_shape)))
657 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
658 v_out.interpolate(u)
659 comm = v_out.function_space.mesh.comm
660 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
661 local_range = V.dofmap.index_map.local_range
662 num_dofs_local = V.dofmap.index_map.size_local
663 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
664 ad = ArrayData(
665 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Point"
666 )
667 logger.debug(
668 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write point data"
669 )
670 backend_cls = get_backend(backend)
671 return backend_cls.write_data(
672 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
673 )
676def write_cell_data(
677 filename: Path | str,
678 u: dolfinx.fem.Function,
679 time: str | float | None,
680 mode: FileMode,
681 backend_args: dict[str, Any] | None,
682 backend: str = "vtkhdf",
683):
684 """Write function to file by interpolating into cell midpoints.
687 Args:
688 filename: Path to file
689 point_data: Data to write to file
690 time: Time stamp
691 mode: Append or write
692 backend_args: The backend arguments
693 """
694 logger.debug(f"Writing cell data to {filename} for function {u.name} at time {time}")
695 V = dolfinx.fem.functionspace(u.function_space.mesh, ("DG", 0, u.ufl_shape))
696 v_out = dolfinx.fem.Function(V, name=u.name, dtype=u.x.array.dtype)
697 v_out.interpolate(u)
698 comm = v_out.function_space.mesh.comm
699 data_shape = (V.dofmap.index_map.size_global, V.dofmap.index_map_bs)
700 local_range = V.dofmap.index_map.local_range
701 num_dofs_local = V.dofmap.index_map.size_local
702 data = v_out.x.array.reshape(-1, V.dofmap.index_map_bs)[:num_dofs_local]
704 ad = ArrayData(
705 name=v_out.name, values=data, global_shape=data_shape, local_range=local_range, type="Cell"
706 )
707 logger.debug(
708 f"Using {backend} backend with arguments {backend_args} and mode {mode} to write cell data"
709 )
710 backend_cls = get_backend(backend)
712 return backend_cls.write_data(
713 filename, comm=comm, mode=mode, time=time, array_data=ad, backend_args=backend_args
714 )