Coverage for C:\src\imod-python\imod\msw\coupler_mapping.py: 100%

1 from typing import Optional

2

3 import numpy as np

4 import pandas as pd

5 import xarray as xr

6

7 from imod.mf6.dis import StructuredDiscretization

8 from imod.mf6.wel import WellDisStructured

9 from imod.msw.fixed_format import VariableMetaData

10 from imod.msw.pkgbase import MetaSwapPackage

11

12

13 class CouplerMapping(MetaSwapPackage):

14 """

15 This contains the data to connect MODFLOW 6 cells to MetaSWAP svats.

16

17 This class is responsible for the file `mod2svat.inp`. It also includes

18 connection to wells.

19

20 Parameters

21 ----------

22 modflow_dis: StructuredDiscretization

23 Modflow 6 structured discretization

24 well: WellDisStructured (optional)

25 If given, this parameter describes sprinkling of SVAT units from MODFLOW

26 cells.

27 """

28

29 _file_name = "mod2svat.inp"

30 _metadata_dict = {

31 "mod_id": VariableMetaData(10, 1, 9999999, int),

32 "free": VariableMetaData(2, None, None, str),

33 "svat": VariableMetaData(10, 1, 9999999, int),

34 "layer": VariableMetaData(5, 0, 9999, int),

35 }

36

37 _with_subunit = ("mod_id",)

38 _without_subunit = ()

39 _to_fill = ("free",)

40

41 def __init__(

42 self,

43 modflow_dis: StructuredDiscretization,

44 well: Optional[WellDisStructured] = None,

45 ):

46 super().__init__()

47

48 self.well = well

49 # Test if equal or larger than 1, to ignore idomain == -1 as well. Don't

50 # assign to self.dataset, as grid extent might differ from svat when

51 # MetaSWAP only covers part of the Modflow grid domain.

52 self.idomain_active = modflow_dis["idomain"] >= 1

53

54 def _create_mod_id_rch(self, svat):

55 """

56 Create modflow indices for the recharge layer, which is where

57 infiltration will take place.

58 """

59 self.dataset["mod_id"] = xr.full_like(svat, fill_value=0, dtype=np.int64)

60 n_subunit = svat["subunit"].size

61 idomain_top_active = self.idomain_active.sel(layer=1, drop=True)

62

63 n_mod_top = idomain_top_active.sum()

64

65 # idomain does not have a subunit dimension, so tile for n_subunits

66 mod_id_1d = np.tile(np.arange(1, n_mod_top + 1), (n_subunit, 1))

67

68 self.dataset["mod_id"].values[:, idomain_top_active.values] = mod_id_1d

69

70 def _render(self, file, index, svat):

71 self._create_mod_id_rch(svat)

72 # package check only possible after calling _create_mod_id_rch

73 self._pkgcheck()

74

75 data_dict = {"svat": svat.values.ravel()[index]}

76

77 data_dict["layer"] = np.full_like(data_dict["svat"], 1)

78

79 for var in self._with_subunit:

80 data_dict[var] = self._index_da(self.dataset[var], index)

81

82 # Get well values

83 if self.well:

84 mod_id_well, svat_well, layer_well = self._create_well_id(svat)

85 data_dict["mod_id"] = np.append(mod_id_well, data_dict["mod_id"])

86 data_dict["svat"] = np.append(svat_well, data_dict["svat"])

87 data_dict["layer"] = np.append(layer_well, data_dict["layer"])

88

89 for var in self._to_fill:

90 data_dict[var] = ""

91

92 dataframe = pd.DataFrame(

93 data=data_dict, columns=list(self._metadata_dict.keys())

94 )

95

96 self._check_range(dataframe)

97

98 return self.write_dataframe_fixed_width(file, dataframe)

99

100 def _create_well_id(self, svat):

101 """

102 Get modflow indices, svats, and layer number for the wells

103 """

104 n_subunit = svat["subunit"].size

105

106 # Convert to Python's 0-based index

107 well_row = self.well["row"] - 1

108 well_column = self.well["column"] - 1

109 well_layer = self.well["layer"] - 1

110

111 n_mod = self.idomain_active.sum()

112 mod_id = xr.full_like(self.idomain_active, 0, dtype=np.int64)

113 mod_id.values[self.idomain_active.values] = np.arange(1, n_mod + 1)

114

115 well_mod_id = mod_id[well_layer, well_row, well_column]

116 well_mod_id = np.tile(well_mod_id, (n_subunit, 1))

117

118 well_svat = svat.values[:, well_row, well_column]

119

120 well_active = well_svat != 0

121

122 well_svat_1d = well_svat[well_active]

123 well_mod_id_1d = well_mod_id[well_active]

124

125 # Tile well_layers for each subunit

126 layer = np.tile(well_layer + 1, (n_subunit, 1))

127 layer_1d = layer[well_active]

128

129 return (well_mod_id_1d, well_svat_1d, layer_1d)