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

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 = compat.dofmap(mesh) 

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() # type: ignore[call-arg] 

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

67 try: 

68 cell_offsets = cell_map.offsets[: num_cells_local + 1] # type: ignore[attr-defined] 

69 if cell_offsets[-1] == 0: 

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

71 else: 

72 cell_array = cell_map.array[: cell_offsets[-1]] # type: ignore[attr-defined] 

73 except AttributeError: 

74 cell_array, cell_offsets = cell_map 

75 cell_offsets = cell_offsets[: num_cells_local + 1] 

76 if cell_offsets[-1] == 0: 

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

78 else: 

79 cell_array = cell_array[: cell_offsets[-1]] 

80 

81 # Compute adjacency with current process as first entry 

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

83 ownership_offset = (cell_offsets + np.arange(len(cell_offsets))).astype(np.int32) 

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

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

86 ownership_array[insert_position] = cell_array 

87 

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

89 ownership_offset += partition_map.local_range[0] 

90 partition_range = partition_map.local_range 

91 partition_global = partition_map.size_global 

92 else: 

93 partition_processes = None 

94 ownership_array = None 

95 ownership_offset = None 

96 partition_range = None 

97 partition_global = None 

98 

99 return MeshData( 

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

101 local_geometry_pos=geometry_range, 

102 num_nodes_global=num_xdofs_global, 

103 local_topology=dofs_out, 

104 local_topology_pos=cell_range, 

105 num_cells_global=num_cells_global, 

106 cell_type=mesh.topology.cell_name(), 

107 degree=cmap.degree, 

108 lagrange_variant=cmap.variant, 

109 store_partition=store_partition_info, 

110 partition_processes=partition_processes, 

111 ownership_array=ownership_array, 

112 ownership_offset=ownership_offset, 

113 partition_range=partition_range, 

114 partition_global=partition_global, 

115 ) 

116 

117 

118def write_mesh( 

119 filename: Path, 

120 comm: MPI.Comm, 

121 mesh_data: MeshData, 

122 time: float = 0.0, 

123 mode: FileMode = FileMode.write, 

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

125 backend: str | None = None, 

126): 

127 """ 

128 Write a mesh to file 

129 

130 Args: 

131 filename: Path to file to write to 

132 comm: MPI communicator used in storage 

133 mesh_data: Internal data structure for the mesh data to save to file 

134 time: Time stamp associated with mesh 

135 mode: File mode to use (write or append) 

136 backend_args: Arguments for the backend 

137 backend: Backend to use 

138 """ 

139 backend_cls = get_backend(backend) 

140 backend_args = backend_cls.get_default_backend_args(backend_args) 

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

142 

143 

144def write_function( 

145 filename: Path, 

146 comm: MPI.Comm, 

147 u: FunctionData, 

148 time: float = 0.0, 

149 mode: FileMode = FileMode.append, 

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

151 backend: str | None = None, 

152): 

153 """ 

154 Write a function to file 

155 

156 Args: 

157 filename: Path to file to write to 

158 comm: MPI communicator used in storage 

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

160 time: Time stamp associated with function 

161 mode: File mode to use (write or append) 

162 backend_args: Arguments for the backend 

163 backend: Backend to use 

164 """ 

165 backend_cls = get_backend(backend) 

166 backend_args = backend_cls.get_default_backend_args(backend_args) 

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