Skip to content

Commit

Permalink
Merge pull request #60 from ImperialCollegeLondon/node_registration
Browse files Browse the repository at this point in the history
Dynamic inclusion of node types
  • Loading branch information
dalonsoa committed Feb 19, 2024
2 parents 3038e6a + 31b6216 commit e3b66fe
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 96 deletions.
2 changes: 1 addition & 1 deletion wsimod/nodes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
65 changes: 31 additions & 34 deletions wsimod/nodes/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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"""

Expand Down Expand Up @@ -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):
""""""

Expand Down
111 changes: 50 additions & 61 deletions wsimod/orchestration/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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 = []
Expand Down Expand Up @@ -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_"]
Expand All @@ -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()]

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit e3b66fe

Please sign in to comment.