The Wavefunction Collapse Algorithm explained very clearly

The Wavefunction Collapse Algorithm explained very clearly
The Wavefunction Collapse Algorithm teaches your computer how to riff. The algorithm takes in an archetypical input, and produces procedurally-generated outputs that look like it.
波函数坍缩算法教会你的电脑如何即兴弹奏。该算法接受一个典型的输入，并产生类似的程序生成的输出。

(Source)

It is most commonly used to create images, but is also capable of building towns, skateparks, and terrible poetry.

be capable of v. 能够
capable 英 /ˈkeɪpəbl/  美 /ˈkeɪpəbl/
adj. 有能力的；有才干的；容许……的；可以做（某事）的；综合性的；有资格的

(Source)

Wavefunction Collapse is a very independent-minded algorithm, and needs almost no outside help or instruction. You feed it an example of the vibe you’re going for, and it figures everything else out for itself. Despite this self-sufficiency, it is surprisingly simple. It doesn’t use any neural networks, random forests, or anything else that sounds like machine learning. This makes it very clean and intuitive once you get the idea.

independent-minded
有主见的
独立思考的

vibe
英 /vaɪb/  美 /vaɪb/
n. 气氛，氛围
vt. <非正式>传递……感觉或氛围
vi. <非正式>享受（音乐等）；和谐相处
n. （Vibe）（挪）维伯（人名）

despite
英 /dɪˈspaɪt/  美 /dɪˈspaɪt/
prep. 即使，尽管
n. （诗/文）侮辱，伤害；轻视，鄙视；憎恨
vt. <废语>故意使烦恼，故意伤害；<古>蔑视，轻蔑

self-sufficiency
英 /ˌself səˈfɪʃnsi/ 美 /ˌself səˈfɪʃnsi/
n. 自给自足；自负

neural
英 /ˈnjʊərəl/  美 /ˈnʊrəl/
adj. 神经的；神经系统的；背的；神经中枢的
n. (Neural)人名；(捷)诺伊拉尔

intuitive
英 /ɪnˈtjuːɪtɪv/  美 /ɪnˈtuːɪtɪv/
adj. 直觉的；凭直觉获知的

Most implementations and explanations of Wavefunction Collapse are of a full, performance-optimized version of the algorithm. These are of course important and necessary, but can be challenging to make sense of from a standing start. This post keeps things clear and simple by focussing on a constrained version of Wavefunction Collapse that I’ve named an Even Simpler Tiled Model. I’ve also put an example implementation of an ESTM on Github. The code is inefficient and slow, but very readable and well commented. Once you fully understand the technology behind an ESTM, explanations of more advanced versions of the algorithm should start to click into place too. If you want to understand the Wavefunction Collapse Algorithm, this is the place to start.

Let’s begin with a story.
Wedding

Imagine that you are planning your wedding. On top of ensuring that you have enough turtle doves and power ballads, you need to design the seating plan for dinner. Your family can be very argumentative and volatile, so this will be difficult. Your dad can’t sit within 2 tables of your mum. Your cousin will get grumpy and lonely if she doesn’t sit with your other cousin. And it’s probably for the best if Uncle Roy doesn’t sit with the environmentalist wing of your partner’s family. With only 5 hours left until the caterers start to arrive, you decide to attack this seemingly intractable problem using the Wavefunction Collapse Algorithm.

You start with a long list of rules and a blank seating plan.

You construct your seating plan’s initial wavefunction. This is a mapping of every seat to a list of the people able to sit in it. For now, any seat can contain any person, including you, your friend who you only ever see at weddings, and your partner’s disconcertingly unattractive ex. Your seating plan’s wavefunction starts in a complete superposition (borrowing from the quantum) of every possible valid layout.

Schrodinger’s Cat was simultaneously both dead and alive until someone opened its box and checked on it; your table plan is simultaneously in every single possible layout until you get your shit together. A complete superposition is a useful theoretical construct, but it won’t help grandma work out where she should be sitting. You need to reduce your seating plan wavefunction to a single, definite state that you can pin onto a notice board and turn into non-quantum namecards.

You begin by collapsing the wavefunction for a single seat. You choose a seat, look at the list of people who are able to sit in it, and randomly assign it to one of them. The seat’s wavefunction is now collapsed.

This choice has consequences that ripple out through the rest of your seating plan’s wavefunction. If Uncle Roy is going to be on Table 2 then cousin Frank and Michele Obama (your partner is a family friend) certainly can’t be. And if Michele isn’t going to be on Table 2 then Barack won’t want to be either. You update the seating plan wavefunction by crossing people off of lists of possible assignees.

Once the ripples have settled you repeat this process. You choose another seat with more than one possible occupant, and collapse its wavefunction by random selecting one of these plausible people to sit in it. Once again, you propagate the ripples of this choice throughout the rest of the seating plan, deleting people from a seat’s wavefunction as they become ineligible to sit in it.

You keep repeating this process until either every seat’s wavefunction is collapsed (meaning that it has exactly 1 person sitting in it), or until you reach a contradiction. A contradiction is a seat in which nobody is able to sit, because they have all been ruled out by your previous choices. A contradiction makes fully collapsing the entire wavefunction impossible.

If you reach a contradiction then the easiest thing to do is to start again. Throw away your work so far, find a new blank table plan, and re-start the algorithm by collapsing the wavefunction for a different random seat. You could also implement a backtracking system that allows you to undo individual choices instead of discarding everything (“well what happens if Shilpa goes in seat 54 instead?”), but the wedding starts in 3 hours, your list of power ballads is just Don’t Wanna Miss A Thing repeated 15 times, and you’ve got no time to be fancy.

After a few false starts you finally reach a fully collapsed state - one in which every seat is assigned to exactly one person and all of the rules are obeyed. You’re done! You rush outside the wedding hall, pin the plan to the noticeboard, and sit back to calmly await the cake, presents, and Aerosmith that await you.
From weddings to bitmaps

This is not a theoretical example. You absolutely could implement a variant of Wavefunction Collapse that produced seating plans for your wedding. It would take in a set of rules and collapse them into a valid dinner layout, unless it ran into so many contradictions trying to deal with your screwball family that it implored you to call the whole thing off.

Of course, in more traditional Wavefunction Collapse we’re not trying to arrange people in a Holiday Inn convention room. We’re trying to arrange pixels in an output image. Nonetheless, the process is remarkably similar. We teach the algorithm a set of rules that its output must obey. We initialize a wavefunction. We collapse one element, and propagate the consequences of this collapse throughout the rest of the wavefunction. And we keep going until the wavefunction is either fully collapsed, or until we reach a contradiction.

Traditional Wavefunction Collapse differs from Wedding Wavefunction Collapse in the way that you teach the algorithm the rules it must obey. In the Wedding version, we had to write down all the rules ourselves. But in the Traditional version we simply give the algorithm an example image, and it figures everything else out from there. It parses the example, analyzes its patterns, and deduces how pixels or tiles are allowed to be arranged.

Let’s start our exploration of real Wavefunction Collapse by considering a simple, special case that ExUtumno (the algorithm’s creator) calls a Simple Tiled Model.
Simple Tiled Model

In a Simple Tiled Model, input and output images are built out of a small number of pre-defined tiles, and each square in the output image is affected and constrained only by its 4 immediate neighbors. For example, suppose we are generating random worlds for a top-down, 2-D game. We might have tiles for land, coast, and sea, and we might have rules like “coast can go next to sea”, “land can go next to coast”, and “sea can go next to other sea”.

A Simple Tiled Model accounts for its tiles’ symmetry and rotation. For example, land can go next to coast, but only in the correct orientation.

This symmetry-handling makes for better output images, but more complex code. To keep things vanilla whilst we’re still learning, let’s consider an even simpler form of Wavefunction Collapse, which I’ll call an Even Simpler Tiled Model.
Even Simpler Tiled Model

An Even Simpler Tiled Model is like a Simple Tiled Model, but its tiles have no symmetry properties. Each tile is a single pixel of a single color, meaning that there is no danger of mismatching their edges.

