mapgen_rivers/rivermapper.py

116 lines
3.4 KiB
Python
Raw Normal View History

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):
2020-04-26 17:13:38 +02:00
"""
Calculates flow direction in D4 (4 choices) for every pixel of the DEM
Also returns an array of lake elevation
"""
(Y, X) = dem.shape
2020-04-26 17:13:38 +02:00
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))
2020-04-26 17:13:38 +02:00
# Make a binary heap
heapq.heapify(borders)
dirs = np.zeros((Y+2, X+2), dtype='u1')
2020-04-26 17:13:38 +02:00
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:
2020-04-26 17:13:38 +02:00
(alt, altmax, y, x) = heapq.heappop(borders) # Take the lowest pixel in the queue
neighbours = dirs[y-1:y+2, x-1:x+2]
2020-04-26 17:13:38 +02:00
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
2020-04-26 17:13:38 +02:00
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())
2020-04-26 17:13:38 +02:00
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):
2020-04-26 17:13:38 +02:00
"""
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]
2020-04-26 17:13:38 +02:00
q = 1 # Consider that every pixel contains a water quantity of 1 by default.
neighbours = dirs_margin[y:y+3, x:x+3]
2020-04-26 17:13:38 +02:00
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:
2020-04-26 17:13:38 +02:00
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):
2020-04-26 17:13:38 +02:00
"""
Calculates flow directions and water quantity
"""
dirs, lakes = flow_dirs_lakes(dem)
return dirs, lakes, accumulate(dirs, dem)