%pylab inline

import numpy as np
# Import pylab to provide scientific Python libraries (NumPy, SciPy, Matplotlib)
%pylab --no-import-all
#import pylab as pl
# import the Image display module
from IPython.display import Image
import math

np.array([1,2,3,4,5,6])

"""
Explanation: II. Numpy and Scipy
Numpy contains core routines for doing fast vector, matrix, and linear algebra-type operations in Python. Scipy contains additional routines for optimization, special functions, and so on. Both contain modules written in C and Fortran so that they're as fast as possible. Together, they give Python roughly the same capability that the Matlab program offers. (In fact, if you're an experienced Matlab user, there a guide to Numpy for Matlab users just for you.)
Making vectors and matrices
Fundamental to both Numpy and Scipy is the ability to work with vectors and matrices. You can create vectors from lists using the array command:
End of explanation
"""


np.array([1,2,3,4,5,6],'d')

np.array([1,2,3,4,5,6],'D')

np.array([1,2,3,4,5,6],'i')

"""
Explanation: You can pass in a second argument to array that gives the numeric type. There are a number of types listed here that your matrix can be. Some of these are aliased to single character codes. The most common ones are 'd' (double precision floating point number), 'D' (double precision complex number), and 'i' (int32). Thus,
End of explanation
"""


np.array([[0,1],[1,0]],'d')

"""
Explanation: To build matrices, you can either use the array command with lists of lists:
End of explanation
"""


np.zeros((3,3),'d')

"""
Explanation: You can also form empty (zero) matrices of arbitrary shape (including vectors, which Numpy treats as vectors with one row), using the zeros command:
End of explanation
"""


np.zeros(3,'d')

np.zeros((1,3),'d')

"""
Explanation: The first argument is a tuple containing the shape of the matrix, and the second is the data type argument, which follows the same conventions as in the array command. Thus, you can make row vectors:
End of explanation
"""


np.zeros((3,1),'d')

"""
Explanation: or column vectors:
End of explanation
"""


np.identity(4,'d')

"""
Explanation: There's also an identity command that behaves as you'd expect:
End of explanation
"""


np.linspace(0,1)

"""
Explanation: as well as a ones command.
Linspace, matrix functions, and plotting
The linspace command makes a linear array of points from a starting to an ending value.
End of explanation
"""


np.linspace(0,1,11)

"""
Explanation: If you provide a third argument, it takes that as the number of points in the space. If you don't provide the argument, it gives a length 50 linear space.
End of explanation
"""


x = np.linspace(0,2*np.pi)
np.sin(x)

"""
Explanation: linspace is an easy way to make coordinates for plotting. Functions in the numpy library (all of which are imported into IPython notebook) can act on an entire vector (or even a matrix) of points at once. Thus,
End of explanation
"""


plot(x,np.sin(x))

"""
Explanation: In conjunction with matplotlib, this is a nice way to plot things:
End of explanation
"""


0.125*identity(3,'d')

"""
Explanation: Matrix operations
Matrix objects act sensibly when multiplied by scalars:
End of explanation
"""


identity(2,'d') + array([[1,1],[1,2]])

"""
Explanation: as well as when you add two matrices together. (However, the matrices have to be the same shape.)
End of explanation
"""


identity(2)*ones((2,2))

"""
Explanation: Something that confuses Matlab users is that the times (*) operator give element-wise multiplication rather than matrix multiplication:
End of explanation
"""


dot(identity(2),ones((2,2)))

"""
Explanation: To get matrix multiplication, you need the dot command:
End of explanation
"""


v = array([3,4],'d')
sqrt(dot(v,v))

"""
Explanation: dot can also do dot products (duh!):
End of explanation
"""


m = array([[1,2],[3,4]])
m.T

"""
Explanation: as well as matrix-vector products.
There are determinant, inverse, and transpose functions that act as you would suppose. Transpose can be abbreviated with ".T" at the end of a matrix object:
End of explanation
"""


diag([1,2,3,4,5])

"""
Explanation: There's also a diag() function that takes a list or a vector and puts it along the diagonal of a square matrix.
End of explanation
"""


A = array([[1,1,1],[0,2,5],[2,5,-1]])
b = array([6,-4,27])
solve(A,b)