The rules for an Even Simpler Tiled Model specify which tiles may be placed next to each other, and at which orientation. Each rule consists of a 3-tuple of 2 tiles and a direction. For example (SEA, COAST, LEFT) means that a SEA tile can be placed to the LEFT of a COAST tile. This rule needs to be accompanied by another rule describing the situation from the COAST’s point of view - (COAST, SEA, RIGHT).

If you want SEA tiles to be permitted to the RIGHT as well as the LEFT of COAST tiles then you’ll need additional rules: (SEA, COAST, RIGHT) and (COAST, SEA, LEFT).

As I mentioned earlier, we don’t have to create the list of all of these rules ourselves. Wavefunction Collapse can produce the ruleset for an Even Simpler Tile Model by parsing an example input image and compiling a list of all the 3-tuples that it contains.

By inspecting the above example image, an Even Simpler Tiled Model observes that sea tiles can only go below or to the side of coast tiles, or anywhere next to other sea tiles. It also notes that coast tiles can go to the side of land, sea or other coast tiles, but only above sea tiles and below land ones. It makes no attempt to infer any more complex rules, like “sea tiles must be adjacent to at least one other sea tile” or “every island must contain at least one land tile”. No tile exerts any influence over which types of tile may or may not be placed 2 or more squares away from it. This is like a wedding plan model in which the only type of rule is “X may sit next to Y”.

When using our to analyze the input image, we also need to record the frequency at which each of its tiles appears. We will later use these numbers as weights when deciding which square’s wavefunction to collapse, and when choosing which tile to assign to a square when it is being collapsed.

Once we know the rules that our output image must adhere to, we are ready to build and collapse our output image’s wavefunction.
Collapse

As in our seating plan example, we start the collapsing process with a wavefunction in which every square in our output image is in a superposition of every type of tile.

We start by choosing the square whose wavefunction we will collapse. In our wedding planning example we made this choice randomly. However, as ExUtumno has observed, this isn’t how humans tend to approach these kinds of problems. Instead, they look for the squares with the lowest entropy. Entropy is a measurement of uncertainty and disorder. In general, a square with high entropy is one with lots of possible tiles remaining in its wavefunction. Which tile it will eventually collapse to is still very uncertain. By contrast, a square with low entropy is one with few possible tiles remaining in its wavefunction. Which tile it will eventually collapse to is already very constrained.

For example, in an Even Simpler Tile Model, a square with no information from its surrounding squares is completely unconstrained and is still able to be any tile. It therefore has very high entropy. But a square with several of its surrounding squares already collapsed might only have 2 tiles that it can possibly take on.

Even though the wavefunction of the centre square in the above diagram has not been fully collapsed, we do know that it can’t be a land tile. It has still been somewhat constrained, and therefore has a lower entropy than the square in the top right, which can still be either land, sea, or coast.

It is low-entropy, restricted tiles that humans tend to focus on when working on Wavefunction Collapse-like problems manually. Even if you’re not rigorously using Wavefunction Collapse to design your wedding seating plan, you will still tend to focus on the areas of the plan that already have the most strictures. You don’t put Dwayne on Table 1, Seat 5, then arbitrarily leap over to putting Kathy on Table 7 (which is currently empty). Instead you seat Dwayne, then figure out who can sit next to him, then who can go next to them, and so on. I haven’t seen this written anywhere else, but my intuition says that following this minimal entropy heuristic probably results in fewer contradictions than randomly choosing squares to collapse.

The entropy formula used in Wavefunction Collapse is Shannon Entropy. It makes use of the tile weights that we parsed from the input image in the previous step:

# Sums are over the weights of each remaining
# allowed tile type for the square whose
# entropy we are calculating.
shannon_entropy_for_square =
  log(sum(weight)) -
  (sum(weight * log(weight)) / sum(weight))

Once we’ve found the square in our wavefunction with the lowest entropy, we collapse its wavefunction. We do this by randomly choosing one of the tiles still available to the square, weighted by the tile weights that we parsed from the example input. We use the weights because they give us a more realistic output image. Suppose that a square’s wavefunction says that it can either be land or coast. We don’t necessarily want to chose each option 50% of the time. If our input image contains more land tiles than coast tiles, we will want to reflect this bias in our output images too. We achieve this by using simple, global weights. If our example image contained 20 land tiles and 10 coast ones, we will collapse our square to land 2/3 of the time and coast the other 1/3.

