Coverage for C:\src\imod-python\imod\mf6\multimodel\exchange_creator.py: 96%
99 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1import abc
2from typing import Dict
4import numpy as np
5import pandas as pd
6import xarray as xr
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
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
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.
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.
34 """
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()
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 )
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
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
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
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
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
97 def __init__(
98 self, submodel_labels: GridDataArray, partition_info: list[PartitionInfo]
99 ):
100 self._submodel_labels = submodel_labels
102 self._global_cell_indices = to_cell_idx(submodel_labels)
104 self._connected_cells = self._find_connected_cells()
106 self.rearrange_connected_cells()
108 self._global_to_local_mapping = (
109 self._create_global_cellidx_to_local_cellid_mapping(partition_info)
110 )
112 self._geometric_information = self._compute_geometric_information()
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
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.
129 Returns
130 -------
131 a list of GWFGWF-exchanges
133 """
134 layers = layers.to_dataframe().filter(["layer"])
136 connected_cells_with_geometric_info = pd.merge(
137 self._connected_cells, self._geometric_information
138 )
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 )
154 exchanges.append(
155 GWFGWF(
156 f"{model_name}_{model_id1}",
157 f"{model_name}_{model_id2}",
158 **connected_cells_dataset,
159 )
160 )
162 return exchanges
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 )
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 )
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 )
199 connected_cells = pd.merge(layers, connected_cells, how="cross")
201 connected_cells_dataset = self._to_xarray(connected_cells)
203 _adjust_gridblock_indexing(connected_cells_dataset)
205 return connected_cells_dataset
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"])
212 connected_cells_with_geometric_info = pd.merge(
213 self._connected_cells, self._geometric_information
214 )
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 )
239 return exchanges
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)
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 )
256 return mapping
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)
265 return to_cell_idx(local_domain)
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
273 local_cell_indices = cls._get_local_cell_indices(submodel_partition_info)
274 local_cell_id = list(np.ndindex(local_cell_indices.shape))
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 )
283 return local_cell_idx_to_id
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
291 label_decreasing = df["cell_label1"] > df["cell_label2"]
293 colnames = ["cell_idx1", "cell_idx2", "cell_label1", "cell_label2"]
294 colnames_reversed = ["cell_idx2", "cell_idx1", "cell_label2", "cell_label1"]
296 decreasing_connections = df.loc[label_decreasing, colnames].values.astype(int)
298 df.loc[label_decreasing, colnames_reversed] = decreasing_connections
300 self._connected_cells = df