279 lines
7.9 KiB
Python
279 lines
7.9 KiB
Python
import numpy as np
|
||
import numpy.random as npr
|
||
from collections import defaultdict
|
||
|
||
# This file provide functions to construct the river tree from an elevation model.
|
||
# Based on a research paper:
|
||
# | Cordonnier, G., Bovy, B., and Braun, J.:
|
||
# | A versatile, linear complexity algorithm for flow routing in topographies with depressions,
|
||
# | Earth Surf. Dynam., 7, 549–562, https://doi.org/10.5194/esurf-7-549-2019, 2019.
|
||
# Big thanks to them for releasing this paper under a free license ! :)
|
||
|
||
# The algorithm here makes use of most of the paper's concepts, including the Planar Boruvka algorithm.
|
||
# Only flow_local and accumulate_flow are custom algorithms.
|
||
|
||
# Define two different method for local flow routing
|
||
def flow_local_steepest(plist):
|
||
vmax = 0.0
|
||
imax = 0.0
|
||
for i, p in enumerate(plist):
|
||
if p > vmax:
|
||
vmax = p
|
||
imax = i
|
||
if vmax > 0.0:
|
||
return imax+1
|
||
return 0
|
||
|
||
def flow_local_semirandom(plist):
|
||
"""
|
||
Determines a flow direction based on denivellation for every neighbouring node.
|
||
Denivellation must be positive for downward and zero for flat or upward:
|
||
dz = max(zref-z, 0)
|
||
"""
|
||
psum = sum(plist)
|
||
if psum == 0:
|
||
return 0
|
||
r = npr.random() * psum
|
||
for i, p in enumerate(plist):
|
||
if r < p:
|
||
return i+1
|
||
r -= p
|
||
|
||
flow_local_methods = {
|
||
'steepest' : flow_local_steepest,
|
||
'semirandom' : flow_local_semirandom,
|
||
}
|
||
|
||
def flow(dem, method='semirandom'):
|
||
if method in flow_local_methods:
|
||
flow_local = flow_local_methods[method]
|
||
else:
|
||
raise KeyError('Flow method \'{}\' does not exist'.format(method))
|
||
|
||
# Flow locally
|
||
dirs1 = np.zeros(dem.shape, dtype=int)
|
||
dirs2 = np.zeros(dem.shape, dtype=int)
|
||
(X, Y) = dem.shape
|
||
Xmax, Ymax = X-1, Y-1
|
||
singular = []
|
||
for x in range(X):
|
||
z0 = z1 = z2 = dem[x,0]
|
||
for y in range(Y):
|
||
z0 = z1
|
||
z1 = z2
|
||
if y < Ymax:
|
||
z2 = dem[x, y+1]
|
||
|
||
plist = [
|
||
max(z1-dem[x+1,y],0) if x<Xmax else 0, # 1: x -> x+1
|
||
max(z1-z2,0), # 2: y -> y+1
|
||
max(z1-dem[x-1,y],0) if x>0 else 0, # 3: x -> x-1
|
||
max(z1-z0,0), # 4: y -> y-1
|
||
]
|
||
|
||
pdir = flow_local(plist)
|
||
dirs2[x,y] = pdir
|
||
if pdir == 0:
|
||
singular.append((x,y))
|
||
elif pdir == 1:
|
||
dirs1[x+1,y] += 1
|
||
elif pdir == 2:
|
||
dirs1[x,y+1] += 2
|
||
elif pdir == 3:
|
||
dirs1[x-1,y] += 4
|
||
elif pdir == 4:
|
||
dirs1[x,y-1] += 8
|
||
|
||
# Compute basins
|
||
basin_id = np.zeros(dem.shape, dtype=int)
|
||
stack = []
|
||
|
||
for i, s in enumerate(singular):
|
||
queue = [s]
|
||
while queue:
|
||
x, y = queue.pop()
|
||
basin_id[x,y] = i
|
||
d = int(dirs1[x,y])
|
||
|
||
if d & 1:
|
||
queue.append((x-1,y))
|
||
if d & 2:
|
||
queue.append((x,y-1))
|
||
if d & 4:
|
||
queue.append((x+1,y))
|
||
if d & 8:
|
||
queue.append((x,y+1))
|
||
|
||
del dirs1
|
||
|
||
# Link basins
|
||
nsing = len(singular)
|
||
links = {}
|
||
def add_link(b0, b1, elev, bound):
|
||
b = (min(b0,b1),max(b0,b1))
|
||
if b not in links or links[b][0] > elev:
|
||
links[b] = (elev, bound)
|
||
|
||
for x in range(X):
|
||
b0 = basin_id[x,0]
|
||
add_link(-1, b0, dem[x,0], (True, x, 0))
|
||
for y in range(1,Y):
|
||
b1 = basin_id[x,y]
|
||
if b0 != b1:
|
||
add_link(b0, b1, max(dem[x,y-1],dem[x,y]), (True, x, y))
|
||
b0 = b1
|
||
add_link(-1, b1, dem[x,Ymax], (True, x, Y))
|
||
for y in range(Y):
|
||
b0 = basin_id[0,y]
|
||
add_link(-1, b0, dem[0,y], (False, 0, y))
|
||
for x in range(1,X):
|
||
b1 = basin_id[x,y]
|
||
if b0 != b1:
|
||
add_link(b0, b1, max(dem[x-1,y],dem[x,y]), (False, x, y))
|
||
b0 = b1
|
||
add_link(-1, b1, dem[Xmax,y], (False, X, y))
|
||
|
||
# Computing basin tree
|
||
graph = planar_boruvka(links)
|
||
|
||
basin_links = defaultdict(dict)
|
||
for elev, b1, b2, bound in graph:
|
||
basin_links[b1][b2] = basin_links[b2][b1] = (elev, bound)
|
||
basins = np.zeros(nsing+1)
|
||
stack = [(-1, float('-inf'))]
|
||
|
||
# Applying basin flowing
|
||
dir_reverse = (0, 3, 4, 1, 2)
|
||
while stack:
|
||
b1, elev1 = stack.pop()
|
||
basins[b1] = elev1
|
||
|
||
for b2, (elev2, bound) in basin_links[b1].items():
|
||
stack.append((b2, max(elev1, elev2)))
|
||
|
||
# Reverse flow direction in b2 (TODO)
|
||
isY, x, y = bound
|
||
backward = True # Whether water will escape the basin in +X/+Y direction
|
||
if not (x < X and y < Y and basin_id[x,y] == b2):
|
||
if isY:
|
||
y -= 1
|
||
else:
|
||
x -= 1
|
||
backward = False
|
||
d = 2*backward + isY + 1
|
||
while d > 0:
|
||
d, dirs2[x,y] = dirs2[x,y], d
|
||
if d == 1:
|
||
x += 1
|
||
elif d == 2:
|
||
y += 1
|
||
elif d == 3:
|
||
x -= 1
|
||
elif d == 4:
|
||
y -= 1
|
||
d = dir_reverse[d]
|
||
|
||
del basin_links[b2][b1]
|
||
del basin_links[b1]
|
||
|
||
# Calculating water quantity
|
||
dirs2[-1,:][dirs2[-1,:]==1] = 0
|
||
dirs2[:,-1][dirs2[:,-1]==2] = 0
|
||
dirs2[0,:][dirs2[0,:]==3] = 0
|
||
dirs2[:,0][dirs2[:,0]==4] = 0
|
||
|
||
waterq = accumulate_flow(dirs2)
|
||
|
||
return dirs2, basins[basin_id], waterq
|
||
|
||
def accumulate_flow(dirs):
|
||
ndonors = np.zeros(dirs.shape, dtype=int)
|
||
ndonors[1:,:] += dirs[:-1,:] == 1
|
||
ndonors[:,1:] += dirs[:,:-1] == 2
|
||
ndonors[:-1,:] += dirs[1:,:] == 3
|
||
ndonors[:,:-1] += dirs[:,1:] == 4
|
||
waterq = np.ones(dirs.shape, dtype=int)
|
||
|
||
(X, Y) = dirs.shape
|
||
rangeX = range(X)
|
||
rangeY = range(Y)
|
||
for x in rangeX:
|
||
for y in rangeY:
|
||
if ndonors[x,y] > 0:
|
||
continue
|
||
xw, yw = x, y
|
||
w = waterq[xw, yw]
|
||
while 1:
|
||
d = dirs[xw, yw]
|
||
if d <= 0:
|
||
break
|
||
elif d == 1:
|
||
xw += 1
|
||
elif d == 2:
|
||
yw += 1
|
||
elif d == 3:
|
||
xw -= 1
|
||
elif d == 4:
|
||
yw -= 1
|
||
|
||
w += waterq[xw, yw]
|
||
waterq[xw, yw] = w
|
||
|
||
if ndonors[xw, yw] > 1:
|
||
ndonors[xw, yw] -= 1
|
||
break
|
||
|
||
return waterq
|
||
|
||
def planar_boruvka(links):
|
||
# Compute basin tree
|
||
|
||
basin_list = defaultdict(dict)
|
||
|
||
for (b1, b2), (elev, bound) in links.items():
|
||
basin_list[b1][b2] = basin_list[b2][b1] = (elev, b1, b2, bound)
|
||
|
||
threshold = 8
|
||
lowlevel = {}
|
||
for k, v in basin_list.items():
|
||
if len(v) <= threshold:
|
||
lowlevel[k] = v
|
||
|
||
basin_graph = []
|
||
n = len(basin_list)
|
||
while n > 1:
|
||
(b1, lnk1) = lowlevel.popitem()
|
||
b2 = min(lnk1, key=lnk1.get)
|
||
lnk2 = basin_list[b2]
|
||
|
||
# Add link to the graph
|
||
basin_graph.append(lnk1[b2])
|
||
|
||
# Union : merge basin 1 into basin 2
|
||
# First, delete the direct link
|
||
del lnk1[b2]
|
||
del lnk2[b1]
|
||
|
||
# Look for basin 1's neighbours, and add them to basin 2 if they have a lower pass
|
||
for k, v in lnk1.items():
|
||
bk = basin_list[k]
|
||
if k in lnk2 and lnk2[k] < v:
|
||
del bk[b1]
|
||
else:
|
||
lnk2[k] = v
|
||
bk[b2] = bk.pop(b1)
|
||
|
||
if k not in lowlevel and len(bk) <= threshold:
|
||
lowlevel[k] = bk
|
||
|
||
if b2 in lowlevel:
|
||
if len(lnk2) > threshold:
|
||
del lowlevel[b2]
|
||
elif len(lnk2) <= threshold:
|
||
lowlevel[b2] = lnk2
|
||
del lnk1
|
||
|
||
n -= 1
|
||
|
||
return basin_graph
|