mapgen_rivers/terrainlib/rivermapper.py

258 lines
7.3 KiB
Python
Raw Normal View History

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, 549562, 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.
def flow_local(plist):
2020-04-26 17:13:38 +02:00
"""
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)
2020-04-26 17:13:38 +02:00
"""
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
2020-04-26 17:13:38 +02:00
def flow(dem):
# 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