diamond-sq.py

# For my untitled roguelike, I want to procedurally generate the 
# map of the surface. After reading/skimming many articles and blog
# posts, I hit on the diamond-square algorithm. Some of the examples
# I'd seen produced nice enough maps, and I was able to understand 
# how it worked better than, say, Perlin Noise. My roguelike is going
# to be done in Rust, but I find Python way easier to prototype in.
# So, here's my attempt to implement the diamond-square algorithm
# (https://en.wikipedia.org/wiki/Diamond-square_algorithm)

import math
from enum import Enum
from random import random
from random import uniform
from PIL import Image

class Terrain(Enum):
	WATER = 0
	BEACH = 1
	GRASS = 2
	DIRT = 3
	FOREST = 4
	MOUNTAIN = 5
	SNOWPEAK = 6
	SHALLOW_WATER = 7

def to_terrain(val):
	if val < -0.50:
		return Terrain.WATER
	elif val < -0.25:
		return Terrain.SHALLOW_WATER
	elif val < -0.0:
		return Terrain.BEACH
	elif val < 0.25:
		return Terrain.DIRT
	elif val < 0.85:
		return Terrain.FOREST
	elif val < 1.5:
		return Terrain.MOUNTAIN
	
	return Terrain.SNOWPEAK

def to_terrain_alt(val):
	if val > 160:	
		return Terrain.MOUNTAIN
	elif val > 120:
		return Terrain.FOREST
	elif val > 100:
		return Terrain.DIRT
	elif val > 90:
		return Terrain.BEACH
	else:
		return Terrain.WATER

def terrain_to_rgb(val):
	if val == Terrain.WATER:
		return (0, 0, 221)
	elif val == Terrain.DIRT:
		return (153, 0, 0)
	elif val == Terrain.BEACH:
		return (255, 178, 127)
	elif val == Terrain.FOREST:
		#return (34, 139, 34)
		return (0, 80, 0)
	elif val == Terrain.MOUNTAIN:
		return (112, 128, 144)
	elif val == Terrain.SNOWPEAK:
		return (255, 255, 255)
	elif val == Terrain.SHALLOW_WATER:
		return (173, 216, 230)

def to_png(grid, size):
	img = Image.new('RGB', (size * 2, size * 2))
	pixels = []	

	for row in grid:
		for val in row:
			terrain = to_terrain(val)
			pixels.append(terrain_to_rgb(terrain))
			pixels.append(terrain_to_rgb(terrain))
		for val in row:
			terrain = to_terrain(val)
			pixels.append(terrain_to_rgb(terrain))
			pixels.append(terrain_to_rgb(terrain))

	img.putdata(pixels)
	img.show()
	
def dump_grid(grid):
	for row in grid:
		s = ""
		for col in row:
			s += f'{col:.2f}  '	
		print(s)

# All amount of fuzziness to add to the averages we
# calculate
def fuzz():
	return uniform(-0.75, 0.75)
	#return (random() - 0.5) / 1

def fuzz2(size):
	return (uniform(0.0, 1.0) * 2 - 1) * size * (1 / size)

