Coverage for C:\src\imod-python\imod\mf6\multimodel\exchange_creator.py: 96%

99 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import abc 

2from typing import Dict 

3 

4import numpy as np 

5import pandas as pd 

6import xarray as xr 

7 

8from imod.mf6.gwfgwf import GWFGWF 

9from imod.mf6.gwtgwt import GWTGWT 

10from imod.mf6.multimodel.modelsplitter import PartitionInfo 

11from imod.mf6.utilities.grid import get_active_domain_slice, to_cell_idx 

12from imod.typing import GridDataArray 

13 

14 

15def _adjust_gridblock_indexing(connected_cells: xr.Dataset) -> xr.Dataset: 

16 """ 

17 adjusts the gridblock numbering from 0-based to 1-based. 

18 """ 

19 connected_cells["cell_id1"].values = connected_cells["cell_id1"].values + 1 

20 connected_cells["cell_id2"].values = connected_cells["cell_id2"].values + 1 

21 return connected_cells 

22 

23 

24class ExchangeCreator(abc.ABC): 

25 """ 

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

27 and a PartitionInfo object. This file contains the cell indices of coupled cells. With coupled cells we mean 

28 cells that are adjacent but that are located in different subdomains. At the moment only structured grids are 

29 supported, for unstructured grids the geometric information is still set to default values. 

30 

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

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

33 

34 """ 

35 

36 def _to_xarray(self, connected_cells: pd.DataFrame) -> xr.Dataset: 

37 """ 

38 converts a panda dataframe with exchange data to a xarray dataset. The 

39 dataframe must have columns called cell_id1 and cell_id2: indices of 

40 cells that are part of the exchange boundary (the subdomain boundary, on 

41 both sides of the boundary) 

42 """ 

43 dataset = connected_cells.to_xarray() 

44 

45 coord_names1 = np.char.add(self._coordinate_names, "_1") 

46 dataset["cell_id1"] = xr.DataArray( 

47 np.array(list(zip(*connected_cells["cell_id1"]))).T, 

48 dims=("index", "cell_dims1"), 

49 coords={"cell_dims1": coord_names1}, 

50 ) 

51 

52 coord_names2 = np.char.add(self._coordinate_names, "_2") 

53 dataset["cell_id2"] = xr.DataArray( 

54 np.array(list(zip(*connected_cells["cell_id2"]))).T, 

55 dims=("index", "cell_dims2"), 

56 coords={"cell_dims2": coord_names2}, 

57 ) 

58 return dataset 

59 

60 @property 

61 @abc.abstractmethod 

62 def _coordinate_names(self): 

63 """ 

64 abstract method that returns the names of the coordinates e.g. ["row", "column"] for structured grids or 

65 ["cell_index"] for usntructured grids 

66 """ 

67 raise NotImplementedError 

68 

69 @abc.abstractmethod 

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

71 """ 

72 abstract method that creates a dataframe with the cell indices and subdomain indices of cells that are on the 

73 boundary of subdomains. The cell index is a number assigned to each cell on the unpartitioned domain 

74 For structured grids it replaces the row and column index which are usually used to identify a cell. 

75 """ 

76 raise NotImplementedError 

77 

78 @abc.abstractmethod 

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

80 """ 

81 abstract method that computes the geometric information needed for the gwfgwf files such as 

82 ihc, cl1, cl2, hwva ( see modflow documentation ) 

83 """ 

84 raise NotImplementedError 

85 

86 @classmethod 

87 @abc.abstractmethod 

88 def _create_global_to_local_idx( 

89 cls, partition_info: list[PartitionInfo], global_cell_indices: GridDataArray 

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

91 """ 

92 abstract method that creates for each partition a mapping from global cell indices to local cells in that 

93 partition. 

94 """ 

95 raise NotImplementedError 

96 

97 def __init__( 

98 self, submodel_labels: GridDataArray, partition_info: list[PartitionInfo] 

99 ): 

100 self._submodel_labels = submodel_labels 

101 

102 self._global_cell_indices = to_cell_idx(submodel_labels) 

103 

104 self._connected_cells = self._find_connected_cells() 

105 

106 self.rearrange_connected_cells() 

107 

108 self._global_to_local_mapping = ( 

109 self._create_global_cellidx_to_local_cellid_mapping(partition_info) 

110 ) 

111 

112 self._geometric_information = self._compute_geometric_information() 

113 

114 def create_gwfgwf_exchanges( 

115 self, model_name: str, layers: GridDataArray 

116 ) -> list[GWFGWF]: 

117 """ 

118 Create GroundWaterFlow-GroundWaterFlow exchanges based on the submodel_labels array provided in the class 

119 constructor. The layer parameter is used to extrude the cell connection through all the layers. An exchange 

120 contains: 

121 - the model names of the connected submodel 

122 - the local cell id of the first model 

123 - the local cell id of the second model 

124 - the layer on which the connected cells reside 

125 

126 For each connection between submodels only a single exchange is created. So if an exchange of model1 to 

127 model2 is created then no exchange for model2 to model1 will be created. 

128 

129 Returns 

130 ------- 

131 a list of GWFGWF-exchanges 

132 

133 """ 

134 layers = layers.to_dataframe().filter(["layer"]) 

135 

136 connected_cells_with_geometric_info = pd.merge( 

137 self._connected_cells, self._geometric_information 

138 ) 

139 

140 exchanges = [] 

141 for ( 

142 model_id1, 

143 grouped_connected_models, 

144 ) in connected_cells_with_geometric_info.groupby("cell_label1"): 

145 for model_id2, connected_domain_pair in grouped_connected_models.groupby( 

146 "cell_label2" 

147 ): 

148 connected_cells_dataset = ( 

149 self._collect_geometric_constants_connected_cells( 

150 model_id1, model_id2, connected_domain_pair, layers 

151 ) 

152 ) 

153 

154 exchanges.append( 

155 GWFGWF( 

156 f"{model_name}_{model_id1}", 

157 f"{model_name}_{model_id2}", 

158 **connected_cells_dataset, 

159 ) 

160 ) 

161 

162 return exchanges 

163 

164 def _collect_geometric_constants_connected_cells( 

165 self, 

166 model_id1: int, 

167 model_id2: int, 

168 connected_domain_pair: pd.DataFrame, 

169 layers: GridDataArray, 

170 ) -> xr.Dataset: 

171 mapping1 = ( 

172 self._global_to_local_mapping[model_id1] 

173 .drop(columns=["local_idx"]) 

174 .rename(columns={"global_idx": "cell_idx1", "local_cell_id": "cell_id1"}) 

175 ) 

176 

177 mapping2 = ( 

178 self._global_to_local_mapping[model_id2] 

179 .drop(columns=["local_idx"]) 

180 .rename(columns={"global_idx": "cell_idx2", "local_cell_id": "cell_id2"}) 

181 ) 

182 

183 connected_cells = ( 

184 connected_domain_pair.merge(mapping1) 

185 .merge(mapping2) 

186 .filter( 

187 [ 

188 "cell_id1", 

189 "cell_id2", 

190 "cl1", 

191 "cl2", 

192 "hwva", 

193 "angldegx", 

194 "cdist", 

195 ] 

196 ) 

197 ) 

198 

199 connected_cells = pd.merge(layers, connected_cells, how="cross") 

200 

201 connected_cells_dataset = self._to_xarray(connected_cells) 

202 

203 _adjust_gridblock_indexing(connected_cells_dataset) 

204 

205 return connected_cells_dataset 

206 

207 def create_gwtgwt_exchanges( 

208 self, transport_model_name: str, flow_model_name: str, layers: GridDataArray 

209 ) -> list[GWTGWT]: 

210 layers = layers.to_dataframe().filter(["layer"]) 

211 

212 connected_cells_with_geometric_info = pd.merge( 

213 self._connected_cells, self._geometric_information 

214 ) 

215 

216 exchanges = [] 

217 for ( 

218 model_id1, 

219 grouped_connected_models, 

220 ) in connected_cells_with_geometric_info.groupby("cell_label1"): 

221 for model_id2, connected_domain_pair in grouped_connected_models.groupby( 

222 "cell_label2" 

223 ): 

224 connected_cells_dataset = ( 

225 self._collect_geometric_constants_connected_cells( 

226 model_id1, model_id2, connected_domain_pair, layers 

227 ) 

228 ) 

229 exchanges.append( 

230 GWTGWT( 

231 f"{transport_model_name}_{model_id1}", 

232 f"{transport_model_name}_{model_id2}", 

233 f"{flow_model_name}_{model_id1}", 

234 f"{flow_model_name}_{model_id2}", 

235 **connected_cells_dataset, 

236 ) 

237 ) 

238 

239 return exchanges 

240 

241 def _create_global_cellidx_to_local_cellid_mapping( 

242 self, partition_info: list[PartitionInfo] 

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

244 global_to_local_idx = self._create_global_to_local_idx( 

245 partition_info, self._global_cell_indices 

246 ) 

247 local_cell_idx_to_id = self._local_cell_idx_to_id(partition_info) 

248 

249 mapping = {} 

250 for submodel_partition_info in partition_info: 

251 model_id = submodel_partition_info.id 

252 mapping[model_id] = pd.merge( 

253 global_to_local_idx[model_id], local_cell_idx_to_id[model_id] 

254 ) 

255 

256 return mapping 

257 

258 @classmethod 

259 def _get_local_cell_indices( 

260 cls, submodel_partition_info: PartitionInfo 

261 ) -> xr.DataArray: 

262 domain_slice = get_active_domain_slice(submodel_partition_info.active_domain) 

263 local_domain = submodel_partition_info.active_domain.sel(domain_slice) 

264 

265 return to_cell_idx(local_domain) 

266 

267 @classmethod 

268 def _local_cell_idx_to_id(cls, partition_info) -> Dict[int, pd.DataFrame]: 

269 local_cell_idx_to_id = {} 

270 for submodel_partition_info in partition_info: 

271 model_id = submodel_partition_info.id 

272 

273 local_cell_indices = cls._get_local_cell_indices(submodel_partition_info) 

274 local_cell_id = list(np.ndindex(local_cell_indices.shape)) 

275 

276 local_cell_idx_to_id[model_id] = pd.DataFrame( 

277 { 

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

279 "local_cell_id": local_cell_id, 

280 } 

281 ) 

282 

283 return local_cell_idx_to_id 

284 

285 def rearrange_connected_cells(self): 

286 """ 

287 connections will be shuffled so that the lowest partition number comes first. 

288 """ 

289 df = self._connected_cells 

290 

291 label_decreasing = df["cell_label1"] > df["cell_label2"] 

292 

293 colnames = ["cell_idx1", "cell_idx2", "cell_label1", "cell_label2"] 

294 colnames_reversed = ["cell_idx2", "cell_idx1", "cell_label2", "cell_label1"] 

295 

296 decreasing_connections = df.loc[label_decreasing, colnames].values.astype(int) 

297 

298 df.loc[label_decreasing, colnames_reversed] = decreasing_connections 

299 

300 self._connected_cells = df