Moved Python files inside a folder (package), except the 2 that are directly executable

terrainlib/__init__.py

@@ -0,0 +1,7 @@
+# Load packages and provide easy access to important functions
+
+from .settings import read_config_file
+from .erosion import EvolutionModel
+from .save import save
+from .bounds import make_bounds, twist, get_fixed
+from .view import stats, update, plot

terrainlib/bounds.py

@@ -0,0 +1,74 @@
+import numpy as np
+
+def make_bounds(dirs, rivers):
+    """
+    Give an array of all horizontal and vertical bounds
+    """
+
+    (Y, X) = dirs.shape
+    bounds_h = np.zeros((Y, X-1), dtype='i4')
+    bounds_v = np.zeros((Y-1, X), dtype='i4')
+
+    bounds_v += (rivers * (dirs==1))[:-1,:]
+    bounds_h += (rivers * (dirs==2))[:,:-1]
+    bounds_v -= (rivers * (dirs==3))[1:,:]
+    bounds_h -= (rivers * (dirs==4))[:,1:]
+
+    return bounds_h, bounds_v
+
+def get_fixed(dirs):
+    """
+    Give the list of points that should not be twisted
+    """
+
+    borders = np.zeros(dirs.shape, dtype='?')
+    borders[-1,:] |= dirs[-1,:]==1
+    borders[:,-1] |= dirs[:,-1]==2
+    borders[0,:] |= dirs[0,:]==3
+    borders[:,0] |= dirs[:,0]==4
+
+    donors = np.zeros(dirs.shape, dtype='?')
+    donors[1:,:] |= dirs[:-1,:]==1
+    donors[:,1:] |= dirs[:,:-1]==2
+    donors[:-1,:] |= dirs[1:,:]==3
+    donors[:,:-1] |= dirs[:,1:]==4
+    return borders | ~donors
+
+def twist(bounds_x, bounds_y, fixed, d=0.1, n=5):
+    """
+    Twist the grid (define an offset for every node). Model river bounds as if they were elastics.
+    Smoothes preferentially big rivers.
+    """
+
+    moveable = ~fixed
+
+    (Y, X) = fixed.shape
+    offset_x = np.zeros((Y, X))
+    offset_y = np.zeros((Y, X))
+
+    for i in range(n):
+        force_long = np.abs(bounds_x) * (1+np.diff(offset_x, axis=1))
+        force_trans = np.abs(bounds_y) * np.diff(offset_x, axis=0)
+
+        force_x = np.zeros((Y, X))
+        force_x[:,:-1] = force_long
+        force_x[:,1:] -= force_long
+        force_x[:-1,:]+= force_trans
+        force_x[1:,:] -= force_trans
+
+        force_long = np.abs(bounds_y) * (1+np.diff(offset_y, axis=0))
+        force_trans = np.abs(bounds_x) * np.diff(offset_y, axis=1)
+
+        force_y = np.zeros((Y, X))
+        force_y[:-1,:] = force_long
+        force_y[1:,:] -= force_long
+        force_y[:,:-1]+= force_trans
+        force_y[:,1:] -= force_trans
+
+        length = np.hypot(force_x, force_y)
+        length[length==0] = 1
+        coeff = d / length * moveable # Normalize, take into account the direction only
+        offset_x += force_x * coeff
+        offset_y += force_y * coeff
+
+    return offset_x, offset_y

terrainlib/erosion.py

