diff --git a/wsimod/nodes/__init__.py b/wsimod/nodes/__init__.py index b67eb56..a8de732 100644 --- a/wsimod/nodes/__init__.py +++ b/wsimod/nodes/__init__.py @@ -3,7 +3,7 @@ from wsimod.nodes.demand import Demand, NonResidentialDemand, ResidentialDemand from wsimod.nodes.distribution import Distribution, UnlimitedDistribution from wsimod.nodes.land import Land -from wsimod.nodes.nodes import Node +from wsimod.nodes.nodes import NODES_REGISTRY, Node from wsimod.nodes.sewer import EnfieldFoulSewer, Sewer from wsimod.nodes.storage import ( Groundwater, diff --git a/wsimod/nodes/nodes.py b/wsimod/nodes/nodes.py index b8ba231..fdfa77a 100644 --- a/wsimod/nodes/nodes.py +++ b/wsimod/nodes/nodes.py @@ -5,15 +5,24 @@ Converted to totals on Thur Apr 21 2022 """ +import logging + from wsimod.arcs.arcs import AltQueueArc, DecayArcAlt from wsimod.core import constants from wsimod.core.core import DecayObj, WSIObj -from wsimod.nodes import nodes class Node(WSIObj): """""" + def __init_subclass__(cls, **kwargs): + """Adds all subclasses to the nodes registry.""" + super().__init_subclass__(**kwargs) + if cls.__name__ in NODES_REGISTRY: + logging.warning(f"Overwriting {cls.__name__} in NODES_REGISTRY with {cls}") + + NODES_REGISTRY[cls.__name__] = cls + def __init__(self, name, data_input_dict=None): """Base class for CWSD nodes. Constructs all the necessary attributes for the node object. @@ -34,22 +43,7 @@ def __init__(self, name, data_input_dict=None): Input data and parameter requirements: - All nodes require a `name` """ - - # Get node types - def all_subclasses(cls): - """ - - Args: - cls: - - Returns: - - """ - return set(cls.__subclasses__()).union( - [s for c in cls.__subclasses__() for s in all_subclasses(c)] - ) - - node_types = [x.__name__ for x in all_subclasses(nodes.Node)] + ["Node"] + node_types = list(NODES_REGISTRY.keys()) # Default essential parameters # Dictionary of arcs @@ -632,8 +626,8 @@ def reinit(self): """ This is an attempt to generalise the behaviour of pull/push_distributed It doesn't yet work... - - def general_distribute(self, vqip, of_type = None, tag = 'default', direction = + + def general_distribute(self, vqip, of_type = None, tag = 'default', direction = None): if direction == 'push': arcs = self.out_arcs @@ -649,49 +643,49 @@ def general_distribute(self, vqip, of_type = None, tag = 'default', direction = values()} else: print('No direction') - + if len(arcs) == 1: - if (of_type == None) | any([x in of_type for x, y in arcs_type.items() if + if (of_type == None) | any([x in of_type for x, y in arcs_type.items() if len(y) > 0]): arc = next(iter(arcs.keys())) return requests[arc](vqip) else: #No viable arcs return tracker - + connected = self.get_connected(direction = direction, of_type = of_type, tag = tag) - + iter_ = 0 - + target = self.copy_vqip(vqip) - #Iterate over sending nodes until deficit met - while (((target['volume'] > constants.FLOAT_ACCURACY) & - (connected['avail'] > constants.FLOAT_ACCURACY)) & + #Iterate over sending nodes until deficit met + while (((target['volume'] > constants.FLOAT_ACCURACY) & + (connected['avail'] > constants.FLOAT_ACCURACY)) & (iter_ < constants.MAXITER)): - - amount = min(connected['avail'], target['volume']) #Deficit or amount + + amount = min(connected['avail'], target['volume']) #Deficit or amount still to push replies = self.empty_vqip() - + for key, allocation in connected['allocation'].items(): to_request = amount * allocation / connected['priority'] to_request = self.v_change_vqip(target, to_request) reply = requests[key](to_request) replies = self.sum_vqip(replies, reply) - + if direction == 'pull': target = self.extract_vqip(target, replies) elif direction == 'push': target = replies - + connected = self.get_connected(direction = direction, of_type = of_type, tag = tag) - iter_ += 1 + iter_ += 1 - if iter_ == constants.MAXITER: + if iter_ == constants.MAXITER: print('Maxiter reached') return target""" @@ -739,6 +733,9 @@ def general_distribute(self, vqip, of_type = None, tag = 'default', direction = # self.__dict__.update(newnode.__dict__) +NODES_REGISTRY: dict[str, type[Node]] = {Node.__name__: Node} + + class Tank(WSIObj): """""" diff --git a/wsimod/orchestration/model.py b/wsimod/orchestration/model.py index 6aca8ff..8e89017 100644 --- a/wsimod/orchestration/model.py +++ b/wsimod/orchestration/model.py @@ -15,12 +15,11 @@ import yaml from tqdm import tqdm -from wsimod import nodes from wsimod.arcs import arcs as arcs_mod from wsimod.core import constants from wsimod.core.core import WSIObj from wsimod.nodes.land import ImperviousSurface -from wsimod.nodes.nodes import Node, QueueTank, ResidenceTank, Tank +from wsimod.nodes.nodes import NODES_REGISTRY, QueueTank, ResidenceTank, Tank os.environ["USE_PYGEOS"] = "0" @@ -141,25 +140,6 @@ def __init__(self): self.nodes = {} self.nodes_type = {} - def all_subclasses(cls): - """ - - Args: - cls: - - Returns: - - """ - return set(cls.__subclasses__()).union( - [s for c in cls.__subclasses__() for s in all_subclasses(c)] - ) - - self.nodes_type = [x.__name__ for x in all_subclasses(Node)] + ["Node"] - self.nodes_type = set( - getattr(nodes, x)(name="").__class__.__name__ for x in self.nodes_type - ).union(["Foul"]) - self.nodes_type = {x: {} for x in self.nodes_type} - def get_init_args(self, cls): """Get the arguments of the __init__ method for a class and its superclasses.""" init_args = [] @@ -414,19 +394,6 @@ def add_nodes(self, nodelist): nodelist (list): List of dicts, where a dict is a node """ - def all_subclasses(cls): - """ - - Args: - cls: - - Returns: - - """ - return set(cls.__subclasses__()).union( - [s for c in cls.__subclasses__() for s in all_subclasses(c)] - ) - for data in nodelist: name = data["name"] type_ = data["type_"] @@ -441,7 +408,14 @@ def all_subclasses(cls): if "geometry" in data.keys(): del data["geometry"] del data["type_"] - self.nodes_type[type_][name] = getattr(nodes, node_type)(**dict(data)) + + if node_type not in NODES_REGISTRY.keys(): + raise ValueError(f"Node type {type_} not recognised") + + if type_ not in self.nodes_type.keys(): + self.nodes_type[type_] = {} + + self.nodes_type[type_][name] = NODES_REGISTRY[node_type](**dict(data)) self.nodes[name] = self.nodes_type[type_][name] self.nodelist = [x for x in self.nodes.values()] @@ -455,7 +429,10 @@ def add_instantiated_nodes(self, nodelist): self.nodelist = nodelist self.nodes = {x.name: x for x in nodelist} for x in nodelist: - self.nodes_type[x.__class__.__name__][x.name] = x + type_ = x.__class__.__name__ + if type_ not in self.nodes_type.keys(): + self.nodes_type[type_] = {} + self.nodes_type[type_][x.name] = x def add_arcs(self, arclist): """Add nodes to the model object from a list of dicts, where each dict contains @@ -488,15 +465,20 @@ def add_arcs(self, arclist): river_arcs[name] = self.arcs[name] if any(river_arcs): - upstreamness = {x: 0 for x in self.nodes_type["Waste"].keys()} + upstreamness = ( + {x: 0 for x in self.nodes_type["Waste"].keys()} + if "Waste" in self.nodes_type + else {} + ) upstreamness = self.assign_upstream(river_arcs, upstreamness) self.river_discharge_order = [] - for node in sorted( - upstreamness.items(), key=lambda item: item[1], reverse=True - ): - if node[0] in self.nodes_type["River"].keys(): - self.river_discharge_order.append(node[0]) + if "River" in self.nodes_type: + for node in sorted( + upstreamness.items(), key=lambda item: item[1], reverse=True + ): + if node[0] in self.nodes_type["River"]: + self.river_discharge_order.append(node[0]) def add_instantiated_arcs(self, arclist): """Add arcs to the model object from a list of objects, where each object is an @@ -522,16 +504,23 @@ def add_instantiated_arcs(self, arclist): "Reservoir", ]: river_arcs[arc.name] = arc + + upstreamness = ( + {x: 0 for x in self.nodes_type["Waste"].keys()} + if "Waste" in self.nodes_type + else {} + ) upstreamness = {x: 0 for x in self.nodes_type["Waste"].keys()} upstreamness = self.assign_upstream(river_arcs, upstreamness) self.river_discharge_order = [] - for node in sorted( - upstreamness.items(), key=lambda item: item[1], reverse=True - ): - if node[0] in self.nodes_type["River"].keys(): - self.river_discharge_order.append(node[0]) + if "River" in self.nodes_type: + for node in sorted( + upstreamness.items(), key=lambda item: item[1], reverse=True + ): + if node[0] in self.nodes_type["River"]: + self.river_discharge_order.append(node[0]) def assign_upstream(self, arcs, upstreamness): """Recursive function to trace upstream up arcs to determine which are the most @@ -726,53 +715,53 @@ def enablePrint(stdout): node.monthyear = date.to_period("M") # Run FWTW - for node in self.nodes_type["FWTW"].values(): + for node in self.nodes_type.get("FWTW", {}).values(): node.treat_water() # Create demand (gets pushed to sewers) - for node in self.nodes_type["Demand"].values(): + for node in self.nodes_type.get("Demand", {}).values(): node.create_demand() # Create runoff (impervious gets pushed to sewers, pervious to groundwater) - for node in self.nodes_type["Land"].values(): + for node in self.nodes_type.get("Land", {}).values(): node.run() # Infiltrate GW - for node in self.nodes_type["Groundwater"].values(): + for node in self.nodes_type.get("Groundwater", {}).values(): node.infiltrate() # Discharge sewers (pushed to other sewers or WWTW) - for node in self.nodes_type["Sewer"].values(): + for node in self.nodes_type.get("Sewer", {}).values(): node.make_discharge() # Foul second so that it can discharge any misconnection - for node in self.nodes_type["Foul"].values(): + for node in self.nodes_type.get("Foul", {}).values(): node.make_discharge() # Discharge WWTW - for node in self.nodes_type["WWTW"].values(): + for node in self.nodes_type.get("WWTW", {}).values(): node.calculate_discharge() # Discharge GW - for node in self.nodes_type["Groundwater"].values(): + for node in self.nodes_type.get("Groundwater", {}).values(): node.distribute() # river - for node in self.nodes_type["River"].values(): + for node in self.nodes_type.get("River", {}).values(): node.calculate_discharge() # Abstract - for node in self.nodes_type["Reservoir"].values(): + for node in self.nodes_type.get("Reservoir", {}).values(): node.make_abstractions() - for node in self.nodes_type["Land"].values(): + for node in self.nodes_type.get("Land", {}).values(): node.apply_irrigation() - for node in self.nodes_type["WWTW"].values(): + for node in self.nodes_type.get("WWTW", {}).values(): node.make_discharge() # Catchment routing - for node in self.nodes_type["Catchment"].values(): + for node in self.nodes_type.get("Catchment", {}).values(): node.route() # river @@ -905,7 +894,7 @@ def enablePrint(stdout): for pol in constants.POLLUTANTS: tanks[-1][pol] = prop.storage[pol] - for name, node in self.nodes_type["Land"].items(): + for name, node in self.nodes_type.get("Land", {}).items(): for surface in node.surfaces: if not isinstance(surface, ImperviousSurface): surfaces.append(