Note
This page was generated from a jupyter notebook.
Profiling and Scaling Analysis of the NetworkSedimentTransporter¶
Planning to run the NetworkSedimentTransporter (NST) on a large river network, or for a long time? You might wonder: how efficient is this component?
Most field-scale applications of the NST would likley have 100-500 links in the network and be initialized with at least 100 parcels per link (e.g., up to 50k+ parcels). This number of parcels per link are commonly need to capture the full grain size distribution.
Because the component may use many parcels on many links, we will want to assess two parts of the components efficiency. First, we will profile the code to see what happens each time the component runs. This is useful to see which parts of the computation take the most time and occur the most often. If code is running too slowly, it is usefull to profile it in order to identify which parts may yield the biggest speedups.
The second goal is scaling, or seeing how changes in the number of parcels and links changes run time. If doubling the number of links doubled the run time, we would say that runtime for the component “scales linearly” with the number of links. In computer science this is often called the time complexity of the code/algorithm.
The goal of this notebook is to explore profiling and scaling with the NST. We will look at how changing the grid nodes and number of parcels per link impacts the initialization and run time of this component. In order to do this we will need to make some python functions to create variably sized grids. We will also use some common python tools for profiling code. The notebook is organized into three parts:
Part 1: Create generic, variable sized grids and parcels.
Part 2: Profiling the code.
Part 3: Scaling analysis.
We begin by importaing all needed python modules.
[ ]:
import cProfile
import io
import pstats
import time
import warnings
from pstats import SortKey
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from landlab.components import FlowDirectorSteepest, NetworkSedimentTransporter
from landlab.data_record import DataRecord
from landlab.grid.network import NetworkModelGrid
from landlab.plot import graph
warnings.filterwarnings("ignore")
Part 1: Variably-sized grids and parcels¶
Part 1a: Grid topology¶
First, we need the ability to create different sizes of network model grids. We will use these different size grids to assess how long it takes the NST to run with different grid sizes.
A simple approach is to create a generic grid in which each node has two recievers. Lets start by writing a function that creates the x and y node coordinates and the linking structure for a given number of layers.
While these grids do not look like typical river networks, they are topologically similar.
We create a function create_node_xy_and_links which takes the following inputs
n_layersx0y0x_percdy
and returns three arrays: - x_of_node - y_of_node - nodes_at_link
These inputs and outputs are defined in the function docstring below. The function was designed to produce output that could be directly provided to the `NetworkModelGrid <https://landlab.readthedocs.io/en/master/reference/grid/network_api.html#module-landlab.grid.network>`__ init function.
[ ]:
def create_node_xy_and_links(n_layers, x0=0.0, y0=0.0, xperc=0.9, dy=1.0):
r"""Create node and link structure of a branching binary tree.
The tree can have an arbitrary number of layers. For example,
a tree with one layer has three nodes and two links:
::
* *
\ /
*
The lowest of the nodes is the "origin" node, and has coordinates
of `(x0, y0)`. The y spacing between layers is given by `dy`. Finally,
in order to ensure that links do not cross and nodes are not co-located
a shrinking factor, `xperc` that must be less than 1.0 is specified.
Each layer has 2^{layer} nodes: Layer 0 has 1 node, layer 1 has 2 nodes,
layer 2 has 4 nodes.
A tree with three layers has seven nodes and six links:
::
* * * *
\ / \ /
* *
\ /
*
Parameters
----------
n_layers : int
Number of layers of the binary tree
x0 : float
x coordinate position of the origin node. Default of 0.
y0=0. : float
y coordinate position of the origin node. Default of 0.
xperc : float
x direction shrinkage factor to prevent co-location of nodes
and corssing links. Must be between 0.0 and 1.0 noninclusive.
Default of 0.9.
dy : float
y direction spacing between layers. Default of 1.
Returns
------
x_of_node : list
Node x coordinates.
y_of_node : list
Node y coordinates.
nodes_at_link : list of (head, tail) tuples
Nodes at link tail and head.
"""
assert xperc < 1.0
assert xperc > 0.0
nodes_per_layer = np.power(2, np.arange(n_layers + 1))
np.sum(nodes_per_layer)
x_of_node = [x0]
y_of_node = [y0]
nodes_at_link = []
id_start_layer = 0
for nl in np.arange(1, n_layers + 1):
nodes_last_layer = np.power(2, nl - 1)
nodes_this_layer = np.power(2, nl)
dx = xperc * (dy) * (0.5 ** (nl - 1))
for ni in range(nodes_last_layer):
head_id = id_start_layer + ni
tail_id = len(x_of_node)
x = x_of_node[head_id]
y = y_of_node[head_id]
x_of_node.extend([x - dx, x + dx])
y_of_node.extend([y + dy, y + dy])
nodes_at_link.extend([(head_id, tail_id), (head_id, tail_id + 1)])
id_start_layer = len(x_of_node) - nodes_this_layer
return x_of_node, y_of_node, nodes_at_link
Next, let’s demonstrate the different sorts of grids we get with different numbers of layers. We’ll look at grids with between 3 and 1023 nodes.
[ ]:
example_layers = [1, 3, 5, 7, 9]
nodes = []
for i, n_layers in enumerate(example_layers):
x_of_node, y_of_node, nodes_at_link = create_node_xy_and_links(n_layers)
grid = NetworkModelGrid((y_of_node, x_of_node), nodes_at_link)
graph.plot_graph(grid, at="node,link", with_id=False)
nodes.append(grid.number_of_nodes)
We can see the number of nodes grows exponentially with the number of layers.
[ ]:
plt.plot(example_layers, nodes)
plt.xlabel("Number of Layers")
plt.ylabel("Number of Nodes")
plt.show()
1b: Generic grid.¶
In order to run the NST, our grids need fields related to channel geometry and flow characteristics. Here, we’ll create a function that takes only the number of layers and uses the function create_node_xy_and_links that we just created, creates a grid, adds, those fields and populate them with generic values. The NST also needs a
`FlowDirectorSteepest <https://landlab.readthedocs.io/en/master/reference/components/flow_director.html?highlight=FlowDirectorSteepest#module-landlab.components.flow_director.flow_director_steepest>`__ instance, so we’ll create that too.
Thus, with these two functions we can specifiy the desired number of layers and get both of these objects we need to instantiate the NST.
[ ]:
def create_nmg_and_fd(n_layers):
"""Create a generic NetworkModelGrid and FlowDirectorSteepest.
This function will also add the following fields to the NetworkModelGrid.
- topographic__elevation at node
- bedrock__elevation at node
- flow_depth at link
- reach_length at link
- channel_width at link.
Parameters
----------
n_layers : int
Number of layers of the binary tree
Returns
-------
grid : NetworkModelGrid instance
fd : FlowDirectorSteepest instance
"""
x_of_node, y_of_node, nodes_at_link = create_node_xy_and_links(n_layers)
grid = NetworkModelGrid((y_of_node, x_of_node), nodes_at_link)
_ = grid.add_field("topographic__elevation", grid.y_of_node.copy(), at="node")
_ = grid.add_field("bedrock__elevation", grid.y_of_node.copy(), at="node")
_ = grid.add_field(
"flow_depth", 2.5 * np.ones(grid.number_of_links), at="link"
) # m
_ = grid.add_field(
"reach_length", 200.0 * np.ones(grid.number_of_links), at="link"
) # m
_ = grid.add_field(
"channel_width", 1.0 * np.ones(grid.number_of_links), at="link"
) # m
fd = FlowDirectorSteepest(grid)
fd.run_one_step()
return grid, fd
1c: Generic sets of parcels¶
Our second to last step in writing functions to create a generic model grid and set of parcels is to create a function that will create our parcels. The parcels are associated with the grid because parcels must be located on a grid element, and thus we will always need the grid to create the parcels.
[ ]:
def create_parcels(grid, parcels_per_link=5):
"""Create a generic set of parcels.
The NetworkSedimentTransporter requires a set of parcels with some
specific attributes (e.g., density) that are used in order to
calculate travel distances. This function creates the parcels in
the correct format and populates all necessary attributes. Specifically
it creates the following attributes:
- "abrasion_rate"
- "density"
- "time_arrival_in_link"
- "active_layer"
- "location_in_link"
- "D"
- "volume"
Parameters
----------
grid : NetworkModelGrid
parcels_per_link : int
Number of parcels to create for each link. Default of 5.
Returns
-------
parcels : DataRecord
"""
# element_id is the link on which the parcel begins.
element_id = np.repeat(np.arange(grid.number_of_links), parcels_per_link)
element_id = np.expand_dims(element_id, axis=1)
# scale volume with parcels per link so we end up with a similar quantity of sediment.
volume = (1.0 / parcels_per_link) * np.ones(np.shape(element_id)) # (m3)
active_layer = np.zeros(np.shape(element_id)) # 1= active, 0 = inactive
density = 2650 * np.ones(np.size(element_id)) # (kg/m3)
abrasion_rate = 0.0 * np.ones(np.size(element_id)) # (mass loss /m)
# Lognormal GSD
medianD = 0.085 # m
mu = np.log(medianD)
sigma = np.log(2) # assume that D84 = sigma*D50
np.random.seed(0)
D = np.random.lognormal(
mu, sigma, np.shape(element_id)
) # (m) the diameter of grains in each parcel
time_arrival_in_link = np.random.rand(np.size(element_id), 1)
location_in_link = np.random.rand(np.size(element_id), 1)
variables = {
"abrasion_rate": (["item_id"], abrasion_rate),
"density": (["item_id"], density),
"time_arrival_in_link": (["item_id", "time"], time_arrival_in_link),
"active_layer": (["item_id", "time"], active_layer),
"location_in_link": (["item_id", "time"], location_in_link),
"D": (["item_id", "time"], D),
"volume": (["item_id", "time"], volume),
}
items = {"grid_element": "link", "element_id": element_id}
parcels = DataRecord(
grid,
items=items,
time=[0.0],
data_vars=variables,
dummy_elements={"link": [NetworkSedimentTransporter.OUT_OF_NETWORK]},
)
return parcels
Part 2: Profile the code¶
Next we will create a NST, run it, and profile it using the python built-in profiling tool cProfile.
We create our grid and parcels, instantiate the profile, run the model for 60 timesteps, and then compile the results of profiling.
We will profile the code in two parts, the instantiation and running.
[ ]:
# feel free to change these parameters and see
# how it impacts the results
nlayer = 5
timesteps = 50
parcels_per_link = 50
# calculate dt and set seed.
dt = 60 * 60 * 24 * 12 # length of timestep (seconds)
np.random.seed(1234)
pr = cProfile.Profile()
pr.enable()
grid, fd = create_nmg_and_fd(nlayer)
parcels = create_parcels(grid, parcels_per_link=parcels_per_link)
nst = NetworkSedimentTransporter(
grid,
parcels,
fd,
bed_porosity=0.3,
g=9.81,
fluid_density=1000,
transport_method="WilcockCrowe",
)
pr.disable()
s = io.StringIO()
sortby = SortKey.CUMULATIVE
ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
ps.print_stats()
print(s.getvalue())
We can see that the majority of the time it took to create the grid, parcels, and the NST component occurs in xarray functions. This is good, as those are likely already efficient.
Next we profile the code as we use the run_one_step function that advances the component forward in time.
[ ]:
pr = cProfile.Profile()
pr.enable()
for t in range(0, (timesteps * dt), dt):
nst.run_one_step(dt)
pr.disable()
s = io.StringIO()
sortby = SortKey.CUMULATIVE
ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
ps.print_stats()
print(s.getvalue())
We can see that the majority of the time it takes to run the component happens in a function called _partition_active_and_storage_layers. This function is used to figure out which parcels are moving and which are not active.
Part 3: Scaling¶
Next we look at how initialization and moving the model forward in time change as we increase the number of nodes and number of parcels on a given link.
3a: Create a timing function.¶
To do this we create one final convenience function that will create everything we need for our scaling analysis, run the model, and report how long it took to initialize and run.
[ ]:
def time_code(nlayer=3, parcels_per_link=100, timesteps=10):
"""Time the initializiation and runtime.
Parameters
----------
n_layers : int
Number of layers of the binary tree used to create the
NetworkModelGrid. Default of 3
parcels_per_link : int
Number of parcels to create for each link. Default of 100.
timesteps : int
Number of timesteps. Default of 10.
Returns
-------
(number_of_nodes, parcels_per_link, total_parcels) : tuple
Tuple indicating the key inputs in our scaling analysis. The number
of nodes, the number of parcels per link, the total number of parcels.
init_duration : float
Duration of the initiation step, in seconds.
r1s_per : float
Duration of the average run_one_step call, in seconds
"""
init_start = time.time()
grid, fd = create_nmg_and_fd(nlayer)
parcels = create_parcels(grid, parcels_per_link=parcels_per_link)
dt = 60 * 60 * 24 * 12 # length of timestep (seconds)
nst = NetworkSedimentTransporter(
grid,
parcels,
fd,
bed_porosity=0.3,
g=9.81,
fluid_density=1000,
transport_method="WilcockCrowe",
)
init_duration = time.time() - init_start
if timesteps > 0:
r1s_start = time.time()
for t in range(timesteps):
nst.run_one_step(dt)
r1s_per = (time.time() - r1s_start) / timesteps
else:
r1s_per = 0.0
return (grid.number_of_nodes, parcels_per_link), init_duration, r1s_per
Next, we use or new time_code function with a few different grid sizes, a few different parcels per link, and for 10 seconds. Feel free to experiment and change these values. Some of these values have been reduced to ensure that this notebook always works in the Landlab continuous integration.
3b: Time the code while systematically varying inputs¶
We save the results in a list structured to easily make a `pandas.Dataframe <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html>`__.
[ ]:
np.random.seed(345)
out = []
# this range for i in reduced for testing.
for i in range(2, 5):
for j in [10, 20, 50, 100, 200, 500]:
print(i, j)
(nn, ppl), init, r1s_per = time_code(nlayer=i, parcels_per_link=j, timesteps=10)
out.append((nn, ppl, init, r1s_per))
We make a dataframe and investigate the contents with df.head. We’ll use some shorthand for the column and axis names:
nnodes = number of nodes in the NetworkModelGrid
ppl = parcels per link
init = duration of time (in seconds) of instantiating the grid, parcels, and NST
r1s_per = duration of time (in seconds) of running the model forward one step in time.
[ ]:
df = pd.DataFrame(out, columns=["nnodes", "ppl", "init", "r1s_per"])
df = df.pivot(index="nnodes", columns="ppl", values=["init", "r1s_per"])
df.head()
3c: Plot the output to see how the code scales.¶
First, lets plot the duration of init (left) and run one step (right) as a function of the number of nodes and parcels per link. We will put number of nodes on the x axis and make a line for each different nubmer of parcels per link. For each As is common, we will make these plots in log-log space.
[ ]:
fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, dpi=300)
df["init"].plot(loglog=True, ax=ax[0], title="init duration")
ax[0].set_ylabel("duration (s)")
ax[0].set_xlabel("Number of Nodes")
df["r1s_per"].plot(loglog=True, ax=ax[1], title="run one step duration")
ax[1].set_ylabel("duration (s)")
ax[1].set_xlabel("Number of Nodes")
# plt.savefig("scaling1.png")
plt.show()