primes.py

#http://pages.physics.cornell.edu/~sethna/StatMech/ComputerExercises/PythonSoftware/Primes.py

"""
Routines to calculate prime numbers and plot their probability distribution. 
An introduction to writing functions and making histograms in Python.
""" 

"""
Import libraries.
"""
import scipy, pylab

def isEven(n):
    """
    Function which will "return True" if n is even, False otherwise.
    Uses the modulo operator: n%m is 0 if n is divisible by m, and 1 if n
    is not divisible by m. So n%3==0 evaluates to True if n is divisible by
    three.
    """
    return n%2==0

def isPrime(n):
    """
    Function which returns True if the integer n is prime. Tests integers
    d from two up to Dmax = scipy.sqrt(n), stopping if any are divisors of n
    (or, test if n is even and then test odd divisors). This is most naturally
    done using the "while" command, 
      while n%d != 0 and d <= Dmax:
          d+=1		[or 2]
    What condition will d satisfy after the while loop if n is prime?
    """
    Dmax = scipy.sqrt(n)
    if n == 2:
        return True
    if isEven(n):
        return False
    d = 3
    while n%d != 0 and d <= Dmax:
        d += 2
    return d > Dmax

def primeList(nMax):
    """
    Returns a list of all prime numbers less than nMax.
    You can use isPrime to generate a list of primes using the nice 
    Python feature of "List comprehensions". For example, the squares of the
    even numbers between seven and nineteen can be generated by
        [n**2 for n in scipy.arange(7,19) if isEven(n)]
    List comprehensions return a list using the elements generated by the 
    "for" loop that satisfy the (optional) if expression.
    """
    return [n for n in scipy.arange(2,nMax) if isPrime(n)]

def PlotPrimeDensity(primes, bins=100):
    """
    Uses pylab.hist to plot a histogram of the
    probability distribution of the numbers in the list "primes".
    (In Python, one can define "optional arguments" to functions; 
        pylab.hist(primes) 
    will create a histogram with a number of bins chosen
    automatically, while
        pylab.hist(primes, bins=100) 
    will generate a histogram with 100 bins.)
    Try plotting the primes up to nMax=100000.
    """
    pylab.hist(primes, bins=bins)
    pylab.show()

def nthPrimeTheory(n):
    """
    p(n) is the nth prime.
    An amazing result in mathematics gives an asympotic expansion of p(n)
    valid at large n:
        p(n) ~ n (log n + log log n - 1) + ...
    see, e.g., http://primes.utm.edu/howmany.shtml.
    Notice that nthPrimeTheory, if written in terms of scipy functions, can be
    applied to a whole scipy array (just like scipy.sin(y)).
    """
    return n * (scipy.log(n) + scipy.log(scipy.log(n)) - 1.)

def PrimeDensityTheory(n):
    """
    We want to return the density of primes rho(p)=dn/dp. Knowing p(n),
    rho = 1/(dp/dn) = 1/(log(n) + log(log(n)) + 1/log(n)). 
    Notice that this is a function of n, not p!
    """
    return 1./(scipy.log(n) + scipy.log(scipy.log(n)) + 1./scipy.log(n))

def PlotPrimeDensityTheory(primes, bins=100, nPoints=1000):
    """
    Plots the theoretical estimate for the histogram, C*rho(n) versus p(n), 
    for suitable normalization C, at nPoints=1000 points
    n between 2 and len(primes). (Use scipy.arange to get a numpy array for n.)
    Notice:
      (a) this is a parametric curve: the horizontal coordinate p is 
      	  a function of n
      (b) we start at n=3 because nthPrimeTheory becomes negative below n~2.72.
    The histogram in PlotPrimeDensity shows the number of counts per bin,
    corresponding to integrating the density over the bin-size. 
    Hence C is the bin size 
       (biggest prime - smallest prime)/bins = (primes[-1]-2.0)/bins.
    Notice 
       (c) the decimal point for the smallest prime 2, avoiding integer
       division,
       (d) the notation primes[-1], giving the LAST element of the list 
       primes.
    Also notice
       (e) pylab.hist, followed by pylab.plot, followed by pylab.show(),
       will give both the histogram and the theory curve on the same plot,
       (f) pylab.plot can take an argument color='r' to make the curve red,
       and an argument linewidth=2 to make a thicker line.
    """
    pylab.hist(primes, bins=bins)
    nMax = len(primes)
    n = scipy.arange(3,nMax,nMax/nPoints)
    p = nthPrimeTheory(n)
    rho = PrimeDensityTheory(n)
    C = (primes[-1]-primes[0])/bins
    pylab.plot(p, rho*C, color='r', linewidth=2)
    pylab.show()

def nextPrime(x):
    i=1
    while True:
        if isPrime(x+i): return x+i
        i+=1

def demo():
    """Demonstrates solution for exercise: example of usage"""
    print "Prime number generation demo"
    primes = primeList(10000)
    PlotPrimeDensityTheory(primes)

if __name__ == '__main__':
    demo()