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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1from typing import Dict, List
3import numpy as np
4import pandas as pd
5import xarray as xr
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
12NOT_CONNECTED_VALUE = -999
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.
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.
27 """
29 def __init__(
30 self, submodel_labels: GridDataArray, partition_info: List[PartitionInfo]
31 ):
32 super().__init__(submodel_labels, partition_info)
34 @property
35 def _coordinate_names(self):
36 return ["row", "column"]
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")
42 return pd.merge(connected_cells_along_x, connected_cells_along_y, how="outer")
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")
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)
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)
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 ]
74 return connected_cell_info
76 def _compute_geometric_information(self) -> pd.DataFrame:
77 geometric_grid_info = create_geometric_grid_info(self._global_cell_indices)
79 cell1_df = geometric_grid_info.take(self._connected_cells["cell_idx1"])
80 cell2_df = geometric_grid_info.take(self._connected_cells["cell_idx2"])
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)
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)
94 cdist = np.where(is_x_connection, distance_x, distance_y)
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
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 )
114 return geometric_information
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)
124 overlap = xr.merge(
125 (global_cell_indices, local_cell_indices),
126 join="inner",
127 fill_value=np.nan,
128 compat="override",
129 )["idomain"]
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 )
139 return global_to_local_idx