ulams-spiral.py

#!/usr/bin/python
# -*- coding: utf8 -*-

# http://mathworld.wolfram.com/PrimeSpiral.html

import math
import gmpy
import numpy
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import matplotlib.cm as cm

RIGHT = 1
UP = 2
LEFT = 3
DOWN = 4


class UlamSpiral:

    def __init__(self, size):
        self.size = size
        self.matrix = numpy.zeros((size, size), dtype=bool)

    def step(self, x, y, step):

        if step == RIGHT:
            y += 1
        if step == UP:
            x -= 1
        if step == LEFT:
            y -= 1
        if step == DOWN:
            x += 1

        return (x, y)

    def next_step(self, step):

        if step == RIGHT:
            return UP
        if step == UP:
            return LEFT
        if step == LEFT:
            return DOWN
        if step == DOWN:
            return RIGHT

    def is_prime(self, n):
        if not gmpy.is_prime(n):  # Rabin-Miller test
            return False
        i = 2
        while i <= math.sqrt(n):
            if n % i == 0:
                return False
            i += 1
        return True

    def generate_spiral(self):

        x = y = self.size / 2

        n = 1
        self.matrix[x][y] = False
        step = RIGHT
        n += 1

        for i in xrange(1, self.size):

            for j in xrange(i):
                (x, y) = self.step(x, y, step)
                self.matrix[x][y] = self.is_prime(n)
                n += 1

            step = self.next_step(step)

            for j in xrange(i):
                (x, y) = self.step(x, y, step)
                self.matrix[x][y] = self.is_prime(n)
                n += 1

            step = self.next_step(step)

        for j in xrange(i):
            (x, y) = self.step(x, y, step)
            self.matrix[x][y] = self.is_prime(n)
            n += 1

        plt.imshow(self.matrix.astype('uint8'), cmap=cm.Greys)
        plt.show()


if __name__ == "__main__":
    s = UlamSpiral(501)
    s.generate_spiral()