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

1import numpy as np 

2import pandas as pd 

3import xarray as xr 

4 

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 

9 

10 

11class CouplerMapping(MetaSwapPackage): 

12 """ 

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

14 

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

16 connection to wells. 

17 

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 """ 

26 

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 } 

34 

35 _with_subunit = ("mod_id",) 

36 _without_subunit = () 

37 _to_fill = ("free",) 

38 

39 def __init__( 

40 self, 

41 modflow_dis: StructuredDiscretization, 

42 well: WellDisStructured = None, 

43 ): 

44 super().__init__() 

45 

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 

51 

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) 

60 

61 n_mod_top = idomain_top_active.sum() 

62 

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)) 

65 

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

67 

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() 

72 

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

74 

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

76 

77 for var in self._with_subunit: 

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

79 

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"]) 

86 

87 for var in self._to_fill: 

88 data_dict[var] = "" 

89 

90 dataframe = pd.DataFrame( 

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

92 ) 

93 

94 self._check_range(dataframe) 

95 

96 return self.write_dataframe_fixed_width(file, dataframe) 

97 

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 

103 

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 

108 

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) 

112 

113 well_mod_id = mod_id[well_layer, well_row, well_column] 

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

115 

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

117 

118 well_active = well_svat != 0 

119 

120 well_svat_1d = well_svat[well_active] 

121 well_mod_id_1d = well_mod_id[well_active] 

122 

123 # Tile well_layers for each subunit 

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

125 layer_1d = layer[well_active] 

126 

127 return (well_mod_id_1d, well_svat_1d, layer_1d)