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

1 import abc

2 from typing import Dict

3

4 import numpy as np

5 import pandas as pd

6 import xarray as xr

7

8 from imod.mf6.gwfgwf import GWFGWF

9 from imod.mf6.gwtgwt import GWTGWT

10 from imod.mf6.multimodel.modelsplitter import PartitionInfo

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

12 from imod.typing import GridDataArray

13

14

15 def _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

24 class 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