Wavefunction collapse in SmallBASIC

wavefunction_collapse.bas

#
# 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)