@@ -0,0 +1,97 @@
+import numpy as np
+import scipy.ndimage as im
+from .rivermapper import flow
+
+def advection(dem, dirs, rivers, time, K=1, m=0.5, sea_level=0):
+    """
+    Simulate erosion by rivers.
+    This models erosion as an upstream advection of elevations ("erosion waves").
+    Advection speed depends on water flux and parameters:
+
+    v = K * flux^m
+    """
+
+    dirs = dirs.copy()
+    dirs[0,:] = 0
+    dirs[-1,:] = 0
+    dirs[:,0] = 0
+    dirs[:,-1] = 0
+
+    adv_time = 1 / (K*rivers**m) # For every pixel, calculate the time an "erosion wave" will need to cross it.
+    dem = np.maximum(dem, sea_level)
+    dem_new = np.zeros(dem.shape)
+
+    for y in range(dirs.shape[0]):
+        for x in range(dirs.shape[1]):
+            # Elevations propagate upstream, so for every pixel we seek the downstream pixel whose erosion wave just reached the current pixel.
+            # This means summing the advection times downstream until we reach the erosion time.
+            x0, y0 = x, y
+            x1, y1 = x, y
+            remaining = time
+            while True:
+                # Move one pixel downstream
+                flow_dir = dirs[y0,x0]
+                if flow_dir == 0:
+                    remaining = 0
+                    break
+                elif flow_dir == 1:
+                    y1 += 1
+                elif flow_dir == 2:
+                    x1 += 1
+                elif flow_dir == 3:
+                    y1 -= 1
+                elif flow_dir == 4:
+                    x1 -= 1
+
+                if remaining <= adv_time[y0,x0]: # Time is over, we found it.
+                    break
+                remaining -= adv_time[y0,x0]
+                x0, y0 = x1, y1
+
+            c = remaining / adv_time[y0,x0]
+            dem_new[y,x] = c*dem[y1,x1] + (1-c)*dem[y0,x0] # If between 2 pixels, perform linear interpolation.
+
+    return np.minimum(dem, dem_new)
+
+def diffusion(dem, time, d=1):
+    radius = d * time**.5
+    return im.gaussian_filter(dem, radius, mode='reflect') # Diffusive erosion is a simple Gaussian blur
+
+class EvolutionModel:
+    def __init__(self, dem, K=1, m=0.5, d=1, sea_level=0, flow=False, flex_radius=100):
+        self.dem = dem
+        #self.bedrock = dem
+        self.K = K
+        self.m = m
+        self.d = d
+        self.sea_level = sea_level
+        self.flex_radius = flex_radius
+        self.define_isostasy()
+        if flow:
+            self.calculate_flow()
+        else:
+            self.lakes = dem
+            self.dirs = np.zeros(dem.shape, dtype='u1')
+            self.rivers = np.zeros(dem.shape, dtype='u4')
+            self.flow_uptodate = False
+
+    def calculate_flow(self):
+        self.dirs, self.lakes, self.rivers = flow(self.dem)
+        self.flow_uptodate = True
+
+    def advection(self, time):
+        dem = advection(self.lakes, self.dirs, self.rivers, time, K=self.K, m=self.m, sea_level=self.sea_level)
+        self.dem = np.minimum(dem, self.dem)
+        self.flow_uptodate = False
+
+    def diffusion(self, time):
+        self.dem = diffusion(self.dem, time, d=self.d)
+        self.flow_uptodate = False
+
+    def define_isostasy(self):
+        self.ref_isostasy = im.gaussian_filter(self.dem, self.flex_radius, mode='reflect') # Define a blurred version of the DEM that will be considered as the reference isostatic elevation.
+
+    def adjust_isostasy(self, rate=1):
+        isostasy = im.gaussian_filter(self.dem, self.flex_radius, mode='reflect') # Calculate blurred DEM
+        correction = (self.ref_isostasy - isostasy) * rate # Compare it with the reference isostasy
+        self.dem = self.dem + correction # Adjust

terrainlib/rivermapper.py

@@ -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)

terrainlib/save.py

@@ -0,0 +1,13 @@
+import numpy as np
+import zlib
+
+def save(data, fname, dtype=None):
+    if dtype is not None:
+        data = data.astype(dtype)
+
+    bin_data = data.tobytes()
+    bin_data_comp = zlib.compress(bin_data, 9)
+    if len(bin_data_comp) < len(bin_data):
+        bin_data = bin_data_comp
+    with open(fname, 'wb') as f:
+        f.write(bin_data)

terrainlib/settings.py

@@ -0,0 +1,16 @@
+import os.path
+
+def read_config_file(fname):
+    settings = {}
+
+    if not os.path.isfile(fname):
+        return settings
+
+    with open(fname, 'r') as f:
+        for line in f:
+            slist = line.split('=', 1)
+            if len(slist) >= 2:
+                prefix, suffix = slist
+                settings[prefix.strip()] = suffix.strip()
+
+    return settings

