mirror of
https://gitlab.com/gaelysam/mapgen_rivers.git
synced 2025-07-07 10:50:43 +02:00
Moved Python files inside a folder (package), except the 2 that are directly executable
This commit is contained in:
115
terrainlib/rivermapper.py
Normal file
115
terrainlib/rivermapper.py
Normal file
@ -0,0 +1,115 @@
|
||||
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)
|
Reference in New Issue
Block a user