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

1 from typing import Dict, List, Tuple

2

3 import numpy as np

4 import pandas as pd

5 import xarray as xr

6

7 from imod.mf6.multimodel.exchange_creator import ExchangeCreator

8 from imod.mf6.multimodel.modelsplitter import PartitionInfo

9 from imod.typing import GridDataArray

10

11

12 class 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) -> Tuple[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