From 336f9bceb04aa9c8107ac96cb43ff059de11194b Mon Sep 17 00:00:00 2001 From: Daniel Cummins Date: Tue, 1 Oct 2024 11:25:16 +0100 Subject: [PATCH 1/2] Moved tank class and subclasses from nodes module to their own module --- wsimod/nodes/nodes.py | 622 +---------------------------------------- wsimod/nodes/tanks.py | 624 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 625 insertions(+), 621 deletions(-) create mode 100644 wsimod/nodes/tanks.py diff --git a/wsimod/nodes/nodes.py b/wsimod/nodes/nodes.py index 2e125af9..ba7ade47 100644 --- a/wsimod/nodes/nodes.py +++ b/wsimod/nodes/nodes.py @@ -8,9 +8,8 @@ import logging from typing import Any, Dict -from wsimod.arcs.arcs import AltQueueArc, DecayArcAlt from wsimod.core import constants -from wsimod.core.core import DecayObj, WSIObj +from wsimod.core.core import WSIObj class Node(WSIObj): @@ -747,622 +746,3 @@ def general_distribute(self, vqip, of_type = None, tag = 'default', direction = NODES_REGISTRY: dict[str, type[Node]] = {Node.__name__: Node} - - -class Tank(WSIObj): - """""" - - def __init__(self, capacity=0, area=1, datum=10, initial_storage=0): - """A standard storage object. - - Args: - capacity (float, optional): Volumetric tank capacity. Defaults to 0. - area (float, optional): Area of tank. Defaults to 1. - datum (float, optional): Datum of tank base (not currently used in any - functions). Defaults to 10. - initial_storage (optional): Initial storage for tank. - float: Tank will be initialised with zero pollutants and the float - as volume - dict: Tank will be initialised with this VQIP - Defaults to 0 (i.e., no volume, no pollutants). - """ - # Set parameters - self.capacity = capacity - self.area = area - self.datum = datum - self.initial_storage = initial_storage - - WSIObj.__init__(self) # Not sure why I do this rather than super() - - # TODO I don't think the outer if statement is needed - if "initial_storage" in dir(self): - if isinstance(self.initial_storage, dict): - # Assume dict is VQIP describing storage - self.storage = self.copy_vqip(self.initial_storage) - self.storage_ = self.copy_vqip( - self.initial_storage - ) # Lagged storage for mass balance - else: - # Assume number describes initial stroage - self.storage = self.v_change_vqip( - self.empty_vqip(), self.initial_storage - ) - self.storage_ = self.v_change_vqip( - self.empty_vqip(), self.initial_storage - ) # Lagged storage for mass balance - else: - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() # Lagged storage for mass balance - - def ds(self): - """Should be called by parent object to get change in storage. - - Returns: - (dict): Change in storage - """ - return self.ds_vqip(self.storage, self.storage_) - - def pull_ponded(self): - """Pull any volume that is above the tank's capacity. - - Returns: - ponded (vqip): Amount of ponded water that has been removed from the - tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, - 'phosphate' : 0.2}) - >>> print(my_tank.storage) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.pull_ponded()) - {'volume' : 1, 'phosphate' : 0.02} - >>> print(my_tank.storage) - {'volume' : 9, 'phosphate' : 0.18} - """ - # Get amount - ponded = max(self.storage["volume"] - self.capacity, 0) - # Pull from tank - ponded = self.pull_storage({"volume": ponded}) - return ponded - - def get_avail(self, vqip=None): - """Get minimum of the amount of water in storage and vqip (if provided). - - Args: - vqip (dict, optional): Maximum water required (only 'volume' is used). - Defaults to None. - - Returns: - reply (dict): Water available - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, - 'phosphate' : 0.2}) - >>> print(my_tank.storage) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.get_avail()) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.get_avail({'volume' : 1})) - {'volume' : 1, 'phosphate' : 0.02} - """ - reply = self.copy_vqip(self.storage) - if vqip is None: - # Return storage - return reply - else: - # Adjust storage pollutants to match volume in vqip - reply = self.v_change_vqip(reply, min(reply["volume"], vqip["volume"])) - return reply - - def get_excess(self, vqip=None): - """Get difference between current storage and tank capacity. - - Args: - vqip (dict, optional): Maximum capacity required (only 'volume' is - used). Defaults to None. - - Returns: - (dict): Difference available - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.get_excess()) - {'volume' : 4, 'phosphate' : 0.16} - >>> print(my_tank.get_excess({'volume' : 2})) - {'volume' : 2, 'phosphate' : 0.08} - """ - vol = max(self.capacity - self.storage["volume"], 0) - if vqip is not None: - vol = min(vqip["volume"], vol) - - # Adjust storage pollutants to match volume in vqip - # TODO the v_change_vqip in the reply here is a weird default (if a VQIP is not - # provided) - return self.v_change_vqip(self.storage, vol) - - def push_storage(self, vqip, force=False): - """Push water into tank, updating the storage VQIP. Force argument can be used - to ignore tank capacity. - - Args: - vqip (dict): VQIP amount to be pushed - force (bool, optional): Argument used to cause function to ignore tank - capacity, possibly resulting in pooling. Defaults to False. - - Returns: - reply (dict): A VQIP of water not successfully pushed to the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> constants.POLLUTANTS = ['phosphate'] - >>> constants.NON_ADDITIVE_POLLUTANTS = [] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> my_push = {'volume' : 10, 'phosphate' : 0.5} - >>> reply = my_tank.push_storage(my_push) - >>> print(reply) - {'volume' : 6, 'phosphate' : 0.3} - >>> print(my_tank.storage) - {'volume': 9.0, 'phosphate': 0.4} - >>> print(my_tank.push_storage(reply, force = True)) - {'phosphate': 0, 'volume': 0} - >>> print(my_tank.storage) - {'volume': 15.0, 'phosphate': 0.7} - """ - if force: - # Directly add request to storage - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - # Check whether request can be met - excess = self.get_excess()["volume"] - - # Adjust accordingly - reply = max(vqip["volume"] - excess, 0) - reply = self.v_change_vqip(vqip, reply) - entered = self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) - - # Update storage - self.storage = self.sum_vqip(self.storage, entered) - - return reply - - def pull_storage(self, vqip): - """Pull water from tank, updating the storage VQIP. Pollutants are removed from - tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). - - Args: - vqip (dict): VQIP amount to be pulled, (only 'volume' key is needed) - - Returns: - reply (dict): A VQIP water successfully pulled from the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.pull_storage({'volume' : 6})) - {'volume': 5.0, 'phosphate': 0.2} - >>> print(my_tank.storage) - {'volume': 0, 'phosphate': 0} - """ - # Pull from Tank by volume (taking pollutants in proportion) - if self.storage["volume"] == 0: - return self.empty_vqip() - - # Adjust based on available volume - reply = min(vqip["volume"], self.storage["volume"]) - - # Update reply to vqip (in proportion to concentration in storage) - reply = self.v_change_vqip(self.storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - - return reply - - def pull_pollutants(self, vqip): - """Pull water from tank, updating the storage VQIP. Pollutants are removed from - tank in according to their values in vqip. - - Args: - vqip (dict): VQIP amount to be pulled - - Returns: - vqip (dict): A VQIP water successfully pulled from the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.pull_pollutants({'volume' : 2, 'phosphate' : 0.15})) - {'volume': 2.0, 'phosphate': 0.15} - >>> print(my_tank.storage) - {'volume': 3, 'phosphate': 0.05} - """ - # Adjust based on available mass - for pol in constants.ADDITIVE_POLLUTANTS + ["volume"]: - vqip[pol] = min(self.storage[pol], vqip[pol]) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, vqip) - return vqip - - def get_head(self, datum=None, non_head_storage=0): - """Area volume calculation for head calcuations. Datum and storage that does not - contribute to head can be specified. - - Args: - datum (float, optional): Value to add to pressure head in tank. - Defaults to None. - non_head_storage (float, optional): Amount of storage that does - not contribute to generation of head. The tank must exceed - this value to generate any pressure head. Defaults to 0. - - Returns: - head (float): Total head in tank - - Examples: - >>> my_tank = Tank(datum = 10, initial_storage = 5, capacity = 10, area = 2) - >>> print(my_tank.get_head()) - 12.5 - >>> print(my_tank.get_head(non_head_storage = 1)) - 12 - >>> print(my_tank.get_head(non_head_storage = 1, datum = 0)) - 2 - """ - # If datum not provided use object datum - if datum is None: - datum = self.datum - - # Calculate pressure head generating storage - head_storage = max(self.storage["volume"] - non_head_storage, 0) - - # Perform head calculation - head = head_storage / self.area + datum - - return head - - def evaporate(self, evap): - """Wrapper for v_distill_vqip to apply a volumetric subtraction from tank - storage. Volume removed from storage and no change in pollutant values. - - Args: - evap (float): Volume to evaporate - - Returns: - evap (float): Volumetric amount of evaporation successfully removed - """ - avail = self.get_avail()["volume"] - - evap = min(evap, avail) - self.storage = self.v_distill_vqip(self.storage, evap) - return evap - - ##Old function no longer needed (check it is not used anywhere and remove) - def push_total(self, vqip): - """ - - Args: - vqip: - - Returns: - - """ - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - ##Old function no longer needed (check it is not used anywhere and remove) - def push_total_c(self, vqip): - """ - - Args: - vqip: - - Returns: - - """ - # Push vqip to storage where pollutants are given as a concentration rather - # than storage - vqip = self.concentration_to_total(self.vqip) - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - def end_timestep(self): - """Function to be called by parent object, tracks previously timestep's - storage.""" - self.storage_ = self.copy_vqip(self.storage) - - def reinit(self): - """Set storage to an empty VQIP.""" - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() - - -class ResidenceTank(Tank): - """""" - - def __init__(self, residence_time=2, **kwargs): - """A tank that has a residence time property that limits storage pulled from the - 'pull_outflow' function. - - Args: - residence_time (float, optional): Residence time, in theory given - in timesteps, in practice it just means that storage / - residence time can be pulled each time pull_outflow is called. - Defaults to 2. - """ - self.residence_time = residence_time - super().__init__(**kwargs) - - def pull_outflow(self): - """Pull storage by residence time from the tank, updating tank storage. - - Returns: - outflow (dict): A VQIP with volume of pulled volume and pollutants - proportionate to the tank's pollutants - """ - # Calculate outflow - outflow = self.storage["volume"] / self.residence_time - # Update pollutant amounts - outflow = self.v_change_vqip(self.storage, outflow) - # Remove from tank - outflow = self.pull_storage(outflow) - return outflow - - -class DecayTank(Tank, DecayObj): - """""" - - def __init__(self, decays={}, parent=None, **kwargs): - """A tank that has DecayObj functions. Decay occurs in end_timestep, after - updating state variables. In this sense, decay is occurring at the very - beginning of the timestep. - - Args: - decays (dict): A dict of dicts containing a key for each pollutant that - decays and, within that, a key for each parameter (a constant and - exponent) - parent (object): An object that can be used to read temperature data from - """ - # Store parameters - self.parent = parent - - # Initialise Tank - Tank.__init__(self, **kwargs) - - # Initialise decay object - DecayObj.__init__(self, decays) - - # Update timestep and ds functions - self.end_timestep = self.end_timestep_decay - self.ds = self.decay_ds - - def end_timestep_decay(self): - """Update state variables and call make_decay.""" - self.total_decayed = self.empty_vqip() - self.storage_ = self.copy_vqip(self.storage) - - self.storage = self.make_decay(self.storage) - - def decay_ds(self): - """Track storage and amount decayed. - - Returns: - ds (dict): A VQIP of change in storage and total decayed - """ - ds = self.ds_vqip(self.storage, self.storage_) - ds = self.sum_vqip(ds, self.total_decayed) - return ds - - -class QueueTank(Tank): - """""" - - def __init__(self, number_of_timesteps=0, **kwargs): - """A tank with an internal queue arc, whose queue must be completed before - storage is available for use. The storage that has completed the queue is under - the 'active_storage' property. - - Args: - number_of_timesteps (int, optional): Built in delay for the internal - queue - it is always added to the queue time, although delay can be - provided with pushes only. Defaults to 0. - """ - # Set parameters - self.number_of_timesteps = number_of_timesteps - - super().__init__(**kwargs) - self.end_timestep = self._end_timestep - self.active_storage = self.copy_vqip(self.storage) - - # TODO enable queue to be initialised not empty - self.out_arcs = {} - self.in_arcs = {} - # Create internal queue arc - self.internal_arc = AltQueueArc( - in_port=self, out_port=self, number_of_timesteps=self.number_of_timesteps - ) - # TODO should mass balance call internal arc (is this arc called in arc mass - # balance?) - - def get_avail(self): - """Return the active_storage of the tank. - - Returns: - (dict): VQIP of active_storage - """ - return self.copy_vqip(self.active_storage) - - def push_storage(self, vqip, time=0, force=False): - """Push storage into QueueTank, applying travel time, unless forced. - - Args: - vqip (dict): A VQIP of the amount to push - time (int, optional): Number of timesteps to spend in queue, in addition - to number_of_timesteps property of internal_arc. Defaults to 0. - force (bool, optional): Force property that will ignore tank capacity - and ignore travel time. Defaults to False. - - Returns: - reply (dict): A VQIP of water that could not be received by the tank - """ - if force: - # Directly add request to storage, skipping queue - self.storage = self.sum_vqip(self.storage, vqip) - self.active_storage = self.sum_vqip(self.active_storage, vqip) - return self.empty_vqip() - - # Push to QueueTank - reply = self.internal_arc.send_push_request(vqip, force=force, time=time) - # Update storage - # TODO storage won't be accurately tracking temperature.. - self.storage = self.sum_vqip( - self.storage, self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) - ) - return reply - - def pull_storage(self, vqip): - """Pull storage from the QueueTank, only water in active_storage is available. - Returning water pulled and updating tank states. Pollutants are removed from - tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). - - Args: - vqip (dict): VQIP amount to pull, only 'volume' property is used - - Returns: - reply (dict): VQIP amount that was pulled - """ - # Adjust based on available volume - reply = min(vqip["volume"], self.active_storage["volume"]) - - # Update reply to vqip - reply = self.v_change_vqip(self.active_storage, reply) - - # Extract from active_storage - self.active_storage = self.extract_vqip(self.active_storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - - return reply - - def pull_storage_exact(self, vqip): - """Pull storage from the QueueTank, only water in active_storage is available. - Pollutants are removed from tank in according to their values in vqip. - - Args: - vqip (dict): A VQIP amount to pull - - Returns: - reply (dict): A VQIP amount successfully pulled - """ - # Adjust based on available - reply = self.copy_vqip(vqip) - for pol in ["volume"] + constants.ADDITIVE_POLLUTANTS: - reply[pol] = min(reply[pol], self.active_storage[pol]) - - # Pull from QueueTank - self.active_storage = self.extract_vqip(self.active_storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - return reply - - def push_check(self, vqip=None, tag="default"): - """Wrapper for get_excess but applies comparison to volume in VQIP. - Needed to enable use of internal_arc, which assumes it is connecting nodes . - rather than tanks. - NOTE: this is intended only for use with the internal_arc. Pushing to - QueueTanks should use 'push_storage'. - - Args: - vqip (dict, optional): VQIP amount to push. Defaults to None. - tag (str, optional): Tag, see Node, don't think it should actually be - used for a QueueTank since there are no handlers. Defaults to - 'default'. - - Returns: - excess (dict): a VQIP amount of excess capacity - """ - # TODO does behaviour for volume = None need to be defined? - excess = self.get_excess() - if vqip is not None: - excess["volume"] = min(vqip["volume"], excess["volume"]) - return excess - - def push_set(self, vqip, tag="default"): - """Behaves differently from normal push setting, it assumes sufficient tank - capacity and receives VQIPs that have reached the END of the internal_arc. - NOTE: this is intended only for use with the internal_arc. Pushing to - QueueTanks should use 'push_storage'. - - Args: - vqip (dict): VQIP amount to push - tag (str, optional): Tag, see Node, don't think it should actually be - used for a QueueTank since there are no handlers. Defaults to - 'default'. - - Returns: - (dict): Returns empty VQIP, indicating all water received (since it - assumes capacity was checked before entering the internal arc) - """ - # Update active_storage (since it has reached the end of the internal_arc) - self.active_storage = self.sum_vqip(self.active_storage, vqip) - - return self.empty_vqip() - - def _end_timestep(self): - """Wrapper for end_timestep that also ends the timestep in the internal_arc.""" - self.internal_arc.end_timestep() - self.internal_arc.update_queue() - self.storage_ = self.copy_vqip(self.storage) - - def reinit(self): - """Zeros storages and arc.""" - self.internal_arc.reinit() - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() - self.active_storage = self.empty_vqip() - - -class DecayQueueTank(QueueTank): - """""" - - def __init__(self, decays={}, parent=None, number_of_timesteps=1, **kwargs): - """Adds a DecayAltArc in QueueTank to enable decay to occur within the - internal_arc queue. - - Args: - decays (dict): A dict of dicts containing a key for each pollutant and, - within that, a key for each parameter (a constant and exponent) - parent (object): An object that can be used to read temperature data from - number_of_timesteps (int, optional): Built in delay for the internal - queue - it is always added to the queue time, although delay can be - provided with pushes only. Defaults to 0. - """ - # Initialise QueueTank - super().__init__(number_of_timesteps=number_of_timesteps, **kwargs) - # Replace internal_arc with a DecayArcAlt - self.internal_arc = DecayArcAlt( - in_port=self, - out_port=self, - number_of_timesteps=number_of_timesteps, - parent=parent, - decays=decays, - ) - - self.end_timestep = self._end_timestep - - def _end_timestep(self): - """End timestep wrapper that removes decayed pollutants and calls internal - arc.""" - # TODO Should the active storage decay if decays are given (probably.. though - # that sounds like a nightmare)? - self.storage = self.extract_vqip(self.storage, self.internal_arc.total_decayed) - self.storage_ = self.copy_vqip(self.storage) - self.internal_arc.end_timestep() diff --git a/wsimod/nodes/tanks.py b/wsimod/nodes/tanks.py new file mode 100644 index 00000000..b2036edb --- /dev/null +++ b/wsimod/nodes/tanks.py @@ -0,0 +1,624 @@ +"""Module for defining tanks.""" + +from wsimod.arcs import AltQueueArc, DecayArcAlt +from wsimod.core import constants +from wsimod.core.core import DecayObj, WSIObj + + +class Tank(WSIObj): + """""" + + def __init__(self, capacity=0, area=1, datum=10, initial_storage=0): + """A standard storage object. + + Args: + capacity (float, optional): Volumetric tank capacity. Defaults to 0. + area (float, optional): Area of tank. Defaults to 1. + datum (float, optional): Datum of tank base (not currently used in any + functions). Defaults to 10. + initial_storage (optional): Initial storage for tank. + float: Tank will be initialised with zero pollutants and the float + as volume + dict: Tank will be initialised with this VQIP + Defaults to 0 (i.e., no volume, no pollutants). + """ + # Set parameters + self.capacity = capacity + self.area = area + self.datum = datum + self.initial_storage = initial_storage + + WSIObj.__init__(self) # Not sure why I do this rather than super() + + # TODO I don't think the outer if statement is needed + if "initial_storage" in dir(self): + if isinstance(self.initial_storage, dict): + # Assume dict is VQIP describing storage + self.storage = self.copy_vqip(self.initial_storage) + self.storage_ = self.copy_vqip( + self.initial_storage + ) # Lagged storage for mass balance + else: + # Assume number describes initial stroage + self.storage = self.v_change_vqip( + self.empty_vqip(), self.initial_storage + ) + self.storage_ = self.v_change_vqip( + self.empty_vqip(), self.initial_storage + ) # Lagged storage for mass balance + else: + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() # Lagged storage for mass balance + + def ds(self): + """Should be called by parent object to get change in storage. + + Returns: + (dict): Change in storage + """ + return self.ds_vqip(self.storage, self.storage_) + + def pull_ponded(self): + """Pull any volume that is above the tank's capacity. + + Returns: + ponded (vqip): Amount of ponded water that has been removed from the + tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, + 'phosphate' : 0.2}) + >>> print(my_tank.storage) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.pull_ponded()) + {'volume' : 1, 'phosphate' : 0.02} + >>> print(my_tank.storage) + {'volume' : 9, 'phosphate' : 0.18} + """ + # Get amount + ponded = max(self.storage["volume"] - self.capacity, 0) + # Pull from tank + ponded = self.pull_storage({"volume": ponded}) + return ponded + + def get_avail(self, vqip=None): + """Get minimum of the amount of water in storage and vqip (if provided). + + Args: + vqip (dict, optional): Maximum water required (only 'volume' is used). + Defaults to None. + + Returns: + reply (dict): Water available + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, + 'phosphate' : 0.2}) + >>> print(my_tank.storage) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.get_avail()) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.get_avail({'volume' : 1})) + {'volume' : 1, 'phosphate' : 0.02} + """ + reply = self.copy_vqip(self.storage) + if vqip is None: + # Return storage + return reply + else: + # Adjust storage pollutants to match volume in vqip + reply = self.v_change_vqip(reply, min(reply["volume"], vqip["volume"])) + return reply + + def get_excess(self, vqip=None): + """Get difference between current storage and tank capacity. + + Args: + vqip (dict, optional): Maximum capacity required (only 'volume' is + used). Defaults to None. + + Returns: + (dict): Difference available + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.get_excess()) + {'volume' : 4, 'phosphate' : 0.16} + >>> print(my_tank.get_excess({'volume' : 2})) + {'volume' : 2, 'phosphate' : 0.08} + """ + vol = max(self.capacity - self.storage["volume"], 0) + if vqip is not None: + vol = min(vqip["volume"], vol) + + # Adjust storage pollutants to match volume in vqip + # TODO the v_change_vqip in the reply here is a weird default (if a VQIP is not + # provided) + return self.v_change_vqip(self.storage, vol) + + def push_storage(self, vqip, force=False): + """Push water into tank, updating the storage VQIP. Force argument can be used + to ignore tank capacity. + + Args: + vqip (dict): VQIP amount to be pushed + force (bool, optional): Argument used to cause function to ignore tank + capacity, possibly resulting in pooling. Defaults to False. + + Returns: + reply (dict): A VQIP of water not successfully pushed to the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> constants.POLLUTANTS = ['phosphate'] + >>> constants.NON_ADDITIVE_POLLUTANTS = [] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> my_push = {'volume' : 10, 'phosphate' : 0.5} + >>> reply = my_tank.push_storage(my_push) + >>> print(reply) + {'volume' : 6, 'phosphate' : 0.3} + >>> print(my_tank.storage) + {'volume': 9.0, 'phosphate': 0.4} + >>> print(my_tank.push_storage(reply, force = True)) + {'phosphate': 0, 'volume': 0} + >>> print(my_tank.storage) + {'volume': 15.0, 'phosphate': 0.7} + """ + if force: + # Directly add request to storage + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + # Check whether request can be met + excess = self.get_excess()["volume"] + + # Adjust accordingly + reply = max(vqip["volume"] - excess, 0) + reply = self.v_change_vqip(vqip, reply) + entered = self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) + + # Update storage + self.storage = self.sum_vqip(self.storage, entered) + + return reply + + def pull_storage(self, vqip): + """Pull water from tank, updating the storage VQIP. Pollutants are removed from + tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). + + Args: + vqip (dict): VQIP amount to be pulled, (only 'volume' key is needed) + + Returns: + reply (dict): A VQIP water successfully pulled from the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.pull_storage({'volume' : 6})) + {'volume': 5.0, 'phosphate': 0.2} + >>> print(my_tank.storage) + {'volume': 0, 'phosphate': 0} + """ + # Pull from Tank by volume (taking pollutants in proportion) + if self.storage["volume"] == 0: + return self.empty_vqip() + + # Adjust based on available volume + reply = min(vqip["volume"], self.storage["volume"]) + + # Update reply to vqip (in proportion to concentration in storage) + reply = self.v_change_vqip(self.storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + + return reply + + def pull_pollutants(self, vqip): + """Pull water from tank, updating the storage VQIP. Pollutants are removed from + tank in according to their values in vqip. + + Args: + vqip (dict): VQIP amount to be pulled + + Returns: + vqip (dict): A VQIP water successfully pulled from the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.pull_pollutants({'volume' : 2, 'phosphate' : 0.15})) + {'volume': 2.0, 'phosphate': 0.15} + >>> print(my_tank.storage) + {'volume': 3, 'phosphate': 0.05} + """ + # Adjust based on available mass + for pol in constants.ADDITIVE_POLLUTANTS + ["volume"]: + vqip[pol] = min(self.storage[pol], vqip[pol]) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, vqip) + return vqip + + def get_head(self, datum=None, non_head_storage=0): + """Area volume calculation for head calcuations. Datum and storage that does not + contribute to head can be specified. + + Args: + datum (float, optional): Value to add to pressure head in tank. + Defaults to None. + non_head_storage (float, optional): Amount of storage that does + not contribute to generation of head. The tank must exceed + this value to generate any pressure head. Defaults to 0. + + Returns: + head (float): Total head in tank + + Examples: + >>> my_tank = Tank(datum = 10, initial_storage = 5, capacity = 10, area = 2) + >>> print(my_tank.get_head()) + 12.5 + >>> print(my_tank.get_head(non_head_storage = 1)) + 12 + >>> print(my_tank.get_head(non_head_storage = 1, datum = 0)) + 2 + """ + # If datum not provided use object datum + if datum is None: + datum = self.datum + + # Calculate pressure head generating storage + head_storage = max(self.storage["volume"] - non_head_storage, 0) + + # Perform head calculation + head = head_storage / self.area + datum + + return head + + def evaporate(self, evap): + """Wrapper for v_distill_vqip to apply a volumetric subtraction from tank + storage. Volume removed from storage and no change in pollutant values. + + Args: + evap (float): Volume to evaporate + + Returns: + evap (float): Volumetric amount of evaporation successfully removed + """ + avail = self.get_avail()["volume"] + + evap = min(evap, avail) + self.storage = self.v_distill_vqip(self.storage, evap) + return evap + + ##Old function no longer needed (check it is not used anywhere and remove) + def push_total(self, vqip): + """ + + Args: + vqip: + + Returns: + + """ + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + ##Old function no longer needed (check it is not used anywhere and remove) + def push_total_c(self, vqip): + """ + + Args: + vqip: + + Returns: + + """ + # Push vqip to storage where pollutants are given as a concentration rather + # than storage + vqip = self.concentration_to_total(self.vqip) + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + def end_timestep(self): + """Function to be called by parent object, tracks previously timestep's + storage.""" + self.storage_ = self.copy_vqip(self.storage) + + def reinit(self): + """Set storage to an empty VQIP.""" + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() + + +class ResidenceTank(Tank): + """""" + + def __init__(self, residence_time=2, **kwargs): + """A tank that has a residence time property that limits storage pulled from the + 'pull_outflow' function. + + Args: + residence_time (float, optional): Residence time, in theory given + in timesteps, in practice it just means that storage / + residence time can be pulled each time pull_outflow is called. + Defaults to 2. + """ + self.residence_time = residence_time + super().__init__(**kwargs) + + def pull_outflow(self): + """Pull storage by residence time from the tank, updating tank storage. + + Returns: + outflow (dict): A VQIP with volume of pulled volume and pollutants + proportionate to the tank's pollutants + """ + # Calculate outflow + outflow = self.storage["volume"] / self.residence_time + # Update pollutant amounts + outflow = self.v_change_vqip(self.storage, outflow) + # Remove from tank + outflow = self.pull_storage(outflow) + return outflow + + +class DecayTank(Tank, DecayObj): + """""" + + def __init__(self, decays={}, parent=None, **kwargs): + """A tank that has DecayObj functions. Decay occurs in end_timestep, after + updating state variables. In this sense, decay is occurring at the very + beginning of the timestep. + + Args: + decays (dict): A dict of dicts containing a key for each pollutant that + decays and, within that, a key for each parameter (a constant and + exponent) + parent (object): An object that can be used to read temperature data from + """ + # Store parameters + self.parent = parent + + # Initialise Tank + Tank.__init__(self, **kwargs) + + # Initialise decay object + DecayObj.__init__(self, decays) + + # Update timestep and ds functions + self.end_timestep = self.end_timestep_decay + self.ds = self.decay_ds + + def end_timestep_decay(self): + """Update state variables and call make_decay.""" + self.total_decayed = self.empty_vqip() + self.storage_ = self.copy_vqip(self.storage) + + self.storage = self.make_decay(self.storage) + + def decay_ds(self): + """Track storage and amount decayed. + + Returns: + ds (dict): A VQIP of change in storage and total decayed + """ + ds = self.ds_vqip(self.storage, self.storage_) + ds = self.sum_vqip(ds, self.total_decayed) + return ds + + +class QueueTank(Tank): + """""" + + def __init__(self, number_of_timesteps=0, **kwargs): + """A tank with an internal queue arc, whose queue must be completed before + storage is available for use. The storage that has completed the queue is under + the 'active_storage' property. + + Args: + number_of_timesteps (int, optional): Built in delay for the internal + queue - it is always added to the queue time, although delay can be + provided with pushes only. Defaults to 0. + """ + # Set parameters + self.number_of_timesteps = number_of_timesteps + + super().__init__(**kwargs) + self.end_timestep = self._end_timestep + self.active_storage = self.copy_vqip(self.storage) + + # TODO enable queue to be initialised not empty + self.out_arcs = {} + self.in_arcs = {} + # Create internal queue arc + self.internal_arc = AltQueueArc( + in_port=self, out_port=self, number_of_timesteps=self.number_of_timesteps + ) + # TODO should mass balance call internal arc (is this arc called in arc mass + # balance?) + + def get_avail(self): + """Return the active_storage of the tank. + + Returns: + (dict): VQIP of active_storage + """ + return self.copy_vqip(self.active_storage) + + def push_storage(self, vqip, time=0, force=False): + """Push storage into QueueTank, applying travel time, unless forced. + + Args: + vqip (dict): A VQIP of the amount to push + time (int, optional): Number of timesteps to spend in queue, in addition + to number_of_timesteps property of internal_arc. Defaults to 0. + force (bool, optional): Force property that will ignore tank capacity + and ignore travel time. Defaults to False. + + Returns: + reply (dict): A VQIP of water that could not be received by the tank + """ + if force: + # Directly add request to storage, skipping queue + self.storage = self.sum_vqip(self.storage, vqip) + self.active_storage = self.sum_vqip(self.active_storage, vqip) + return self.empty_vqip() + + # Push to QueueTank + reply = self.internal_arc.send_push_request(vqip, force=force, time=time) + # Update storage + # TODO storage won't be accurately tracking temperature.. + self.storage = self.sum_vqip( + self.storage, self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) + ) + return reply + + def pull_storage(self, vqip): + """Pull storage from the QueueTank, only water in active_storage is available. + Returning water pulled and updating tank states. Pollutants are removed from + tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). + + Args: + vqip (dict): VQIP amount to pull, only 'volume' property is used + + Returns: + reply (dict): VQIP amount that was pulled + """ + # Adjust based on available volume + reply = min(vqip["volume"], self.active_storage["volume"]) + + # Update reply to vqip + reply = self.v_change_vqip(self.active_storage, reply) + + # Extract from active_storage + self.active_storage = self.extract_vqip(self.active_storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + + return reply + + def pull_storage_exact(self, vqip): + """Pull storage from the QueueTank, only water in active_storage is available. + Pollutants are removed from tank in according to their values in vqip. + + Args: + vqip (dict): A VQIP amount to pull + + Returns: + reply (dict): A VQIP amount successfully pulled + """ + # Adjust based on available + reply = self.copy_vqip(vqip) + for pol in ["volume"] + constants.ADDITIVE_POLLUTANTS: + reply[pol] = min(reply[pol], self.active_storage[pol]) + + # Pull from QueueTank + self.active_storage = self.extract_vqip(self.active_storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + return reply + + def push_check(self, vqip=None, tag="default"): + """Wrapper for get_excess but applies comparison to volume in VQIP. + Needed to enable use of internal_arc, which assumes it is connecting nodes . + rather than tanks. + NOTE: this is intended only for use with the internal_arc. Pushing to + QueueTanks should use 'push_storage'. + + Args: + vqip (dict, optional): VQIP amount to push. Defaults to None. + tag (str, optional): Tag, see Node, don't think it should actually be + used for a QueueTank since there are no handlers. Defaults to + 'default'. + + Returns: + excess (dict): a VQIP amount of excess capacity + """ + # TODO does behaviour for volume = None need to be defined? + excess = self.get_excess() + if vqip is not None: + excess["volume"] = min(vqip["volume"], excess["volume"]) + return excess + + def push_set(self, vqip, tag="default"): + """Behaves differently from normal push setting, it assumes sufficient tank + capacity and receives VQIPs that have reached the END of the internal_arc. + NOTE: this is intended only for use with the internal_arc. Pushing to + QueueTanks should use 'push_storage'. + + Args: + vqip (dict): VQIP amount to push + tag (str, optional): Tag, see Node, don't think it should actually be + used for a QueueTank since there are no handlers. Defaults to + 'default'. + + Returns: + (dict): Returns empty VQIP, indicating all water received (since it + assumes capacity was checked before entering the internal arc) + """ + # Update active_storage (since it has reached the end of the internal_arc) + self.active_storage = self.sum_vqip(self.active_storage, vqip) + + return self.empty_vqip() + + def _end_timestep(self): + """Wrapper for end_timestep that also ends the timestep in the internal_arc.""" + self.internal_arc.end_timestep() + self.internal_arc.update_queue() + self.storage_ = self.copy_vqip(self.storage) + + def reinit(self): + """Zeros storages and arc.""" + self.internal_arc.reinit() + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() + self.active_storage = self.empty_vqip() + + +class DecayQueueTank(QueueTank): + """""" + + def __init__(self, decays={}, parent=None, number_of_timesteps=1, **kwargs): + """Adds a DecayAltArc in QueueTank to enable decay to occur within the + internal_arc queue. + + Args: + decays (dict): A dict of dicts containing a key for each pollutant and, + within that, a key for each parameter (a constant and exponent) + parent (object): An object that can be used to read temperature data from + number_of_timesteps (int, optional): Built in delay for the internal + queue - it is always added to the queue time, although delay can be + provided with pushes only. Defaults to 0. + """ + # Initialise QueueTank + super().__init__(number_of_timesteps=number_of_timesteps, **kwargs) + # Replace internal_arc with a DecayArcAlt + self.internal_arc = DecayArcAlt( + in_port=self, + out_port=self, + number_of_timesteps=number_of_timesteps, + parent=parent, + decays=decays, + ) + + self.end_timestep = self._end_timestep + + def _end_timestep(self): + """End timestep wrapper that removes decayed pollutants and calls internal + arc.""" + # TODO Should the active storage decay if decays are given (probably.. though + # that sounds like a nightmare)? + self.storage = self.extract_vqip(self.storage, self.internal_arc.total_decayed) + self.storage_ = self.copy_vqip(self.storage) + self.internal_arc.end_timestep() From d86ed58ccb85244a7e46c91573d92407674fc40e Mon Sep 17 00:00:00 2001 From: Daniel Cummins Date: Tue, 1 Oct 2024 11:59:22 +0100 Subject: [PATCH 2/2] Fix node and tank imports --- tests/test_nodes.py | 245 +------------------------------ tests/test_tanks.py | 268 ++++++++++++++++++++++++++++++++++ wsimod/nodes/land.py | 3 +- wsimod/nodes/sewer.py | 3 +- wsimod/nodes/storage.py | 3 +- wsimod/nodes/tanks.py | 2 +- wsimod/nodes/wtw.py | 3 +- wsimod/orchestration/model.py | 3 +- 8 files changed, 280 insertions(+), 250 deletions(-) create mode 100644 tests/test_tanks.py diff --git a/tests/test_nodes.py b/tests/test_nodes.py index b285f712..576ef7b9 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -11,14 +11,7 @@ from wsimod.arcs.arcs import Arc from wsimod.core import constants -from wsimod.nodes.nodes import ( - DecayQueueTank, - DecayTank, - Node, - QueueTank, - ResidenceTank, - Tank, -) +from wsimod.nodes.nodes import Node from wsimod.nodes.storage import Storage from wsimod.nodes.waste import Waste @@ -424,242 +417,6 @@ def test_data_read(self): self.assertEqual(15, node.get_data_input("temperature")) - def test_tank_ds(self): - tank = Tank( - capacity=10, - initial_storage={"volume": 5, "phosphate": 0.4, "temperature": 10}, - ) - tank.end_timestep() - - d1 = {"volume": 2, "phosphate": 0.01, "temperature": 15} - - _ = tank.push_storage(d1) - - diff = tank.ds() - - d2 = {"volume": 2, "phosphate": 0.01, "temperature": 0} - - self.assertDictAlmostEqual(d2, diff, 16) - - def test_ponded(self): - tank = Tank( - capacity=10, - initial_storage={"volume": 15, "phosphate": 0.4, "temperature": 10}, - ) - d1 = {"volume": 5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.pull_ponded() - self.assertDictAlmostEqual(d1, reply) - - def test_tank_get_avail(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - reply = tank.get_avail() - self.assertDictAlmostEqual(d1, reply) - - reply = tank.get_avail({"volume": 2.5}) - d2 = {"volume": 2.5, "phosphate": 0.2, "temperature": 10} - self.assertDictAlmostEqual(d2, reply) - - reply = tank.get_avail({"volume": 10}) - self.assertDictAlmostEqual(d1, reply) - - def test_tank_get_excess(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.get_excess() - self.assertDictAlmostEqual(d2, reply) - - d2 = {"volume": 1, "phosphate": 0.4 * 1 / 7.5, "temperature": 10} - reply = tank.get_excess({"volume": 1}) - self.assertDictAlmostEqual(d2, reply) - - def test_tank_push_storage(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.4, "temperature": 15} - - d3 = {"volume": 2.5, "phosphate": 0.2, "temperature": 15} - reply = tank.push_storage(d2) - self.assertDictAlmostEqual(d3, reply) - - d4 = {"volume": 0, "phosphate": 0, "temperature": 0} - reply = tank.push_storage(d2, force=True) - self.assertDictAlmostEqual(d4, reply) - - def test_tank_pull_storage(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.4 * 5 / 7.5, "temperature": 10} - - reply = tank.pull_storage({"volume": 5}) - self.assertDictAlmostEqual(d2, reply) - - d3 = {"volume": 2.5, "phosphate": 0.4 * 2.5 / 7.5, "temperature": 10} - - reply = tank.pull_storage({"volume": 5}) - - self.assertDictAlmostEqual(d3, reply, 15) - - def test_tank_pull_pollutants(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.1, "temperature": 10} - - reply = tank.pull_pollutants(d2) - self.assertDictAlmostEqual(d2, reply) - - reply = tank.pull_pollutants(d2) - d3 = {"volume": 2.5, "phosphate": 0.1, "temperature": 10} - self.assertDictAlmostEqual(d3, reply, 15) - - def test_tank_head(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1, datum=5, area=2.5) - - reply = tank.get_head() - self.assertEqual(8, reply) - - reply = tank.get_head(datum=-1) - self.assertEqual(2, reply) - - reply = tank.get_head(non_head_storage=2) - self.assertEqual(7.2, reply) - - reply = tank.get_head(non_head_storage=10) - self.assertEqual(5, reply) - - def test_evap(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 0, "phosphate": 0.4, "temperature": 10} - - reply = tank.evaporate(10) - self.assertEqual(7.5, reply) - self.assertDictAlmostEqual(d2, tank.storage) - - def test_residence_tank(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = ResidenceTank(residence_time=3, initial_storage=d1) - - d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.pull_outflow() - self.assertDictAlmostEqual(d2, reply) - - def test_decay_tank(self): - node = Node(name="", data_input_dict={("temperature", 1): 15}) - node.t = 1 - d1 = {"volume": 8, "phosphate": 0.4, "temperature": 10} - - tank = DecayTank( - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - initial_storage=d1, - parent=node, - ) - _ = tank.pull_storage({"volume": 2}) - - d3 = {"volume": -2, "phosphate": -0.1, "temperature": 0} - - diff = tank.decay_ds() - self.assertDictAlmostEqual(d3, diff, 16) - - tank.end_timestep_decay() - - d2 = { - "volume": 6, - "phosphate": 0.3 - 0.3 * 0.001 * 1.005 ** (15 - 20), - "temperature": 10, - } - - self.assertDictAlmostEqual(d2, tank.storage, 16) - - self.assertAlmostEqual( - 0.3 * 0.001 * 1.005 ** (15 - 20), tank.total_decayed["phosphate"] - ) - - def test_queue_push(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - - tank.push_storage(d2) - - d3 = {"volume": 6, "phosphate": 0.5, "temperature": (5 * 10 + 15) / 6} - - self.assertDictAlmostEqual(d3, tank.storage) - self.assertDictAlmostEqual(d1, tank.active_storage) - self.assertDictAlmostEqual(d2, tank.internal_arc.queue[1]) - - tank.push_storage(d2, force=True) - self.assertDictAlmostEqual(d3, tank.active_storage) - - tank.end_timestep() - - d4 = {"volume": 7, "phosphate": 0.6, "temperature": ((5 * 10) + (15 * 2)) / 7} - self.assertDictAlmostEqual(d4, tank.active_storage) - - def test_queue_pull(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - reply = tank.push_storage(d2) - - reply = tank.pull_storage({"volume": 6}) - self.assertDictAlmostEqual(d1, reply) - tank.end_timestep() - self.assertDictAlmostEqual(d2, tank.active_storage) - - def test_queue_pull_exact(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - reply = tank.push_storage(d2) - - reply = tank.pull_storage_exact( - {"volume": 6, "phosphate": 0.1, "temperature": 10} - ) - - d3 = {"volume": 5, "phosphate": 0.1, "temperature": 10} - self.assertDictAlmostEqual(d3, reply) - - reply = tank.pull_storage_exact( - {"volume": 0, "phosphate": 0.6, "temperature": 10} - ) - d4 = {"volume": 0, "phosphate": 0.3, "temperature": 10} - self.assertDictAlmostEqual(d4, reply, 16) - - def test_decay_queue(self): - node = Node(name="", data_input_dict={("temperature", 1): 15}) - node.t = 1 - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = DecayQueueTank( - number_of_timesteps=1, - capacity=10, - initial_storage=d1, - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - parent=node, - ) - - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - - _ = tank.push_storage(d2) - - tank.end_timestep() - - d4 = { - "volume": 6, - "phosphate": 0.4 + 0.1 * (1 - 0.001 * 1.005 ** (15 - 20)), - "temperature": ((5 * 10) + (15 * 1)) / 6, - } - self.assertDictAlmostEqual(d4, tank.storage, 15) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_tanks.py b/tests/test_tanks.py new file mode 100644 index 00000000..5c6c759d --- /dev/null +++ b/tests/test_tanks.py @@ -0,0 +1,268 @@ +"""Tests for the tanks module.""" + +import unittest +from unittest import TestCase + +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import ( + DecayQueueTank, + DecayTank, + QueueTank, + ResidenceTank, + Tank, +) + + +class MyTestClass(TestCase): + def assertDictAlmostEqual(self, d1, d2, accuracy=19): + """ + + Args: + d1: + d2: + accuracy: + """ + for d in [d1, d2]: + for key, item in d.items(): + d[key] = round(item, accuracy) + self.assertDictEqual(d1, d2) + + def test_tank_ds(self): + tank = Tank( + capacity=10, + initial_storage={"volume": 5, "phosphate": 0.4, "temperature": 10}, + ) + tank.end_timestep() + + d1 = {"volume": 2, "phosphate": 0.01, "temperature": 15} + + _ = tank.push_storage(d1) + + diff = tank.ds() + + d2 = {"volume": 2, "phosphate": 0.01, "temperature": 0} + + self.assertDictAlmostEqual(d2, diff, 16) + + def test_ponded(self): + tank = Tank( + capacity=10, + initial_storage={"volume": 15, "phosphate": 0.4, "temperature": 10}, + ) + d1 = {"volume": 5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.pull_ponded() + self.assertDictAlmostEqual(d1, reply) + + def test_tank_get_avail(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + reply = tank.get_avail() + self.assertDictAlmostEqual(d1, reply) + + reply = tank.get_avail({"volume": 2.5}) + d2 = {"volume": 2.5, "phosphate": 0.2, "temperature": 10} + self.assertDictAlmostEqual(d2, reply) + + reply = tank.get_avail({"volume": 10}) + self.assertDictAlmostEqual(d1, reply) + + def test_tank_get_excess(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.get_excess() + self.assertDictAlmostEqual(d2, reply) + + d2 = {"volume": 1, "phosphate": 0.4 * 1 / 7.5, "temperature": 10} + reply = tank.get_excess({"volume": 1}) + self.assertDictAlmostEqual(d2, reply) + + def test_tank_push_storage(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.4, "temperature": 15} + + d3 = {"volume": 2.5, "phosphate": 0.2, "temperature": 15} + reply = tank.push_storage(d2) + self.assertDictAlmostEqual(d3, reply) + + d4 = {"volume": 0, "phosphate": 0, "temperature": 0} + reply = tank.push_storage(d2, force=True) + self.assertDictAlmostEqual(d4, reply) + + def test_tank_pull_storage(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.4 * 5 / 7.5, "temperature": 10} + + reply = tank.pull_storage({"volume": 5}) + self.assertDictAlmostEqual(d2, reply) + + d3 = {"volume": 2.5, "phosphate": 0.4 * 2.5 / 7.5, "temperature": 10} + + reply = tank.pull_storage({"volume": 5}) + + self.assertDictAlmostEqual(d3, reply, 15) + + def test_tank_pull_pollutants(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.1, "temperature": 10} + + reply = tank.pull_pollutants(d2) + self.assertDictAlmostEqual(d2, reply) + + reply = tank.pull_pollutants(d2) + d3 = {"volume": 2.5, "phosphate": 0.1, "temperature": 10} + self.assertDictAlmostEqual(d3, reply, 15) + + def test_tank_head(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1, datum=5, area=2.5) + + reply = tank.get_head() + self.assertEqual(8, reply) + + reply = tank.get_head(datum=-1) + self.assertEqual(2, reply) + + reply = tank.get_head(non_head_storage=2) + self.assertEqual(7.2, reply) + + reply = tank.get_head(non_head_storage=10) + self.assertEqual(5, reply) + + def test_evap(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 0, "phosphate": 0.4, "temperature": 10} + + reply = tank.evaporate(10) + self.assertEqual(7.5, reply) + self.assertDictAlmostEqual(d2, tank.storage) + + def test_residence_tank(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = ResidenceTank(residence_time=3, initial_storage=d1) + + d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.pull_outflow() + self.assertDictAlmostEqual(d2, reply) + + def test_decay_tank(self): + node = Node(name="", data_input_dict={("temperature", 1): 15}) + node.t = 1 + d1 = {"volume": 8, "phosphate": 0.4, "temperature": 10} + + tank = DecayTank( + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + initial_storage=d1, + parent=node, + ) + _ = tank.pull_storage({"volume": 2}) + + d3 = {"volume": -2, "phosphate": -0.1, "temperature": 0} + + diff = tank.decay_ds() + self.assertDictAlmostEqual(d3, diff, 16) + + tank.end_timestep_decay() + + d2 = { + "volume": 6, + "phosphate": 0.3 - 0.3 * 0.001 * 1.005 ** (15 - 20), + "temperature": 10, + } + + self.assertDictAlmostEqual(d2, tank.storage, 16) + + self.assertAlmostEqual( + 0.3 * 0.001 * 1.005 ** (15 - 20), tank.total_decayed["phosphate"] + ) + + def test_queue_push(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + + tank.push_storage(d2) + + d3 = {"volume": 6, "phosphate": 0.5, "temperature": (5 * 10 + 15) / 6} + + self.assertDictAlmostEqual(d3, tank.storage) + self.assertDictAlmostEqual(d1, tank.active_storage) + self.assertDictAlmostEqual(d2, tank.internal_arc.queue[1]) + + tank.push_storage(d2, force=True) + self.assertDictAlmostEqual(d3, tank.active_storage) + + tank.end_timestep() + + d4 = {"volume": 7, "phosphate": 0.6, "temperature": ((5 * 10) + (15 * 2)) / 7} + self.assertDictAlmostEqual(d4, tank.active_storage) + + def test_queue_pull(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + reply = tank.push_storage(d2) + + reply = tank.pull_storage({"volume": 6}) + self.assertDictAlmostEqual(d1, reply) + tank.end_timestep() + self.assertDictAlmostEqual(d2, tank.active_storage) + + def test_queue_pull_exact(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + reply = tank.push_storage(d2) + + reply = tank.pull_storage_exact( + {"volume": 6, "phosphate": 0.1, "temperature": 10} + ) + + d3 = {"volume": 5, "phosphate": 0.1, "temperature": 10} + self.assertDictAlmostEqual(d3, reply) + + reply = tank.pull_storage_exact( + {"volume": 0, "phosphate": 0.6, "temperature": 10} + ) + d4 = {"volume": 0, "phosphate": 0.3, "temperature": 10} + self.assertDictAlmostEqual(d4, reply, 16) + + def test_decay_queue(self): + node = Node(name="", data_input_dict={("temperature", 1): 15}) + node.t = 1 + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = DecayQueueTank( + number_of_timesteps=1, + capacity=10, + initial_storage=d1, + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + parent=node, + ) + + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + + _ = tank.push_storage(d2) + + tank.end_timestep() + + d4 = { + "volume": 6, + "phosphate": 0.4 + 0.1 * (1 - 0.001 * 1.005 ** (15 - 20)), + "temperature": ((5 * 10) + (15 * 1)) / 6, + } + self.assertDictAlmostEqual(d4, tank.storage, 15) + + +if __name__ == "__main__": + unittest.main() diff --git a/wsimod/nodes/land.py b/wsimod/nodes/land.py index f34fb750..7bcaf86a 100644 --- a/wsimod/nodes/land.py +++ b/wsimod/nodes/land.py @@ -8,8 +8,9 @@ from math import exp, log, log10, sin from wsimod.core import constants -from wsimod.nodes.nodes import DecayTank, Node, ResidenceTank +from wsimod.nodes.nodes import Node from wsimod.nodes.nutrient_pool import NutrientPool +from wsimod.nodes.tanks import DecayTank, ResidenceTank class Land(Node): diff --git a/wsimod/nodes/sewer.py b/wsimod/nodes/sewer.py index af566658..545495c1 100644 --- a/wsimod/nodes/sewer.py +++ b/wsimod/nodes/sewer.py @@ -5,7 +5,8 @@ Converted to totals on 2022-05-03 """ from wsimod.core import constants -from wsimod.nodes.nodes import Node, QueueTank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import QueueTank class Sewer(Node): diff --git a/wsimod/nodes/storage.py b/wsimod/nodes/storage.py index eab68665..2cd75990 100644 --- a/wsimod/nodes/storage.py +++ b/wsimod/nodes/storage.py @@ -6,7 +6,8 @@ from math import exp from wsimod.core import constants -from wsimod.nodes.nodes import DecayQueueTank, DecayTank, Node, QueueTank, Tank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import DecayQueueTank, DecayTank, QueueTank, Tank class Storage(Node): diff --git a/wsimod/nodes/tanks.py b/wsimod/nodes/tanks.py index b2036edb..4806d6ae 100644 --- a/wsimod/nodes/tanks.py +++ b/wsimod/nodes/tanks.py @@ -1,6 +1,6 @@ """Module for defining tanks.""" -from wsimod.arcs import AltQueueArc, DecayArcAlt +from wsimod.arcs.arcs import AltQueueArc, DecayArcAlt from wsimod.core import constants from wsimod.core.core import DecayObj, WSIObj diff --git a/wsimod/nodes/wtw.py b/wsimod/nodes/wtw.py index e40f0a5d..afd6c6a5 100644 --- a/wsimod/nodes/wtw.py +++ b/wsimod/nodes/wtw.py @@ -7,7 +7,8 @@ from typing import Any, Dict from wsimod.core import constants -from wsimod.nodes.nodes import Node, Tank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import Tank class WTW(Node): diff --git a/wsimod/orchestration/model.py b/wsimod/orchestration/model.py index 8dc4e917..4026a325 100644 --- a/wsimod/orchestration/model.py +++ b/wsimod/orchestration/model.py @@ -19,7 +19,8 @@ from wsimod.core import constants from wsimod.core.core import WSIObj from wsimod.nodes.land import ImperviousSurface -from wsimod.nodes.nodes import NODES_REGISTRY, QueueTank, ResidenceTank, Tank +from wsimod.nodes.nodes import NODES_REGISTRY +from wsimod.nodes.tanks import QueueTank, ResidenceTank, Tank os.environ["USE_PYGEOS"] = "0"