"""
Explanation: We'll find this useful later on.
Matrix Solvers
You can solve systems of linear equations using the solve command:
End of explanation
"""


A = array([[13,-4],[-4,7]],'d')
eigvalsh(A)

eigh(A)

"""
Explanation: There are a number of routines to compute eigenvalues and eigenvectors

eigvals returns the eigenvalues of a matrix
eigvalsh returns the eigenvalues of a Hermitian matrix
eig returns the eigenvalues and eigenvectors of a matrix
eigh returns the eigenvalues and eigenvectors of a Hermitian matrix.
End of explanation
"""


def nderiv(y,x):
    "Finite difference derivative of the function f"
    n = len(y)
    d = zeros(n,'d') # assume double
    # Use centered differences for the interior points, one-sided differences for the ends
    for i in range(1,n-1):
        d[i] = (y[i+1]-y[i-1])/(x[i+1]-x[i-1])
    d[0] = (y[1]-y[0])/(x[1]-x[0])
    d[n-1] = (y[n-1]-y[n-2])/(x[n-1]-x[n-2])
    return d

"""
Explanation: Example: Finite Differences
Now that we have these tools in our toolbox, we can start to do some cool stuff with it. Many of the equations we want to solve in Physics involve differential equations. We want to be able to compute the derivative of functions:
$$ y' = \frac{y(x+h)-y(x)}{h} $$
by discretizing the function $y(x)$ on an evenly spaced set of points $x_0, x_1, \dots, x_n$, yielding $y_0, y_1, \dots, y_n$. Using the discretization, we can approximate the derivative by
$$ y_i' \approx \frac{y_{i+1}-y_{i-1}}{x_{i+1}-x_{i-1}} $$
We can write a derivative function in Python via
End of explanation
"""


x = linspace(0,2*pi)
dsin = nderiv(sin(x),x)
plot(x,dsin,label='numerical')
plot(x,cos(x),label='analytical')
title("Comparison of numerical and analytical derivatives of sin(x)")
legend()

"""
Explanation: Let's see whether this works for our sin example from above:
End of explanation
"""


def Laplacian(x):
    h = x[1]-x[0] # assume uniformly spaced points
    n = len(x)
    M = -2*identity(n,'d')
    for i in range(1,n):
        M[i,i-1] = M[i-1,i] = 1
    return M/h**2

x = linspace(-3,3)
m = 1.0
ohm = 1.0
T = (-0.5/m)*Laplacian(x)
V = 0.5*(ohm**2)*(x**2)
H =  T + diag(V)
E,U = eigh(H)
h = x[1]-x[0]

# Plot the Harmonic potential
plot(x,V,color='k')

for i in range(4):
    # For each of the first few solutions, plot the energy level:
    axhline(y=E[i],color='k',ls=":")
    # as well as the eigenfunction, displaced by the energy level so they don't
    # all pile up on each other:
    plot(x,-U[:,i]/sqrt(h)+E[i])
title("Eigenfunctions of the Quantum Harmonic Oscillator")
xlabel("Displacement (bohr)")
ylabel("Energy (hartree)")

"""
Explanation: Pretty close!
One-Dimensional Harmonic Oscillator using Finite Difference
Now that we've convinced ourselves that finite differences aren't a terrible approximation, let's see if we can use this to solve the one-dimensional harmonic oscillator.
We want to solve the time-independent Schrodinger equation
$$ -\frac{\hbar^2}{2m}\frac{\partial^2\psi(x)}{\partial x^2} + V(x)\psi(x) = E\psi(x)$$
for $\psi(x)$ when $V(x)=\frac{1}{2}m\omega^2x^2$ is the harmonic oscillator potential. We're going to use the standard trick to transform the differential equation into a matrix equation by multiplying both sides by $\psi^*(x)$ and integrating over $x$. This yields
$$ -\frac{\hbar}{2m}\int\psi(x)\frac{\partial^2}{\partial x^2}\psi(x)dx + \int\psi(x)V(x)\psi(x)dx = E$$
We will again use the finite difference approximation. The finite difference formula for the second derivative is
$$ y'' = \frac{y_{i+1}-2y_i+y_{i-1}}{x_{i+1}-x_{i-1}} $$
We can think of the first term in the Schrodinger equation as the overlap of the wave function $\psi(x)$ with the second derivative of the wave function $\frac{\partial^2}{\partial x^2}\psi(x)$. Given the above expression for the second derivative, we can see if we take the overlap of the states $y_1,\dots,y_n$ with the second derivative, we will only have three points where the overlap is nonzero, at $y_{i-1}$, $y_i$, and $y_{i+1}$. In matrix form, this leads to the tridiagonal Laplacian matrix, which has -2's along the diagonals, and 1's along the diagonals above and below the main diagonal.
The second term turns leads to a diagonal matrix with $V(x_i)$ on the diagonal elements. Putting all of these pieces together, we get:
End of explanation
"""


