Withdrawing water from a canal: an agent-based model using Python
Agent-based modeling (ABM) é a powerful technique used in a bunch of areas, including hydrology. ABMs allow us to simulate the complex behavior of natural and social systems, so scientists can better understand the relationship between different elements in a system. In hydrology, ABM has been used to simulate processes in rivers, watersheds and other hydro-systems.
In this article, we will explore the importance of ABM in hydrology and we will present an agent-based model that simulates water withdrawing from a canal. We will only use Python, which will give us the freedom to implement personalized and flexible functions. In this way, we will not be dependent on any software.
As a demonstration, we will use an example of users who withdraw water from a linear water canal. Among the functionalities that we will see are: agent creation, implementation of interaction rules, graph-based model (upstream users affect the water quantity available for downstream users), and a personalized agent activation rule. You can check this link to see the Jupyter Notebook of this tutorial. I also published this article in Portuguese.
Setup
We will use the Mesa Python package to build our model. Mesa is used to implement ABMs. It has base classes to build the main ABM functionalities.
If you do not have Mesa installed, follow the documentation. Also, I recommend using Anaconda to manage the Python packages.
In short, create a new virtual environment at Anaconda and activate it.
#Anaconda
conda activate mesaenv
#Terminal
source bin/activateThen, install all four packages we are going to use here: mesa, matplotlib, networkx, and jupyter (to run the notebook).
python3 -m pip install mesa
python3 -m pip install matplotlib
python3 -m pip install networkx
python3 -m pip install jupyterAgent and Model Classes
We start creating a simple model, then we are going to implement more functions into it. We create two classes, one represents the users that withdraw water from the canal, while the other handles global model variables. They are represented by the classesWaterUserAgent and WaterModel.
from mesa import Agent, Model
class WaterUserAgent(Agent):
"""An agent that interacts with the water canal by withdrawing water for he/she own use."""
_last_id = 0
def __init__(self, model):
super().__init__(WaterUserAgent._last_id+1, model)
WaterUserAgent._last_id += 1
self.water_to_withdraw = 1
class WaterModel(Model):
"""
Create a new Irrigation model with the given parameters.
Args:
N: int number that represent the number of water users to create per year
"""
def __init__(self, N=10):
self.num_agents = N
# Create agents
for i in range(self.num_agents):
w = WaterUserAgent(self)First, we import gent and Model from mesa. At the init method in WaterUserAgent we define that each water user has an attribute called water_to_withdraw that represents the water volume the user is going to withdraw. In WaterModel, we receive the number of agents that are going to be created and create them (but we did not add them into the model yet!).
Setting up space
As we explained before, in a water canal actions from upstream users affect downstream ones. For example, in the extreme case that the first water user withdraws all the available water from the canal, all other users won’t have any water to use. To model this behavior we are going to use the NetworkGrid, which is a graph-based way to define space. We are going to need help from NetworkX, a python library to create graphs (a kind of network). We are going to create a linear graph. In this way, each node is connected to the previous and next node only.
Based on a linear graph we create the logic of a canal. See that by changing the graph relations we can simulate not only canals but entire watersheds!
Let’s create our linear graph:
import networkx as nx
import matplotlib.pyplot as plt
G = nx.path_graph(100)
nx.draw_spectral(G, with_labels=True, font_size=8)
Due to the number of nodes, the visualization is not pretty, but it works like a charm. We create a graph with 100 nodes. Each node represents a possible place that an agent can allocate themselves. Only one agent can be allocated in a node. So, the number of nodes restricts the maximum number of agents in the model (keep that in mind when you’re building yours). In this example, we are going to simulate a small water canal. So, 100 is fine.
Our linear graph is ready, so we are going to update the WaterModel. We also are going to create a function that allocates water users in a random node.
from mesa.space import NetworkGrid
import random
class WaterModel(Model):
"""
Create a new Irrigation model with the given parameters.
Args:
N: int number that represent the number of water users to create per year
G: generated linear graph using networkx
"""
def __init__(self, N=2):
self.num_agents = N
self.grid = NetworkGrid(G)
self.number_of_farmers_to_create = 10
# Create agents
for i in range(self.num_agents):
w = WaterUserAgent(self)
def create_farmers_random_position(self):
"""Create water users at random position in a linear graph."""
while (i < self.number_of_farmers_to_create):
random_node = random.sample(list(G.nodes()), 1)
if (len(self.grid.get_cell_list_contents(random_node)) == 0):
f = WaterUserAgent(self)
self.schedule.add(f)
self.grid.place_agent(f, random_node[0])
i += 1
def step(self):
""" Execute the step for all agents, one at a time. At the end, advance model by one step """
# Preparation
self.create_farmers_random_position()
# Run step
self.schedule.step()Above, we created some agents and we positioned them in a random node. By executing this behavior in the step method we make sure that this is executed at each model iteration.
Reordering Agents Activation (Custom Schedule)
After creating agents, we need to create the sequence of steps that are going to be executed when they are activated. To do so, we are going to create another stepmethod. This time, it is going to be placed in the Agents class. We need to execute step in the Schedule to execute the functions of this method to all agents, one at a time. The Schedule is part of the time module of mesa. It controls when each agent executes their actions. As the order of agent activation impacts the results, it needs to be expressed.
We want agents to withdraw water from the canal (this is the action they will perform). By default, the scheduler activates agents randomly. However, in a water canal, this is not how it is supposed to work. Upstream users affect downstream users regarding water availability. Therefore, upstream agents need to be activated first. To redefine the activation order based on their position in the canal, we need to create a custom Scheduler.
Let’s put the creation of the step method on hold and focus on Scheduler.
First, we are going to create a custom scheduler that inherits the class BaseSchedule. The BaseSchedule is a base class that activates agents one by one and assumes that they have a step method. Then, we are going to create e method that receives the correct order to activate agents, ordering them by their position in the model. Finally, we create a method that reorders the agent's activation.
Mesa uses an Ordered Dictionary, which is a subclass of the traditional dictionary in Python. It “remembers” the order that items were added to it using a key to each dictionary.
In case you have multiple agents, you can also redefine the order which each agent class is activated. See this link for more details.
Below we create our custom scheduler.
from mesa.time import BaseScheduler
from collections import OrderedDict
class CanalScheduler(BaseScheduler):
def __init__(self, model):
super().__init__(model)
def sort_agents(self):
sorted_list_of_agents = sorted(self.agents, key=lambda x: x.pos)
new_queue = OrderedDict(dict((i,j) for i,j in enumerate(sorted_list_of_agents)))
return new_queue
def step(self):
self._agents = self.sort_agents()
super().step()The scheduler needs to have a step method. This method is called when we run self.schedule.step() at the WaterModel. It runs sort_agentsthat reorders agents based on their position.
Now that we have our custom scheduler, we finally can create the run_model method that calls the step method from the scheduler. To test, let’s create a step method so each agent prints a message when they are activated. In this way, we can test if our implementation is correct.
class WaterUserAgent(Agent):
"""An agent that interacts with the water canal by withdrawing water for he/she own use."""
_last_id = 0
def __init__(self, model):
super().__init__(WaterUserAgent._last_id+1, model)
WaterUserAgent._last_id += 1
self.water_to_withdraw = 1
def step(self):
print("Olá, Eu sou o Water Agent nº {} e estou na posição {}.".format(
self.unique_id, self.pos))
class WaterModel(Model):
"""
Create a new Irrigation model with the given parameters.
Args:
N: int number that represent the number of water users to create per year
G: generated linear graph using networkx
"""
def __init__(self, N=2):
self.num_agents = N
self.grid = NetworkGrid(G)
self.number_of_farmers_to_create = 2
# Define scheduler
self.schedule = CanalScheduler(
WaterModel)
# Create agents
for i in range(self.num_agents):
w = WaterUserAgent(self)
def create_farmers_random_position(self):
"""Create water users at random position in a linear graph."""
i = 0
while (i < self.number_of_farmers_to_create):
random_node = random.sample(list(G.nodes()), 1)
if (len(self.grid.get_cell_list_contents(random_node)) == 0):
f = WaterUserAgent(self)
self.schedule.add(f)
self.grid.place_agent(f, random_node[0])
i += 1
def step(self):
""" Execute the step for all agents, one at a time. At the end, advance model by one step """
# Preparation
self.create_farmers_random_position()
# Run step
self.schedule.step()
def run_model(self, step_count=3):
for i in range(1,step_count+1):
print("-------------- \n" +
"Initiating year n. " + str(i) + "\n" +
"--------------")
self.step()model.run_model()
--------------
Initiating year n. 1
--------------
Hi, I am Water Agent n. 5 and I am in position 26. I am withdrawing 4.245385761173964 m³/year
Available Water: 95.75461423882604
Hi, I am Water Agent n. 4 and I am in position 50. I am withdrawing 7.98141842703323 m³/year
Available Water: 87.7731958117928
Hi, I am Water Agent n. 6 and I am in position 68. I am withdrawing 2.270441679844418 m³/year
Available Water: 85.50275413194838
--------------
Initiating year n. 2
--------------
Hi, I am Water Agent n. 7 and I am in position 12. I am withdrawing 4.615901240437739 m³/year
Available Water: 95.38409875956226
Hi, I am Water Agent n. 5 and I am in position 26. I am withdrawing 4.245385761173964 m³/year
Available Water: 91.1387129983883
Hi, I am Water Agent n. 4 and I am in position 50. I am withdrawing 7.98141842703323 m³/year
Available Water: 83.15729457135507
Hi, I am Water Agent n. 8 and I am in position 51. I am withdrawing 5.55253531624829 m³/year
Available Water: 77.60475925510677
Hi, I am Water Agent n. 9 and I am in position 64. I am withdrawing 6.911028324209107 m³/year
Available Water: 70.69373093089766
Hi, I am Water Agent n. 6 and I am in position 68. I am withdrawing 2.270441679844418 m³/year
Available Water: 68.42328925105323
...As agents’ positions are ordered, it seems our model has been implemented correctly so far. And this means that our model is done! Before we move on to visualization, I would like to address a few things:
- We implemented a random draw from a uniform distribution in
water_to_withdrawjust to show this example. In a real case, I suggest you try to fit a power-law, log-normal or gama distributions, because they will considerate the assimetry in water withdrawal data. - We defined the step size as 1 year because it simplifies the water balance. You can always choose to use a finer discretization if it fits your model purpose.
- At most countries, a government agency manages water withdraws. In this case, you can define multiple agent classes by multiple water purposes (irrigation, human supply, …) and new water structures like the presence of reservoirs.
Collecting data and plotting them
To get data from our ABM, we are going to use an internal function from mesa called DataCollector. Including data collect, the final version of our ABM is going to look like this:
from mesa.datacollection import DataCollector
class WaterUserAgent(Agent):
"""An agent that interacts with the water canal by withdrawing water for he/she own use."""
_last_id = 0
def __init__(self, model):
super().__init__(WaterUserAgent._last_id+1, model)
WaterUserAgent._last_id += 1
self.water_to_withdraw = random.uniform(1, 9)
def withdraw_water(self):
model.water_available -= self.water_to_withdraw
def step(self):
print("Hi, I am Water Agent n. {} and I am in position {}. I am withdrawing {} m³/year".format(
self.unique_id, self.pos, self.water_to_withdraw))
self.withdraw_water()
print("Available Water: {}".format(model.water_available))
class WaterModel(Model):
"""
Create a new Irrigation model with the given parameters.
Args:
N: int number that represent the number of water users to create per year
G: generated linear graph using networkx
water_available: Total water to supply the canal (in m³/year)
"""
def __init__(self,
N,
G,
water_available):
self.num_agents = N
self.grid = NetworkGrid(G)
self.number_of_farmers_to_create = N
self.water_available_to_reset = water_available
# Define scheduler
self.schedule = CanalScheduler(WaterModel)
# Create agents
for i in range(self.num_agents):
w = WaterUserAgent(self)
# Data Collector
self.datacollector = DataCollector(
agent_reporters={
"Position":
lambda x: x.pos,
"Withdrawn water (m³/year)":
lambda x: x.water_to_withdraw,
},
)
self.datacollector.collect(self)
def create_farmers_random_position(self):
"""Create water users at random position in a linear graph."""
i = 0
while (i < self.number_of_farmers_to_create):
random_node = random.sample(list(G.nodes()), 1)
if (len(self.grid.get_cell_list_contents(random_node)) == 0):
f = WaterUserAgent(self)
self.schedule.add(f)
self.grid.place_agent(f, random_node[0])
i += 1
def step(self):
""" Execute the step for all agents, one at a time. At the end, advance model by one step """
# Preparation
self.create_farmers_random_position()
self.water_available = self.water_available_to_reset
# Run step
self.schedule.step()
# Save data
self.datacollector.collect(self)
def run_model(self, step_count=3):
for i in range(1,step_count+1):
print("-------------- \n" +
"Initiating year n. " + str(i) + "\n" +
"--------------")
self.step()water_available = 100
model = WaterModel(random.randint(1, 5), G, water_available)
model.run_model(10)
agents_results = model.datacollector.get_agent_vars_dataframe()import matplotlib.pyplot as plt
def water_balance_in_canal(agent_results):
index = []
water_volume = []
for i in range(agent_results.index.get_level_values(0).min(), agent_results.index.get_level_values(0).max()+1):
results_current_step = agent_results.xs(i, level=0)
index.append(i)
water_volume.append(results_current_step['Withdrawn water (m³/year)'].sum())
water_balance = [water_available - x for x in water_volume]
fig, ax = plt.subplots(dpi=100)
ax.axhline(0, linestyle='--', color='red')
ax.plot(index, water_balance, color='blue')
ax.set_xlabel('Step')
ax.set_ylabel('Balanço hídrico anual')
water_balance_in_canal(agents_results)
Which modifications would you make and which methods would you implement to get even closer to the real case?
