Coverage for C:\src\imod-python\imod\mf6\multimodel\exchange_creator_unstructured.py: 100%

76 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1from typing import Dict, List 

2 

3import numpy as np 

4import pandas as pd 

5import xarray as xr 

6 

7from imod.mf6.multimodel.exchange_creator import ExchangeCreator 

8from imod.mf6.multimodel.modelsplitter import PartitionInfo 

9from imod.typing import GridDataArray 

10 

11 

12class ExchangeCreator_Unstructured(ExchangeCreator): 

13 """ 

14 Creates the GroundWaterFlow to GroundWaterFlow exchange package (gwfgwf) as a function of a submodel label array and a 

15 PartitionInfo object. This file contains the cell indices of coupled cells. With coupled cells we mean cells that are adjacent but 

16 that are located in different subdomains. At the moment only structured grids are supported, for unstructured grids the geometric information 

17 is still set to default values. 

18 

19 The submodel_labels array should have the same topology as the domain being partitioned. The array will be used 

20 to determine the connectivity of the submodels after the split operation has been performed. 

21 

22 """ 

23 

24 def __init__( 

25 self, submodel_labels: GridDataArray, partition_info: List[PartitionInfo] 

26 ): 

27 self._connected_cell_edge_indices = ( 

28 self._find_subdomain_connection_edge_indices(submodel_labels) 

29 ) 

30 super().__init__(submodel_labels, partition_info) 

31 

32 @property 

33 def _coordinate_names(self): 

34 return ["cell_index"] 

35 

36 def _find_connected_cells(self) -> pd.DataFrame: 

37 edge_face_connectivity = self._submodel_labels.ugrid.grid.edge_face_connectivity 

38 

39 face1 = edge_face_connectivity[self._connected_cell_edge_indices, 0] 

40 face2 = edge_face_connectivity[self._connected_cell_edge_indices, 1] 

41 

42 label_of_face1 = self._submodel_labels.values[face1] 

43 label_of_face2 = self._submodel_labels.values[face2] 

44 

45 connected_cell_info = pd.DataFrame( 

46 { 

47 "cell_idx1": face1, 

48 "cell_idx2": face2, 

49 "cell_label1": label_of_face1, 

50 "cell_label2": label_of_face2, 

51 } 

52 ) 

53 

54 return connected_cell_info 

55 

56 def _compute_geometric_information(self) -> pd.DataFrame: 

57 grid = self._submodel_labels.ugrid.grid 

58 face1, face2 = self._get_partition_sorted_connected_faces() 

59 centroid_1 = grid.centroids[face1] 

60 centroid_2 = grid.centroids[face2] 

61 

62 cdist = np.linalg.norm(centroid_2 - centroid_1, axis=1) 

63 

64 edge_coordinates = grid.edge_node_coordinates[self._connected_cell_edge_indices] 

65 

66 U = np.diff(edge_coordinates, axis=1)[:, 0] 

67 # Compute vector of first cell centroid to first edge vertex 

68 Vi = centroid_1 - edge_coordinates[:, 0] 

69 # Compute vector of second cell centroid to first edge vertex 

70 Vj = centroid_2 - edge_coordinates[:, 0] 

71 length = np.linalg.norm(U, axis=1) 

72 

73 # get the normal to the cell edge from U 

74 dx = U[:, 0] 

75 dy = U[:, 1] 

76 normal = np.array((dy[:], -dx[:]), dtype=np.float_).T 

77 

78 # If the inner product of the normal with a vector on the edge to the face centroid is positive 

79 # then the normal vector points inwards 

80 inward_vector_mask = np.sum(normal * Vi, axis=-1) > 0 

81 normal[inward_vector_mask] = -normal[inward_vector_mask] 

82 

83 angle = np.degrees(np.arctan2(normal[:, 1], normal[:, 0])) 

84 

85 df = pd.DataFrame( 

86 { 

87 "cell_idx1": self._connected_cells["cell_idx1"].values, 

88 "cell_idx2": self._connected_cells["cell_idx2"].values, 

89 "cl1": np.abs(np.cross(U, Vi)) / length, 

90 "cl2": np.abs(np.cross(U, Vj)) / length, 

91 "hwva": length, 

92 "angldegx": angle, 

93 "cdist": cdist, 

94 } 

95 ) 

96 return df 

97 

98 @classmethod 

99 def _create_global_to_local_idx( 

100 cls, partition_info: List[PartitionInfo], global_cell_indices: GridDataArray 

101 ) -> Dict[int, pd.DataFrame]: 

102 global_to_local_idx = {} 

103 for submodel_partition_info in partition_info: 

104 local_cell_indices = cls._get_local_cell_indices(submodel_partition_info) 

105 

106 global_cell_indices_partition = global_cell_indices.where( 

107 submodel_partition_info.active_domain == 1 

108 ) 

109 global_cell_indices_partition = global_cell_indices_partition.dropna( 

110 "mesh2d_nFaces", how="all" 

111 ) 

112 

113 global_cell_indices_df = global_cell_indices_partition.to_dataframe() 

114 global_cell_indices_da = xr.Dataset.from_dataframe(global_cell_indices_df) 

115 

116 overlap = xr.merge( 

117 (global_cell_indices_da, xr.DataArray(local_cell_indices)), 

118 join="inner", 

119 fill_value=np.nan, 

120 compat="override", 

121 )["idomain"] 

122 

123 model_id = submodel_partition_info.id 

124 global_to_local_idx[model_id] = pd.DataFrame( 

125 { 

126 "global_idx": overlap.values.flatten(), 

127 "local_idx": local_cell_indices.values.flatten(), 

128 } 

129 ) 

130 

131 return global_to_local_idx 

132 

133 @staticmethod 

134 def _find_subdomain_connection_edge_indices(submodel_labels): 

135 edge_face_connectivity = submodel_labels.ugrid.grid.edge_face_connectivity 

136 

137 face1 = edge_face_connectivity[:, 0] 

138 face2 = edge_face_connectivity[:, 1] 

139 

140 label_of_face1 = submodel_labels.values[face1] 

141 label_of_face2 = submodel_labels.values[face2] 

142 

143 is_internal_edge = label_of_face1 - label_of_face2 == 0 

144 is_external_boundary_edge = np.any((face1 == -1, face2 == -1), axis=0) 

145 is_inter_domain_edge = ~is_internal_edge & ~is_external_boundary_edge 

146 

147 return is_inter_domain_edge 

148 

149 def _get_partition_numbers(self, face: np.ndarray) -> np.ndarray: 

150 return self._submodel_labels.data[face] 

151 

152 def _get_partition_sorted_connected_faces(self) -> (np.ndarray, np.ndarray): 

153 grid = self._submodel_labels.ugrid.grid 

154 edge_face_connectivity = grid.edge_face_connectivity 

155 

156 unordered_face1, unordered_face2 = edge_face_connectivity[ 

157 self._connected_cell_edge_indices 

158 ].T 

159 

160 # Obtain the cellface indices on both sides of each edge. 

161 # They should be ordered by giving the cellface with the lowest partition number first. 

162 face_partition_1 = self._get_partition_numbers(unordered_face1) 

163 face_partition_2 = self._get_partition_numbers(unordered_face2) 

164 face1 = np.where( 

165 face_partition_1 > face_partition_2, unordered_face2, unordered_face1 

166 ) 

167 face2 = np.where( 

168 face_partition_1 > face_partition_2, unordered_face1, unordered_face2 

169 ) 

170 return (face1, face2)