minimalistic dfs search for cellular automata

search.py

from dataclasses import dataclass
from abc import ABC, abstractmethod

"""Perform a depth-first search for patterns in a cellular automaton.

Some jargon specific for this file
State: Possible values for a `Variable`.
Variable: A place where one can assign a `State` to.
Address: An (x, y, phase) tuple that uniquely identifies a `Variable`.
"""


def conwaylife(
        nw, nc, ne,
        cw, cc, ce,
        sw, sc, se):
    """
    >>> conwaylife(1, 0, 0, 0, 0, 0, 0, 0, 0)
    0
    >>> conwaylife(1, 0, 0, 0, 1, 0, 0, 0, 0)
    0

    >>> conwaylife(1, 1, 0, 0, 0, 0, 0, 0, 0)
    0
    >>> conwaylife(1, 1, 0, 0, 1, 0, 0, 0, 0)
    1

    >>> conwaylife(1, 1, 1, 0, 0, 0, 0, 0, 0)
    1
    >>> conwaylife(1, 1, 1, 0, 1, 0, 0, 0, 0)
    1

    >>> conwaylife(1, 1, 1, 1, 0, 0, 0, 0, 0)
    0
    >>> conwaylife(1, 1, 1, 1, 1, 0, 0, 0, 0)
    0
    """
    neighbor_count = nw + nc + ne + cw + ce + sw + sc + se
    satisfies_b3 = not cc and neighbor_count == 3
    satisfies_s2 = cc and neighbor_count == 2
    satisfies_s3 = cc and neighbor_count == 3
    return int(satisfies_b3 or satisfies_s2 or satisfies_s3)


class State:
    unknown = -1
    off = 0
    on = 1


@dataclass
class Address:
    x: int
    y: int
    phase: int


class Variable(ABC):
    @abstractmethod
    def try_assign(self, value) -> bool:
        """Try to assign value to variable and return whether the assignment was successful."""
        pass

    @abstractmethod
    def unassign(self):
        pass


class MutableVariable(Variable):
    def __init__(self):
        self.value = State.unknown

    def try_assign(self, value) -> bool:
        self.value = value
        return True

    def unassign(self):
        self.value = State.unknown


class ImmutableVariable(Variable):
    def __init__(self, value):
        self.value = value

    def try_assign(self, value):
        return value == self.value

    def unassign(self):
        pass


class Search:

    def __init__(self, /, width, height, period, rule=conwaylife, padding=2):
        self.width = width
        self.height = height
        self.padding = padding

        self.period = period
        self.rule = rule
        self.search_order = [
            Address(x=x, y=y, phase=phase)
            for y in range(height + padding)
            for x in range(width + padding)
            for phase in range(period)]
        self.variables = []

        # Create a grid of `off` cells
        for phase in range(self.period):
            plane = []
            for y in range(self.get_board_height()):
                row = []
                for x in range(self.get_board_width()):
                    row.append(ImmutableVariable(State.off))
                plane.append(row)
            self.variables.append(plane)

        # leave the padding `off`,
        # but replace the middle cells with `unknown` variables
        for phase in range(self.period):
            for y in range(self.height):
                for x in range(self.width):
                    self.variables[phase][y + padding][x + padding] = (
                        MutableVariable())

    def get_board_width(self):
        return self.width + self.padding * 2

    def get_board_height(self):
        return self.height + self.padding * 2

    @staticmethod
    def char_for_state(state):
        if state == State.unknown:
            return "?"
        if state == State.off:
            return "."
        if state == State.on:
            return "*"
        raise ValueError("Invalid state!")

    def print_search_state(self):
        chars = []
        for phase in range(self.period):
            for y in range(self.get_board_height()):
                for x in range(self.get_board_width()):
                    c = Search.char_for_state(
                        self.variables[phase][y][x].value)
                    chars.append(c)
                chars.append("\n")
            chars.append("\n")
        search_state_string = "".join(chars)
        print(search_state_string)

    def print_solution(self):
        chars = []
        phase = 0
        for y in range(self.get_board_height()):
            for x in range(self.get_board_width()):
                c = Search.char_for_state(self.variables[phase][y][x].value)
                chars.append(c)
            chars.append("\n")
        chars.append("\n")
        solution_string = "".join(chars)
        print(solution_string)

    def get_variable_at(self, address):
        return self.variables[address.phase][address.y + self.padding][address.x + self.padding]

    def get_state_at(self, address):
        return self.get_variable_at(address).value

    def get_cell_translated_address(self, address, dx, dy):
        return Address(address.x + dx, address.y + dy, address.phase)

    def next_phase(self, address):
        return (
            Address(address.x, address.y, 0)
            if address.phase == self.period - 1
            else Address(address.x, address.y, address.phase + 1))

    def is_assignment_consistent(self, address):
        """Check if cell at coordinate `(x, y)` is consistent at phase `phase`.

        This assumes that neighbors' states at `phase` are not `unknown`
        and the cell's state is also determined at phase `phase + 1`.
        """
        [
            nw, nc, ne,
            cw, cc, ce,
            sw, sc, se,
        ] = [
            self.get_state_at(
                self.get_cell_translated_address(address, dx, dy))
            for dy in [-1, 0, 1]
            for dx in [-1, 0, 1]
        ]
        cc_next = self.get_state_at(self.next_phase(address))
        if any(state == State.unknown for state in (nw, nc, ne, cw, cc, ce, sw, sc, se, cc_next)):
            raise ValueError("Expected all known cells, got unknown cell")
        return self.rule(nw, nc, ne, cw, cc, ce, sw, sc, se) == cc_next

    def start_search_at_ith_address(self, i: int):
        try:
            address = self.search_order[i]
        except IndexError:
            self.print_solution()
            return
        for state in [State.off, State.on]:
            if not self.get_variable_at(address).try_assign(state):
                continue
            # When assigning to a cell at phase,
            # check consistency at the northwest cell at that phase.
            if not self.is_assignment_consistent(self.get_cell_translated_address(address, -1, -1)):
                self.get_variable_at(address).unassign()
                continue
            self.start_search_at_ith_address(i + 1)

    def search(self):
        self.start_search_at_ith_address(0)


if __name__ == "__main__":
    s = Search(width=3, height=3, period=2)
    s.search()