Cayley graph of the free group F2

f2.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["Pillow>=11.3", "numpy"]
# ///

"""
Generate the Cayley graph of the free group F2.
"""

from __future__ import annotations
from typing import Iterator, NamedTuple
from argparse import ArgumentParser

import numpy as np
from PIL import Image


class Direction(NamedTuple):
    dx: int
    dy: int

    def opposite(self) -> Direction:
        return Direction(-self.dx, -self.dy)


DIRECTIONS: tuple[Direction, ...] = (
    Direction(0, -1),
    Direction(1, 0),
    Direction(0, 1),
    Direction(-1, 0),
)


class Position(NamedTuple):
    x: int
    y: int

    def move(self, direction: Direction) -> Position:
        dx, dy = direction
        return Position(self.x + dx, self.y + dy)


class Point(NamedTuple):
    position: Position
    step: int


def recursive(position: Position, direction: Direction, step: int) -> Iterator[Point]:
    if step <= 0:
        return
    step -= 1
    for _ in range(2**step):
        position = position.move(direction)
        yield Point(position, step)
    for next_direction in DIRECTIONS:
        if next_direction != direction.opposite():
            yield from recursive(position, next_direction, step)


def generate_points(center: Position, steps: int) -> Iterator[Point]:
    yield Point(center, steps)
    for direction in DIRECTIONS:
        yield from recursive(center, direction, steps)


def make_array(steps: int):
    width = 2 ** (steps + 1) + 1
    grid = np.zeros((width, width), dtype=np.uint8)
    center = Position(width // 2, width // 2)
    for point in generate_points(center, steps):
        color = int(round(255 * (point.step + 1) / (steps + 1)))
        grid[point.position] = color
    return grid


def main():
    parser = ArgumentParser(
        description="Generate the Cayley graph of the free group F2."
    )
    parser.add_argument(
        "-s",
        "--steps",
        type=int,
        default=9,
        help="Number of steps for the fractal tree (default: 9)",
    )
    args = parser.parse_args()
    array = make_array(args.steps)
    img = Image.fromarray(array)
    img.show()


if __name__ == "__main__":
    main()