terrainlib/view.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+import numpy as np
+
+has_matplotlib = True
+try:
+    import matplotlib.colors as mcl
+    import matplotlib.pyplot as plt
+    try:
+        import colorcet as cc
+        cmap1 = cc.cm.CET_L11
+        cmap2 = cc.cm.CET_L12
+    except ImportError: # No module colorcet
+        import matplotlib.cm as cm
+        cmap1 = cm.summer
+        cmap2 = cm.Blues
+except ImportError: # No module matplotlib
+    has_matplotlib = False
+
+if has_matplotlib:
+    def view_map(dem, lakes, scale=1, title=None):
+        if not has_matplotlib:
+            return
+        lakes_sea = np.maximum(lakes, 0)
+        water = np.maximum(lakes_sea - dem, 0)
+        max_elev = lakes_sea.max()
+        max_depth = water.max()
+
+        ls = mcl.LightSource(azdeg=315, altdeg=45)
+        rgb = ls.shade(lakes_sea, cmap=cmap1, vert_exag=1/scale, blend_mode='soft', vmin=0, vmax=max_elev)
+
+        (X, Y) = dem.shape
+        extent = (0, Y*scale, 0, X*scale)
+        plt.imshow(np.flipud(rgb), extent=extent, interpolation='antialiased')
+        alpha = (water > 0).astype('u1')
+        plt.imshow(np.flipud(water), alpha=np.flipud(alpha), cmap=cmap2, extent=extent, vmin=0, vmax=max_depth, interpolation='antialiased')
+
+        sm1 = plt.cm.ScalarMappable(cmap=cmap1, norm=plt.Normalize(vmin=0, vmax=max_elev))
+        plt.colorbar(sm1).set_label('Elevation')
+
+        sm2 = plt.cm.ScalarMappable(cmap=cmap2, norm=plt.Normalize(vmin=0, vmax=max_depth))
+        plt.colorbar(sm2).set_label('Water depth')
+
+        plt.xlabel('X')
+        plt.ylabel('Z')
+
+        if title is not None:
+            plt.title(title, fontweight='bold')
+
+    def update(*args, t=0.01, **kwargs):
+        plt.clf()
+        view_map(*args, **kwargs)
+        plt.pause(t)
+
+    def plot(*args, **kwargs):
+        plt.clf()
+        view_map(*args, **kwargs)
+        plt.show()
+
+else:
+    def update(*args, **kwargs):
+        pass
+    def plot(*args, **kwargs):
+        pass
+
+def stats(dem, lake_dem, scale=1):
+    surface = dem.size
+
+    continent = lake_dem >= 0
+    continent_surface = continent.sum()
+
+    lake = continent & (lake_dem>dem)
+    lake_surface = lake.sum()
+
+    print('---   General    ---')
+    print('Grid size:    {:5d}x{:5d}'.format(dem.shape[0], dem.shape[1]))
+    if scale > 1:
+        print('Map size:     {:5d}x{:5d}'.format(int(dem.shape[0]*scale), int(dem.shape[1]*scale)))
+    print()
+    print('---   Surfaces   ---')
+    print('Continents:        {:6.2%}'.format(continent_surface/surface))
+    print('-> Ground:         {:6.2%}'.format((continent_surface-lake_surface)/surface))
+    print('-> Lakes:          {:6.2%}'.format(lake_surface/surface))
+    print('Oceans:            {:6.2%}'.format(1-continent_surface/surface))
+    print()
+    print('---  Elevations  ---')
+    print('Mean elevation:      {:4.0f}'.format(dem.mean()))
+    print('Mean ocean depth:    {:4.0f}'.format((dem*~continent).sum()/(surface-continent_surface)))
+    print('Mean continent elev: {:4.0f}'.format((dem*continent).sum()/continent_surface))
+    print('Lowest elevation:    {:4.0f}'.format(dem.min()))
+    print('Highest elevation:   {:4.0f}'.format(dem.max()))