from numpy.polynomial.hermite import Hermite
def ho_evec(x,n,m,ohm):
    vec = [0]*9
    vec[n] = 1
    Hn = Hermite(vec)
    return (1/sqrt(2**n*math.factorial(n)))*pow(m*ohm/pi,0.25)*exp(-0.5*m*ohm*x**2)*Hn(x*sqrt(m*ohm))

"""
Explanation: We've made a couple of hacks here to get the orbitals the way we want them. First, I inserted a -1 factor before the wave functions, to fix the phase of the lowest state. The phase (sign) of a quantum wave function doesn't hold any information, only the square of the wave function does, so this doesn't really change anything. 
But the eigenfunctions as we generate them aren't properly normalized. The reason is that finite difference isn't a real basis in the quantum mechanical sense. It's a basis of Dirac δ functions at each point; we interpret the space betwen the points as being "filled" by the wave function, but the finite difference basis only has the solution being at the points themselves. We can fix this by dividing the eigenfunctions of our finite difference Hamiltonian by the square root of the spacing, and this gives properly normalized functions.
Special Functions
The solutions to the Harmonic Oscillator are supposed to be Hermite polynomials. The Wikipedia page has the HO states given by
$$\psi_n(x) = \frac{1}{\sqrt{2^n n!}}
\left(\frac{m\omega}{\pi\hbar}\right)^{1/4}
\exp\left(-\frac{m\omega x^2}{2\hbar}\right)
H_n\left(\sqrt{\frac{m\omega}{\hbar}}x\right)$$
Let's see whether they look like those. There are some special functions in the Numpy library, and some more in Scipy. Hermite Polynomials are in Numpy:
End of explanation
"""


plot(x,ho_evec(x,0,1,1),label="Analytic")
plot(x,-U[:,0]/sqrt(h),label="Numeric")
xlabel('x (bohr)')
ylabel(r'$\psi(x)$')
title("Comparison of numeric and analytic solutions to the Harmonic Oscillator")
legend()

"""
Explanation: Let's compare the first function to our solution.
End of explanation
"""


phase_correction = [-1,1,1,-1,-1,1]
for i in range(6):
    subplot(2,3,i+1)
    plot(x,ho_evec(x,i,1,1),label="Analytic")
    plot(x,phase_correction[i]*U[:,i]/sqrt(h),label="Numeric")

"""
Explanation: The agreement is almost exact.
We can use the subplot command to put multiple comparisons in different panes on a single plot:
End of explanation
"""


from scipy.special import airy,jn,eval_chebyt,eval_legendre
subplot(2,2,1)
x = linspace(-1,1)
Ai,Aip,Bi,Bip = airy(x)
plot(x,Ai)
plot(x,Aip)
plot(x,Bi)
plot(x,Bip)
title("Airy functions")

subplot(2,2,2)
x = linspace(0,10)
for i in range(4):
    plot(x,jn(i,x))
title("Bessel functions")

subplot(2,2,3)
x = linspace(-1,1)
for i in range(6):
    plot(x,eval_chebyt(i,x))
title("Chebyshev polynomials of the first kind")

subplot(2,2,4)
x = linspace(-1,1)
for i in range(6):
    plot(x,eval_legendre(i,x))
title("Legendre polynomials")

"""
Explanation: Other than phase errors (which I've corrected with a little hack: can you find it?), the agreement is pretty good, although it gets worse the higher in energy we get, in part because we used only 50 points.
The Scipy module has many more special functions:
End of explanation
"""


raw_data = """\
3.1905781584582433,0.028208609537968457
4.346895074946466,0.007160804747670053
5.374732334047101,0.0046962988461934805
8.201284796573875,0.0004614473299618756
10.899357601713055,0.00005038370219939726
16.295503211991434,4.377451812785309e-7
21.82012847965739,3.0799922117601088e-9
32.48394004282656,1.524776208284536e-13
43.53319057815846,5.5012073588707224e-18"""

"""
Explanation: As well as Jacobi, Laguerre, Hermite polynomials, Hypergeometric functions, and many others. There's a full listing at the Scipy Special Functions Page.
Least squares fitting
Very often we deal with some data that we want to fit to some sort of expected behavior. Say we have the following:
End of explanation
"""


data = []
for line in raw_data.splitlines():
    words = line.split(',')
    data.append(map(float,words))
data = array(data)

title("Raw Data")
xlabel("Distance")
plot(data[:,0],data[:,1],'bo')

"""
Explanation: There's a section below on parsing CSV data. We'll steal the parser from that. For an explanation, skip ahead to that section. Otherwise, just assume that this is a way to parse that text into a numpy array that we can plot and do other analyses with.
End of explanation
"""


title("Raw Data")
xlabel("Distance")
semilogy(data[:,0],data[:,1],'bo')

"""
Explanation: Since we expect the data to have an exponential decay, we can plot it using a semi-log plot.
End of explanation
"""


params = polyfit(data[:,0],log(data[:,1]),1)
a = params[0]
A = exp(params[1])

"""
Explanation: For a pure exponential decay like this, we can fit the log of the data to a straight line. The above plot suggests this is a good approximation. Given a function
$$ y = Ae^{-ax} $$
$$ \log(y) = \log(A) - ax$$
Thus, if we fit the log of the data versus x, we should get a straight line with slope $a$, and an intercept that gives the constant $A$.
There's a numpy function called polyfit that will fit data to a polynomial form. We'll use this to fit to a straight line (a polynomial of order 1)
End of explanation
"""


x = linspace(1,45)
title("Raw Data")
xlabel("Distance")
semilogy(data[:,0],data[:,1],'bo')
semilogy(x,A*exp(a*x),'b-')

"""
Explanation: Let's see whether this curve fits the data.
End of explanation
"""


gauss_data = """\
-0.9902286902286903,1.4065274110372852e-19
-0.7566104566104566,2.2504438576596563e-18
-0.5117810117810118,1.9459459459459454
-0.31887271887271884,10.621621621621626
-0.250997150997151,15.891891891891893
-0.1463309463309464,23.756756756756754
-0.07267267267267263,28.135135135135133
-0.04426734426734419,29.02702702702703
-0.0015939015939017698,29.675675675675677
0.04689304689304685,29.10810810810811
0.0840994840994842,27.324324324324326
0.1700546700546699,22.216216216216214
0.370878570878571,7.540540540540545
0.5338338338338338,1.621621621621618
0.722014322014322,0.08108108108108068
0.9926849926849926,-0.08108108108108646"""

data = []
for line in gauss_data.splitlines():
    words = line.split(',')
    data.append(map(float,words))
data = array(data)

plot(data[:,0],data[:,1],'bo')

"""
Explanation: If we have more complicated functions, we may not be able to get away with fitting to a simple polynomial. Consider the following data:
End of explanation
"""


def gauss(x,A,a): return A*exp(a*x**2)

"""
Explanation: This data looks more Gaussian than exponential. If we wanted to, we could use polyfit for this as well, but let's use the curve_fit function from Scipy, which can fit to arbitrary functions. You can learn more using help(curve_fit).
First define a general Gaussian function to fit to.
End of explanation
"""


from scipy.optimize import curve_fit

params,conv = curve_fit(gauss,data[:,0],data[:,1])
x = linspace(-1,1)
plot(data[:,0],data[:,1],'bo')
A,a = params
plot(x,gauss(x,A,a),'b-')

"""
Explanation: Now fit to it using curve_fit:
End of explanation
"""


from random import random
rands = []
for i in range(100):
    rands.append(random())
plot(rands)

"""
Explanation: The curve_fit routine we just used is built on top of a very good general minimization capability in Scipy. You can learn more at the scipy documentation pages.
Monte Carlo, random numbers, and computing $\pi$
Many methods in scientific computing rely on Monte Carlo integration, where a sequence of (pseudo) random numbers are used to approximate the integral of a function. Python has good random number generators in the standard library. The random() function gives pseudorandom numbers uniformly distributed between 0 and 1:
End of explanation
"""


from random import gauss
grands = []
for i in range(100):
    grands.append(gauss(0,1))
plot(grands)

"""
Explanation: random() uses the Mersenne Twister algorithm, which is a highly regarded pseudorandom number generator. There are also functions to generate random integers, to randomly shuffle a list, and functions to pick random numbers from a particular distribution, like the normal distribution:
End of explanation
"""


plot(rand(100))

"""
Explanation: It is generally more efficient to generate a list of random numbers all at once, particularly if you're drawing from a non-uniform distribution. Numpy has functions to generate vectors and matrices of particular types of random distributions.
End of explanation
"""


npts = 5000
xs = 2*rand(npts)-1
ys = 2*rand(npts)-1
r = xs**2+ys**2
ninside = (r<1).sum()
figsize(6,6) # make the figure square
title("Approximation to pi = %f" % (4*ninside/float(npts)))
plot(xs[r<1],ys[r<1],'b.')
plot(xs[r>1],ys[r>1],'r.')
figsize(8,6) # change the figsize back to 4x3 for the rest of the notebook

"""
Explanation: One of the first programs I ever wrote was a program to compute $\pi$ by taking random numbers as x and y coordinates, and counting how many of them were in the unit circle. For example:
End of explanation
"""


n = 100
total = 0
for k in range(n):
    total += pow(-1,k)/(2*k+1.0)
print 4*total

"""
Explanation: The idea behind the program is that the ratio of the area of the unit circle to the square that inscribes it is $\pi/4$, so by counting the fraction of the random points in the square that are inside the circle, we get increasingly good estimates to $\pi$. 
The above code uses some higher level Numpy tricks to compute the radius of each point in a single line, to count how many radii are below one in a single line, and to filter the x,y points based on their radii. To be honest, I rarely write code like this: I find some of these Numpy tricks a little too cute to remember them, and I'm more likely to use a list comprehension (see below) to filter the points I want, since I can remember that.
As methods of computing $\pi$ go, this is among the worst. A much better method is to use Leibniz's expansion of arctan(1):
$$\frac{\pi}{4} = \sum_k \frac{(-1)^k}{2*k+1}$$
End of explanation
"""


from numpy import sqrt
def f(x): return exp(-x)
x = linspace(0,10)
plot(x,exp(-x))

"""
Explanation: If you're interested a great method, check out Ramanujan's method. This converges so fast you really need arbitrary precision math to display enough decimal places. You can do this with the Python decimal module, if you're interested.
Numerical Integration
Integration can be hard, and sometimes it's easier to work out a definite integral using an approximation. For example, suppose we wanted to figure out the integral:
$$\int_0^\infty\exp(-x)dx=1$$
End of explanation
"""


from scipy.integrate import quad
quad(f,0,inf)

"""
Explanation: Scipy has a numerical integration routine quad (since sometimes numerical integration is called quadrature), that we can use for this:
End of explanation
"""


from scipy.fftpack import fft,fftfreq

npts = 4000
nplot = npts/10
t = linspace(0,120,npts)
def acc(t): return 10*sin(2*pi*2.0*t) + 5*sin(2*pi*8.0*t) + 2*rand(npts)

signal = acc(t)

FFT = abs(fft(signal))
freqs = fftfreq(npts, t[1]-t[0])

subplot(211)
plot(t[:nplot], signal[:nplot])
subplot(212)
plot(freqs,20*log10(FFT),',')
show()

"""
Explanation: There are also 2d and 3d numerical integrators in Scipy. See the docs for more information.
Fast Fourier Transform and Signal Processing
Very often we want to use FFT techniques to help obtain the signal from noisy data. Scipy has several different options for this.
End of explanation
"""


