Skip to content

Instantly share code, notes, and snippets.

@nitori

nitori/wfc.py

Last active August 17, 2024 23:59
Show Gist options
  • Save nitori/4697667a35f3beecb12cca04481c01b6 to your computer and use it in GitHub Desktop.
Save nitori/4697667a35f3beecb12cca04481c01b6 to your computer and use it in GitHub Desktop.
Download ZIP
very simple wave function collapse thingy
Raw
wfc.py
from collections.abc import Iterable
import math
import heapq
import random
from collections import deque
import time
type WeightedOptions = list[tuple[int, int]]
UP = 0b1000
RIGHT = 0b0100
DOWN = 0b0010
LEFT = 0b0001
def entropy(weights: Iterable[int]) -> float:
ws = sum(weights)
weights = (w / ws for w in weights)
return -sum(p * math.log2(p) if p > 0 else 0 for p in weights)
class Cell:
options: WeightedOptions
def __init__(self, index, options: WeightedOptions):
self.index = index
self._options = options
self._entropy = None
self.collapsed = False
@property
def options(self):
return self._options
@options.setter
def options(self, value):
self._options = value
self._entropy = None
@property
def entropy(self) -> float:
if self._entropy is None:
self._entropy = entropy([w for _, w in self.options])
return self._entropy
def collapse(self, wave: 'Wave'):
self.collapsed = True
options, weights = zip(*self.options)
option = wave.rng.choices(options, k=1, weights=weights)[0]
self.options = [(option, 1)]
def __lt__(self, other: 'Cell'):
return self.entropy < other.entropy
def __repr__(self):
return (f'<{self.__class__.__name__} collapsed={self.collapsed} '
f'index={self.index} entropy={self.entropy:.4f}>')
class Wave:
# entropy, index, possible tiles
grid: list[Cell]
# maps index to cell.
grid_map: dict[int, Cell]
def __init__(self, width: int, height: int, options: WeightedOptions):
self.width = width
self.height = height
self.options = options
self.reset()
self.rng = random.Random()
def reset(self):
self.grid = [Cell(i, self.options[:]) for i in range(self.width * self.height)]
self.grid_map = {c.index: c for c in self.grid}
heapq.heapify(self.grid)
def get_minimum_entropy_tile(self) -> Cell | None:
minvalues = []
others = []
heapq.heapify(self.grid)
while self.grid:
cell: Cell = heapq.heappop(self.grid)
if cell.collapsed:
others.append(cell)
elif not minvalues:
minvalues.append(cell)
elif math.isclose(minvalues[0].entropy, cell.entropy):
minvalues.append(cell)
else:
others.append(cell)
break
for cell in others:
heapq.heappush(self.grid, cell)
for cell in minvalues:
heapq.heappush(self.grid, cell)
if not minvalues:
return None
return self.rng.choice(minvalues)
def collapse(self):
while True:
cell = self.get_minimum_entropy_tile()
if cell is None:
return
if not self.collapse_cell(cell):
return
def collapse_cell(self, cell: Cell):
cell.collapse(self)
return self.propagate(cell)
def propagate(self, cell: Cell):
stack = deque([cell])
while stack:
cell = stack.popleft()
for nexti, direction in self.neighbours(cell.index):
other_cell = self.get(nexti)
if other_cell.collapsed:
continue
options_before = other_cell.options[:]
options_after = self.reduce(cell, other_cell, direction)
if options_before != options_after:
other_cell.options = options_after
if not other_cell.options:
return False
stack.append(other_cell)
return True
def neighbours(self, index: int):
y, x = divmod(index, self.width)
for dx, dy, direction in [(0, -1, UP), (1, 0, RIGHT), (0, 1, DOWN), (-1, 0, LEFT)]:
nx, ny = dx + x, dy + y
if self.in_range(nx, ny):
ni = ny * self.width + nx
yield ni, direction
def in_range(self, x, y):
return 0 <= x < self.width and 0 <= y < self.height
def get(self, x_or_i: int, y: int | None = None) -> Cell:
if y is not None:
index = y * self.width + x_or_i
else:
index = x_or_i
return self.grid_map[index]
@staticmethod
def reduce(a: Cell, b: Cell, direction: int) -> WeightedOptions:
"""
direction is from a to b. e.g. "UP" means "b" is above "a", so we
have to compare a's "UP" side, and b's "DOWN" side.
return new list of weighted options for "b".
"""
if direction not in (UP, RIGHT, DOWN, LEFT):
raise ValueError('Invalid value for direction.')
other_direction = direction >> 2 if direction > 2 else direction << 2
keep = []
for tp_b, w in b.options:
for tp_a, _ in a.options:
if bool(tp_a & direction) == bool(tp_b & other_direction):
keep.append((tp_b, w))
break
return keep
def dbg(lst: WeightedOptions):
def _dbg_opt(opt):
v, w = opt
vals = []
for s, l in ((UP, 'UP'), (RIGHT, 'RIGHT'), (DOWN, 'DOWN'), (LEFT, 'LEFT')):
if v & s:
vals.append(l)
return '(' + ' | '.join(vals) + f', {w})'
return f'[' + ', '.join(_dbg_opt(ow) for ow in lst) + ']'
def output(w: Wave):
boxes: list[list[None | list[str]]] = [[None] * w.width for _ in range(w.height)]
a = '\033[33;43m##\033[0m'
b = '\033[32;42m \033[0m'
for y in range(w.height):
for x in range(w.width):
cell = w.get(x, y)
if cell.collapsed:
opt = cell.options[0][0]
boxes[y][x] = [
b + a + b if opt & UP else b + b + b,
''.join([
a if opt & LEFT else b,
a if opt else b,
a if opt & RIGHT else b,
]),
b + a + b if opt & DOWN else b + b + b,
]
else:
boxes[y][x] = [' ', f'{len(cell.options):^6d}', ' ']
for y in range(w.height):
for i in range(3):
for x in range(w.width):
print(boxes[y][x][i], end='')
print()
def main() -> None:
options = [
(0, 10),
(UP | DOWN, 10),
(LEFT | RIGHT, 10),
(LEFT | DOWN, 10),
(LEFT | UP, 10),
(RIGHT | DOWN, 10),
(RIGHT | UP, 10),
# T-cross roads (3-way)
(UP | RIGHT | DOWN, 5),
(RIGHT | DOWN | LEFT, 5),
(DOWN | LEFT | UP, 5),
(LEFT | UP | RIGHT, 5),
# 4-way intersection
(LEFT | RIGHT | UP | DOWN, 2),
]
w = Wave(25, 15, options)
print('\033[2J')
while True:
print('\033[H')
w.reset()
x1 = w.rng.choice(range(1, w.width - 1))
x2 = w.rng.choice(range(1, w.width - 1))
for x in range(w.width):
c = w.get(x, 0)
c.options = [(0 if x1 != x else UP | DOWN, 1)]
w.collapse_cell(c)
c = w.get(x, w.height - 1)
c.options = [(0 if x2 != x else UP | DOWN, 1)]
w.collapse_cell(c)
for y in range(w.height):
c = w.get(0, y)
c.options = [(0, 1)]
w.collapse_cell(c)
c = w.get(w.width - 1, y)
c.options = [(0, 1)]
w.collapse_cell(c)
w.collapse()
output(w)
time.sleep(1)
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment