Completely change flow routing algorithm,
use a local flow calculation, determine depressions, and link them using a minimum spanning tree (Boruvka's algorithm). This is based on a paper by Cordonnier et al, 2019.
This commit is contained in:
Gael-de-Sailly
parent
40098d6be3
commit
ecd2c7d3f9
|
|
@ -1,115 +1,247 @@
|
|||
import numpy as np
|
||||
import heapq
|
||||
import sys
|
||||
import numpy.random as npr
|
||||
from collections import defaultdict
|
||||
|
||||
# Conventions:
|
||||
# 1 = South (+Y)
|
||||
# 2 = East (+X)
|
||||
# 3 = North (-Y)
|
||||
# 4 = West (-X)
|
||||
|
||||
sys.setrecursionlimit(65536)
|
||||
|
||||
neighbours_dirs = np.array([
|
||||
[0,1,0],
|
||||
[2,0,4],
|
||||
[0,3,0],
|
||||
], dtype=int)
|
||||
|
||||
neighbours_pattern = neighbours_dirs > 0
|
||||
|
||||
def flow_dirs_lakes(dem, random=0):
|
||||
def flow_local(plist):
|
||||
"""
|
||||
Calculates flow direction in D4 (4 choices) for every pixel of the DEM
|
||||
Also returns an array of lake elevation
|
||||
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)
|
||||
"""
|
||||
|
||||
(Y, X) = dem.shape
|
||||
|
||||
dem_margin = np.zeros((Y+2, X+2)) # We need a margin of one pixel at every edge, to prevent crashes when scanning the neighbour pixels on the borders
|
||||
dem_margin[1:-1,1:-1] = dem
|
||||
if random > 0:
|
||||
dem_margin += np.random.random(dem_margin.shape) * random
|
||||
|
||||
# Initialize: list map borders
|
||||
borders = []
|
||||
|
||||
for x in range(1,X+1):
|
||||
dem_north = dem_margin[1,x]
|
||||
borders.append((dem_north, dem_north, 1, x))
|
||||
dem_south = dem_margin[Y,x]
|
||||
borders.append((dem_south, dem_south, Y, x))
|
||||
|
||||
for y in range(2,Y):
|
||||
dem_west = dem_margin[y,1]
|
||||
borders.append((dem_west, dem_west, y, 1))
|
||||
dem_east = dem_margin[y,X]
|
||||
borders.append((dem_east, dem_east, y, X))
|
||||
|
||||
# Make a binary heap
|
||||
heapq.heapify(borders)
|
||||
|
||||
dirs = np.zeros((Y+2, X+2), dtype=int)
|
||||
dirs[-2:,:] = 1 # Border pixels flow outside the map
|
||||
dirs[:,-2:] = 2
|
||||
dirs[ :2,:] = 3
|
||||
dirs[:, :2] = 4
|
||||
|
||||
lakes = np.zeros((Y, X))
|
||||
|
||||
def add_point(y, x, altmax):
|
||||
alt = dem_margin[y, x]
|
||||
heapq.heappush(borders, (alt, altmax, y, x))
|
||||
|
||||
while len(borders) > 0:
|
||||
(alt, altmax, y, x) = heapq.heappop(borders) # Take the lowest pixel in the queue
|
||||
neighbours = dirs[y-1:y+2, x-1:x+2]
|
||||
empty_neighbours = (neighbours == 0) * neighbours_pattern # Find the neighbours whose flow direction is not yet defined
|
||||
neighbours += empty_neighbours * neighbours_dirs # They flow into the pixel being studied
|
||||
|
||||
lake = max(alt, altmax) # Set lake elevation to the maximal height of the downstream section.
|
||||
lakes[y-1,x-1] = lake
|
||||
|
||||
coords = np.transpose(empty_neighbours.nonzero())
|
||||
for (dy,dx) in coords-1: # Add these neighbours into the queue
|
||||
add_point(y+dy, x+dx, lake)
|
||||
|
||||
return dirs[1:-1,1:-1], lakes
|
||||
|
||||
def accumulate(dirs, dem=None):
|
||||
"""
|
||||
Calculates the quantity of water that accumulates at every pixel,
|
||||
following flow directions.
|
||||
"""
|
||||
|
||||
(Y, X) = dirs.shape
|
||||
dirs_margin = np.zeros((Y+2,X+2))-1
|
||||
dirs_margin[1:-1,1:-1] = dirs
|
||||
quantity = np.zeros((Y, X), dtype=int)
|
||||
|
||||
def calculate_quantity(y, x):
|
||||
if quantity[y,x] > 0:
|
||||
return quantity[y,x]
|
||||
q = 1 # Consider that every pixel contains a water quantity of 1 by default.
|
||||
neighbours = dirs_margin[y:y+3, x:x+3]
|
||||
donors = neighbours == neighbours_dirs # Identify neighbours that give their water to the pixel being studied
|
||||
|
||||
coords = np.transpose(donors.nonzero())
|
||||
for (dy,dx) in coords-1:
|
||||
q += calculate_quantity(y+dy, x+dx) # Add water quantity of the donors pixels (this triggers calculation for these pixels, recursively)
|
||||
quantity[y, x] = q
|
||||
return q
|
||||
|
||||
for x in range(X):
|
||||
for y in range(Y):
|
||||
calculate_quantity(y, x)
|
||||
|
||||
return quantity
|
||||
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):
|
||||
"""
|
||||
Calculates flow directions and water quantity
|
||||
"""
|
||||
|
||||
dirs, lakes = flow_dirs_lakes(dem)
|
||||
return dirs, lakes, accumulate(dirs, 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, np.maximum(basins[basin_id], dem), 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
|
||||
|
|
|
|||
Loading…
Reference in New Issue