Coverage for C:\src\imod-python\imod\msw\coupler_mapping.py: 29%
58 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
1import numpy as np
2import pandas as pd
3import xarray as xr
5from imod.mf6.dis import StructuredDiscretization
6from imod.mf6.wel import WellDisStructured
7from imod.msw.fixed_format import VariableMetaData
8from imod.msw.pkgbase import MetaSwapPackage
11class CouplerMapping(MetaSwapPackage):
12 """
13 This contains the data to connect MODFLOW 6 cells to MetaSWAP svats.
15 This class is responsible for the file `mod2svat.inp`. It also includes
16 connection to wells.
18 Parameters
19 ----------
20 modflow_dis: StructuredDiscretization
21 Modflow 6 structured discretization
22 well: WellDisStructured (optional)
23 If given, this parameter describes sprinkling of SVAT units from MODFLOW
24 cells.
25 """
27 _file_name = "mod2svat.inp"
28 _metadata_dict = {
29 "mod_id": VariableMetaData(10, 1, 9999999, int),
30 "free": VariableMetaData(2, None, None, str),
31 "svat": VariableMetaData(10, 1, 9999999, int),
32 "layer": VariableMetaData(5, 0, 9999, int),
33 }
35 _with_subunit = ("mod_id",)
36 _without_subunit = ()
37 _to_fill = ("free",)
39 def __init__(
40 self,
41 modflow_dis: StructuredDiscretization,
42 well: WellDisStructured = None,
43 ):
44 super().__init__()
46 self.well = well
47 # Test if equal or larger than 1, to ignore idomain == -1 as well. Don't
48 # assign to self.dataset, as grid extent might differ from svat when
49 # MetaSWAP only covers part of the Modflow grid domain.
50 self.idomain_active = modflow_dis["idomain"] >= 1
52 def _create_mod_id_rch(self, svat):
53 """
54 Create modflow indices for the recharge layer, which is where
55 infiltration will take place.
56 """
57 self.dataset["mod_id"] = xr.full_like(svat, fill_value=0, dtype=np.int64)
58 n_subunit = svat["subunit"].size
59 idomain_top_active = self.idomain_active.sel(layer=1, drop=True)
61 n_mod_top = idomain_top_active.sum()
63 # idomain does not have a subunit dimension, so tile for n_subunits
64 mod_id_1d = np.tile(np.arange(1, n_mod_top + 1), (n_subunit, 1))
66 self.dataset["mod_id"].values[:, idomain_top_active.values] = mod_id_1d
68 def _render(self, file, index, svat):
69 self._create_mod_id_rch(svat)
70 # package check only possible after calling _create_mod_id_rch
71 self._pkgcheck()
73 data_dict = {"svat": svat.values.ravel()[index]}
75 data_dict["layer"] = np.full_like(data_dict["svat"], 1)
77 for var in self._with_subunit:
78 data_dict[var] = self._index_da(self.dataset[var], index)
80 # Get well values
81 if self.well:
82 mod_id_well, svat_well, layer_well = self._create_well_id(svat)
83 data_dict["mod_id"] = np.append(mod_id_well, data_dict["mod_id"])
84 data_dict["svat"] = np.append(svat_well, data_dict["svat"])
85 data_dict["layer"] = np.append(layer_well, data_dict["layer"])
87 for var in self._to_fill:
88 data_dict[var] = ""
90 dataframe = pd.DataFrame(
91 data=data_dict, columns=list(self._metadata_dict.keys())
92 )
94 self._check_range(dataframe)
96 return self.write_dataframe_fixed_width(file, dataframe)
98 def _create_well_id(self, svat):
99 """
100 Get modflow indices, svats, and layer number for the wells
101 """
102 n_subunit = svat["subunit"].size
104 # Convert to Python's 0-based index
105 well_row = self.well["row"] - 1
106 well_column = self.well["column"] - 1
107 well_layer = self.well["layer"] - 1
109 n_mod = self.idomain_active.sum()
110 mod_id = xr.full_like(self.idomain_active, 0, dtype=np.int64)
111 mod_id.values[self.idomain_active.values] = np.arange(1, n_mod + 1)
113 well_mod_id = mod_id[well_layer, well_row, well_column]
114 well_mod_id = np.tile(well_mod_id, (n_subunit, 1))
116 well_svat = svat.values[:, well_row, well_column]
118 well_active = well_svat != 0
120 well_svat_1d = well_svat[well_active]
121 well_mod_id_1d = well_mod_id[well_active]
123 # Tile well_layers for each subunit
124 layer = np.tile(well_layer + 1, (n_subunit, 1))
125 layer_1d = layer[well_active]
127 return (well_mod_id_1d, well_svat_1d, layer_1d)