forked from kavinbharathii/wave-function-collapse
-
Notifications
You must be signed in to change notification settings - Fork 0
/
grid.py
111 lines (90 loc) · 4.75 KB
/
grid.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# --------------------------------------------------------------------------------- #
# required libraries
import copy
import random
from cell import Cell
# --------------------------------------------------------------------------------- #
# Grid class
class Grid:
def __init__(self, width, height, rez, options):
self.width = width
self.height = height
self.rez = rez
self.w = self.width // self.rez
self.h = self.height // self.rez
self.grid = []
self.options = options
# initiate each spot in the grid with a cell object
def initiate(self):
for i in range(self.w):
self.grid.append([])
for j in range(self.h):
cell = Cell(i, j, self.rez, self.options)
self.grid[i].append(cell)
# draw each cell in the grid
def draw(self, win):
for row in self.grid:
for cell in row:
cell.draw(win)
# randomly pick a cell using [entropy heuristic]
def heuristic_pick(self):
# shallow copy of a grid
grid_copy = [i for row in self.grid for i in row]
grid_copy.sort(key = lambda x:x.entropy())
filtered_grid = list(filter(lambda x:x.entropy() > 1, grid_copy))
if filtered_grid == []:
return None
initial = filtered_grid[0]
filtered_grid = list(filter(lambda x:x.entropy()==initial.entropy(), filtered_grid))
# return a pick if filtered copy os not empty
pick = random.choice(filtered_grid)
return pick
# [WAVE FUNCTION COLLAPSE] algorithm
def collapse(self):
# pick a random cell using entropy heuristic
pick = self.heuristic_pick()
if pick:
self.grid[pick.x][pick.y].options
self.grid[pick.x][pick.y].observe()
else:
return
# shallow copy of the gris
next_grid = copy.copy(self.grid)
# update the entropy values and superpositions of each cell in the grid
for i in range(len(self.grid)):
for j in range(len(self.grid[0])):
if self.grid[i][j].collapsed:
next_grid[i][j] = self.grid[i][j]
else:
# cumulative_valid_options will hold the options that will satisfy the "down", "right", "up", "left"
# conditions of each cell in the grid. The cumulative_valid_options is computed by,
cumulative_valid_options = self.options
# check above cell
cell_above = self.grid[(i - 1) % self.w][j]
valid_options = [] # holds the valid options for the current cell to fit with the above cell
for option in cell_above.options:
valid_options.extend(option.down)
cumulative_valid_options = [option for option in cumulative_valid_options if option in valid_options]
# check right cell
cell_right = self.grid[i][(j + 1) % self.h]
valid_options = [] # holds the valid options for the current cell to fit with the right cell
for option in cell_right.options:
valid_options.extend(option.left)
cumulative_valid_options = [option for option in cumulative_valid_options if option in valid_options]
# check down cell
cell_down = self.grid[(i + 1) % self.w][j]
valid_options = [] # holds the valid options for the current cell to fit with the down cell
for option in cell_down.options:
valid_options.extend(option.up)
cumulative_valid_options = [option for option in cumulative_valid_options if option in valid_options]
# check left cell
cell_left = self.grid[i][(j - 1) % self.h]
valid_options = [] # holds the valid options for the current cell to fit with the left cell
for option in cell_left.options:
valid_options.extend(option.right)
cumulative_valid_options = [option for option in cumulative_valid_options if option in valid_options]
# finally assign the cumulative_valid_options options to be the current cells valid options
next_grid[i][j].options = cumulative_valid_options
next_grid[i][j].update()
# re-assign the grid value after cell evaluation
self.grid = copy.copy(next_grid)