Coverage for C:\src\imod-python\imod\wq\evt.py: 88%

65 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import abc 

2 

3import jinja2 

4import numpy as np 

5 

6from imod.wq.pkgbase import BoundaryCondition 

7 

8 

9class Evapotranspiration(BoundaryCondition, abc.ABC): 

10 _pkg_id = "evt" 

11 

12 _mapping = ( 

13 ("evtr", "maximum_rate"), 

14 ("surf", "surface"), 

15 ("exdp", "extinction_depth"), 

16 ) 

17 

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

19 

20 _template = jinja2.Template( 

21 "[evt]\n" 

22 " nevtop = {{recharge_option}}\n" 

23 " ievtcb = {{save_budget}}\n" 

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

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

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

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

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

29 " {%- endfor -%}" 

30 " {%- endfor -%}" 

31 ) 

32 

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

34 d = { 

35 "mapping": self._mapping, 

36 "save_budget": self["save_budget"].values, 

37 "recharge_option": self._option, 

38 } 

39 self._replace_keyword(d, "save_budget") 

40 

41 dicts = {} 

42 for _, name in self._mapping: 

43 dicts[name] = self._compose_values_timelayer( 

44 name, globaltimes, directory, nlayer=nlayer 

45 ) 

46 d["dicts"] = dicts 

47 

48 return self._template.render(d) 

49 

50 def _pkgcheck(self, ibound=None): 

51 pass 

52 

53 def repeat_stress( 

54 self, surface=None, maximum_rate=None, extinction_depth=None, use_cftime=False 

55 ): 

56 varnames = ["surface", "maximum_rate", "extinction_depth"] 

57 values = [surface, maximum_rate, extinction_depth] 

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

59 self._repeat_stress(varname, value, use_cftime) 

60 

61 

62class EvapotranspirationTopLayer(Evapotranspiration): 

63 _option = 1 

64 

65 def __init__( 

66 self, 

67 maximum_rate, 

68 surface, 

69 extinction_depth, 

70 concentration=0.0, 

71 save_budget=False, 

72 ): 

73 super().__init__() 

74 self["maximum_rate"] = maximum_rate 

75 self["surface"] = surface 

76 self["extinction_depth"] = extinction_depth 

77 self["concentration"] = concentration 

78 self["save_budget"] = save_budget 

79 

80 def _set_ssm_layers(self, ibound): 

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

82 

83 

84class EvapotranspirationLayers(Evapotranspiration): 

85 _option = 2 

86 

87 _mapping = ( 

88 ("evtr", "maximum_rate"), 

89 ("surf", "surface"), 

90 ("exdp", "extinction_depth"), 

91 ("ievt", "evapotranspiration_layer"), 

92 ) 

93 

94 def __init__( 

95 self, 

96 maximum_rate, 

97 surface, 

98 extinction_depth, 

99 evapotranspiration_layer, 

100 concentration=0.0, 

101 save_budget=False, 

102 ): 

103 super().__init__() 

104 self["maximum_rate"] = maximum_rate 

105 self["surface"] = surface 

106 self["extinction_depth"] = extinction_depth 

107 self["evapotranspiration_layer"] = evapotranspiration_layer 

108 self["concentration"] = concentration 

109 self["save_budget"] = save_budget 

110 

111 def _set_ssm_layers(self, ibound): 

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

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

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

115 

116 

117class EvapotranspirationHighestActive(Evapotranspiration): 

118 _option = 3 

119 

120 def __init__( 

121 self, 

122 maximum_rate, 

123 surface, 

124 extinction_depth, 

125 concentration=0.0, 

126 save_budget=False, 

127 ): 

128 super().__init__() 

129 self["maximum_rate"] = maximum_rate 

130 self["surface"] = surface 

131 self["extinction_depth"] = extinction_depth 

132 self["concentration"] = concentration 

133 self["save_budget"] = save_budget 

134 

135 def _set_ssm_layers(self, ibound): 

136 top_layer = ibound["layer"].where(ibound > 0).min("layer") 

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

138 unique_layers = np.unique(top_layer.values) 

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

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