# coding: utf-8

## Warm-up Exercise: Prime Numbers

### BACKGROUND

# We introduce several useful features of the Python language and associated libraries by writing a program to generate prime numbers. We also test mathematical predictions of the asymptotic density of primes.

### LEARNING GOALS

# <i>Science</i>: There are no real science goals for this simple warmup exercise. If you're interested in the mathematics of prime numbers, feel free to work through the material on the asymptotic density of primes. Since this is really an exercise to get familiar with Python, however, the PrimesHints.py file already implements much of that.
# 
# <i>Computation</i>: As part of <a href="http://pages.physics.cornell.edu/~myers/teaching/ComputationalMethods/GettingStarted.html"> Getting Started with Python</a>, this exercise will introduce you to modifying Hints files, writing Python functions, using the editor and the ipython interpreter, and simple plotting commands.

### PROCEDURE 

# Consult the instructions in the section "External editors, %run, and Hints files" in the Getting Started with Python page.
# 

# In[2]:

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

# Import the necessary libraries.
import scipy, pylab

def isEven(n):
    """
    Function that will "return True" if n is even, False otherwise.
    Uses the modulo operator, which returns the remainder of the
    division of n by m:  e.g., n%m is 0 if n is divisible by m.
    """
    return n%2==0

def isPrime(n):
    """
    Function that 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, d += 2 if testing only odds

    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 
    Python list comprehension syntax. For example, the squares of the
    even numbers between seven and nineteen can be generated by

    [n**2 for n in range(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 range(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.linspace to get a scipy array for n.
    (Type help(scipy.linspace) to get information on usage.)
    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.linspace(3,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 demo():
    """Demonstrates solution for exercise: example of usage"""
    print("Prime number generation demo")
    primes = primeList(10000)
    PlotPrimeDensityTheory(primes)


# Copyright (C) Cornell University
# All rights reserved.
# Apache License, Version 2.0