116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
import numpy as np
|
|
import heapq
|
|
import sys
|
|
|
|
# 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='u1')
|
|
|
|
neighbours_pattern = neighbours_dirs > 0
|
|
|
|
def flow_dirs_lakes(dem, random=0):
|
|
"""
|
|
Calculates flow direction in D4 (4 choices) for every pixel of the DEM
|
|
Also returns an array of lake elevation
|
|
"""
|
|
|
|
(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='u1')
|
|
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='i4')
|
|
|
|
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
|
|
|
|
def flow(dem):
|
|
"""
|
|
Calculates flow directions and water quantity
|
|
"""
|
|
|
|
dirs, lakes = flow_dirs_lakes(dem)
|
|
return dirs, lakes, accumulate(dirs, dem)
|