Extending Pywr with custom Nodes

Nodes and subclasses thereof provide the basic network structure in Pywr. There are several different types of node available. Two major categories exist: flow nodes and storage nodes. Flow nodes are the typical nodes that represent rivers, pipes and other features from, through or to which a resource can flow. These nodes are typically characterised by minimum and maximum flow rates. Storage nodes provide the ability to store resource from one time-step to another, and are characterised by minimum, maximum and initial volumes.

There are three fundamental types of node in Pywr:

Input nodes add water to the system
Output nodes remove water from the system
Link nodes do not add or remove water from the system

There is a fourth node type, Storage, which can be considered fundamental because there are special rules for it’s behaviour in the linear programme used to solve the water balance:

Storage nodes can carry water from one timestep to the next

All other node types in Pywr are subclasses of these base types. For example, the pywr.nodes.Catchment node type is a special case of Input where the min_flow and max_flow properties are equal.

The most common way to create a new node type is using a compound node. A compound node contains one or more existing nodes and is used to manage common or more complex arrangements of the basic node types. An example of a compound node is the PiecewiseLink. It is composed of a link (OUT_1) which receives water from upstream and an link (IN_1) which conveys water downstream, connected by a set of links in parallel (LNK_1 … LNK_N) each with a different max_flow and cost, illustrated below:

                           /-->-- LNK_1 -->--\
UPSTREAM -->-- OUT_1 -->--|--->--  ...  -->---|-->-- IN_1 -->-- DOWNSTREAM
                           \-->-- LNK_N -->--/

Let’s look at an example to create a new node type that represents a leaky pipe. To remove water from the system we need to use an output node (LEAK), with two links representing the boundaries of the compound node (INFLOW and OUTFLOW):

UPSTREAM -->-- INFLOW -->-- OUTFLOW -->-- DOWNSTREAM
                 |
                 \------>--  LEAK

This is a simple structure which represents leakage as a demand with a maximum value and a benefit to be supplied. It is slightly flawed as the leakage volume does not vary proportionally to flow through the link, but is sufficient as an example:

from pywr.nodes import Node, Link, Output

class LeakyPipe(Node):
    def __init__(self, leakage, leakage_cost=-99999, *args, **kwargs):
        self.allow_isolated = True  # Required for compound nodes

        super(LeakyPipe, self).__init__(*args, **kwargs)

        # Define the internal nodes. The parent of the nodes is defined to identify them as sub-nodes.
        self.inflow = Link(self.model, name='{} In'.format(self.name), parent=self)
        self.outflow = Link(self.model, name='{} Out'.format(self.name), parent=self)
        self.leak = Output(self.model, name='{} Leak'.format(self.name), parent=self)

        # Connect the internal nodes
        self.inflow.connect(self.outflow)
        self.inflow.connect(self.leak)

        # Define the properties of the leak (demand and benefit)
        self.leak.max_flow = leakage
        self.leak.cost = leakage_cost

    def iter_slots(self, slot_name=None, is_connector=True):
        # This is required so that connecting to this node actually connects to the outflow sub-node, and
        # connecting from this node actually connects to the input sub-node
        if is_connector:
            yield self.outflow
        else:
            yield self.inflow

    def after(self, timestep):
        # Make the flow on the compound node appear as the flow _after_ the leak
        self.commit_all(self.outflow.flow)
        # Make sure save is done after setting aggregated flow
        super(LeakyPipe, self).after(timestep)

    @classmethod
    def load(cls, data, model):
        del(data["type"])
        leakage = data.pop("leakage")
        leakage_cost = data.pop("leakage_cost", None)
        return cls(model, leakage, leakage_cost, **data)

The custom node does not need to be “registered”, unlike Parameters, as this is done automatically using metaclasses. The new node type can be referenced from a JSON model provided that the class has been imported before the JSON is loaded:

from pywr.model import Model
import leakypipe

model = Model.load("leaky_pipe_model.sjon")

{
    "type": "leakypipe",
    "leakage": "1.0"
}

The allow_isolated attribute identifies nodes of this type as compound nodes. Without this the model would raise an error that the node is not connected to the rest of the network, as the connections are actually to its sub-nodes.

The after method is not required but is useful so that recorders can be attached to the compound node. Without this the flow would appear to be zero as the flow doesn’t actually pass through the compound node.

The iter_slots method is required so that connecting to/from the node (e.g. upstream.connect(leaky)) creates connections to the sub-nodes.

A more advanced representation of the leaky pipe could use an additional AggregatedNode to constrain the ratio of flow through the OUTFLOW and LEAK nodes. [*]