Coverage for  / dolfinx-env / lib / python3.12 / site-packages / io4dolfinx / writers.py: 97%

1# Copyright (C) 2024-2026 Jørgen Schartum Dokken 

2# 

3# This file is part of io4dolfinx 

4# 

5# SPDX-License-Identifier: MIT 

6 

7from pathlib import Path 

8from typing import Any 

9 

10from mpi4py import MPI 

11 

12import dolfinx 

13import numpy as np 

14from packaging.version import Version 

15 

16from . import compat 

17from .backends import FileMode, get_backend 

18from .structures import FunctionData, MeshData 

19 

20 

21def prepare_meshdata_for_storage(mesh: dolfinx.mesh.Mesh, store_partition_info: bool) -> MeshData: 

22 """ 

23 Helper function for extracting the required data from a distributed 

24 {py:class}`dolfinx.mesh.Mesh`. 

25 

26 Args: 

27 mesh: The mesh 

28 store_partition_info: If one should store the partitioning info 

29 Returns: 

30 Data-container with the info that should be stored. 

31 """ 

32 

33 num_xdofs_local = mesh.geometry.index_map().size_local 

34 num_xdofs_global = mesh.geometry.index_map().size_global 

35 geometry_range = mesh.geometry.index_map().local_range 

36 gdim = mesh.geometry.dim 

37 

38 # Convert local connectivity to globa l connectivity 

39 g_imap = mesh.geometry.index_map() 

40 g_dmap = mesh.geometry.dofmap 

41 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local 

42 num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global 

43 cell_range = mesh.topology.index_map(mesh.topology.dim).local_range 

44 cmap = compat.cmap(mesh) 

45 

46 geom_layout = cmap.create_dof_layout() 

47 if hasattr(geom_layout, "num_entity_closure_dofs"): 

48 num_dofs_per_cell = geom_layout.num_entity_closure_dofs(mesh.topology.dim) 

49 else: 

50 num_dofs_per_cell = len(geom_layout.entity_closure_dofs(mesh.topology.dim, 0)) 

51 dofs_out = np.zeros((num_cells_local, num_dofs_per_cell), dtype=np.int64) 

52 assert g_dmap.shape[1] == num_dofs_per_cell 

53 dofs_out[:, :] = np.asarray( 

54 g_imap.local_to_global(g_dmap[:num_cells_local, :].reshape(-1)) 

55 ).reshape(dofs_out.shape) 

56 

57 if store_partition_info: 

58 partition_processes = mesh.comm.size 

59 

60 # Get partitioning 

61 if Version(dolfinx.__version__) > Version("0.9.0"): 

62 consensus_tag = 1202 

63 cell_map = mesh.topology.index_map(mesh.topology.dim).index_to_dest_ranks(consensus_tag) 

64 else: 

65 cell_map = mesh.topology.index_map(mesh.topology.dim).index_to_dest_ranks() 

66 num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local 

67 cell_offsets = cell_map.offsets[: num_cells_local + 1] 

68 if cell_offsets[-1] == 0: 

69 cell_array = np.empty(0, dtype=np.int32) 

70 else: 

71 cell_array = cell_map.array[: cell_offsets[-1]] 

72 

73 # Compute adjacency with current process as first entry 

74 ownership_array = np.full(num_cells_local + cell_offsets[-1], -1, dtype=np.int32) 

75 ownership_offset = cell_offsets + np.arange(len(cell_offsets), dtype=np.int32) 

76 ownership_array[ownership_offset[:-1]] = mesh.comm.rank 

77 insert_position = np.flatnonzero(ownership_array == -1) 

78 ownership_array[insert_position] = cell_array 

79 

80 partition_map = dolfinx.common.IndexMap(mesh.comm, ownership_array.size) 

81 ownership_offset += partition_map.local_range[0] 

82 partition_range = partition_map.local_range 

83 partition_global = partition_map.size_global 

84 else: 

85 partition_processes = None 

86 ownership_array = None 

87 ownership_offset = None 

88 partition_range = None 

89 partition_global = None 

90 

91 return MeshData( 

92 local_geometry=mesh.geometry.x[:num_xdofs_local, :gdim].copy(), 

93 local_geometry_pos=geometry_range, 

94 num_nodes_global=num_xdofs_global, 

95 local_topology=dofs_out, 

96 local_topology_pos=cell_range, 

97 num_cells_global=num_cells_global, 

98 cell_type=mesh.topology.cell_name(), 

99 degree=cmap.degree, 

100 lagrange_variant=cmap.variant, 

101 store_partition=store_partition_info, 

102 partition_processes=partition_processes, 

103 ownership_array=ownership_array, 

104 ownership_offset=ownership_offset, 

105 partition_range=partition_range, 

106 partition_global=partition_global, 

107 ) 

108 

109 

110def write_mesh( 

111 filename: Path, 

112 comm: MPI.Intracomm, 

113 mesh_data: MeshData, 

114 time: float = 0.0, 

115 mode: FileMode = FileMode.write, 

116 backend_args: dict[str, Any] | None = None, 

117 backend: str = "adios2", 

118): 

119 """ 

120 Write a mesh to file using ADIOS2 

121 

122 Args: 

123 comm: MPI communicator used in storage 

124 mesh: Internal data structure for the mesh data to save to file 

125 filename: Path to file to write to 

126 engine: ADIOS2 engine to use 

127 mode: ADIOS2 mode to use (write or append) 

128 io_name: Internal name used for the ADIOS IO object 

129 """ 

130 backend_cls = get_backend(backend) 

131 backend_args = backend_cls.get_default_backend_args(backend_args) 

132 backend_cls.write_mesh(filename, comm, mesh_data, backend_args, mode, time) 

133 

134 

135def write_function( 

136 filename: Path, 

137 comm: MPI.Intracomm, 

138 u: FunctionData, 

139 time: float = 0.0, 

140 mode: FileMode = FileMode.append, 

141 backend_args: dict[str, Any] | None = None, 

142 backend: str = "adios2", 

143): 

144 """ 

145 Write a function to file using ADIOS2 

146 

147 Args: 

148 comm: MPI communicator used in storage 

149 u: Internal data structure for the function data to save to file 

150 filename: Path to file to write to 

151 engine: ADIOS2 engine to use 

152 mode: ADIOS2 mode to use (write or append) 

153 time: Time stamp associated with function 

154 io_name: Internal name used for the ADIOS IO object 

155 """ 

156 backend_cls = get_backend(backend) 

157 backend_args = backend_cls.get_default_backend_args(backend_args) 

158 backend_cls.write_function(filename, comm, u=u, time=time, backend_args=backend_args, mode=mode)