Coverage for wsimod/nodes/distribution.py: 24%

1# -*- coding: utf-8 -*- 

2"""Created on Sun Aug 14 16:27:14 2022. 

3 

4@author: bdobson 

5""" 

6 

7from wsimod.core import constants 

8from wsimod.nodes.nodes import Node 

9 

10 

11def decorate_leakage_set(self, f): 

12 """Decorator to extend the functionality of `f` by introducing leakage. This is 

13 achieved by adjusting the volume of the request (vqip) to include anticipated 

14 leakage, calling the original function `f`, and then distributing the leaked amount 

15 to groundwater. 

16 

17 Args: 

18 self (instance of Distribution class): The Distribution object to be 

19 extended 

20 f (function): The function to be extended. Expected to be the 

21 Distribution object's pull_set function. 

22 

23 Returns: 

24 pull_set (function): The decorated function which includes the 

25 original functionality of `f` and additional leakage operations. 

26 """ 

27 

28 def pull_set(vqip, **kwargs): 

29 """ 

30 

31 Args: 

32 vqip: 

33 **kwargs: 

34 

35 Returns: 

36 

37 """ 

38 vqip["volume"] /= 1 - self.leakage 

39 

40 reply = f(vqip, **kwargs) 

41 

42 amount_leaked = self.v_change_vqip(reply, reply["volume"] * self.leakage) 

43 

44 reply = self.extract_vqip(reply, amount_leaked) 

45 

46 unsuccessful_leakage = self.push_distributed( 

47 amount_leaked, of_type="Groundwater" 

48 ) 

49 if unsuccessful_leakage["volume"] > constants.FLOAT_ACCURACY: 

50 print( 

51 "warning, distribution leakage not going to GW in {0} at {1}".format( 

52 self.name, self.t 

53 ) 

54 ) 

55 reply = self.sum_vqip(reply, unsuccessful_leakage) 

56 

57 return reply 

58 

59 return pull_set 

60 

61 

62def decorate_leakage_check(self, f): 

63 """Decorator to extend the functionality of `f` by introducing leakage. This is 

64 achieved by adjusting the volume of the request (vqip) to include anticipated 

65 leakage and then calling the original function `f`. 

66 

67 Args: 

68 self (instance of Distribution class): The Distribution object to be 

69 extended 

70 f (function): The function to be extended. Expected to be the 

71 Distribution object's pull_set function. 

72 

73 Returns: 

74 pull_check (function): The decorated function which includes the 

75 original functionality of `f` and additional leakage operations. 

76 """ 

77 

78 def pull_check(vqip, **kwargs): 

79 """ 

80 

81 Args: 

82 vqip: 

83 **kwargs: 

84 

85 Returns: 

86 

87 """ 

88 if vqip is not None: 

89 vqip["volume"] /= 1 - self.leakage 

90 reply = f(vqip, **kwargs) 

91 amount_leaked = self.v_change_vqip(reply, reply["volume"] * self.leakage) 

92 

93 reply = self.extract_vqip(reply, amount_leaked) 

94 return reply 

95 

96 return pull_check 

97 

98 

99class Distribution(Node): 

100 """""" 

101 

102 def __init__(self, leakage=0, **kwargs): 

103 """A Node that cannot be pushed to. Intended to pass calls to FWTW - though this 

104 currently relies on the user to connect it properly. 

105 

106 Args: 

107 leakage (float, optional): 1 > float >= 0 to express how much 

108 water should be leaked to any attached groundwater nodes. This 

109 number represents the proportion of total flow through the node 

110 that should be leaked. 

111 Defaults to 0. 

112 

113 Functions intended to call in orchestration: 

114 None 

115 

116 Key assumptions: 

117 - No distribution processes yet represented, this class is just 

118 for conveyance. 

119 

120 Input data and parameter requirements: 

121 - None 

122 """ 

123 self.leakage = leakage 

124 super().__init__(**kwargs) 

125 # Update handlers 

126 self.push_set_handler["default"] = self.push_set_deny 

127 self.push_check_handler["default"] = self.push_check_deny 

128 

129 if leakage > 0: 

130 self.pull_set_handler["default"] = decorate_leakage_set( 

131 self, self.pull_set_handler["default"] 

132 ) 

133 self.pull_check_handler["default"] = decorate_leakage_check( 

134 self, self.pull_check_handler["default"] 

135 ) 

136 

137 

138class UnlimitedDistribution(Distribution): 

139 """""" 

140 

141 def __init__(self, **kwargs): 

142 """A distribution node that provides unlimited water while tracking pass 

143 balance. 

144 

145 Functions intended to call in orchestration: 

146 None 

147 

148 Key assumptions: 

149 - Water demand is always satisfied. 

150 

151 Input data and parameter requirements: 

152 - None 

153 """ 

154 super().__init__(**kwargs) 

155 # Update handlers 

156 self.pull_set_handler["default"] = self.pull_set_unlimited 

157 self.pull_check_handler["default"] = lambda x: self.v_change_vqip( 

158 self.empty_vqip(), constants.UNBOUNDED_CAPACITY 

159 ) 

160 

161 # States 

162 self.supplied = self.empty_vqip() 

163 

164 self.mass_balance_in.append(lambda: self.supplied) 

165 

166 def pull_set_unlimited(self, vqip): 

167 """Respond that VQIP was fulfilled and update state variables for mass balance. 

168 

169 Args: 

170 vqip (dict): A VQIP amount to request 

171 

172 Returns: 

173 vqip (dict): A VQIP amount that was supplied 

174 """ 

175 # TODO maybe need some pollutant concentrations? 

176 vqip = self.v_change_vqip(self.empty_vqip(), vqip["volume"]) 

177 self.supplied = self.sum_vqip(self.supplied, vqip) 

178 return vqip 

179 

180 def end_timestep(self): 

181 """Update state variables.""" 

182 self.supplied = self.empty_vqip()