├── .gitignore ├── LICENSE ├── README.md ├── example.png ├── main.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | vendor/* 2 | *.swp 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | This is free and unencumbered software released into the public domain. 2 | 3 | Anyone is free to copy, modify, publish, use, compile, sell, or 4 | distribute this software, either in source code form or as a compiled 5 | binary, for any purpose, commercial or non-commercial, and by any 6 | means. 7 | 8 | In jurisdictions that recognize copyright laws, the author or authors 9 | of this software dedicate any and all copyright interest in the 10 | software to the public domain. We make this dedication for the benefit 11 | of the public at large and to the detriment of our heirs and 12 | successors. We intend this dedication to be an overt act of 13 | relinquishment in perpetuity of all present and future rights to this 14 | software under copyright law. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 19 | IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | For more information, please refer to 25 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Wavefunction Collapse 2 | 3 | This is the example code for my blog post, ["The Wavefunction Collapse Algorithm explained very clearly"](https://robertheaton.com/2018/12/17/wavefunction-collapse-algorithm/). To install dependencies: 4 | 5 | ``` 6 | virtualenv vendor 7 | source vendor/bin/activate 8 | pip install -r requirements.txt 9 | ``` 10 | 11 | Then to run: 12 | 13 | ``` 14 | python main.py 15 | ``` 16 | 17 | You will get a nice picture of a coastline. 18 | 19 | ![WFC Example](example.png) 20 | -------------------------------------------------------------------------------- /example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/robert/wavefunction-collapse/21b2e5d95ec7db6057382bfb61ba4557cdd436f4/example.png -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import random 2 | import math 3 | from typing import Literal 4 | import colorama 5 | 6 | # Types 7 | Tile = str 8 | Coordinates = tuple[int, int] 9 | Up = tuple[Literal[1], Literal[0]]; Down = tuple[Literal[-1], Literal[0]] 10 | Left = tuple[Literal[0], Literal[-1]]; Right = tuple[Literal[0], Literal[1]] 11 | Direction = Up | Down | Left | Right 12 | Compatibility = tuple[Tile, Tile, Direction] 13 | Weights = dict[Tile, int] 14 | Coefficients = set[Tile] 15 | CoefficientMatrix = list[list[Coefficients]] 16 | 17 | UP = (1, 0) 18 | DOWN = (-1, 0) 19 | LEFT = (0, -1) 20 | RIGHT = (0, 1) 21 | DIRS = [UP, DOWN, LEFT, RIGHT] 22 | 23 | 24 | class CompatibilityOracle(object): 25 | 26 | """The CompatibilityOracle class is responsible for telling us 27 | which combinations of tiles and directions are compatible. It's 28 | so simple that it perhaps doesn't need to be a class, but I think 29 | it helps keep things clear. 30 | """ 31 | 32 | def __init__(self, data: set[Compatibility]): 33 | self.data = data 34 | 35 | def check(self, tile1: Tile, tile2: Tile, direction: Direction) -> bool: 36 | return (tile1, tile2, direction) in self.data 37 | 38 | 39 | class Wavefunction(object): 40 | 41 | """The Wavefunction class is responsible for storing which tiles 42 | are permitted and forbidden in each location of an output image. 43 | """ 44 | 45 | @staticmethod 46 | def mk(size: tuple[int, int], weights: Weights): 47 | """Initialize a new Wavefunction for a grid of `size`, 48 | where the different tiles have overall weights `weights`. 49 | 50 | Arguments: 51 | size -- a 2-tuple of (width, height) 52 | weights -- a dict of tile -> weight of tile 53 | """ 54 | coefficient_matrix = Wavefunction.init_coefficient_matrix(size, list(weights.keys())) 55 | return Wavefunction(coefficient_matrix, weights) 56 | 57 | @staticmethod 58 | def init_coefficient_matrix(size: tuple[int, int], tiles: list[Tile]) -> CoefficientMatrix: 59 | """Initializes a 2-D wavefunction matrix of coefficients. 60 | The matrix has size `size`, and each element of the matrix 61 | starts with all tiles as possible. No tile is forbidden yet. 62 | 63 | NOTE: coefficients is a slight misnomer, since they are a 64 | set of possible tiles instead of a tile -> number/bool dict. This 65 | makes the code a little simpler. We keep the name `coefficients` 66 | for consistency with other descriptions of Wavefunction Collapse. 67 | 68 | Arguments: 69 | size -- a 2-tuple of (width, height) 70 | tiles -- a set of all the possible tiles 71 | 72 | Returns: 73 | A 2-D matrix in which each element is a set 74 | """ 75 | coefficient_matrix: CoefficientMatrix = [] 76 | 77 | for _ in range(size[1]): 78 | row: list[Coefficients] = [] 79 | for _ in range(size[0]): 80 | row.append(set(tiles)) 81 | coefficient_matrix.append(row) 82 | 83 | return coefficient_matrix 84 | 85 | def __init__(self, coefficient_matrix: CoefficientMatrix, weights: Weights): 86 | self.coefficient_matrix = coefficient_matrix 87 | self.weights = weights 88 | 89 | def get(self, co_ords: Coordinates) -> Coefficients: 90 | """Fetches the set of possible tiles at `co_ords`. 91 | 92 | Arguments: 93 | co_ords -- Tuple representing 2D co-ordinates in the format (y, x). 94 | 95 | Returns: 96 | The set of possible tiles. 97 | """ 98 | y, x = co_ords 99 | return self.coefficient_matrix[y][x] 100 | 101 | def get_collapsed(self, co_ords: Coordinates) -> Tile: 102 | """Returns the only remaining possible tile at `co_ords`. 103 | If there is not exactly 1 remaining possible tile then 104 | this method raises an exception. 105 | """ 106 | opts = self.get(co_ords) 107 | assert(len(opts) == 1) 108 | return next(iter(opts)) 109 | 110 | def get_all_collapsed(self) -> list[list[Tile]]: 111 | """Returns a 2-D matrix of the only remaining possible 112 | tiles at each location in the wavefunction. If any location 113 | does not have exactly 1 remaining possible tile then 114 | this method raises an exception. 115 | """ 116 | height = len(self.coefficient_matrix) 117 | width = len(self.coefficient_matrix[0]) 118 | 119 | collapsed: list[list[Tile]] = [] 120 | for y in range(height): 121 | row: list[Tile] = [] 122 | for x in range(width): 123 | row.append(self.get_collapsed((y, x))) 124 | collapsed.append(row) 125 | 126 | return collapsed 127 | 128 | def shannon_entropy(self, co_ords: Coordinates) -> float: 129 | """Calculates the Shannon Entropy of the wavefunction at 130 | `co_ords`. 131 | """ 132 | y, x = co_ords 133 | 134 | sum_of_weights = 0 135 | sum_of_weight_log_weights = 0 136 | for opt in self.coefficient_matrix[y][x]: 137 | weight = self.weights[opt] 138 | sum_of_weights += weight 139 | sum_of_weight_log_weights += weight * math.log(weight) 140 | 141 | return math.log(sum_of_weights) - (sum_of_weight_log_weights / sum_of_weights) 142 | 143 | 144 | def is_fully_collapsed(self) -> bool: 145 | """Returns true if every element in Wavefunction is fully 146 | collapsed, and false otherwise. 147 | """ 148 | for row in self.coefficient_matrix: 149 | for sq in row: 150 | if len(sq) > 1: 151 | return False 152 | 153 | return True 154 | 155 | def collapse(self, co_ords: Coordinates) -> None: 156 | """Collapses the wavefunction at `co_ords` to a single, definite 157 | tile. The tile is chosen randomly from the remaining possible tiles 158 | at `co_ords`, weighted according to the Wavefunction's global 159 | `weights`. 160 | 161 | This method mutates the Wavefunction, and does not return anything. 162 | """ 163 | y, x = co_ords 164 | opts = self.coefficient_matrix[y][x] 165 | filtered_tiles_with_weights = [(tile, weight) for tile, weight in self.weights.items() if tile in opts] 166 | 167 | total_weights = sum([weight for _, weight in filtered_tiles_with_weights]) 168 | rnd = random.random() * total_weights 169 | 170 | chosen = filtered_tiles_with_weights[0][0] 171 | for tile, weight in filtered_tiles_with_weights: 172 | rnd -= weight 173 | if rnd < 0: 174 | chosen = tile 175 | break 176 | 177 | self.coefficient_matrix[y][x] = set([chosen]) 178 | 179 | def constrain(self, co_ords: Coordinates, forbidden_tile: Tile) -> None: 180 | """Removes `forbidden_tile` from the list of possible tiles 181 | at `co_ords`. 182 | 183 | This method mutates the Wavefunction, and does not return anything. 184 | """ 185 | y, x = co_ords 186 | self.coefficient_matrix[y][x].remove(forbidden_tile) 187 | 188 | 189 | class Model(object): 190 | 191 | """The Model class is responsible for orchestrating the 192 | Wavefunction Collapse algorithm. 193 | """ 194 | 195 | def __init__(self, output_size: tuple[int, int], weights: Weights, compatibility_oracle: CompatibilityOracle): 196 | self.output_size = output_size 197 | self.compatibility_oracle = compatibility_oracle 198 | 199 | self.wavefunction = Wavefunction.mk(output_size, weights) 200 | 201 | def run(self) -> list[list[Tile]]: 202 | """Collapses the Wavefunction until it is fully collapsed, 203 | then returns a 2-D matrix of the final, collapsed state. 204 | """ 205 | while not self.wavefunction.is_fully_collapsed(): 206 | self.iterate() 207 | 208 | return self.wavefunction.get_all_collapsed() 209 | 210 | def iterate(self) -> None: 211 | """Performs a single iteration of the Wavefunction Collapse 212 | Algorithm. 213 | """ 214 | # 1. Find the co-ordinates of minimum entropy 215 | co_ords = self.min_entropy_co_ords() 216 | # 2. Collapse the wavefunction at these co-ordinates 217 | self.wavefunction.collapse(co_ords) 218 | # 3. Propagate the consequences of this collapse 219 | self.propagate(co_ords) 220 | 221 | def propagate(self, co_ords: Coordinates) -> None: 222 | """Propagates the consequences of the wavefunction at `co_ords` 223 | collapsing. If the wavefunction at (y, x) collapses to a fixed tile, 224 | then some tiles may not longer be theoretically possible at 225 | surrounding locations. 226 | 227 | This method keeps propagating the consequences of the consequences, 228 | and so on until no consequences remain. 229 | """ 230 | stack = [co_ords] 231 | 232 | while len(stack) > 0: 233 | cur_co_ords = stack.pop() 234 | # Get the set of all possible tiles at the current location 235 | cur_possible_tiles = self.wavefunction.get(cur_co_ords) 236 | 237 | # Iterate through each location immediately adjacent to the 238 | # current location. 239 | for d in valid_dirs(cur_co_ords, self.output_size): 240 | other_co_ords = (cur_co_ords[0] + d[0], cur_co_ords[1] + d[1]) 241 | 242 | # Iterate through each possible tile in the adjacent location's 243 | # wavefunction. 244 | for other_tile in set(self.wavefunction.get(other_co_ords)): 245 | # Check whether the tile is compatible with any tile in 246 | # the current location's wavefunction. 247 | other_tile_is_possible = any([ 248 | self.compatibility_oracle.check(cur_tile, other_tile, d) for cur_tile in cur_possible_tiles 249 | ]) 250 | # If the tile is not compatible with any of the tiles in 251 | # the current location's wavefunction then it is impossible 252 | # for it to ever get chosen. We therefore remove it from 253 | # the other location's wavefunction. 254 | if not other_tile_is_possible: 255 | self.wavefunction.constrain(other_co_ords, other_tile) 256 | stack.append(other_co_ords) 257 | 258 | def min_entropy_co_ords(self) -> Coordinates: 259 | """Returns the co-ords of the location whose wavefunction has 260 | the lowest entropy. 261 | """ 262 | min_entropy = None 263 | min_entropy_co_ords: Coordinates = (0, 0) 264 | 265 | width, height = self.output_size 266 | for y in range(height): 267 | for x in range(width): 268 | if len(self.wavefunction.get((y, x))) == 1: 269 | continue 270 | 271 | entropy = self.wavefunction.shannon_entropy((y, x)) 272 | # Add some noise to mix things up a little 273 | entropy_plus_noise = entropy - (random.random() / 1000) 274 | if min_entropy is None or entropy_plus_noise < min_entropy: 275 | min_entropy = entropy_plus_noise 276 | min_entropy_co_ords = (y, x) 277 | 278 | return min_entropy_co_ords 279 | 280 | 281 | def render_colors(matrix: list[list[Tile]], colors: dict[str, str]) -> None: 282 | """Render the fully collapsed `matrix` using the given `colors. 283 | 284 | Arguments: 285 | matrix -- 2-D matrix of tiles 286 | colors -- dict of tile -> `colorama` color 287 | """ 288 | for row in matrix: 289 | output_row: list[str] = [] 290 | for val in row: 291 | color = colors[val] 292 | output_row.append(color + val + colorama.Style.RESET_ALL) 293 | 294 | print("".join(output_row)) 295 | 296 | 297 | def valid_dirs(cur_co_ords: Coordinates, matrix_size: tuple[int, int]) -> list[Direction]: 298 | """Returns the valid directions from `cur_co_ord` in a matrix 299 | of `matrix_size`. Ensures that we don't try to take step to the 300 | left when we are already on the left edge of the matrix. 301 | """ 302 | y, x = cur_co_ords 303 | width, height = matrix_size 304 | dirs: list[Direction] = [] 305 | 306 | if y < height-1: dirs.append(UP) 307 | if y > 0: dirs.append(DOWN) 308 | if x > 0: dirs.append(LEFT) 309 | if x < width-1: dirs.append(RIGHT) 310 | 311 | return dirs 312 | 313 | 314 | def parse_example_matrix(matrix: list[list[Tile]]) -> tuple[set[Compatibility], Weights]: 315 | """Parses an example `matrix`. Extracts: 316 | 317 | 1. Tile compatibilities - which pairs of tiles can be placed next 318 | to each other and in which directions 319 | 2. Tile weights - how common different tiles are 320 | 321 | Arguments: 322 | matrix -- a 2-D matrix of tiles 323 | 324 | Returns: 325 | A tuple of: 326 | * A set of compatibile tile combinations, where each combination is of 327 | the form (tile1, tile2, direction) 328 | * A dict of weights of the form tile -> weight 329 | """ 330 | compatibilities: set[Compatibility] = set() 331 | matrix_height = len(matrix) 332 | matrix_width = len(matrix[0]) 333 | 334 | weights: Weights = {} 335 | 336 | for y, row in enumerate(matrix): 337 | for x, cur_tile in enumerate(row): 338 | if cur_tile not in weights: 339 | weights[cur_tile] = 0 340 | weights[cur_tile] += 1 341 | 342 | for d in valid_dirs((y, x), (matrix_width, matrix_height)): 343 | other_tile = matrix[y+d[0]][x+d[1]] 344 | compatibilities.add((cur_tile, other_tile, d)) 345 | 346 | return compatibilities, weights 347 | 348 | 349 | input_matrix = [ 350 | ['L','L','L','L'], 351 | ['L','L','L','L'], 352 | ['L','L','L','L'], 353 | ['L','C','C','L'], 354 | ['C','S','S','C'], 355 | ['S','S','S','S'], 356 | ['S','S','S','S'], 357 | ] 358 | input_matrix2 = [ 359 | ['A','A','A','A'], 360 | ['A','A','A','A'], 361 | ['A','A','A','A'], 362 | ['A','C','C','A'], 363 | ['C','B','B','C'], 364 | ['C','B','B','C'], 365 | ['A','C','C','A'], 366 | ] 367 | 368 | compatibilities, weights = parse_example_matrix(input_matrix) 369 | compatibility_oracle = CompatibilityOracle(compatibilities) 370 | model = Model((50, 10), weights, compatibility_oracle) 371 | output = model.run() 372 | 373 | colors = { 374 | 'L': colorama.Fore.GREEN, 375 | 'S': colorama.Fore.BLUE, 376 | 'C': colorama.Fore.YELLOW, 377 | 'A': colorama.Fore.CYAN, 378 | 'B': colorama.Fore.MAGENTA, 379 | } 380 | 381 | render_colors(output, colors) 382 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | colorama 2 | --------------------------------------------------------------------------------