From ecd2c7d3f9564c67871ae8e23d0f869d328edb7b Mon Sep 17 00:00:00 2001 From: Gael-de-Sailly Date: Fri, 4 Dec 2020 01:03:03 +0100 Subject: [PATCH] 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. --- terrainlib/rivermapper.py | 348 ++++++++++++++++++++++++++------------ 1 file changed, 240 insertions(+), 108 deletions(-) diff --git a/terrainlib/rivermapper.py b/terrainlib/rivermapper.py index fcd738e..687c14d 100644 --- a/terrainlib/rivermapper.py +++ b/terrainlib/rivermapper.py @@ -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 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