Commit 59fc8e08 authored by gcordonnier's avatar gcordonnier

add d4/d8 option

parent 19c8a630
......@@ -14,14 +14,23 @@ def reset_receivers(receivers, nnodes):
@numba.njit
def compute_receivers_d8(receivers, dist2receivers, elevation,
def compute_receivers(receivers, dist2receivers, elevation,
active_nodes,
nx, ny, dx, dy):
# queen (D8) neighbor lookup
dr = np.array([-1, -1, -1, 0, 0, 1, 1, 1], dtype=np.intp)
dc = np.array([-1, 0, 1, -1, 1, -1, 0, 1], dtype=np.intp)
nx, ny, dx, dy, directions = 8):
if directions == 8:
# queen (D8) neighbor lookup
dr = (-1, -1, -1, 0, 0, 1, 1, 1)
dc = (-1, 0, 1, -1, 1, -1, 0, 1)
else:
directions = 4
# king (D4) neighbor lookup
dr = (-1, 0, 0, 1, 0, 0, 0, 0)
dc = ( 0, -1, 1, 0, 0, 0, 0, 0)
length = [math.sqrt((dy * dr[i])**2 + (dx * dc[i])**2) for i in range(8)]
length = np.sqrt((dy * dr)**2 + (dx * dc)**2)
tiny = np.finfo(elevation.dtype).tiny
......@@ -34,7 +43,7 @@ def compute_receivers_d8(receivers, dist2receivers, elevation,
slope_max = tiny
for k in range(8):
for k in range(directions):
if r + dr[k] >=0 and r + dr[k]<ny and c + dc[k] >=0 and c + dc[k] <nx:
ineighbor = (r + dr[k]) * nx + (c + dc[k])
slope = (elevation[inode] - elevation[ineighbor]) / length[k]
......@@ -124,7 +133,7 @@ def compute_pits(pits, outlets, active_nodes, nbasins):
@numba.njit
def _connect_basins(conn_basins, conn_nodes, conn_weights,
nbasins, basins, outlets, receivers, stack,
active_nodes, elevation, nx, ny):
active_nodes, elevation, nx, ny, directions = 8):
"""Connect adjacent basins together through their lowest pass.
Creates an (undirected) graph of basins and their connections.
......@@ -173,8 +182,8 @@ def _connect_basins(conn_basins, conn_nodes, conn_weights,
iactive = False
# king (D4) neighbor lookup
dr = (0, -1, 1, 0)
dc = (-1, 0, 0, 1)
dr = (0, -1, 1, 0, 1, 1, -1, -1)
dc = (-1, 0, 0, 1, 1, -1, 1, -1)
for istack in stack:
irec = receivers[istack]
......@@ -201,7 +210,7 @@ def _connect_basins(conn_basins, conn_nodes, conn_weights,
r = istack // nx
c = istack % nx
for k in range(4):
for k in range(directions):
kr = r + dr[k]
kc = c + dc[k]
......@@ -437,7 +446,7 @@ def _update_pits_receivers(receivers, dist2receivers, outlets,
@numba.njit
def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy):
def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy, directions = 8):
""" Parse the basins from sea to crests, filling each of the basin
by looking at each of the basin nodes, sorted by distance to the sill
......@@ -475,7 +484,7 @@ def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elev
parse_begin +=1
for nb_dir_i in range (8):
for nb_dir_i in range (directions):
kr = i_node // nx + dr[nb_dir_i]
kc = i_node % nx + dc[nb_dir_i]
......@@ -505,7 +514,9 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
active_nodes, elevation, nx, ny, dx, dy,
method='mst_linear', sediment_lakes = False,
sediment_slope = 1e-10, water_height=None,
reorder = True):
reorder = True,
randomize = True,
directions = 8):
"""Ensure that no flow is captured in sinks.
If needed, update `receivers`, `dist2receivers`, `ndonors`,
......@@ -527,14 +538,14 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
nconn, basin0 = _connect_basins(
conn_basins, conn_nodes, conn_weights,
nbasins, basins, outlets, receivers, stack,
active_nodes, elevation, nx, ny)
active_nodes, elevation, nx, ny, directions)
#resize conn arrays, applying a random shuffle:
# In case of basins with equal height sills (flat area or 1 sized localminima),
# the resulting connection order in the mst will depend on the order of conn_basins
# Another way to see that is that MST is not unique when some edges has the same weigth,
# in this case we chose the most 'natural' random one
perm = np.random.permutation(nconn)
perm = np.random.permutation(nconn) if randomize else np.arange(nconn)
conn_basins = conn_basins[perm]
conn_nodes = conn_nodes[perm]
conn_weights = conn_weights[perm]
......@@ -567,17 +578,18 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
water_height = np.empty_like(elevation)
water_height[:] = elevation[:]
_sediment_lakes(sills, basin_stack[:basin_stack_size], basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy)
_sediment_lakes(sills, basin_stack[:basin_stack_size], basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy, directions)
# the water height obtained from _sediment_lakes is a very regular geometric shape.
# we apply some ramdomness to obtin a more natural drainage pattern,
# without adding any other local minima
water_height += 0.99 * sediment_slope * np.random.rand(nx*ny) * (water_height > elevation)
if randomize:
water_height += 0.99 * sediment_slope * np.random.rand(nx*ny) * (water_height > elevation)
if reorder:
reset_receivers(receivers, nnodes)
compute_receivers_d8(receivers, dist2receivers,
water_height, active_nodes, nx, ny, dx, dy)
compute_receivers(receivers, dist2receivers,
water_height, active_nodes, nx, ny, dx, dy, directions)
else:
_update_pits_receivers(receivers, dist2receivers, outlets,
......@@ -590,7 +602,10 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes = None, method='mst_linear', sediment_slope = 1e-10):
def compute_water_level(elevation, nx, ny, dx, dy, water_height,
active_nodes = None, method='mst_linear',
sediment_slope = 1e-10, randomize = True,
directions = 8):
""" Gives, in 'water_height', the elevation of water for a dem 'elevation' of size (nx, ny).
The water level can be given a small slope 'sediment_slope' from the sill, in which case
......@@ -620,8 +635,8 @@ def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes =
active_nodes = mask.flatten()
reset_receivers(receivers, nnodes)
compute_receivers_d8(receivers, dist2receivers,
elevation, active_nodes, nx, ny, dx, dy)
compute_receivers(receivers, dist2receivers,
elevation, active_nodes, nx, ny, dx, dy, directions)
compute_donors(ndonors, donors, receivers, nnodes)
compute_stack(stack, ndonors, donors, receivers, nnodes)
......@@ -636,4 +651,8 @@ def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes =
method=method, sediment_lakes = True,
sediment_slope = sediment_slope,
water_height = water_height,
reorder = False)
reorder = False,
randomize = randomize,
directions = directions)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment