149 lines
3.5 KiB
Python
149 lines
3.5 KiB
Python
from heapq import heappush, heappop, heapify
|
||
import sys
|
||
import numpy as np
|
||
sys.setrecursionlimit(65536)
|
||
|
||
def generate_rivermap(heightmap, sea_level=128, river_limit=1000, max_river_hdiff=40, river_power=0.25):
|
||
print("Generating rivermap")
|
||
print("[rivers] Finding start points")
|
||
|
||
(Y, X) = heightmap.shape
|
||
visited = np.zeros((Y,X), dtype=bool)
|
||
|
||
start_points = []
|
||
|
||
def add_start_point(y,x):
|
||
start_points.append((heightmap[y, x] + np.random.random(), y, x))
|
||
visited[y, x] = True
|
||
|
||
def find_start_points(t, x=1, y=1):
|
||
sy, sx = t.shape
|
||
if t.all() or not t.any():
|
||
return
|
||
if max(sx, sy) == 3:
|
||
if (not t[1,1]) and (t[0,1] or t[1,0] or t[1,2] or t[2,1]):
|
||
add_start_point(y,x)
|
||
return
|
||
if sx < sy:
|
||
cut = sy//2
|
||
find_start_points(t[:cut+1,:], x=x, y=y)
|
||
find_start_points(t[cut-1:,:], x=x, y=y+cut-1)
|
||
else:
|
||
cut = sx//2
|
||
find_start_points(t[:,:cut+1], x=x, y=y)
|
||
find_start_points(t[:,cut-1:], x=x+cut-1, y=y)
|
||
|
||
seas = heightmap <= sea_level
|
||
find_start_points(seas)
|
||
|
||
to_explore = X * Y - np.count_nonzero(seas)
|
||
|
||
for x in np.flatnonzero(~seas[0,:]):
|
||
add_start_point(0, x)
|
||
for x in np.flatnonzero(~seas[-1,:]):
|
||
add_start_point(Y-1, x)
|
||
for y in np.flatnonzero(~seas[1:-1,0]):
|
||
add_start_point(y+1, 0)
|
||
for y in np.flatnonzero(~seas[1:-1,-1]):
|
||
add_start_point(y+1, X-1)
|
||
|
||
del seas
|
||
|
||
print("[rivers] Found", str(len(start_points)), "start points")
|
||
|
||
heap = start_points[:]
|
||
heapify(heap)
|
||
|
||
print("[rivers] Building river trees:", str(to_explore), "points to visit")
|
||
|
||
flow_dirs = np.zeros((Y, X), dtype=np.int8)
|
||
|
||
# Directions:
|
||
# 1: +x
|
||
# 2: +y
|
||
# 4: -x
|
||
# 8: -y
|
||
|
||
def try_push(y, x): # try_push does 2 things at once: returning whether water can flow, and push the upward position in heap if yes.
|
||
if not visited[y, x]:
|
||
h = heightmap[y, x]
|
||
if h > sea_level:
|
||
heappush(heap, (h + np.random.random(), y, x))
|
||
visited[y, x] = True
|
||
return True
|
||
return False
|
||
|
||
def process_neighbors(y, x):
|
||
dirs = 0
|
||
if x > 0 and try_push(y, x-1):
|
||
dirs+= 1
|
||
if y > 0 and try_push(y-1, x):
|
||
dirs += 2
|
||
if x < X-1 and try_push(y, x+1):
|
||
dirs += 4
|
||
if y < Y-1 and try_push(y+1, x):
|
||
dirs += 8
|
||
flow_dirs[y, x] = dirs
|
||
|
||
while len(heap) > 0:
|
||
t = heappop(heap)
|
||
to_explore -= 1
|
||
if to_explore % 1000000 == 0:
|
||
print("[rivers]", str(to_explore // 1000000), "× 10⁶ points remaining", "Altitude:", int(t[0]), "Queue:", len(heap))
|
||
process_neighbors(t[1], t[2])
|
||
|
||
visited = None
|
||
|
||
print("[rivers] Calculating water quantity")
|
||
|
||
waterq = np.ones((Y, X))
|
||
river_array = np.zeros((Y, X), dtype=bool)
|
||
|
||
def draw_river(x, y, q):
|
||
if q >= river_limit:
|
||
rsize = int((q / river_limit)**river_power)
|
||
if rsize > 1:
|
||
hmax = heightmap[y,x] + max_river_hdiff
|
||
rsize -= 1
|
||
xmin = max(x-rsize, 0)
|
||
xmax = min(x+rsize+1, X)
|
||
ymin = max(y-rsize, 0)
|
||
ymax = min(y+rsize+1,Y)
|
||
river_array[y,xmin:xmax] += heightmap[y,xmin:xmax] <= hmax
|
||
river_array[ymin:ymax,x] += heightmap[ymin:ymax,x] <= hmax
|
||
else:
|
||
river_array[y,x] = True
|
||
|
||
def set_water(y, x):
|
||
water = 1
|
||
dirs = flow_dirs[y, x]
|
||
|
||
if dirs % 2 == 1:
|
||
water += set_water(y, x-1)
|
||
dirs //= 2
|
||
if dirs % 2 == 1:
|
||
water += set_water(y-1, x)
|
||
dirs //= 2
|
||
if dirs % 2 == 1:
|
||
water += set_water(y, x+1)
|
||
dirs //= 2
|
||
if dirs % 2 == 1:
|
||
water += set_water(y+1, x)
|
||
waterq[y, x] = water
|
||
|
||
if water >= river_limit:
|
||
draw_river(x, y, water)
|
||
return water
|
||
|
||
maxwater = 0
|
||
for start in start_points:
|
||
water = set_water(start[1], start[2])
|
||
if water > maxwater:
|
||
maxwater = water
|
||
|
||
print("[rivers] Maximal water quantity:", str(maxwater))
|
||
|
||
flow_dirs = None
|
||
|
||
return river_array
|