<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 <a href="http
Step2:  Part 1
Step3:  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.
Step4:  We can see the number of nodes grows exponentially with the number of layers.
Step6:  1b
Step8:  1c
Step9:  Part 2
Step10:  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.
Step12:  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
Step13:  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
Step14:  We make a dataframe and investigate the contents with df.head. We'll use some shorthand for the column and axis names
Step15:  3c
<ASSISTANT_TASK:>
Python Code:
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
import xarray as xr
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")
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../landlab_header.png"></a>
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.
End of explanation
def create_node_xy_and_links(n_layers, x0=0.0, y0=0.0, xperc=0.9, dy=1.0):
    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))
    nnodes = 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
Explanation: 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_layers
x0
y0
x_perc
dy
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 init function.
End of explanation
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)
Explanation: 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.
End of explanation
plt.plot(example_layers, nodes)
plt.xlabel("Number of Layers")
plt.ylabel("Number of Nodes")
plt.show()
Explanation: We can see the number of nodes grows exponentially with the number of layers.
End of explanation
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
Explanation: 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 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.
End of explanation
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
Explanation: 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.
End of explanation
# 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())
Explanation: 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.
End of explanation
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())
Explanation: 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.
End of explanation
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
Explanation: 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.
End of explanation
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))
Explanation: 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.
End of explanation
df = pd.DataFrame(out, columns=["nnodes", "ppl", "init", "r1s_per"])
df = df.pivot(index="nnodes", columns="ppl", values=["init", "r1s_per"])
df.head()
Explanation: 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.
End of explanation
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()
Explanation: 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.
End of explanation
<END_TASK>