We propagate the consequences of our choice through the rest of our output’s wavefunction (“if that tile is sea then this one can’t be land, which means that this one can’t be coast”). Once all the after-tremors have quietened down we repeat this process, using the minimal entropy heuristic to choose which tile to collapse next. We repeat this collapse-propagate loop until either our output’s wavefunction is completely collapsed and we can return a result, or we reach a contradiction and return an error.

We have made a world (or an error).
Beyond the Even Simpler Tiled Model

Now that you understand the Even Simpler Tiled Model, you’re ready to climb the ladder of power and complexity. Start with the Simple Tiled Model, which we touched on at the start of this post, then move on to the full Overlapping Model. In an Overlapping Model, tiles or pixels influence each other from further away. If you’re into this kind of thing ExUtumno notes that the Simple Tiled Model is analagous to an order-1 Markov chain, and more complex models are like those of higher order.

Wavefunction Collapse can even accomodate additional constraints, like “this tile must be sea”, or “this pixel must be red”, or “there can only be one monster in the output.” This is discussed in the main project’s README. You can also investigate the performance optimizations made in the full implementation. There’s no need to recalculate every square’s entropy on every iteration, and information propagation throughout the wavefunction can be made substantially faster. These considerations become more important the larger your output image becomes.

Wavefunction Collapse is a pleasing and powerful tool to have in your toolbox. Next time you’re planning a wedding or procedurally generating a world, keep it in mind.

import random
import math
import colorama

UP = (0, 1)
LEFT = (-1, 0)
DOWN = (0, -1)
RIGHT = (1, 0)
DIRS = [UP, DOWN, LEFT, RIGHT]


class CompatibilityOracle(object):

    """The CompatibilityOracle class is responsible for telling us
    which combinations of tiles and directions are compatible. It's
    so simple that it perhaps doesn't need to be a class, but I think
    it helps keep things clear.
    """

    def __init__(self, data):
        self.data = data

    def check(self, tile1, tile2, direction):
        return (tile1, tile2, direction) in self.data


