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 .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 {py: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 = mesh.geometry.dofmap 

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 = mesh.geometry.cmap 

44 geom_layout = cmap.create_dof_layout() 

45 if hasattr(geom_layout, "num_entity_closure_dofs"): 

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

47 else: 

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

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

50 assert g_dmap.shape[1] == num_dofs_per_cell 

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

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

53 ).reshape(dofs_out.shape) 

54 

55 if store_partition_info: 

56 partition_processes = mesh.comm.size 

57 

58 # Get partitioning 

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

60 consensus_tag = 1202 

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

62 else: 

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

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

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

66 if cell_offsets[-1] == 0: 

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

68 else: 

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

70 

71 # Compute adjacency with current process as first entry 

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

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

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

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

76 ownership_array[insert_position] = cell_array 

77 

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

79 ownership_offset += partition_map.local_range[0] 

80 partition_range = partition_map.local_range 

81 partition_global = partition_map.size_global 

82 else: 

83 partition_processes = None 

84 ownership_array = None 

85 ownership_offset = None 

86 partition_range = None 

87 partition_global = None 

88 

89 return MeshData( 

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

91 local_geometry_pos=geometry_range, 

92 num_nodes_global=num_xdofs_global, 

93 local_topology=dofs_out, 

94 local_topology_pos=cell_range, 

95 num_cells_global=num_cells_global, 

96 cell_type=mesh.topology.cell_name(), 

97 degree=mesh.geometry.cmap.degree, 

98 lagrange_variant=mesh.geometry.cmap.variant, 

99 store_partition=store_partition_info, 

100 partition_processes=partition_processes, 

101 ownership_array=ownership_array, 

102 ownership_offset=ownership_offset, 

103 partition_range=partition_range, 

104 partition_global=partition_global, 

105 ) 

106 

107 

108def write_mesh( 

109 filename: Path, 

110 comm: MPI.Intracomm, 

111 mesh_data: MeshData, 

112 time: float = 0.0, 

113 mode: FileMode = FileMode.write, 

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

115 backend: str = "adios2", 

116): 

117 """ 

118 Write a mesh to file using ADIOS2 

119 

120 Args: 

121 comm: MPI communicator used in storage 

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

123 filename: Path to file to write to 

124 engine: ADIOS2 engine to use 

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

126 io_name: Internal name used for the ADIOS IO object 

127 """ 

128 backend_cls = get_backend(backend) 

129 backend_args = backend_cls.get_default_backend_args(backend_args) 

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

131 

132 

133def write_function( 

134 filename: Path, 

135 comm: MPI.Intracomm, 

136 u: FunctionData, 

137 time: float = 0.0, 

138 mode: FileMode = FileMode.append, 

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

140 backend: str = "adios2", 

141): 

142 """ 

143 Write a function to file using ADIOS2 

144 

145 Args: 

146 comm: MPI communicator used in storage 

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

148 filename: Path to file to write to 

149 engine: ADIOS2 engine to use 

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

151 time: Time stamp associated with function 

152 io_name: Internal name used for the ADIOS IO object 

153 """ 

154 backend_cls = get_backend(backend) 

155 backend_args = backend_cls.get_default_backend_args(backend_args) 

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