248 lines
6.7 KiB
Python
248 lines
6.7 KiB
Python
import numpy as np
|
|
import numpy.random as npr
|
|
from collections import defaultdict
|
|
|
|
def flow_local(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
|
|
|
|
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
|