ulam_spiral.py

from itertools import count, cycle, repeat, islice
from PIL import Image, ImageDraw

class ArrayCursor:
    position = 0, 0

    def move_north(self):
        self.position = self.position[0], self.position[1] + 1

    def move_south(self):
        self.position = self.position[0], self.position[1] - 1

    def move_east(self):
        self.position = self.position[0] + 1, self.position[1]

    def move_west(self):
        self.position = self.position[0] - 1, self.position[1]

    def move(self, dir: str):
        {
            'n': self.move_north,
            's': self.move_south,
            'e': self.move_east,
            'w': self.move_west,
        }[dir]()


def spiral_moves():
    steps = count(start=1)
    step_count = 0

    for i, d in enumerate(cycle("enws")):
        if (i + 1) % 2:
            step_count = next(steps)

        for _ in range(step_count):
            yield d


def double_up(iter):
    for step_count in iter:
        yield step_count
        yield step_count


def square_spiral_step_directions(directions="enws"):
    for step_count, dir in zip(double_up(count(start=1)), cycle(directions)):
        for x in repeat(dir, step_count):
            yield x


def spiral_positions(directions):
    cursor = ArrayCursor()
    yield cursor.position
    for x in square_spiral_step_directions(directions):
        cursor.move(x)
        yield cursor.position


def is_prime(num):
    for n in range(2, int(num ** 1 / 2) + 1):
        if num % n == 0:
            return False
    return True


def is_prime_fast(n):
    def is_prime_fast_inner(n):
        if n & 1 == 0:
            return 2
        d = 3
        while d * d <= n:
            if n % d == 0:
                return d
            d = d + 2
        return 0
    return not is_prime_fast_inner(n)

def filtered_positions(filter):
    for number, position in enumerate(spiral_positions("enws"), start=1):
        if filter(number):
            yield number, position


def get_array_bounds(positions):
    min_x, min_y, max_x, max_y = 0, 0, 0, 0
    for x, y in positions:
        min_x = min(min_x, x)
        max_x = max(max_x, x)
        min_y = min(min_y, y)
        max_y = max(max_y, y)
    return min_x, min_y, max_x, max_y


def draw_points(points, size):
    image = Image.new("RGB", size, "white")
    draw = ImageDraw.Draw(image)
    for point in points:
        draw.point(point, fill="black")
    image.save("test.png")


if __name__ == '__main__':
    n = 500000
    top_n = list(islice(filtered_positions(is_prime_fast), n))
    min_x, min_y, max_x, max_y = get_array_bounds((position for _, position in top_n))
    bounds = max_x - min_x, max_y - min_y
    offset = min_x * -1, min_y * -1
    points = ((position[0] + offset[0], position[1] + offset[1]) for number, position in top_n)
    draw_points(points, bounds)