class Wavefunction(object):

    """The Wavefunction class is responsible for storing which tiles
    are permitted and forbidden in each location of an output image.
    """

    @staticmethod
    def mk(size, weights):
        """Initialize a new Wavefunction for a grid of `size`,
        where the different tiles have overall weights `weights`.

        Arguments:
        size -- a 2-tuple of (width, height)
        weights -- a dict of tile -> weight of tile
        """
        coefficients = Wavefunction.init_coefficients(size, weights.keys())
        return Wavefunction(coefficients, weights)

    @staticmethod
    def init_coefficients(size, tiles):
        """Initializes a 2-D wavefunction matrix of coefficients.
        The matrix has size `size`, and each element of the matrix
        starts with all tiles as possible. No tile is forbidden yet.

        NOTE: coefficients is a slight misnomer, since they are a
        set of possible tiles instead of a tile -> number/bool dict. This
        makes the code a little simpler. We keep the name `coefficients`
        for consistency with other descriptions of Wavefunction Collapse.

        Arguments:
        size -- a 2-tuple of (width, height)
        tiles -- a set of all the possible tiles

        Returns:
        A 2-D matrix in which each element is a set
        """
        coefficients = []

        for x in range(size[0]):
            row = []
            for y in range(size[1]):
                row.append(set(tiles))
            coefficients.append(row)

        return coefficients

    def __init__(self, coefficients, weights):
        self.coefficients = coefficients
        self.weights = weights

    def get(self, co_ords):
        """Returns the set of possible tiles at `co_ords`"""
        x, y = co_ords
        return self.coefficients[x][y]

    def get_collapsed(self, co_ords):
        """Returns the only remaining possible tile at `co_ords`.
        If there is not exactly 1 remaining possible tile then
        this method raises an exception.
        """
        opts = self.get(co_ords)
        assert(len(opts) == 1)
        return next(iter(opts))

    def get_all_collapsed(self):
        """Returns a 2-D matrix of the only remaining possible
        tiles at each location in the wavefunction. If any location
        does not have exactly 1 remaining possible tile then
        this method raises an exception.
        """
        width = len(self.coefficients)
        height = len(self.coefficients[0])

        collapsed = []
        for x in range(width):
            row = []
            for y in range(height):
                row.append(self.get_collapsed((x,y)))
            collapsed.append(row)

        return collapsed

    def shannon_entropy(self, co_ords):
        """Calculates the Shannon Entropy of the wavefunction at
        `co_ords`.
        """
        x, y = co_ords

        sum_of_weights = 0
        sum_of_weight_log_weights = 0
        for opt in self.coefficients[x][y]:
            weight = self.weights[opt]
            sum_of_weights += weight
            sum_of_weight_log_weights += weight * math.log(weight)

        return math.log(sum_of_weights) - (sum_of_weight_log_weights / sum_of_weights)


    def is_fully_collapsed(self):
        """Returns true if every element in Wavefunction is fully
        collapsed, and false otherwise.
        """
        for x, row in enumerate(self.coefficients):
            for y, sq in enumerate(row):
                if len(sq) > 1:
                    return False

        return True

    def collapse(self, co_ords):
        """Collapses the wavefunction at `co_ords` to a single, definite
        tile. The tile is chosen randomly from the remaining possible tiles
        at `co_ords`, weighted according to the Wavefunction's global
        `weights`.

        This method mutates the Wavefunction, and does not return anything.
        """
        x, y = co_ords
        opts = self.coefficients[x][y]
        valid_weights = {tile: weight for tile, weight in self.weights.items() if tile in opts}

        total_weights = sum(valid_weights.values())
        rnd = random.random() * total_weights

        chosen = None
        for tile, weight in valid_weights.items():
            rnd -= weight
            if rnd < 0:
                chosen = tile
                break

        self.coefficients[x][y] = set(chosen)

    def constrain(self, co_ords, forbidden_tile):
        """Removes `forbidden_tile` from the list of possible tiles
        at `co_ords`.

        This method mutates the Wavefunction, and does not return anything.
        """
        x, y = co_ords
        self.coefficients[x][y].remove(forbidden_tile)


class Model(object):

    """The Model class is responsible for orchestrating the
    Wavefunction Collapse algorithm.
    """

    def __init__(self, output_size, weights, compatibility_oracle):
        self.output_size = output_size
        self.compatibility_oracle = compatibility_oracle

        self.wavefunction = Wavefunction.mk(output_size, weights)

    def run(self):
        """Collapses the Wavefunction until it is fully collapsed,
        then returns a 2-D matrix of the final, collapsed state.
        """
        while not self.wavefunction.is_fully_collapsed():
            self.iterate()

        return self.wavefunction.get_all_collapsed()

    def iterate(self):
        """Performs a single iteration of the Wavefunction Collapse
        Algorithm.
        """
        # 1. Find the co-ordinates of minimum entropy
        co_ords = self.min_entropy_co_ords()
        # 2. Collapse the wavefunction at these co-ordinates
        self.wavefunction.collapse(co_ords)
        # 3. Propagate the consequences of this collapse
        self.propagate(co_ords)

    def propagate(self, co_ords):
        """Propagates the consequences of the wavefunction at `co_ords`
        collapsing. If the wavefunction at (x,y) collapses to a fixed tile,
        then some tiles may not longer be theoretically possible at
        surrounding locations.

        This method keeps propagating the consequences of the consequences,
        and so on until no consequences remain.
        """
        stack = [co_ords]

        while len(stack) > 0:
            cur_coords = stack.pop()
            # Get the set of all possible tiles at the current location
            cur_possible_tiles = self.wavefunction.get(cur_coords)

            # Iterate through each location immediately adjacent to the
            # current location.
            for d in valid_dirs(cur_coords, self.output_size):
                other_coords = (cur_coords[0] + d[0], cur_coords[1] + d[1])

                # Iterate through each possible tile in the adjacent location's
                # wavefunction.
                for other_tile in set(self.wavefunction.get(other_coords)):
                    # Check whether the tile is compatible with any tile in
                    # the current location's wavefunction.
                    other_tile_is_possible = any([
                        self.compatibility_oracle.check(cur_tile, other_tile, d) for cur_tile in cur_possible_tiles
                    ])
                    # If the tile is not compatible with any of the tiles in
                    # the current location's wavefunction then it is impossible
                    # for it to ever get chosen. We therefore remove it from
                    # the other location's wavefunction.
                    if not other_tile_is_possible:
                        self.wavefunction.constrain(other_coords, other_tile)
                        stack.append(other_coords)

    def min_entropy_co_ords(self):
        """Returns the co-ords of the location whose wavefunction has
        the lowest entropy.
        """
        min_entropy = None
        min_entropy_coords = None

        width, height = self.output_size
        for x in range(width):
            for y in range(height):
                if len(self.wavefunction.get((x,y))) == 1:
                    continue

                entropy = self.wavefunction.shannon_entropy((x, y))
                # Add some noise to mix things up a little
                entropy_plus_noise = entropy - (random.random() / 1000)
                if min_entropy is None or entropy_plus_noise < min_entropy:
                    min_entropy = entropy_plus_noise
                    min_entropy_coords = (x, y)

        return min_entropy_coords


