🧠 #186 — Wave Function Collapse: Overlapping Model Explained with Code & Formulas

H = log(\sum w_i) - \frac{\sum w_i \cdot log(w_i)}{\sum w_i}

wave[x][y] = [True, False, True, True, ...]

import numpy as np
from PIL import Image
from collections import defaultdict
from random import choice

N = 3  # pattern size
output_width = 20
output_height = 20

def get_patterns(img, N):
    patterns = defaultdict(int)
    w, h = img.shape
    for y in range(h - N + 1):
        for x in range(w - N + 1):
            patch = tuple(img[y:y+N, x:x+N].flatten())
            patterns[patch] += 1
    return list(patterns.items())

def pattern_match(a, b, direction):
    N = int(len(a)**0.5)
    a = np.array(a).reshape(N, N)
    b = np.array(b).reshape(N, N)
    if direction == 'right':
        return np.array_equal(a[:, 1:], b[:, :-1])
    if direction == 'down':
        return np.array_equal(a[1:, :], b[:-1, :])
    return False

def get_compatibility(patterns):
    compat = {i: {'right': set(), 'down': set()} for i in range(len(patterns))}
    for i, (a, _) in enumerate(patterns):
        for j, (b, _) in enumerate(patterns):
            if pattern_match(a, b, 'right'):
                compat[i]['right'].add(j)
            if pattern_match(a, b, 'down'):
                compat[i]['down'].add(j)
    return compat

def collapse(wave, entropies):
    # Pick the cell with the lowest entropy
    xs, ys = np.where(entropies == np.min(entropies[entropies > 0]))
    x, y = choice(list(zip(xs, ys)))

    options = [i for i, valid in enumerate(wave[x][y]) if valid]
    chosen = choice(options)
    wave[x][y] = [i == chosen for i in range(len(wave[x][y]))]
    return x, y

def propagate(wave, compat, x, y, W, H):
    changed = [(x, y)]
    while changed:
        cx, cy = changed.pop()
        for dx, dy, dir in [(-1, 0, 'right'), (1, 0, 'right'), (0, -1, 'down'), (0, 1, 'down')]:
            nx, ny = cx + dx, cy + dy
            if 0 <= nx < W and 0 <= ny < H:
                for p in range(len(wave[nx][ny])):
                    if wave[nx][ny][p]:
                        valid = False
                        for q in range(len(wave[cx][cy])):
                            if wave[cx][cy][q] and p in compat[q][dir]:
                                valid = True
                                break
                        if not valid:
                            wave[nx][ny][p] = False
                            changed.append((nx, ny))

def run_wfc(patterns, compat, W, H):
    wave = [[[True]*len(patterns) for _ in range(H)] for _ in range(W)]
    entropies = np.full((W, H), len(patterns), dtype=float)

    for _ in range(W * H):
        x, y = collapse(wave, entropies)
        propagate(wave, compat, x, y, W, H)
        # Update entropy
        for i in range(W):
            for j in range(H):
                possibilities = [p for p in wave[i][j] if p]
                if possibilities:
                    entropies[i][j] = len(possibilities)
                else:
                    entropies[i][j] = 0
    return wave

def visualize(wave, patterns, N, W, H):
    tile_size = N
    output = np.zeros((H + N - 1, W + N - 1), dtype=np.uint8)
    for x in range(W):
        for y in range(H):
            idx = wave[x][y].index(True)
            pattern = np.array(patterns[idx][0]).reshape(N, N)
            output[y:y+N, x:x+N] = pattern
    return output

# Example: Use a grayscale image
input_img = Image.open("input.png").convert("L")
img_array = np.array(input_img)

patterns = get_patterns(img_array, N)
compat = get_compatibility(patterns)
wave = run_wfc(patterns, compat, output_width, output_height)
result = visualize(wave, patterns, N, output_width, output_height)

Image.fromarray(result).save("output.png")

w_p = \frac{f(p)}{\sum_{q} f(q)}