# perform the diamond step described in the wiki article. Average
# the four corners of the square defined by (r, c) as the upper left
# corner of the square, and size as the length of the sides of the square
# (I agree with a blogger I read who said the names of the diamond step
# and square steps seem reversed to me...
def diamond_step(grid, r, c, size):
	avg = grid[r][c]
	avg += grid[r][c + size - 1]
	avg += grid[r + size - 1][c]
	avg += grid[r + size - 1][c + size - 1]
	avg /= 4

	grid[r + size // 2][c + size // 2] = avg + fuzz2(size)

# calculate the average of the 4 corners of the diamond
# with (r, c) as the centre. Note that it's possible for
# one of those points to be outside the grid. We'll ignore those
# points.
def calc_diamond_avg(grid, r, c, size):
	count = 0
	avg = 0.0
	if c - size >= 0:
		avg += grid[r][c - size]
		count += 1
	if c + size < len(grid):
		avg += grid[r][c + size]
		count += 1
	if r - size >= 0:
		avg += grid[r - size][c]
		count += 1
	if r + size < len(grid):
		avg += grid[r + size][c]
		count += 1

	grid[r][c] = avg / count + fuzz2(size)

# in square_step(), we set the four points on the grid that are
# midway between the center of the diamond (r, c) and its four 
# corners (the diagram of the square step in the wiki article is
# handy to see that this actually means. Note that in many cases,
# one of the corners of the diamond will be outside the matrix.
# The two solutions I've seen are to just average 3 points in those
# instances, or treat the matrix as though it wraps around. I'm going
# to start off just taking 3 poins and see how my maps turn out.
# of which (r, c) is the centre. Those 
def square_step(grid, r, c, size):
	half_size = size // 2
	
	# the four points we want to set are grid[r - half_size[c],
	# grid[r][c - half_size], grid[r][c + half_size] and
	# grid[r + half_size][c]
	calc_diamond_avg(grid, r - half_size, c, half_size)
	calc_diamond_avg(grid, r + half_size, c, half_size)
	calc_diamond_avg(grid, r, c - half_size, half_size)
	calc_diamond_avg(grid, r, c + half_size, half_size)

# Average each point with its neighbours to smooth things out
def smooth_map(grid, size):
	size
	for r in range(0, size):
		for c in range(0, size):
			avg = grid[r][c] 
			count = 1
			if r - 1 >= 0: 
				if c - 1 >= 0: 
					avg += grid[r - 1][c - 1]	
					count += 1
				avg += grid[r - 1][c]
				count += 1
				if c + 1 < size:
					avg += grid[r - 1][c + 1]
					count += 1
			if c - 1 >= 0:
				avg += grid[r][c - 1]
				count += 1
			if c + 1 < size:
				avg == grid[r][c + 1]
				count += 1
			if r + 1 < size: 
				if c - 1 >= 0: 
					avg += grid[r + 1][c - 1]	
					count += 1
				avg += grid[r + 1][c]
				count += 1
				if c + 1 < size:
					avg += grid[r + 1][c + 1]
					count += 1
			grid[r][c] = avg / count

def diamond_sq(grid, r, c, size, depth):
	diamond_step(grid, r, c, size)
	half_size = size // 2
	square_step(grid, r + half_size, c + half_size, size)

	if half_size == 1:
		return

	new_size = size // 2
	diamond_sq(grid, r, c, new_size + 1, depth + 1)
	diamond_sq(grid, r, c + new_size, new_size + 1, depth + 1)
	diamond_sq(grid, r + new_size, c, new_size + 1, depth + 1)
	diamond_sq(grid, r + new_size, c + new_size, new_size + 1, depth + 1)

def to_island(grid, size, shift_y):
	for r in range(size):
		for c in range(size):
			xd = c / (size - 1.0) * 2 - 1.0
			yd = r / (size - shift_y) * 2 - 1.0
			island_size = 1.15
			grid[r][c] = grid[r][c] + island_size - math.sqrt(xd*xd + yd*yd) * 3.0
  
SIZE = 65 # This alg. requires a square matrix whose side length is 2^n + 1
grid = []

for _ in range(0, SIZE):
	row = [0.0] * SIZE
	grid.append(row)

# seed the four corners of the matrix
#grid[0][0] = -0.15 #random()
#grid[0][SIZE - 1] = abs(random()) + 0.5 
#grid[SIZE - 1][0] = abs(random()) + 0.5 
#grid[SIZE - 1][SIZE - 1] = abs(random()) + 1.25

# This, without the to_island() call is making maps I think would be
# good for Rogue Village
grid[0][0] = -2.0 #fuzz()
grid[0][SIZE - 1] = 1.0 #fuzz() # abs(random()) + 0.5 
grid[SIZE - 1][0] = 1.5 # fuzz() # abs(random()) + 0.5 
grid[SIZE - 1][SIZE - 1] = 2.0 # fuzz() # abs(random()) + 1.25

try:
	diamond_sq(grid, 0, 0, SIZE, 1)
	smooth_map(grid, SIZE)
	smooth_map(grid, SIZE)
	to_island(grid, SIZE, -0.0)

	#for r in range(0, SIZE):
	#	for c in range(0, SIZE):
	#		grid[r][c] = grid[r][c] * 60 + 120
except IndexError:
	print("boo :(")	

#print(grid[4][0:10])
to_png(grid, SIZE)