This notebook contains a sketch and examples of an approach to calculating flux divergence, using a simple diffusion problem as an example.
Let's suppose we want to create a diffusion model with a raster grid, in which we compute gradients in a field, then flux based on gradient, and finally divergence of flux. Start by creating a small raster grid:
In [1]:
from landlab import RasterModelGrid
import numpy as np
In [2]:
rg = RasterModelGrid(3, 4, 10.0)
z = rg.add_zeros('node', 'topographic__elevation')
rg.set_closed_boundaries_at_grid_edges(True, True, False, False)
#rg.set_closed_nodes([3, 7, 8, 9, 10, 11])
In [3]:
z[5] = 50.
z[6] = 36.
print(z)
There are several ways we can go about running the model, and we're going to explore each in turn. Here are some of the factors at play. We have a field of elevation at all nodes---that's what we want to track. Some of those nodes, however, are boundary nodes, and should be held fixed (or manipulated independently). Some of the links are inactive. Only the core nodes are subject to change due to fluxes in and out. In any event, we can only really do a mass balance on a cell; a node is just a point. As to fluxes: these depend, of course, on the gradient in z. We can compute gradients on links or on faces. The two are nearly identical---both are gradients between values at the nodes on either side of a link---but they differ in array length. A calculation done on links is NL (number of links) long (17 in the case of a 3x4 grid); one done on faces is NF (number of faces) long (7 for a 3x4 grid). Note that the number of faces is equal to the number of interior links.
The upshot, then, is that we'll start with an example that calculates gradient on faces.
Let's test it. The gradient array should be:
[5, 3.6, 5, -1.4, -3.6, -5, -3.6]
In [4]:
grad_at_faces = rg.calculate_gradients_at_faces(z)
grad_at_faces
Out[4]:
So now we have an NF-long array of gradients. Let's make a unit flux out of this, by multiplying gradient by -1 (stuff flows down hill):
In [5]:
unit_flux_at_faces = -grad_at_faces
unit_flux_at_faces
Out[5]:
We might alternatively only want to calculate gradients and fluxes at active faces, which are faces crossed by an active link. How to get this?
In [6]:
rg.active_faces
Out[6]:
The unit flux at active faces should be the same as unit_flux_at_faces but with 3 entries zeroed out:
In [7]:
unit_flux_at_active_faces = np.zeros(rg.number_of_faces)
unit_flux_at_active_faces[rg.active_faces] = unit_flux_at_faces[rg.active_faces]
unit_flux_at_active_faces
Out[7]:
Next, let's find the net face fluxes at cells. This is defined as the sum of total inflows and outflows. In order to do this, we need two data structures: faces_at_cell and link_dirs_at_node. The first is an M x N array of int that contains the IDs of faces at each cell, in counter-clockwise orientation starting with east (raster) or nearest CCW to east (others). Here M is the maximum number of links at a node, which is 4 for raster, 6 for hex, and arbitrary for Voronoi.
To get faces_at_cell, we first need to create links_at_node and link_dirs_at_node. It turns out that links_at_node already exists for raster (not for unstructured). Here's what it looks like for our raster grid:
In [8]:
rg.links_at_node
Out[8]:
(Below under LEFTOVERS are links_at_node and link_dirs_at_node functions for raster; the code now has something like this, so these are no longer needed)
Here's a version of the above that makes link dirs but not links (since that already exists), and in the right array orientation:
In our 3x4 raster grid, the links_at_node array should look like the following. There are N (# nodes) columns and four rows; each row represents a cardinal direction (east, north, west, and south). For example, node 5 is one of only two interior nodes; it is connected to links 8, 11, 7, and 4, on the east, north, west, and south sides, respectively. (Note that unstructured grids will not use cardinal directions, as illustrated later).
array([[ 0, 3, -1, -1],
[ 1, 4, 0, -1],
[ 2, 5, 1, -1],
[-1, 6, 2, -1],
[ 7, 10, -1, 3],
[ 8, 11, 7, 4],
[ 9, 12, 8, 5],
[-1, 13, 9, 6],
[14, -1, -1, 10],
[15, -1, 14, 11],
[16, -1, 15, 12],
[-1, -1, 16, 13]])
In [9]:
rg.links_at_node
Out[9]:
The link_dirs_at_node array indicates the direction of each of a node's links relative to the node. For each link and node, an entry of 1 means that link is entering the node; in other words, the node is that link's head. An entry of -1 means that link is leaving the node; the node is that link's tail. An entry of zero means that there is no link at this position. For example, node 0 (lower-left corner) has outgoing links on the east and north, but no links on the west or south. It should look like this:
array([[-1, -1, 0, 0],
[-1, -1, 1, 0],
[-1, -1, 1, 0],
[ 0, -1, 1, 0],
[-1, -1, 0, 1],
[-1, -1, 1, 1],
[-1, -1, 1, 1],
[ 0, -1, 1, 1],
[-1, 0, 0, 1],
[-1, 0, 1, 1],
[-1, 0, 1, 1],
[ 0, 0, 1, 1]], dtype=int8)
The _link_dirs_at_node setup has now been encoded into raster.py, and base.py now includes a link_dirs_at_node property to return it, so it should be built in:
In [10]:
rg._link_dirs_at_node
Out[10]:
Before we can take the next step, we also need faces_at_cell: an M x NC array containing the face IDs at each cell, ordered CCW.
It turns out that RasterModelGrid has this, though at the moment it's a method rather than a property (Eric plans to fix this soon):
In [11]:
rg.faces_at_cell()
Out[11]:
Now we're ready to define the net_face_flux_at_cells method. This is one of a series of similar methods. We assume the caller has given us an array of length NF (number of faces) containing flux per unit face width. The first step is to convert this to total flux by multiplying by face width.
Then, for each node, we want to add up the incoming fluxes and subtract the outgoing fluxes. The algorithm does this by iterating over the 4 cardinal directions. In the first iteration, we take the IDs of all nodes' east links (which is row 0 of links_at_node. The flux corresponding to these links is total_flux[links_at_node[0]].
Wait a minute, I thought we wanted faces, not links? We do, but here's the thing: every face has a corresponding link. We simply need to find those links corresponding to the faces.
What if there is no east link at a particular node? Then its entry in links_at_node will be -1, which means it will take the value of the last entry in the flux array. That's the wrong flux value, of course, but it does not matter: we multiply each link flux by the corresponding entry in link_dirs_at_node. If there is no such link, this value is zero---so whatever
In [12]:
def net_face_flux_at_cells(grid, unit_flux_at_faces):
total_flux = unit_flux_at_faces * grid.face_width
nffc = np.zeros(grid.number_of_cells)
fac = grid.faces_at_cell() # inefficient because faces_at_cell must be created repeatedly
for r in range(4):
nffc += total_flux[fac[:,r]] \
* grid.link_dirs_at_node[grid.node_at_cell,r]
return nffc
Now we should have everything needed to run the net_face_flux_at_cells function. It ought to return an array with the following values:
At cell 0, we have fluxes of: -14 (east), -50 (north), -50 (west), -50 (south), which sums to -164.
At cell 1, we have fluxes of: -36 (east), -36 (north), +14 (west), -36 (south), which sums to -94.
In [13]:
unit_flux_at_faces
Out[13]:
In [14]:
rg.face_width
Out[14]:
In [15]:
net_face_flux_at_cells(rg, unit_flux_at_faces) ## WHY IS IT POS NOT NEG?
Out[15]:
If we wanted to work just with active links (treating closed boundaries as walls), we would instead use unit_flux_at_active_faces. In this case, with the east and north sides closed, we should get:
cell 0: -14 (east) + 0 (north) - 50 (west) - 50 (south) = -114
cell 1: 0 (east) + 0 (north) + 14 (west) - 36 (south) = -22
In [16]:
net_face_flux_at_cells(rg, unit_flux_at_active_faces)
Out[16]:
Let's pause for some definitions:
Let unit flux, $q$, be defined as the flow of mass, volume, energy, momentum, etc., per unit time per unit face width. For example, if we're working with volume fluxes, $q$ has dimensions of length squared per time. When defined at a link, it is positive in the direction of the link; when defined at a face, it is positive in the direction of the associated link.
Total flux, $Q$, is then simply $Q = q W$, where $W$ is the width of a face.
Flux divergence, $\nabla q$, at a cell is defined as $\nabla q = -\Sigma Q / A_c$, where $\Sigma$ represents a summation over all the cell's faces (or active faces, as the case may be), and $A_c$ is the surface area of the cell. Notice that if $q$ has dimensions of $L^2/T$, then $\nabla q$ has dimensions of $L/T$.
The minus sign on flux divergence keeps us consistent with standard mathematical usage.
So, to find flux divergence at cells, we simply find the total flux and divide by the cell area (and multiply by -1). Here's a function to do this:
In [17]:
def face_flux_divergence_at_cells(grid, unit_flux_at_faces):
return -net_face_flux_at_cells(grid, unit_flux_at_faces) / grid.area_of_cell
In the following example, the result should be 1.64 at cell 0 and 0.94 at cell 1:
In [18]:
face_flux_divergence_at_cells(rg, unit_flux_at_faces)
Out[18]:
In [19]:
face_flux_divergence_at_cells(rg, unit_flux_at_active_faces)
Out[19]:
So now we'll test out using these algorithms repeatedly in a diffusion calculation. Same grid, same initial conditions.
In [20]:
D = 0.01
niter = 10000
dt = 0.0001 # way tinier than it needs to be, just for performance testing
In [21]:
import time
print(z)
start_time = time.time()
for i in range(niter):
g = rg.calculate_gradients_at_faces(z)
q = -g
dq = face_flux_divergence_at_cells(rg, q)
z[rg.node_at_cell] -= dq*dt
end_time = time.time()
print(z)
print(1000*(end_time-start_time)/niter)
Last time I ran this it took about 0.07 ms per iteration. For completeness, we'll try an active-face version:
In [22]:
z[5] = 50.
z[6] = 36.
import time
start_time = time.time()
q = rg.zeros(centering='face')
fal = rg.face_at_link # would be easier if this were a property
af = fal[rg.active_links] # we apparently don't have active_faces
for i in range(niter):
g = rg.calculate_gradients_at_faces(z)
q[af] = -g[af]
dq = face_flux_divergence_at_cells(rg, q)
z[rg.node_at_cell] -= dq*dt
end_time = time.time()
print(z)
print(1000*(end_time-start_time)/niter)
So, based on the last time I ran this example, the two versions have about the same speed.
This next section repeats the previous one but for a hex grid (as an example of a Voronoi).
In [23]:
from landlab import HexModelGrid
In [24]:
hg = HexModelGrid(3, 3, 10.0)
z = rg.add_zeros('node', 'topographic__elevation')
z[4] = 50.
z[5] = 36.
hg.set_closed_nodes([6, 7, 8, 9])
print(z)
The gradients should be:
[5, 5, 3.6, 3.6, 5, -1.4, -3.6, -5, -5, -3.6, -3.6]
In [25]:
g = hg.calculate_gradients_at_faces(z)
g
Out[25]:
In [26]:
hg.set_closed_nodes([6, 7, 8, 9])
In [27]:
unit_flux_at_faces = -g
unit_flux_at_faces
Out[27]:
Test: our (3, 3) hex grid should now have links_at_node that looks like this:
array([[ 0, 3, 2, -1, -1, -1],
[ 1, 5, 4, 0, -1, -1],
[ 7, 6, 1, -1, -1, -1],
[ 8, 11, 2, -1, -1, -1],
[ 9, 13, 12, 8, 3, 4],
[10, 15, 14, 9, 5, 6],
[16, 10, 7, -1, -1, -1],
[17, 11, 12, -1, -1, -1],
[18, 17, 13, 14, -1, -1],
[18, 15, 16, -1, -1, -1]], dtype=int32)
In [49]:
hg.links_at_node
Out[49]:
Link directions should look like this:
array([[-1, -1, -1, 0, 0, 0],
[-1, -1, -1, 1, 0, 0],
[-1, -1, 1, 0, 0, 0],
[-1, -1, 1, 0, 0, 0],
[-1, -1, -1, 1, 1, 1],
[-1, -1, -1, 1, 1, 1],
[-1, 1, 1, 0, 0, 0],
[-1, 1, 1, 0, 0, 0],
[-1, 1, 1, 1, 0, 0],
[ 1, 1, 1, 0, 0, 0]], dtype=int8)
In [50]:
hg.link_dirs_at_node
Out[50]:
In [51]:
hg.number_of_links
Out[51]:
In [52]:
hg.faces_at_cell
Out[52]:
Test 1: for node 4 of our 3x3 hex grid, there should be 6 links with angles 0, 60, 120, 180, 240, and 300 degrees.
In [53]:
180*hg.link_angle(hg.links_at_node[4,:], hg.link_dirs_at_node[4,:]) / np.pi
Out[53]:
Test 2: angle should be $2\pi$ where there is no link (indicated by dir==0). This serves to force no-link entries to the end of the array.
In [54]:
180*hg.link_angle(hg.links_at_node[6,:], hg.link_dirs_at_node[6,:]) / np.pi
Out[54]:
Redefine this function to use faces_at_cell as property rather than fn.
In [64]:
def net_face_flux_at_cells(grid, unit_flux_at_faces):
total_flux = unit_flux_at_faces * grid.face_width
nffc = np.zeros(grid.number_of_cells)
fac = grid.faces_at_cell
for r in range(fac.shape[1]):
nffc += total_flux[fac[:,r]] \
* grid.link_dirs_at_node[grid.node_at_cell,r]
return nffc
For cell 0, the net flux should be:
outflux on 5 faces at rate 5 x 5.77350269 = 144.3 plus outflux to cell 1 at rate 1.4 x 5.77350269 = 8.08, for a total of -152.4.
For cell 1, the net flux should be:
outflux on 5 faces at rate 3.6 x 5.77 = 103.9.
influx from cell 0 at 8.1.
total = -57.1
In [72]:
nffc = net_face_flux_at_cells(hg, unit_flux_at_faces)
nffc
3.6*5.77*5-8.
Out[72]:
NEXT STEPS:
LEFTOVERS:
In [37]:
def make_faces_at_cell(grid):
max_num_faces = grid.gt_links_at_node.shape[1]
grid.gt_faces_at_cell = np.zeros((grid.number_of_cells, max_num_faces), dtype=np.int)
fal = grid.face_at_link
print 'Face at link', fal
nac = grid.node_at_cell
print 'Node at cell:', nac
for r in range(max_num_faces):
lanr = grid.links_at_node[nac,r]
print 'Row', r, 'links at these nodes', lanr
grid.gt_faces_at_cell[:,r] = fal[lanr]
print 'Face', r, 'at these cells', fal[lanr]
print grid.gt_faces_at_cell
In [38]:
def make_links_at_node_array_raster(grid):
"""Make array with links at each node
Examples
--------
>>> from landlab import RasterModelGrid
>>> rmg = RasterModelGrid(3, 4)
>>> make_links_at_node_array_raster(rmg)
>>> rmg.gt_links_at_node
array([[ 0, 1, 2, -1, 7, 8, 9, -1, 14, 15, 16, -1],
[ 3, 4, 5, 6, 10, 11, 12, 13, -1, -1, -1, -1],
[-1, 0, 1, 2, -1, 7, 8, 9, -1, 14, 15, 16],
[-1, -1, -1, -1, 3, 4, 5, 6, 10, 11, 12, 13]], dtype=int32)
>>> rmg._link_dirs_at_node
array([[-1, -1, -1, 0, -1, -1, -1, 0, -1, -1, -1, 0],
[-1, -1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0],
[ 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1],
[ 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=int8)
"""
# Create arrays for link-at-node information
grid.gt_links_at_node = -np.ones((grid.number_of_nodes, 4), dtype=np.int32)
grid.gt_link_dirs_at_node = np.zeros((grid.number_of_nodes, 4), dtype=np.int8)
grid.gt_active_link_dirs_at_node = np.zeros((grid.number_of_nodes, 4), dtype=np.int8)
grid.gt_num_links_at_node = np.zeros(grid.number_of_nodes, dtype=np.uint8) # assume <256 links at any node
grid.gt_num_active_links_at_node = np.zeros(grid.number_of_nodes, dtype=np.uint8) # assume <256 links at any node
num_links_per_row = (grid.number_of_node_columns * 2) - 1
# Sweep over all links
for lk in xrange(grid.number_of_links):
# Find the orientation
is_horiz = (lk % num_links_per_row) < (grid.number_of_node_columns - 1)
# Find the IDs of the tail and head nodes
t = grid.node_at_link_tail[lk]
h = grid.node_at_link_head[lk]
# If the link is horizontal, the index (row) in the links_at_node array
# should be 0 (east) for the tail node, and 2 (west) for the head node.
# If vertical, the index should be 1 (north) for the tail node and
# 3 (south) for the head node.
if is_horiz:
tail_index = 0
head_index = 2
else:
tail_index = 1
head_index = 3
# Add this link to the list for this node, set the direction (outgoing,
# indicated by -1), and increment the number found so far
grid.gt_links_at_node[tail_index][t] = lk
grid.gt_links_at_node[head_index][h] = lk
grid.gt_link_dirs_at_node[tail_index][t] = -1
grid.gt_link_dirs_at_node[head_index][h] = 1
grid.gt_num_links_at_node[h] += 1
grid.gt_num_links_at_node[t] += 1
In [39]:
def make_link_dirs_at_node_array_raster(grid):
"""Make array with link directions at each node
Examples
--------
>>> from landlab import RasterModelGrid
>>> rmg = RasterModelGrid(3, 4)
>>> make_link_dirs_at_node_array_raster(rmg)
>>> rmg.links_at_node
>>> rmg._link_dirs_at_node
"""
# Create arrays for link-at-node information
grid.gt_links_at_node = -np.ones((grid.number_of_nodes, 4), dtype=np.int32)
grid._link_dirs_at_node = np.zeros((grid.number_of_nodes, 4), dtype=np.int8)
grid.gt_num_links_at_node = np.zeros(grid.number_of_nodes, dtype=np.uint8) # assume <256 links at any node
grid.gt_num_active_links_at_node = np.zeros(grid.number_of_nodes, dtype=np.uint8) # assume <256 links at any node
num_links_per_row = (grid.number_of_node_columns * 2) - 1
# Sweep over all links
for lk in xrange(grid.number_of_links):
# Find the orientation
is_horiz = (lk % num_links_per_row) < (grid.number_of_node_columns - 1)
# Find the IDs of the tail and head nodes
t = grid.node_at_link_tail[lk]
h = grid.node_at_link_head[lk]
# If the link is horizontal, the index (row) in the links_at_node array
# should be 0 (east) for the tail node, and 2 (west) for the head node.
# If vertical, the index should be 1 (north) for the tail node and
# 3 (south) for the head node.
if is_horiz:
tail_index = 0
head_index = 2
else:
tail_index = 1
head_index = 3
# Add this link to the list for this node, set the direction (outgoing,
# indicated by -1), and increment the number found so far
grid._link_dirs_at_node[t][tail_index] = -1
grid._link_dirs_at_node[h][head_index] = 1
grid.gt_num_links_at_node[h] += 1
grid.gt_num_links_at_node[t] += 1
In [40]:
def gt_setup_link_and_face_temp(grid):
"""Set up link_at_face and face_at_link arrays.
"""
from landlab import BAD_INDEX_VALUE
grid._face_at_link = np.zeros(grid.number_of_links, dtype=int)
grid._face_at_link[:] = BAD_INDEX_VALUE
grid._num_faces = len(grid.face_width)
grid._link_at_face = np.zeros(grid._num_faces, dtype=int)
face_id = 0
for link in range(grid.number_of_links):
tc = grid.cell_at_node[grid.node_at_link_tail[link]]
hc = grid.cell_at_node[grid.node_at_link_head[link]]
if tc != BAD_INDEX_VALUE or hc != BAD_INDEX_VALUE:
grid._face_at_link[link] = face_id
grid._link_at_face[face_id] = link
face_id += 1
In [41]:
# The numpy arctan2 function tells us the angle of a ray from origin to given point(s)
x = np.array([1.0, 3.0*3.0**0.5, 1.0, 3.0, 0.0])
y = np.array([0.0, 3.0, 1.0, 3.0*3.0**0.5, 1.0])
ang = np.arctan2(-y, x)
180*ang/np.pi
Out[41]:
In [42]:
def link_angle(grid, links, dirs):
"""Find and return the angle of link(s) in given direction.
Parameters
----------
grid : ModelGrid object
reference to the grid
links : 1d numpy array
one or more link IDs
dirs : 1d numpy array (must be same length as links)
direction of links relative to node: +1 means head is origin;
-1 means tail is origin.
Notes
-----
dx and dy are the x and y differences between the link endpoints.
Multiplying this by dirs orients these offsets correctly (i.e.,
the correct node is the origin). The call to arctan2 calculates
the angle in radians. Angles in the lower two quadrants will be
negative and clockwise from the positive x axis. We want them
counter-clockwise, which is what the last couple of lines before
the return statement do.
"""
dx = -dirs * (grid.node_x[grid.node_at_link_head[links]] - grid.node_x[grid.node_at_link_tail[links]])
dy = -dirs * (grid.node_y[grid.node_at_link_head[links]] - grid.node_y[grid.node_at_link_tail[links]])
ang = np.arctan2(dy, dx)
(lower_two_quads, ) = np.where(ang<0.0)
ang[lower_two_quads] = (2 * np.pi) + ang[lower_two_quads]
(no_link, ) = np.where(dirs == 0)
ang[no_link] = 2*np.pi
return ang
In [43]:
def sort_links_at_node_by_angle(grid):
"""Sort the links_at_node and link_dirs_at_node arrays by angle.
"""
for n in range(grid.number_of_nodes):
ang = link_angle(grid, grid.gt_links_at_node[:,n], grid.gt_link_dirs_at_node[:,n])
indices = np.argsort(ang)
grid.gt_links_at_node[:,n] = grid.gt_links_at_node[indices,n]
grid.gt_link_dirs_at_node[:,n] = grid.gt_link_dirs_at_node[indices,n]
In [44]:
def find_number_of_links_per_node(grid):
"""Find and record how many links are attached to each node."""
grid._number_of_links_per_node = np.zeros(grid.number_of_nodes, dtype=np.int)
for ln in range(grid.number_of_links):
grid._number_of_links_per_node[grid.node_at_link_tail[ln]] += 1
grid._number_of_links_per_node[grid.node_at_link_head[ln]] += 1
In [45]:
# when made a member of grid, should be a property
def number_of_links_per_node(grid):
try:
return grid._number_of_links_per_node
except AttributeError:
find_number_of_links_per_node(grid)
return grid._number_of_links_per_node
In [46]:
def make_links_at_node_array_general(grid):
"""Make array with links at each node
"""
# Find maximum number of links per node
nlpn = number_of_links_per_node(grid) #this fn should become member and property
max_num_links = np.amax(nlpn)
nlpn[:] = 0 # we'll zero it out, then rebuild it
#for later examples:
# nlpn should be [3 4 3 3 6 6 3 3 4 3] and max should be 6
# Create arrays for link-at-node information
grid.gt_links_at_node = -np.ones((max_num_links, grid.number_of_nodes), dtype=np.int32)
grid.gt_link_dirs_at_node = np.zeros((max_num_links, grid.number_of_nodes), dtype=np.int8)
grid.gt_active_link_dirs_at_node = np.zeros((max_num_links, grid.number_of_nodes), dtype=np.int8)
grid.gt_num_active_links_at_node = np.zeros(grid.number_of_nodes, dtype=np.uint8) # assume <256 links at any node
# Sweep over all links
for lk in xrange(grid.number_of_links):
# Find the IDs of the tail and head nodes
t = grid.node_at_link_tail[lk]
h = grid.node_at_link_head[lk]
# Add this link to the list for this node, set the direction (outgoing,
# indicated by -1), and increment the number found so far
grid.gt_links_at_node[nlpn[t]][t] = lk
grid.gt_links_at_node[nlpn[h]][h] = lk
grid.gt_link_dirs_at_node[nlpn[t]][t] = -1
grid.gt_link_dirs_at_node[nlpn[h]][h] = 1
nlpn[t] += 1
nlpn[h] += 1