world-ecoregion/map-generator/index.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors, cm
import scipy.interpolate as interpolate
from scipy import ndimage
import math
from io import BytesIO
import base64

parameters = {
    'width': {
        'default': 900,
        'type': 'int',
    },
    'height': {
        'default': 450,
        'type': 'int',
    },
    'mountain_ratio': {
        'default': 0.3,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'sharpness': {
        'default': 0.7,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'max_elevation': {
        'default': 30,
        'type': 'int',
        'min': 0,
        'max': 50,
    },
    'ground_noise': {
        'default': 15,
        'type': 'int',
        'min': 0,
        'max': 50,
    },
    'water_proportion': {
        'default': 0.6,
        'type': 'float',
        'min': 0,
        'max': 0.99,
        'step': 0.01
    },
    'mountain_jaggedness': {
        'default': 1,
        'type': 'int',
        'min': 0,
        'max': 5,
    },
    'mountain_sea_distance': {
        'default': 50,
        'type': 'int',
        'min': 0,
        'max': 200,
    },
    'mountain_sea_threshold': {
        'default': 2,
        'type': 'int',
        'min': 0,
        'max': 5,
    },
    'water_level': {
        'default': 0,
        'type': 'int',
    },
    'mountain_area_elevation': {
        'default': 0.4,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'mountain_area_elevation_n': {
        'default': 5,
        'type': 'int',
        'min': 0,
        'max': 15,
    },
    'mountain_area_elevation_area': {
        'default': 10,
        'type': 'int',
        'min': 0,
        'max': 25,
    },
    'continents': {
        'default': 5,
        'type': 'int',
    },
    'seed': {
        'default': '',
        'type': 'int',
        'description': 'Leave empty for a random seed generated from the current timestamp.'
    },
        
}
p = { k: parameters[k]['default'] for k in parameters }

CONTINENT_MAX_TRIALS = 1e4
SEA_COLOR = np.array((53, 179, 220, 255)) / 255

DIRECTIONS = [(-1, -1), (-1, 0), (-1, 1), (1, 1), (1, 0), (1, -1), (0, -1), (0, 1)]

def s(x):
    return -2 * x**3 + 3 * x**2

def is_ground(value):
    return value > p['water_level']

def in_range(p, m, size):
    x, y = p
    mx, my = m
    return ((x - mx)**2 + (y - my)**2) < size

def max_recursion(fn, max_recursion=0):
    def f(*args, recursion=0, **kwargs):
        if recursion > max_recursion:
            return

        return fn(*args, **kwargs)

def bound_check(ground, point):
    x, y = point
    w, h = ground.shape

    x = max(min(x, w - 1), 0)
    y = max(min(y, h - 1), 0)

    return (x, y)


def continent_agent(ground, position, size):
    if size <= 0: return
    x, y = position

    w, h = ground.shape

    trials = 0

    while True:
        if size <= 0 or trials > CONTINENT_MAX_TRIALS: break

        dx = np.random.randint(2) or -1
        dy = np.random.randint(2) or -1

        r = np.random.randint(3)
        new_point = bound_check(ground, (x + dx, y + dy))
        if r == 0:
            x = new_point[0]
        elif r == 1:
            y = new_point[1]
        else:
            x, y = new_point


        if not is_ground(ground[x, y]) and in_range((x, y), position, size):
            trials = 0
            size -= 1
            ground[x, y] = np.random.randint(1, p['ground_noise'])
        else:
            trials += 1

def neighbours(ground, position, radius):
    x, y = position
    return ground[x-radius:x+radius+1, y-radius:y+radius+1]

def away_from_sea(ground, position, radius=p['mountain_sea_distance']):
    ns = neighbours(ground, position, radius).flatten()
    sea = len([1 for x in ns if not is_ground(x)])

    return sea < p['mountain_sea_threshold']

def random_elevate_agent(ground, position, height, size=p['mountain_area_elevation_n']):
    position = position + np.random.random_integers(-p['mountain_area_elevation_area'], p['mountain_area_elevation_area'], size=2)

    for i in range(size):
        d = DIRECTIONS[np.random.randint(len(DIRECTIONS))]

        change = height * p['mountain_area_elevation'] + np.random.randint(p['mountain_jaggedness'] + 1)
        new_index = bound_check(ground, position + np.array(d))

        if is_ground(ground[new_index]):
            ground[new_index] += change


def mountain_agent(ground, position):
    if not away_from_sea(ground, position):
        return

    x, y = position
    height = np.random.randint(p['max_elevation'])

    ground[x, y] = height

    last_height = height
    for i in range(1, height):
        for d in DIRECTIONS:
            change = np.random.randint(p['mountain_jaggedness'] + 1)
            distance = np.array(d)*i
            new_index = bound_check(ground, position + distance)

            if is_ground(ground[new_index]):
                ground[new_index] = last_height - change

        last_height = last_height - p['mountain_jaggedness']

    random_elevate_agent(ground, position, height)


# takes an initial position and a list of (direction, probability) tuples to walk on
# def split_agent(ground, position, directions):

def constant_filter(a):
    if a[0] > (1 - p['sharpness']):
        return max(1, a[0])
    return 0

def generate_map(**kwargs):
    plt.clf()

    p.update(kwargs)

    np.random.seed(p['seed'] or None)

    width, height = p['width'], p['height']
    continents = p['continents']

    ground = np.zeros((width, height))
    ground_size = width * height * (1 - p['water_proportion'])

    # position = (int(width / 2), int(height / 2))
    # ground_size = width * height * GROUND_PROPORTION
    # continent_agent(ground, position, size=ground_size)
    for continent in range(continents):
        position = (np.random.randint(0, width), np.random.randint(0, height))
        print(position)
        continent_agent(ground, position, size=ground_size)

    ground = ndimage.gaussian_filter(ground, sigma=(1 - p['sharpness']) * 20)

    for i in range(int(ground_size * p['mountain_ratio'] / p['max_elevation']**2)):
        position = (np.random.randint(0, width), np.random.randint(0, height))
        mountain_agent(ground, position)

    norm = colors.Normalize(vmin=1)
    greys = cm.get_cmap('Greys')
    greys.set_under(color=SEA_COLOR)

    ground = ndimage.gaussian_filter(ground, sigma=4)
    ground = ndimage.generic_filter(ground, constant_filter, size=1)
    print(np.min(ground), np.max(ground), p['max_elevation'])
    print(np.unique(ground))

    plt.imshow(ground.T, cmap=greys, norm=norm)

    figfile = BytesIO()
    plt.savefig(figfile, format='png')
    figfile.seek(0)

    return figfile

if __name__ == "__main__":
    generate_map()
    plt.show()