Skip to content

Instantly share code, notes, and snippets.

@chrisws
Created February 5, 2019 10:36
Show Gist options
  • Save chrisws/badd5d0c1a53d86e1020c00878bb4738 to your computer and use it in GitHub Desktop.
Save chrisws/badd5d0c1a53d86e1020c00878bb4738 to your computer and use it in GitHub Desktop.
Wavefunction collapse in SmallBASIC
#
# see https://robertheaton.com/2018/12/17/wavefunction-collapse-algorithm/
# https://raw.githubusercontent.com/robert/wavefunction-collapse/master/main.py
#
const kUP = [0, 1]
const kLEFT = [-1, 0]
const kDOWN = [0, -1]
const kRIGHT = [1, 0]
const kDIRS = [kUP, kDOWN, kLEFT, kRIGHT]
const GREEN = rgb(0,128,0)
const BLUE = rgb(0,0, 128)
const YELLOW = rgb(128,0,0)
#
# 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
#
func CompatibilityOracle(_data)
func check(tile1, tile2, direction)
return [tile1, tile2, direction] in self._data
end
local r = {}
r._data = _data
r.check = @check
return r
end
#
# The Wavefunction class is responsible for storing which tiles
# are permitted and forbidden in each location of an output image.
#
func Wavefunction(size, weights)
#
# 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
#
func init_coefficients(size, tiles)
local w, h, row, result, tile_set
[w, h] = size
result = []
for x = 0 to w - 1
row = []
for y = 0 to h - 1
row << tiles
next y
result << row
next x
return result
end
#
# Returns the set of possible tiles at co_ords
#
func get(co_ords)
local x, y
[x, y] = co_ords
return self.coefficients[x][y]
end
#
# Returns the only remaining possible tile at `co_ords`.
# If there is not exactly 1 remaining possible tile then
# this method raises an exception.
#
func get_collapsed(co_ords)
local opts = get(co_ords)
return opts[0]
end
#
# 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.
#
func get_all_collapsed()
local width = len(self.coefficients) - 1
local height = len(self.coefficients[0]) - 1
local collapsed = []
local x, y, row
for x = 0 to width
row = []
for y = 0 to height
row << get_collapsed([x, y])
next y
collapsed << row
next x
return collapsed
end
#
# Calculates the Shannon Entropy of the wavefunction at co_ords
#
func shannon_entropy(co_ords)
local sum_of_weights = 0
local sum_of_weight_log_weights = 0
local x, y, opt
[x, y] = co_ords
for opt in self.coefficients[x][y]
weight = self.weights[opt]
sum_of_weights += weight
sum_of_weight_log_weights += weight * log(weight)
next opt
if (sum_of_weights == 0) then
return 0
else
return log(sum_of_weights) - (sum_of_weight_log_weights / sum_of_weights)
endif
end
#
# Returns true if every element in Wavefunction is fully
# collapsed, and false otherwise.
#
func is_fully_collapsed()
local row, sq
for row in self.coefficients
for sq in row
if len(sq) > 1 then return False
next sq
next row
return True
end
#
# 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.
#
sub collapse(co_ords)
local opts, x, y, valid_weights, total_weights, weight, _rnd, chosen
[x, y] = co_ords
opts = self.coefficients[x][y]
valid_weights = {}
for tile in self.weights
if tile in opts then
weight = self.weights[tile]
valid_weights[tile] = weight
total_weights += weight
endif
next tile
_rnd = rnd * total_weights
chosen = nil
for tile in valid_weights
weight = valid_weights[tile]
_rnd -= weight
if _rnd < 0 then
chosen = tile
exit for
endif
next tile
# assign the chosen tile
self.coefficients[x][y] = [chosen]
end
#
# Removes `forbidden_tile` from the list of possible tiles at `co_ords`.
# This method mutates the Wavefunction, and does not return anything.
#
sub constrain(co_ords, forbidden_tile)
local x, y, i, tile
[x, y] = co_ords
i = 0
for tile in self.coefficients[x][y]
if tile == forbidden_tile then
delete self.coefficients[x][y], i, 1
exit for
endif
i++
next j
end
func keys(m)
local result, k
dim result
for k in m
result << k
next k
return result
end
#
# Initialize a new Wavefunction for a grid of `size`,
# where the different tiles have overall weights `weights`.
#
local r = {}
r.coefficients = init_coefficients(size, keys(weights))
r.weights = weights
r.check = @check
r.get = @get
r.get_collapsed = @get_collapsed
r.get_all_collapsed = @get_all_collapsed
r.shannon_entropy = @shannon_entropy
r.is_fully_collapsed = @is_fully_collapsed
r.collapse = @collapse
r.constrain = @constrain
return r
end
#
# The Model class is responsible for orchestrating the
# Wavefunction Collapse algorithm.
#
func Model(output_size, weights, compatibility_oracle)
#
# Collapses the Wavefunction until it is fully collapsed,
# then returns a 2-D matrix of the final, collapsed state.
#
func execute()
while not self.wavefunction.is_fully_collapsed()
iterate()
wend
return self.wavefunction.get_all_collapsed()
end
#
# Performs a single iteration of the Wavefunction Collapse Algorithm.
#
sub iterate()
# 1. Find the co-ordinates of minimum entropy
local co_ords = min_entropy_co_ords()
# 2. Collapse the wavefunction at these co-ordinates
self.wavefunction.collapse(co_ords)
# 3. Propagate the consequences of this collapse
propagate(co_ords)
end
#
# 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.
#
sub propagate(co_ords)
local stack = [co_ords]
local cur_coords, other_tile_is_possible
while len(stack) > 0
cur_coords = stack_pop(stack)
# 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.
# note: this required a fix in 0.12.15
for other_tile in self.wavefunction.get(other_coords)
# Check whether the tile is compatible with any tile in
# the current location's wavefunction.
other_tile_is_possible = false
for cur_tile in cur_possible_tiles
other_tile_is_possible = self.compatibility_oracle.check(cur_tile, other_tile, d)
if (other_tile_is_possible) then exit for
next cur_tile
# 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 then
self.wavefunction.constrain(other_coords, other_tile)
stack << other_coords
endif
next other_tile
next d
wend
end
#
# Returns the co-ords of the location whose wavefunction has
# the lowest entropy.
#
func min_entropy_co_ords()
local width, height, entropy, x, y
local min_entropy = Nil
local result = Nil
[width, height] = self.output_size
for x = 0 to width - 1
for y = 0 to height - 1
if len(self.wavefunction.get([x, y])) != 1 then
entropy = self.wavefunction.shannon_entropy([x, y])
# Add some noise to mix things up a little
entropy_plus_noise = entropy - (rnd / 1000)
if min_entropy == Nil or entropy_plus_noise < min_entropy then
min_entropy = entropy_plus_noise
result = [x, y]
endif
endif
next y
next x
return result
end
local r = {}
r.output_size = output_size
r.compatibility_oracle = compatibility_oracle
r.wavefunction = Wavefunction(output_size, weights)
r.execute = @execute
return r
end
#
# Pops the next item from the stack (array)
#
func stack_pop(byref items)
local idx = len(items) - 1
local result = items[idx]
delete items, idx, 1
return result
end
#
# Render the fully collapsed `matrix` using the given colors.
# Arguments
# matrix -- 2-D matrix of tiles
# colors -- dict of tile -> `colorama` color
#
sub render_colors(matrix, colors)
local c, col, row, output_row
for row in matrix
output_row = ""
for c in row
col = colors[c]
output_row += "\033[" + col + "m" + c + "\033[0m"
next c
print output_row
next row
end
#
# 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.
#
func valid_dirs(cur_co_ord, matrix_size)
local x, y, width, height, dirs
[x, y] = cur_co_ord
[width, height] = matrix_size
dirs = []
if x > 0 then dirs << kLEFT
if x < width - 1 then dirs << kRIGHT
if y > 0 then dirs << kDOWN
if y < height - 1 then dirs << kUP
return dirs
end
#
# 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
#
func parse_example_matrix(matrix)
local compatibilities = []
local matrix_width = len(matrix)
local matrix_height = len(matrix[0])
local weights = {}
local items = {}
local x = 0
local y = 0
local row, cur_tile, d, item
for row in matrix
for cur_tile in row
if not cur_tile in weights then
weights[cur_tile] = 0
endif
weights[cur_tile] += 1
for d in valid_dirs([x, y], [matrix_width, matrix_height])
other_tile = matrix[x + d[0]][y + d[1]]
item = [cur_tile, other_tile, d]
if items[item] != 1 then
compatibilities << item
items[item] = 1
endif
next d
y++
next cur_tile
x++
y = 0
next row
return [compatibilities, weights]
end
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"]]
[compatibilities, weights] = parse_example_matrix(input_matrix)
compatibility_oracle = CompatibilityOracle(compatibilities)
model = Model([15, 80], weights, compatibility_oracle)
_output = model.execute()
colors = {}
colors["L"] = 32
colors["S"] = 34
colors["C"] = 33
render_colors(_output, colors)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment