A TypeScript implementation based on Robert Heaton's "Even Simpler Tiled Model" of the Wave Function Collapse algorithm

wfc-simple-tiled.ts

type Tile = number;
type TileOrientation = 0 | 90 | 180 | 270;

export type TileMapping = Record<number, number>;
type Coordinates = [number, number];

enum DirectionKeys {
  UP = 'UP',
  DOWN = 'DOWN',
  LEFT = 'LEFT',
  RIGHT = 'RIGHT',
}

const DIRECTIONS = {
  [DirectionKeys.UP]: [1, 0] as const,
  [DirectionKeys.DOWN]: [-1, 0] as const,
  [DirectionKeys.LEFT]: [0, -1] as const,
  [DirectionKeys.RIGHT]: [0, 1] as const,
};

const UP = DIRECTIONS[DirectionKeys.UP];
const DOWN = DIRECTIONS[DirectionKeys.DOWN];
const LEFT = DIRECTIONS[DirectionKeys.LEFT];
const RIGHT = DIRECTIONS[DirectionKeys.RIGHT];

type Direction = (typeof DIRECTIONS)[keyof typeof DIRECTIONS];

type Compatibility = [Tile, Tile, Direction];
type Weights = Record<Tile, number>;
type Coefficients = Set<Tile>;
type CoefficientMatrix = Coefficients[][];

class CompatibilityOracle {
  constructor(public data: Set<string>) {}

  check(tile1: Tile, tile2: Tile, direction: Direction): boolean {
    const compatibility: Compatibility = [tile1, tile2, direction];
    if (this.data.has(JSON.stringify(compatibility))) {
      return true;
    }

    return false;
  }
}

class Wavefunction {
  constructor(
    public coefficientMatrix: CoefficientMatrix,
    public weights: Weights
  ) {}

  static create(size: [number, number], weights: Weights): Wavefunction {
    const coefficientMatrix = Wavefunction.initCoefficientMatrix(
      size,
      Object.keys(weights).map((key) => parseInt(key))
    );

    return new Wavefunction(coefficientMatrix, weights);
  }

  static initCoefficientMatrix(
    size: [number, number],
    tiles: Tile[]
  ): CoefficientMatrix {
    // initialise the "superposition" of all possible options for each tile in the grid
    const coefficientMatrix: CoefficientMatrix = [];

    for (let i = 0; i < size[1]; i++) {
      const row: Coefficients[] = [];

      for (let j = 0; j < size[0]; j++) {
        row.push(new Set(tiles));
      }

      coefficientMatrix.push(row);
    }

    return coefficientMatrix;
  }

  get(coords: Coordinates): Coefficients {
    const [y, x] = coords;
    return this.coefficientMatrix[y][x];
  }

  getCollapsed(coords: Coordinates): Tile {
    // Return the only possible tile at this coord
    const possibleTiles = this.get(coords);

    if (possibleTiles.size !== 1) {
      throw new Error(`Need just one, actual size ${possibleTiles.size}`);
    }

    return possibleTiles.values().next().value;
  }

  getAllCollapsed(): Tile[][] {
    // returns a matrix of the only remaining tile for each location in the wave function
    const height = this.coefficientMatrix.length;
    const width = this.coefficientMatrix[0].length;

    const collapsed: Tile[][] = [];

    for (let y = 0; y < height; y++) {
      const row: Tile[] = [];

      for (let x = 0; x < width; x++) {
        row.push(this.getCollapsed([y, x]));
      }

      collapsed.push(row);
    }

    return collapsed;
  }

  shannonEntropy(coords: Coordinates): number {
    const [y, x] = coords;

    let sumOfWeights = 0;
    let sumOfWeightLogWeights = 0;

    const possibleTiles = this.coefficientMatrix[y][x];

    for (const tile of possibleTiles) {
      const weight = this.weights[tile];
      sumOfWeights += weight;
      sumOfWeightLogWeights += weight * Math.log(weight);
    }

    return Math.log(sumOfWeights) - sumOfWeightLogWeights / sumOfWeights;
  }

  isFullyCollapsed(): boolean {
    for (const row of this.coefficientMatrix) {
      for (const square of row) {
        if (square.size > 1) {
          return false;
        }
      }
    }

    return true;
  }

  collapse(coords: Coordinates): void {
    // collapse the square at coords to one tile from the remaining options
    const [y, x] = coords;
    const possibleTiles = this.coefficientMatrix[y][x];

    const filteredTilesWithWeights: [Tile, number][] = [];

    for (const tile of possibleTiles) {
      if (this.weights[tile] !== undefined) {
        filteredTilesWithWeights.push([tile, this.weights[tile]]);
      }
    }

    const totalWeights = filteredTilesWithWeights.reduce(
      (sum, [, weight]) => sum + weight,
      0
    );
    let random = Math.random() * totalWeights;

    let chosen = filteredTilesWithWeights[0][0];
    for (const [orientedTile, weight] of filteredTilesWithWeights) {
      random -= weight;

      if (random < 0) {
        chosen = orientedTile;
        break;
      }
    }

    this.coefficientMatrix[y][x] = new Set([chosen]);
  }

  constrain(coords: Coordinates, forbiddenTile: Tile): void {
    const [y, x] = coords;
    this.coefficientMatrix[y][x].delete(forbiddenTile);
  }
}

class Model {
  wavefunction: Wavefunction;

  constructor(
    public outputSize: [number, number],
    public compatibilityOracle: CompatibilityOracle,
    weights: Weights
  ) {
    this.wavefunction = Wavefunction.create(outputSize, weights);
  }

  run(): Tile[][] {
    while (!this.wavefunction.isFullyCollapsed()) {
      this.iterate();
    }

    return this.wavefunction.getAllCollapsed();
  }

  iterate(): void {
    const coords = this.minEntropyCoords();
    this.wavefunction.collapse(coords);
    this.propogate(coords);
  }

  propogate(coords: Coordinates): void {
    const stack: Coordinates[] = [coords];

    while (stack.length > 0) {
      const currentCoords = stack.pop() as Coordinates;
      const currentPossibleTiles = this.wavefunction.get(currentCoords);

      for (const direction of validDirs(currentCoords, this.outputSize)) {
        const otherCoords: Coordinates = [
          currentCoords[0] + direction[0],
          currentCoords[1] + direction[1],
        ];

        for (const otherTile of new Set(this.wavefunction.get(otherCoords))) {
          let otherTileIsPossible = false;

          for (const currentTile of currentPossibleTiles) {
            if (
              this.compatibilityOracle.check(currentTile, otherTile, direction)
            ) {
              otherTileIsPossible = true;
              break;
            }
          }

          if (!otherTileIsPossible) {
            this.wavefunction.constrain(otherCoords, otherTile);
            stack.push(otherCoords);
          }
        }
      }
    }
  }

  minEntropyCoords(): Coordinates {
    let minEntropy: number | null = null;
    let minEntropyCoords: Coordinates = [0, 0];
    const [width, height] = this.outputSize;

    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        if (this.wavefunction.get([y, x]).size === 1) continue;
        const entropy = this.wavefunction.shannonEntropy([y, x]);
        // optional: adds randomness
        const entropyPlusNoise = entropy - Math.random() / 1000;
        if (minEntropy === null || entropyPlusNoise < minEntropy) {
          minEntropy = entropyPlusNoise;
          minEntropyCoords = [y, x];
        }
      }
    }

    return minEntropyCoords;
  }
}

const rotateDirection = (
  direction: Direction,
  rotation: TileOrientation
): Direction => {
  const turns = rotation / 90;

  const rotations: Record<DirectionKeys, Direction[]> = {
    [DirectionKeys.UP]: [
      DIRECTIONS[DirectionKeys.UP],
      DIRECTIONS[DirectionKeys.RIGHT],
      DIRECTIONS[DirectionKeys.DOWN],
      DIRECTIONS[DirectionKeys.LEFT],
    ],
    [DirectionKeys.RIGHT]: [
      DIRECTIONS[DirectionKeys.RIGHT],
      DIRECTIONS[DirectionKeys.DOWN],
      DIRECTIONS[DirectionKeys.LEFT],
      DIRECTIONS[DirectionKeys.UP],
    ],
    [DirectionKeys.DOWN]: [
      DIRECTIONS[DirectionKeys.DOWN],
      DIRECTIONS[DirectionKeys.LEFT],
      DIRECTIONS[DirectionKeys.UP],
      DIRECTIONS[DirectionKeys.RIGHT],
    ],
    [DirectionKeys.LEFT]: [
      DIRECTIONS[DirectionKeys.LEFT],
      DIRECTIONS[DirectionKeys.UP],
      DIRECTIONS[DirectionKeys.RIGHT],
      DIRECTIONS[DirectionKeys.DOWN],
    ],
  };

  const directionKey = Object.keys(DIRECTIONS).find(
    (key) =>
      DIRECTIONS[key as DirectionKeys][0] === direction[0] &&
      DIRECTIONS[key as DirectionKeys][1] === direction[1]
  ) as DirectionKeys;

  return rotations[directionKey][turns];
};

const validDirs = (
  coords: Coordinates,
  matrixSize: [number, number]
): Direction[] => {
  const [y, x] = coords;
  const [width, height] = matrixSize;
  const dirs: Direction[] = [];

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

  return dirs;
};

const parseExampleMatrix = (
  matrix: Tile[][],
  rotationMap: Record<number, number>,
  equivalenceMap: TileMapping
): [Set<string>, Weights] => {
  const compatibilities: Set<string> = new Set();
  const matrixHeight = matrix.length;
  const matrixWidth = matrix[0].length;
  const weights: Weights = {};

  for (let y = 0; y < matrixHeight; y++) {
    for (let x = 0; x < matrixWidth; x++) {
      let currentTile = matrix[y][x];

      // if equivalence exists, use the base tile instead
      currentTile = equivalenceMap[currentTile] ?? currentTile;

      weights[currentTile] = (weights[currentTile] || 0) + 1;

      for (const direction of validDirs([y, x], [matrixWidth, matrixHeight])) {
        const otherTile = matrix[y + direction[0]][x + direction[1]];

        // add main compatibility
        const compatibility: Compatibility = [
          currentTile,
          otherTile,
          direction,
        ];
        compatibilities.add(JSON.stringify(compatibility));

        // if 90 degree rotations exist for both, add that as well
        // rotating the tiles and the compatibility direction
        const currentTileRotation = rotationMap[currentTile];
        const otherTileRotation = rotationMap[otherTile];

        if (currentTileRotation && otherTileRotation) {
          weights[currentTileRotation] =
            (weights[currentTileRotation] || 0) + 1;

          const rotationCompatibility: Compatibility = [
            currentTileRotation,
            otherTileRotation,
            rotateDirection(direction, 90),
          ];

          compatibilities.add(JSON.stringify(rotationCompatibility));
        }
      }
    }
  }

  return [compatibilities, weights];
};

export const generateTerrainWithRotations = (
  inputMatrix: number[][],
  size: [number, number],
  rotationMap: TileMapping = {},
  equivalenceMap: TileMapping = {}
) => {
  const [compatibilities, weights] = parseExampleMatrix(
    inputMatrix,
    rotationMap,
    equivalenceMap
  );
  const compatibilityOracle = new CompatibilityOracle(compatibilities);
  const model = new Model(size, compatibilityOracle, weights);
  const output = model.run();
  return output;
};