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

56 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.mf6.utilities.grid import create_geometric_grid_info 

10from imod.typing import GridDataArray 

11 

12NOT_CONNECTED_VALUE = -999 

13 

14 

15class ExchangeCreator_Structured(ExchangeCreator): 

16 """ 

17 Creates the GroundWaterFlow to GroundWaterFlow exchange package (gwfgwf) as 

18 a function of a submodel label array and a PartitionInfo object. This file 

19 contains the cell indices of coupled cells. With coupled cells we mean cells 

20 that are adjacent but that are located in different subdomains. At the 

21 moment only structured grids are supported, for unstructured grids the 

22 geometric information is still set to default values. 

23 

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

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

26 

27 """ 

28 

29 def __init__( 

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

31 ): 

32 super().__init__(submodel_labels, partition_info) 

33 

34 @property 

35 def _coordinate_names(self): 

36 return ["row", "column"] 

37 

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

39 connected_cells_along_x = self._find_connected_cells_along_axis("x") 

40 connected_cells_along_y = self._find_connected_cells_along_axis("y") 

41 

42 return pd.merge(connected_cells_along_x, connected_cells_along_y, how="outer") 

43 

44 def _find_connected_cells_along_axis(self, axis_label: str) -> pd.DataFrame: 

45 diff1 = self._submodel_labels.diff(f"{axis_label}", label="lower") 

46 diff2 = self._submodel_labels.diff(f"{axis_label}", label="upper") 

47 

48 connected_cells_idx1 = self._global_cell_indices.where( 

49 diff1 != 0, drop=True, other=NOT_CONNECTED_VALUE 

50 ).astype(int) 

51 connected_cells_idx2 = self._global_cell_indices.where( 

52 diff2 != 0, drop=True, other=NOT_CONNECTED_VALUE 

53 ).astype(int) 

54 

55 connected_model_label1 = self._submodel_labels.where( 

56 diff1 != 0, drop=True 

57 ).astype(int) 

58 connected_model_label2 = self._submodel_labels.where( 

59 diff2 != 0, drop=True 

60 ).astype(int) 

61 

62 connected_cell_info = pd.DataFrame( 

63 { 

64 "cell_idx1": connected_cells_idx1.values.flatten(), 

65 "cell_idx2": connected_cells_idx2.values.flatten(), 

66 "cell_label1": connected_model_label1.values.flatten(), 

67 "cell_label2": connected_model_label2.values.flatten(), 

68 } 

69 ) 

70 connected_cell_info = connected_cell_info.loc[ 

71 connected_cell_info.cell_idx1 != NOT_CONNECTED_VALUE 

72 ] 

73 

74 return connected_cell_info 

75 

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

77 geometric_grid_info = create_geometric_grid_info(self._global_cell_indices) 

78 

79 cell1_df = geometric_grid_info.take(self._connected_cells["cell_idx1"]) 

80 cell2_df = geometric_grid_info.take(self._connected_cells["cell_idx2"]) 

81 

82 distance_x = np.abs(cell1_df["x"].values - cell2_df["x"].values) 

83 distance_y = np.abs(cell1_df["y"].values - cell2_df["y"].values) 

84 

85 is_x_connection = distance_x > distance_y 

86 cl1 = 0.5 * np.where( 

87 is_x_connection, cell1_df["dx"].values, cell1_df["dy"].values 

88 ) 

89 cl2 = 0.5 * np.where( 

90 is_x_connection, cell2_df["dx"].values, cell2_df["dy"].values 

91 ) 

92 hwva = np.where(is_x_connection, cell2_df["dy"].values, cell2_df["dx"].values) 

93 

94 cdist = np.where(is_x_connection, distance_x, distance_y) 

95 

96 outward_vector = np.zeros((len(is_x_connection), 2)) 

97 outward_vector[:, 0] = cell2_df["x"].values - cell1_df["x"].values 

98 outward_vector[:, 1] = cell2_df["y"].values - cell1_df["y"].values 

99 anglex = np.arctan2(outward_vector[:, 1], outward_vector[:, 0]) 

100 angledegx = np.degrees(anglex) % 360 

101 

102 geometric_information = pd.DataFrame( 

103 { 

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

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

106 "cl1": cl1, 

107 "cl2": cl2, 

108 "hwva": hwva, 

109 "angldegx": angledegx, 

110 "cdist": cdist, 

111 } 

112 ) 

113 

114 return geometric_information 

115 

116 @classmethod 

117 def _create_global_to_local_idx( 

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

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

120 global_to_local_idx = {} 

121 for submodel_partition_info in partition_info: 

122 local_cell_indices = cls._get_local_cell_indices(submodel_partition_info) 

123 

124 overlap = xr.merge( 

125 (global_cell_indices, local_cell_indices), 

126 join="inner", 

127 fill_value=np.nan, 

128 compat="override", 

129 )["idomain"] 

130 

131 model_id = submodel_partition_info.id 

132 global_to_local_idx[model_id] = pd.DataFrame( 

133 { 

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

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

136 } 

137 ) 

138 

139 return global_to_local_idx