Coverage for wsimod/nodes/catchment.py: 25%

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

2"""Created on Mon Nov 15 14:20:36 2021. 

3 

4@author: bdobson 

5 

6Converted to totals on 2022-05-03 

7""" 

8from wsimod.core import constants 

9from wsimod.nodes.nodes import Node 

10 

11 

12class Catchment(Node): 

13 """""" 

14 

15 def __init__( 

16 self, 

17 name, 

18 data_input_dict={}, 

19 ): 

20 """Node that reads input data to create VQIPs that are pushed downstream and 

21 tracks abstractions made from the node, adjusting pushes accordingly. 

22 

23 Args: 

24 name (str): Node name 

25 data_input_dict (dict, optional): Dictionary of data inputs relevant for 

26 the node. Keys are tuples where first value is the name of the 

27 variable to read from the dict and the second value is the time. 

28 Defaults to {}. 

29 

30 Functions intended to call in orchestration: 

31 route 

32 

33 Key assumptions: 

34 - Flows from `Catchment` nodes are simply read from data, thus 

35 assumed to be unresponsive to wider water cycle changes. 

36 

37 Input data and parameter requirements: 

38 - Flow data in the `data_input_dict` at the model timestep. 

39 _Units_: cubic metres/timestep 

40 - Values for each variable defined in `constants.POLLUTANTS` also 

41 stored in `data_input_dict` at the model timestep. 

42 _Units_: kg/m3/timestep (additive pollutants) 

43 """ 

44 # Update args 

45 super().__init__(name) 

46 self.data_input_dict = data_input_dict 

47 

48 # Update handlers 

49 self.pull_set_handler["default"] = self.pull_set_abstraction 

50 self.pull_check_handler["default"] = self.pull_check_abstraction 

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

52 self.push_check_handler["default"] = self.push_set_deny 

53 self.unrouted_water = self.empty_vqip() 

54 # Mass balance 

55 self.mass_balance_in.append(lambda: self.get_flow()) 

56 self.mass_balance_out.append(lambda: self.unrouted_water) 

57 self.end_timestep = self.end_timestep_ 

58 

59 def get_flow(self): 

60 """Read volume data, read pollutant data, convert additibve pollutants from 

61 kg/m3 to kg. 

62 

63 Returns: 

64 vqip (dict): Return read data as a VQIP 

65 """ 

66 # TODO (if used) - note that if flow is < float accuracy then it won't 

67 # get pushed, and the pollutants will 'disappear', causing a mass balance error 

68 vqip = {"volume": self.data_input_dict[("flow", self.t)]} 

69 for pollutant in constants.POLLUTANTS: 

70 vqip[pollutant] = self.data_input_dict[(pollutant, self.t)] 

71 for pollutant in constants.ADDITIVE_POLLUTANTS: 

72 vqip[pollutant] *= vqip["volume"] 

73 

74 return vqip 

75 

76 def route(self): 

77 """Send any water that has not already been abstracted downstream.""" 

78 # Get amount of water 

79 avail = self.get_avail() 

80 # Route excess flow onwards 

81 reply = self.push_distributed(avail, of_type=["Node", "River", "Waste"]) 

82 self.unrouted_water = self.sum_vqip(self.unrouted_water, reply) 

83 if reply["volume"] > constants.FLOAT_ACCURACY: 

84 pass 

85 # print('Catchment unable to route') 

86 

87 def get_avail(self): 

88 """Water available for abstraction (Read data and subtract pre-existing 

89 abstractions). 

90 

91 Returns: 

92 avail (dict): A VQIP of water available for abstraction 

93 """ 

94 # Get available vqip 

95 avail = self.get_flow() 

96 

97 # Remove abstractions already made 

98 for name, arc in self.out_arcs.items(): 

99 avail = self.v_change_vqip(avail, avail["volume"] - arc.vqip_in["volume"]) 

100 

101 return avail 

102 

103 def pull_check_abstraction(self, vqip=None): 

104 """Check wrapper for get_avail that updates response if VQIP is given. 

105 

106 Args: 

107 vqip (dict, optional): A VQIP that is compared with get_avail and the 

108 minimum is returned. Only the 'volume' key is used. Defaults to None. 

109 

110 Returns: 

111 avail (dict): A VQIP of water available for abstraction 

112 """ 

113 # Respond to abstraction check request 

114 avail = self.get_avail() 

115 

116 if vqip: 

117 avail = self.v_change_vqip(avail, min(avail["volume"], vqip["volume"])) 

118 

119 return avail 

120 

121 def pull_set_abstraction(self, vqip): 

122 """Request set wrapper for get_avail where VQIP is specified. 

123 

124 Args: 

125 vqip (dict): A VQIP of water to pull. Only the 'volume' key is used. 

126 

127 Returns: 

128 avail (dict): A VQIP of water abstracted 

129 """ 

130 # Respond to abstraction set request 

131 avail = self.get_avail() 

132 avail = self.v_change_vqip(avail, min(avail["volume"], vqip["volume"])) 

133 

134 return avail 

135 

136 def end_timestep_(self): 

137 """Reset unrouted water.""" 

138 self.unrouted_water = self.empty_vqip()