Creates an Ulam Spiral as a square matrix of size N.

prime_spiral.py

#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt

def create_number_spiral(N):
    '''Creates a number spiral as
    a square matrix of size N'''

    X = np.ndarray((N,N))
    ctr = 1
    num = 1

    if (N%2 == 0):
        coord = np.array([(N)/2, (N)/2 - 1], dtype=np.int64)
    else:
        coord = np.array([(N)/2, (N)/2], dtype=np.int64)

    vec = np.array([0,1])
    T = np.array([[0, -1], [1, 0]])

    while ctr < N:
        #Do logic here
        k = ctr
        while k > 0:
            X[coord[0]][coord[1]] = num
            #print "Index: ",coord
            #print "Vector: ", vec
            coord += vec
            num += 1
            k -= 1

        vec = np.dot(T, vec)

        k = ctr
        while k > 0:
            X[coord[0]][coord[1]] = num
            #print "Index: ",coord
            #print "Vector: ", vec
            coord += vec
            num += 1
            k -= 1

        vec = np.dot(T, vec)
        #print "Counter: ", ctr
        ctr += 1

    for ctr in xrange(N):
        X[coord[0]][coord[1]] = num
        coord += vec
        num += 1

    return X

def get_primes(n):
    '''Returns a list of all primes below n'''
    sieve = [True] * n
    for i in xrange(3,int(n**0.5)+1,2):
        if sieve[i]:
            sieve[i*i::2*i]=[False]*((n-i*i-1)/(2*i)+1)
    return [2] + [i for i in xrange(3,n,2) if sieve[i]]


def filter_primes(matrix):
    '''Takes a matrix and converts it to 
    a binary array where all primes become
    1 and all non primes become 0'''
    R = matrix.shape[0]
    C = matrix.shape[1]

    m = np.zeros((R,C))
    prime_list = get_primes(R**2 + 1)
    
    for i in xrange(R):
        for j in xrange(C):
            if matrix[i][j] in prime_list:
                m[i][j] = 128
    return m

def prime_spiral_gen(N):
    '''Creates a number spiral matrix and 
    makes all prime numbers 1 and all the
    composite numbers 1'''

    X = create_number_spiral(200)
    Y = filter_primes(X)
    plt.imshow(Y, cmap="hot")
    plt.show()

prime_spiral_gen(200)