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

1 import abc

2

3 import jinja2

4 import numpy as np

5

6 from imod.wq.pkgbase import BoundaryCondition

7

8

9 class 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

62 class 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

84 class 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

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