def render_colors(matrix, colors):
    """Render the fully collapsed `matrix` using the given `colors.

    Arguments:
    matrix -- 2-D matrix of tiles
    colors -- dict of tile -> `colorama` color
    """
    for row in matrix:
        output_row = []
        for val in row:
            color = colors[val]
            output_row.append(color + val + colorama.Style.RESET_ALL)

        print("".join(output_row))


def valid_dirs(cur_co_ord, matrix_size):
    """Returns the valid directions from `cur_co_ord` in a matrix
    of `matrix_size`. Ensures that we don't try to take step to the
    left when we are already on the left edge of the matrix.
    """
    x, y = cur_co_ord
    width, height = matrix_size
    dirs = []

    if x > 0: dirs.append(LEFT)
    if x < width-1: dirs.append(RIGHT)
    if y > 0: dirs.append(DOWN)
    if y < height-1: dirs.append(UP)

    return dirs


def parse_example_matrix(matrix):
    """Parses an example `matrix`. Extracts:
    
    1. Tile compatibilities - which pairs of tiles can be placed next
        to each other and in which directions
    2. Tile weights - how common different tiles are

    Arguments:
    matrix -- a 2-D matrix of tiles

    Returns:
    A tuple of:
    * A set of compatibile tile combinations, where each combination is of
        the form (tile1, tile2, direction)
    * A dict of weights of the form tile -> weight
    """
    compatibilities = set()
    matrix_width = len(matrix)
    matrix_height = len(matrix[0])

    weights = {}

    for x, row in enumerate(matrix):
        for y, cur_tile in enumerate(row):
            if cur_tile not in weights:
                weights[cur_tile] = 0
            weights[cur_tile] += 1

            for d in valid_dirs((x,y), (matrix_width, matrix_height)):
                other_tile = matrix[x+d[0]][y+d[1]]
                compatibilities.add((cur_tile, other_tile, d))

    return compatibilities, weights


input_matrix = [
    ['L','L','L','L'],
    ['L','L','L','L'],
    ['L','L','L','L'],
    ['L','C','C','L'],
    ['C','S','S','C'],
    ['S','S','S','S'],
    ['S','S','S','S'],
]
input_matrix2 = [
    ['A','A','A','A'],
    ['A','A','A','A'],
    ['A','A','A','A'],
    ['A','C','C','A'],
    ['C','B','B','C'],
    ['C','B','B','C'],
    ['A','C','C','A'],
]

compatibilities, weights = parse_example_matrix(input_matrix)
compatibility_oracle = CompatibilityOracle(compatibilities)
model = Model((10, 50), weights, compatibility_oracle)
output = model.run()

colors = {
    'L': colorama.Fore.GREEN,
    'S': colorama.Fore.BLUE,
    'C': colorama.Fore.YELLOW,
    'A': colorama.Fore.CYAN,
    'B': colorama.Fore.MAGENTA,
}

render_colors(output, colors)