Coverage for C:\src\imod-python\imod\wq\rch.py: 34%

77 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 10:26 +0200

1import abc 

2 

3import jinja2 

4import numpy as np 

5 

6from imod.wq.pkgbase import BoundaryCondition 

7 

8 

9class Recharge(BoundaryCondition, abc.ABC): 

10 _pkg_id = "rch" 

11 

12 _mapping = (("rech", "rate"),) 

13 

14 _keywords = {"save_budget": {False: 0, True: 1}} 

15 

16 _template = jinja2.Template( 

17 "[rch]\n" 

18 " nrchop = {{recharge_option}}\n" 

19 " irchcb = {{save_budget}}\n" 

20 " {%- for name, dictname in mapping -%}" 

21 " {%- for time, timedict in dicts[dictname].items() -%}" 

22 " {%- for layer, value in timedict.items() %}\n" 

23 " {{name}}_p{{time}} = {{value}}\n" 

24 " {%- endfor -%}\n" 

25 " {%- endfor -%}" 

26 " {%- endfor -%}" 

27 ) 

28 

29 def _render(self, directory, globaltimes, nlayer, *args, **kwargs): 

30 d = { 

31 "mapping": self._mapping, 

32 "save_budget": self.dataset["save_budget"].values, 

33 "recharge_option": self._option, 

34 } 

35 self._replace_keyword(d, "save_budget") 

36 

37 dicts = {} 

38 for _, name in self._mapping: 

39 dicts[name] = self._compose_values_timelayer( 

40 name, globaltimes, directory, nlayer=nlayer 

41 ) 

42 d["dicts"] = dicts 

43 

44 return self._template.render(d) 

45 

46 def _pkgcheck(self, ibound=None): 

47 pass 

48 

49 def repeat_stress(self, rate=None, concentration=None, use_cftime=False): 

50 varnames = ["rate", "concentration"] 

51 values = [rate, concentration] 

52 for varname, value in zip(varnames, values): 

53 self._repeat_stress(varname, value, use_cftime) 

54 

55 

56class RechargeTopLayer(Recharge): 

57 """ 

58 The Recharge package is used to simulate a specified flux distributed over 

59 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW, 

60 these rates are multiplied by the horizontal area of the cells to which they 

61 are applied to calculate the volumetric flux rates. In this class the 

62 Recharge gets applied to the top grid layer (NRCHOP=1). 

63 

64 Parameters 

65 ---------- 

66 rate: float or xr.DataArray of floats 

67 is the amount of recharge. 

68 concentration: float or xr.DataArray of floats 

69 is the concentration of the recharge 

70 save_budget: bool, optional 

71 flag indicating if the budget needs to be saved. 

72 Default is False. 

73 """ 

74 

75 _option = 1 

76 

77 def __init__(self, rate, concentration, save_budget=False): 

78 super().__init__() 

79 self["rate"] = rate 

80 self["concentration"] = concentration 

81 self["save_budget"] = save_budget 

82 

83 def _set_ssm_layers(self, ibound): 

84 self._ssm_layers = np.array([1]) 

85 

86 

87class RechargeLayers(Recharge): 

88 """ 

89 The Recharge package is used to simulate a specified flux distributed over 

90 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW, 

91 these rates are multiplied by the horizontal area of the cells to which they 

92 are applied to calculate the volumetric flux rates. In this class the 

93 Recharge gets applied to a specific, specified, layer (NRCHOP=2). 

94 

95 Parameters 

96 ---------- 

97 rate: float or xr.DataArray of floats 

98 is the amount of recharge. 

99 recharge_layer: int or xr.DataArray of integers 

100 layer number variable that defines the layer in each vertical column 

101 where recharge is applied 

102 concentration: float or xr.DataArray of floats 

103 is the concentration of the recharge 

104 save_budget: bool, optional 

105 flag indicating if the budget needs to be saved. 

106 Default is False. 

107 """ 

108 

109 _option = 2 

110 

111 _mapping = (("rech", "rate"), ("irch", "recharge_layer")) 

112 

113 def __init__(self, rate, recharge_layer, concentration, save_budget=False): 

114 super().__init__() 

115 self["rate"] = rate 

116 self["recharge_layer"] = recharge_layer 

117 self["concentration"] = concentration 

118 self["save_budget"] = save_budget 

119 

120 def _set_ssm_layers(self, ibound): 

121 unique_layers = np.unique(self.dataset["recharge_layer"].values) 

122 unique_layers = unique_layers[~np.isnan(unique_layers)] 

123 self._ssm_layers = unique_layers.astype(np.int32) 

124 

125 

126class RechargeHighestActive(Recharge): 

127 """ 

128 The Recharge package is used to simulate a specified flux distributed over 

129 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW, 

130 these rates are multiplied by the horizontal area of the cells to which they 

131 are applied to calculate the volumetric flux rates. In this class the 

132 Recharge gets applied to the highest active cell in each vertical column 

133 (NRCHOP=3). 

134 

135 Parameters 

136 ---------- 

137 rate: float or xr.DataArray of floats 

138 is the amount of recharge. 

139 concentration: float or xr.DataArray of floats 

140 is the concentration of the recharge 

141 save_budget: bool, optional 

142 flag indicating if the budget needs to be saved. 

143 Default is False. 

144 """ 

145 

146 _option = 3 

147 

148 def __init__(self, rate, concentration, save_budget=False): 

149 super().__init__() 

150 

151 rate_scalar = np.ndim(rate) == 0 

152 conc_scalar = np.ndim(concentration) == 0 

153 if rate_scalar and (not conc_scalar): 

154 raise ValueError("Rate cannot be scalar if concentration is non-scalar.") 

155 

156 self["rate"] = rate 

157 self["concentration"] = concentration 

158 self["save_budget"] = save_budget 

159 

160 def _set_ssm_layers(self, ibound): 

161 rate = self.dataset["rate"] 

162 rate_idf = ("x" in rate.dims) and ("y" in rate.dims) 

163 conc_scalar = np.ndim(self.dataset["concentration"]) == 0 

164 if rate_idf and conc_scalar: 

165 rch_active = (rate != 0.0) & rate.notnull() 

166 if "time" in rch_active.dims: 

167 rch_active = rch_active.any("time") 

168 rch_active = rch_active & (ibound > 0) 

169 else: 

170 rch_active = ibound > 0 

171 

172 top_layer = ibound["layer"].where(rch_active).min("layer") 

173 top_layer = top_layer.where((ibound > 0).any("layer")) 

174 unique_layers = np.unique(top_layer.values) 

175 unique_layers = unique_layers[~np.isnan(unique_layers)] 

176 self._ssm_layers = unique_layers.astype(np.int32) 

177 

178 def repeat_stress(self, rate=None, concentration=None, use_cftime=False): 

179 varnames = ["rate", "concentration"] 

180 values = [rate, concentration] 

181 for varname, value in zip(varnames, values): 

182 self._repeat_stress(varname, value, use_cftime)