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

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 fenicsx_compat 

14import numpy as np 

15 

16from .backends import FileMode, get_backend 

17from .structures import FunctionData, MeshData 

18 

19 

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

21 """ 

22 Helper function for extracting the required data from a distributed 

23 :class:`dolfinx.mesh.Mesh`. 

24 

25 Args: 

26 mesh: The mesh 

27 store_partition_info: If one should store the partitioning info 

28 Returns: 

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

30 """ 

31 

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

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

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

35 gdim = mesh.geometry.dim 

36 

37 # Convert local connectivity to globa l connectivity 

38 g_imap = mesh.geometry.index_map() 

39 g_dmap = fenicsx_compat.dofmap(mesh) 

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

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

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

43 cmap = fenicsx_compat.cmap(mesh) 

44 

45 geom_layout = cmap.create_dof_layout() 

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

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

48 assert g_dmap.shape[1] == num_dofs_per_cell 

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

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

51 ).reshape(dofs_out.shape) 

52 

53 if store_partition_info: 

54 partition_processes = mesh.comm.size 

55 

56 # Get partitioning 

57 cell_imap = mesh.topology.index_map(mesh.topology.dim) 

58 consensus_tag = 1202 

59 cell_array, cell_offsets = fenicsx_compat.index_to_dest_ranks(cell_imap, consensus_tag) 

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

61 cell_offsets = cell_offsets[: num_cells_local + 1] 

62 if cell_offsets[-1] == 0: 

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

64 else: 

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

66 

67 # Compute adjacency with current process as first entry 

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

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

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

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

72 ownership_array[insert_position] = cell_array 

73 

74 partition_map = fenicsx_compat.create_index_map( 

75 mesh.comm, 

76 ownership_array.size, 

77 np.empty(0, dtype=np.int64), 

78 np.empty(0, dtype=np.int32), 

79 ) 

80 ownership_offset += partition_map.local_range[0] 

81 partition_range = partition_map.local_range 

82 partition_global = partition_map.size_global 

83 else: 

84 partition_processes = None 

85 ownership_array = None 

86 ownership_offset = None 

87 partition_range = None 

88 partition_global = None 

89 

90 return MeshData( 

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

92 local_geometry_pos=geometry_range, 

93 num_nodes_global=num_xdofs_global, 

94 local_topology=dofs_out, 

95 local_topology_pos=cell_range, 

96 num_cells_global=num_cells_global, 

97 cell_type=mesh.topology.cell_name(), 

98 degree=cmap.degree, 

99 lagrange_variant=cmap.variant, 

100 store_partition=store_partition_info, 

101 partition_processes=partition_processes, 

102 ownership_array=ownership_array, 

103 ownership_offset=ownership_offset, 

104 partition_range=partition_range, 

105 partition_global=partition_global, 

106 ) 

107 

108 

109def write_mesh( 

110 filename: Path, 

111 comm: MPI.Comm, 

112 mesh_data: MeshData, 

113 time: float = 0.0, 

114 mode: FileMode = FileMode.write, 

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

116 backend: str | None = None, 

117): 

118 """ 

119 Write a mesh to file 

120 

121 Args: 

122 filename: Path to file to write to 

123 comm: MPI communicator used in storage 

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

125 time: Time stamp associated with mesh 

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

127 backend_args: Arguments for the backend 

128 backend: Backend to use 

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.Comm, 

138 u: FunctionData, 

139 time: float = 0.0, 

140 mode: FileMode = FileMode.append, 

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

142 backend: str | None = None, 

143): 

144 """ 

145 Write a function to file 

146 

147 Args: 

148 filename: Path to file to write to 

149 comm: MPI communicator used in storage 

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

151 time: Time stamp associated with function 

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

153 backend_args: Arguments for the backend 

154 backend: Backend to use 

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)