from __future__ import division, print_function

import numpy as np
import scipy as sp
import matplotlib as mpl
import matplotlib.pyplot as plt

#IPython magic command for inline plotting
%matplotlib inline
#a better plot shape for IPython
mpl.rcParams['figure.figsize']=[15,3]

from scipy import optimize
help(optimize.fmin)

# This code snippet models a falling object using the well-known formula.
a  = -9.8 #m/s^2
v  = 0    #m/s
x0 = 1000 #m
t = 1

y = a*t**2 + v*t + x0
print(y)

# This code snippet models a falling object using the well-known formula.
a  = -9.8 #m/s^2
v  = 0    #m/s
x0 = 1000 #m
t = [0.0, 0.25, 0.5, 0.75, 1.0]

y = a*t**2 + v*t + x0
print(y)

# This code snippet models a falling object using the well-known formula.
a  = -9.8 #m/s^2
v  = 0    #m/s
x0 = 1000 #m
t = np.array([0.0, 0.25, 0.5, 0.75, 1.0])

y = a*t**2 + v*t + x0
print(y)

def y_fall(t, x0, v0):
    a = -9.8
    y = a*t**2 + v0*t + x0
    return y

y_fall(t, x0, v)

t = np.arange(0,10,0.1)
y = y_fall(t, x0, v)

print(y)

plt.plot(t, y_fall(t, x0, v))
plt.show()

from numpy import zeros

dx = 0.1
dy = 0.1
dx2 = dx*dx
dy2 = dy*dy

def py_update(u):
    nx, ny = u.shape
    for i in range(1,nx-1):
        for j in range(1, ny-1):
            u[i,j] = ((u[i+1, j] + u[i-1, j]) * dy2 +
                      (u[i, j+1] + u[i, j-1]) * dx2) / (2*(dx2+dy2))

def calc(N, Niter=100, func=py_update, args=()):
    u = zeros([N, N])
    u[0] = 1
    for i in range(Niter):
        func(u,*args)
    return u

%timeit calc(100, 200, func=py_update)

def num_update(u):
    u[1:-1,1:-1] = ((u[2:,1:-1]+u[:-2,1:-1])*dy2 + 
                    (u[1:-1,2:] + u[1:-1,:-2])*dx2) / (2*(dx2+dy2))

%timeit calc(100, 200, func=num_update)

from scipy import weave

def weave_update(u):
    code = """
    int i, j;
    for (i=1; i<Nu[0]-1; i++) {
       for (j=1; j<Nu[1]-1; j++) {
           U2(i,j) = ((U2(i+1, j) + U2(i-1, j))*dy2 + \
                       (U2(i, j+1) + U2(i, j-1))*dx2) / (2*(dx2+dy2));
       }
    }
    """
    weave.inline(code, ['u', 'dx2', 'dy2'])

%timeit calc(100, 200, func=weave_update)

from IPython.display import Image
Image(filename='numpy.png', width=800, embed=True)

A = np.array([[1,2],[3,4]])
print(A)

B = np.ones([3,2])
print(B)

C = np.zeros((4,5))
print(C)

I = np.eye(3)
print(I)

A + 1

# elementwise multiplication
B * B

# matrix multiplication
B.dot(A)

B - 3

B - 0+3j

A + B

A.dot(B)

# Initialize array elements to one.
np.ones((3,2), dtype=np.int32)

# Initialize array elements to zero.
np.zeros((6,1),dtype=np.complex)

# Do not initialize array elements to any particular values---use with caution!
np.empty((4,5))

print(range(10))
print('range with integers: ', [i for i in range(10)])
print('range with integers: ', [i for i in range(0,10,2)])
print('range with floats: ',   [i for i in range(0,10,0.5)])

print(np.arange(10))
print('arange with integers: ', np.arange(10))
print('arange with integers: ', np.arange(0,10,2))
print('arange with floats: ',   np.arange(0,10,0.5))

# Create a range of floating-point values (useful in for loops).
dx = 0.1
xhi = 2.0
xlo = 1.0
x = np.arange(xlo, xhi+dx, dx)
print(x)

# Create an array with 11 equal spaces from 0 to 1 inclusive.
np.linspace(0, 1, 11)









def y_fall(t, x0, v0):
    a = -9.8
    y = a*t**2 + v0*t + x0
    return y

t = np.arange(0,300,0.1)
print(y_fall(t, x0, v))

x = 7
y = x
x = 8
print x,y

A = np.linspace(0.0, 9.0, 10)
B = A

print("A is B?", A is B)

A = np.linspace(0.0, 9.0, 10)
B = A

print("A =", A)
B[0] = -1 #changes element in A since B is simply another name for A
print("A =", A)

B.shape = 2,5

print("A.shape =", A.shape)
print("A =", A)

print("B.shape =", B.shape)
print("B =", B)

A = np.linspace(0.0, 9.0, 10)
C = A.view()

print("C is A?", C is A)
print("C.base is A?", C.base is A)

A = np.linspace(0.0, 9.0, 10)
C = A.view()

print("A =", A)
C[0] = -1 #changes element in A since C shares the same data as A
print("A =", A)

C.shape = 2,5

print("A.shape =", A.shape)
print("A =", A)

print("C.shape =", C.shape)
print("C =", C)

D = C[1].view()
print(D)

A = np.linspace(0.0, 9.0, 10)
E = A.copy()

print("A is E?", A is E)
print("E.base is A?", E.base is A)

A = np.linspace(0.0, 9.0, 10)
E = A.copy()

print("A =", A)
E[0] = -1 #doesn't changes element in A since E is now a different instance from A
print("A =", A)
print("E =", A)

E.shape = 2,5
print("A.shape =", A.shape)
print("A =", A)
print("E.shape =", E.shape)
print("E =", E)

x = np.linspace(0, 9, 10)
print("x =", x)
print("x[0:5] =", x[0:5])
print("x[0:10:3] =", x[0:10:3])
print("x[1::2] =", x[1::2])
print("x[::] =", x[::])
print("x[::-1] =", x[::-1])

import numpy.random as npr

X = npr.random((3,3))
print("X =\n", X)

print("X[0:2, 0] =", X[0:2, 0])
print("X[1::2] =", X[1::2])
print("X[:, 0] =", X[:, 0])
print("X[1, :] =", X[1, :])
print("X[::-1, ::-1] =\n", X[::-1, ::-1])



A = np.array([[1.5,0.9,4.6,0.1],[0.3,0.8,1.3, 2.7],[2.5,2.5,0.6,3.2]])

print('A = \n', A, '\n')

# np.sort sorts everything in its column or row.
print('Sorted along first index:\n', np.sort(A, 0), '\n')
print('Sorted along second index:\n', np.sort(A, 1), '\n')

# To keep things together, try this:
print('Indices of sorting:\n', np.argsort(A))
print('yielding the matrix columns sorted by the first row:\n', np.take(A, np.argsort(A), 1)[:,0])
print('or the rows sorted by the first column:\n', np.take(A, np.argsort(A, 0), 0)[:,0])

Q = npr.random((10,10))

print('Q = \n', Q)
np.save('random-data.npy', Q)

R = np.load('random-data.npy')
print('R = \n', R)

np.savetxt('random-data.dat', Q)

S = np.genfromtxt('random-data.dat')
print('S = \n', S)

import os

statinfo = os.stat('random-data.npy')
print('Binary file random-data.npy consists of %d bytes.'%statinfo.st_size)

statinfo = os.stat('random-data.dat')
print('Plain-text file random-data.dat consists of %d bytes.'%statinfo.st_size)

from numpy import expm1
def f_scalar(n):
    if n < 0.0:
        return 0.0
    else:
        return expm1(n)

print 'Scalar function on scalar value:', f_scalar(1.0) #this will succeed
print 'Scalar function on vector value:', f_scalar(np.linspace(0,1,6)) #this will fail

np.linspace(0,1,6) < 0.0

from numpy import vectorize
f_vector = vectorize(f_scalar)
print 'Vector function on vector value:', f_vector(np.linspace(0,1,6)) #this will succeed now

from numpy import where
def f_hand(n):
    return where(n < 0.0, 0.0, expm1(n))

print 'Handwritten vector function on vector value:', f_hand(np.linspace(0,1,6))
print 'Handwritten vector function on scalar value:', f_hand(1.0)

n = np.linspace(-5,5,11)
print where(n < 0.0, 0.0, expm1(n))
print where(n < 0.0, 0.0, n)
print where(n < 0.0)

def add_2nd_elems(b, a):
    """adds the second elements of two number containers"""
    print b, a
    return b[1] + a[1]

aa = [[1,2,3,4], [5,6,7,8], [9,10,11,12]]
bb = [[100,200,300,400], [500,600,700,800]]
print add_2nd_elems(bb,aa) #works okay

add_2nd_vec = np.vectorize(add_2nd_elems)
print add_2nd_vec(bb, aa) #breaks---why?



from numpy import linalg as LA
A = np.array([[1.,2],[3,4]])
B = np.array([[0.,-1],[1,0]])
I = np.eye(2)
print A
print B
print I

I*A #Elementwise multiplication

np.dot(I,A) #Conventional matrix multiplication

np.inner(I,A) #Conventional matrix multiplication

LA.eig(A)

LA.eig(B)

A.T #transpose

LA.inv(A) #inverse

# Solve the matrix problem C*x = t
C = np.array([[2.843,-1.326,9.841],[8.673,1.295,-3.215],[0.173,-7.724,2.832]])
t = np.array([[5.643,3.124,1.694]]).T
x = LA.solve(C,t)
print 'x=\n', repr(x)

from scipy.linalg import lu
# Solve the matrix problem C*x = t by LU decomposition.  (Note that SciPy P equals MATLAB P^-1.)
(P, L, U) = lu(C)
print 'P =\n', P
print 'L =\n', L
print 'U =\n', U

#C*x = P*L*U*x = t --> z = P^-1*t, y = L^-1*z, x = U^-1*y
z = LA.inv(P).dot(t)
y = LA.inv(L).dot(z)
x_lu = LA.inv(U).dot(y)
print 'x =\n', x_lu
print 'This should be zero:\nx_solve - x_lu =\n', x - x_lu

print 'Determinant of C =', LA.det(C), '\n'
print 'Matrix norm of C =', LA.norm(C), '\n'
print 'Inverse of C =\n', LA.inv(C), '\n'
print 'Eigenvalues of C =\n', LA.eig(C)[0], '\n'
print 'Eigenvectors of C =\n', LA.eig(C)[1]

T = np.array([[0.5, 1, 0, 0, 0, 0, 0],
              [0.866, 0, 0, 0, 0, 0, 0],
              [-0.5, 0, 0.5, 1, 0, 0, 0],
              [0.866, 0, 0.866, 0, 0, 0, 0],
              [0, -1, -0.5, 0, 0.5, 1, 0],
              [0, 0, 0.866, 0, 0.866, 0, 0],
              [0, 0, 0, -1, -0.5, 0, 0.5]])
print(T)
f1 = 1000
f2 = 2000
f = np.array([f1, -0.433*f1-0.5*f2, -f1, 0, 0, f2, 0])





from scipy.sparse import coo_matrix # COOrdinate format, a common sparse representation
print(coo_matrix(T))

from sys import getsizeof
print('%d bytes for dense T'%getsizeof(T))
print('%d bytes for sparse T'%getsizeof(coo_matrix(T)))

# Banded dense matrix
nx = 4

H = np.zeros((nx,nx), dtype=np.float128)
i,j = np.indices(H.shape)
H[i==j]   =  2.0
H[i==j-1] = -1.0
H[i==j+1] = -1.0

print(H)

# Banded sparse matrix
nx = 4
from scipy.sparse import csr_matrix # Compressed Sparse Row format, a common sparse representation

J = csr_matrix((nx,nx), dtype=np.float128)
i,j = np.indices(J.shape)
J[i==j] = 2.0
J[i==j-1] = -1.0
J[i==j+1] = -1.0

print(J)
print(J.todense())

from numpy.polynomial.polynomial import polyval
poly = [1,4,6,7] # 1 + 4x + 6x**2 + 7x**3
print(polyval(2.0, poly))

from numpy.polynomial import Polynomial

Polynomial([1,2,3]).linspace(25, [0,1])

Polynomial([4,2,-1]).roots()

# Load the energy consumption data from online source.
import urllib.request

url  = urllib.request.urlopen("https://raw.githubusercontent.com/uiuc-cse/python-sp15/gh-pages/lessons/data/eia-data.csv").read(20000).decode('utf-8') # read only 20 000 chars
rows = url.split('\r') # then split it into lines
data_list = [row.split(',') for row in rows]
data_list = [float(pt) for row in data_list for pt in row ]
energy_data = np.array(data_list)
print(energy_data.shape[0])
energy_data = np.reshape(energy_data, (energy_data.shape[0]/2,2))

ax = plt.subplot(111)
ax.plot(energy_data[:,0], energy_data[:,1])
ax.set_title(r'U.S. Energy Consumption in $10^{15}$ BTU')
plt.show()

sample = energy_data[3::12] # sampling in March

ax = plt.subplot(111)
ax.plot(sample[:,0], sample[:,1])
ax.set_title(r'U.S. Energy Consumption in $10^{15}$ BTU')
plt.show()

from numpy.polynomial.polynomial import polyfit

order = 3

x = sample[:,0]
y = sample[:,1]

y_poly = polyfit(x, y, order)
y_fit = polyval(x, y_poly)

ax = plt.subplot(111)
ax.plot(x, y)
ax.plot(x, y_fit, 'r-')
ax.set_title(r'U.S. Energy Consumption in $10^{15}$ BTU, March 1973–March 2014')
plt.show()

y_poly, y_stats = polyfit(x, y, order, full=True)
y_fit = polyval(x, y_poly)

print(y_poly)
print(y_stats)

from numpy.polynomial import Chebyshev as T
from numpy.polynomial import Polynomial as P
x = np.linspace(-1, 1, 100)
for i in range(6):
    ax = plt.plot(x, T.basis(i)(x), lw=2, label=r'$T_%d$'%i)
    print(r'T_%d(x) = %s'%(i, T.basis(i).convert(kind=P)))
plt.legend(loc="upper left")
plt.show()

# Select an arbitrary function to approximate.
from scipy.special import j0 # Bessel function of the first kind, order zero

# Approximate the function f by the Taylor series expansion p (translated to the origin, so p(0) = f(x0)).
from scipy.interpolate import approximate_taylor_polynomial
x0 = 3.5 # location at which we will approximate
order = 3 # <-- tuning knob for order of approximation
series = approximate_taylor_polynomial(j0, x0, order, 0.01)
print(series)

# Plot the approximation(s) versus the original function.
x = np.linspace(x0-0.5, x0+0.5, 200)

mpl.rcParams['figure.figsize']=[15,6]
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x, j0(x), 'r-', lw=2, label=r'$J_0(x)$')
ax.plot(x, series(x-x0))

# Plot residual in subfigure.
ax = fig.add_subplot(222)
ax.plot(x, j0(x)-series(x-x0), lw=2)
plt.plot()

from sympy import Symbol, series
from sympy.functions.special.bessel import besselj

x = Symbol('x')
series(besselj(0,x), x, x0=x0, n=order+1).evalf()

from numpy.random import random
print(random((2,3)))

from numpy.random import uniform
print(uniform(0,1))
print(uniform(0,100))
print(uniform(5))

from numpy.random import normal

n = (1000,1)
x = normal(size=n)
avg = np.mean(x)
std = np.std(x)

x_avg  = np.ones(n)* avg
x_stdl = np.ones(n)*(avg-std)
x_stdh = np.ones(n)*(avg+std)

plt.plot(x,'bx',x_avg,'r-',x_stdl,'r--',x_stdh,'r--')
plt.title('%d Random Gaussian Numbers' % n[0]); plt.xlabel(r'$n$'); plt.ylabel(r'$U(n)$')
plt.show()

plt.hist(x,bins=20)
plt.title('Distribution of %d Random Gaussian Numbers' % n[0]); plt.xlabel(r'$U(n)$'); plt.ylabel(r'Frequency of $U$')
plt.show()

# Generate an array of ordered pairs of independent variable values.
factors = np.array([[i, j] for i in range(0,4) for j in range(0,4)])
print('Initial array:\n', factors)

# Assign a random number to each ordered pair and sort on that basis.  Presto!  A randomized factorial design!
from numpy.random import shuffle
shuffle(factors)
print('Shuffled array:\n', factors)

from scipy.special import jn
sp.info(jn)

from scipy import optimize
sp.source(optimize.fmin)

from scipy import constants
from math import e as ee

print(constants.R) #molar gas constant in J/mol/K
print(constants.physical_constants['molar gas constant'])

print(constants.c)
print(constants.physical_constants['speed of light in vacuum'])

print(constants.physical_constants['neutron mass'])
print(constants.physical_constants['{220} lattice spacing of silicon'])
print(constants.physical_constants['Planck mass energy equivalent in GeV'])

print(constants.femto*constants.gram) #SI equivalent of fg (i.e., in kg)

#(1,000 miles to m) / (speed of sound at 15°C, 1 atm in m/s)
1000 * constants.mile / constants.mach

print(ee) #Euler's number, base of the natural logarithm
print(constants.e) #elementary charge, not Euler's number

from scipy.special import jn, jn_zeros
xlo = 0.0
xhi = jn_zeros(1,4)[3]
x = np.linspace(xlo, xhi, 501)
y = jn(1, x)
plt.plot(x, y, 'r-')
plt.show()

from scipy.integrate import trapz
val = trapz(x, y)
print('The result of the trapezoid rule is %f.'%val)

from scipy.integrate import simps
val = simps(x, y)
print('The result of Simpson\'s rule is %f.'%val)

from scipy.integrate import quadrature
(val, err) = quadrature(lambda x: jn(1, x), xlo, xhi)
print('The result of quadrature is %f with an error of %e.'%(val, err))

from scipy.integrate import quad
(val, err) = quad(lambda x: jn(1, x), xlo, xhi)
print('The result of general integration is %f with an error of %e.'%(val, err))

#from numpy.polynomial import Polynomial as P
from numpy.polynomial.polynomial import polyint
c = [0.667, 4.675, -2.349, 5.602]
print('The base polynomial coefficients are', c)
print('The integrated polynomial coefficients are', polyint(c))

x = np.linspace(-2, 2, 1001)
plt.plot(x, np.polyval(c, x))
plt.plot(x, np.polyval(polyint(c), x))
plt.ylim((0, 20))
plt.show()

from scipy.integrate import odeint

def dydx(y, x):
    if x == 0:
        return float('Inf')
    return (1/(x+x**3.0)) - x/(1+x*x)*y

y0 = np.array((1.0))
dx = 1e-5
x  = np.arange(1e-2, 10.0, dx)
sol= odeint(dydx, y0, x)
plt.plot(x, sol)
plt.show()

def dydx(y, x):
    return y*y

y0 = np.arange(0.2,2.01,0.2)
x  = np.arange(0.0, 0.5, 1e-5)
sol= odeint(dydx, y0, x)

import matplotlib.cm as cm
for i in np.arange(10):
    c = cm.rainbow(i/10.,1)
    plt.plot(x[:], sol[:,i], color=c)
plt.ylim((0,10))
plt.show()

#f = [u, v]
def dfdx(f, x):
    return np.array([f[1], -20*np.sin(x)])

f1 = odeint(dfdx, np.array([0.0, 0.0]), np.linspace(0.0, 1.0, 10000))
f2 = odeint(dfdx, np.array([0.0, 1.0]), np.linspace(0.0, 1.0, 10000))
#Shooting method for converting a BVP to an IVP.
c1 = (1.0 - f2[:,0][-1]) / (f1[:,0][-1] - f2[:,0][-1])
c2 = (f1[:,0][-1] - 1.0) / (f1[:,0][-1] - f2[:,0][-1])
f = c1 * f1[:,0] + c2 * f2[:,0]

plt.plot(np.linspace(0.0, 1.0, 10000), f1[:,0], 'r--',
         np.linspace(0.0, 1.0, 10000), f2[:,0], 'g--',
         np.linspace(0.0, 1.0, 10000), f, 'b-')
plt.title(r'Solution of 1D heat equation $u_{xx}(x) = -20 \sin(x)$')
plt.ylabel(r'$u(x)$')
plt.xlabel(r'$x$')
plt.ylim((-4,2))
plt.show()

from scipy.interpolate import interp1d

xpts = np.linspace(0, 10, 11)
y = np.sin(-xpts**2/8.0)

f = interp1d(xpts, y, kind='linear')
f2 = interp1d(xpts, y, kind='cubic')

x = np.linspace(0, 10, 201)
plt.plot(xpts, y, 'kx',
         x, f(x), 'r--',
         x, f2(x), 'r-',
         x, np.sin(-x**2/8.0), 'k-')
plt.legend(['data', 'linear', 'cubic', 'exact'], loc='upper left', ncol=2)
plt.ylim((-1.5,1.5))
plt.show()

print('f(%f) = %f'%(2.71828, f(2.71828)))

def func(x, y):
    return x*(1-x)*np.cos(4*np.pi*x) * np.sin(2*np.pi*np.sqrt(y))

# Define the basic grid coordinates.
grid_x, grid_y = np.mgrid[0:1:250j, 0:1:250j]

# Define a random subset of the grid for which we will generate data.
pts = np.random.rand(500,2)
vals = func(pts[:,0], pts[:,1])

# Load the methods and generate a grid for each approach.
from scipy.interpolate import griddata
grid_z0 = griddata(pts, vals, (grid_x, grid_y), method='nearest')
grid_z1 = griddata(pts, vals, (grid_x, grid_y), method='linear')
grid_z2 = griddata(pts, vals, (grid_x, grid_y), method='cubic')

plt.subplot(221)
plt.imshow(func(grid_x, grid_y).T, extent=(0,1,0,1), origin='lower', cmap=cm.PiYG)
plt.plot(pts[:,0], pts[:,1], 'k.', ms=1)
plt.title('Original')

plt.subplot(222)
plt.imshow(grid_z0.T, extent=(0,1,0,1), origin='lower', cmap=cm.PiYG)
plt.title('Nearest')

plt.subplot(223)
plt.imshow(grid_z1.T, extent=(0,1,0,1), origin='lower', cmap=cm.PiYG)
plt.title('Linear')

plt.subplot(224)
plt.imshow(grid_z2.T, extent=(0,1,0,1), origin='lower', cmap=cm.PiYG)
plt.title('Cubic')

plt.gcf().set_size_inches(12, 12)
plt.show()

from math import e as ee
def f(x):
    return x ** (1/x)
x = np.linspace(0,4,401)
y = f(x)
mpl.rcParams['figure.figsize']=[15,3]
plt.plot(x, y, 'b-')
plt.axis((0.0, 4.0, 0.0, 1.5))
plt.show()

from scipy.optimize import minimize
def fmod(x):
    return 1.0-f(x)
x0 = 1.0 #a bad guess

res = minimize(fmod, x0, method='nelder-mead', options={'xtol': 1e-8, 'disp': True})
print(res)
print('Absolute error is %e.'%(res.x-ee))
mpl.rcParams['figure.figsize']=[15,3]
plt.plot(x, y, 'b-', [ee, ee], [0.0, 1.5], 'r--', [res.x, res.x], [0.0, 1.5], 'g--')
plt.axis((0.0, 4.0, 0.0, 1.5))
plt.show()

from scipy.special import j0, jn_zeros
x = np.linspace(0, 16, 201)
plt.plot(x, j0(x), 'r-', jn_zeros(0, 5), [0]*5, 'bo')
plt.ylim((-1.0, 1.0))
plt.show()

from scipy.special import airy
(Ai, Aip, Bi, Bip) = airy(1.0)
print('Ai(1.0) = %.4f\nBi(1.0) = %.4f'%(Ai, Bi))

from scipy.special import erf, gamma, eval_legendre
x = np.linspace(0.001, 3, 201)
plt.plot(x, erf(x), 'r-',
         x, gamma(x), 'g-',
         x, eval_legendre(8, x), 'b-')
plt.ylim((-0.5, 2.5))
plt.show()

from scipy.special import jn, jv, j0

print('j0(1.5) = %s\n'%j0(1.5))

# Let's do some time trials to see which is faster.  (This is for 1e6 evaluations, not one.)
import timeit
t = timeit.Timer("jv(0, 1.5)","from scipy.special import jv")
print('jv(0, 1.5) @ %.4s s'%t.timeit())

t = timeit.Timer("jn(0, 1.5)","from scipy.special import jn")
print('jn(0, 1.5) @ %.4s s'%t.timeit())

t = timeit.Timer("j0(1.5)","from scipy.special import j0")
print('j0(1.5) @ %.4s s'%t.timeit())

from scipy.stats import alpha, chi2, maxwell, rayleigh, pareto
x = np.linspace(0.01, 10, 1001)
ax = plt.plot(x, alpha.pdf(x, 1.0), lw=2, label=r'$\alpha$')
ax = plt.plot(x, chi2.pdf(x, 1.0), lw=2, label=r'$\chi^2$')
ax = plt.plot(x, maxwell.pdf(x, 1.0), lw=2, label=r'Maxwell')
ax = plt.plot(x, rayleigh.pdf(x, 1.0), lw=2, label=r'Rayleigh')
ax = plt.plot(x, pareto.pdf(x, 1.0), lw=2, label=r'Pareto')

plt.legend(loc='upper right')
plt.ylim(0.0, 1.2)
plt.show()

from numpy.random import normal
n = (1000,1)
x = normal(size=n)
plt.plot(x, 'bo')

from scipy.stats import tmean, skewtest, histogram
x_mean = np.mean(x)
x_tmean= tmean(x, (-1, 1)) #trimmed mean:  mean of interval [-1, 1]

print('x_mean =\n', x_mean)
print('x_tmean =\n', x_tmean)
print('is skew significant?  the p-value is', skewtest(x)[1])

h = histogram(x, 5)
print('histogram bins size = %f, low value = %f, occupancy = %s'%(h[2], h[1], h[0]))

# Basic parameters
nt = 60
nx = 10
alpha = 0.1 
length = 1.0 
tmax = 0.5

# Derived parameters:  mesh spacing and time step size
dx = length / (nx-1)
dt = tmax / (nt-1)

# Create arrays to save data in process.
x = []
t = []
u = []
for i in range(nt):
    t.append(i*dt)
    ulist = []
    for j in range(nx):
        ulist.append(0.0)
    u.append(ulist)
for i in range(nx):
    x.append(i*dx)

# Set initial and boundary conditions.
from math import sin, pi
for j in range(nx):
    u[0][j] = sin(pi*(j*dx)/length)**2
#boundaries are implicitly set by this initial condition

# Loop through each time step.
r  = alpha * dt / (dx*dx)
s  = 1 - 2*r
for n in range(1, nt):
    for j in range(1, nx-1):
        u[n][j] = r*u[n-1][j-1] + s*u[n-1][j] + r*u[n-1][j+1]

# Plot the results (initial and final conditions).
mpl.rcParams['figure.figsize']=[15,3]
plt.plot(x, u[0], 'k-', lw=2)
plt.plot(x, u[-1],  'b-', lw=2)

plt.ylim((0,1))
plt.xlim((0,1))
plt.show()

# Load parameters from disk.
with open('fd-params.cfg', 'r') as config:
    for line in config:
        variable, value = [word.strip() for word in line.split('=')]
        exec(variable + '=' + value)

# Compute mesh spacing and time step.
dx = length / (nx-1)
dt = tmax / (nt-1)

# Create arrays to save data in process.
x = np.linspace(0, length+1e-15, nx)
t = np.linspace(0, tmax+1e-15, nt)
u = np.zeros([nx, nt])

# Set initial and boundary conditions.
u[:, 0] = np.sin(np.pi*x/length)**2
#boundaries are implicitly set by this initial condition

# Loop through each time step.
r  = alpha * dt / (dx*dx)
s  = 1 - 2*r
for n in range(1, nt):
    for j in range(1, nx-1):
        u[j, n] = r*u[j-1, n-1] + s*u[j, n-1] + r*u[j+1, n-1]

# Plot the results (initial and final conditions).
mpl.rcParams['figure.figsize']=[15,3]
plt.plot(x, u[:,0], 'k-', lw=2)
plt.plot(x, u[:,-1],  'b-', lw=2)

plt.ylim((0,1))
plt.xlim((0,1))
plt.show()

def calc_params(nx, nt, alpha, length, tmax):
    # Compute mesh spacing and time step.
    dx = length / (nx-1)
    dt = tmax / (nt-1)
    return (dx, dt)

def create_arrays(nx, nt):
    # Create arrays to save data in process.
    x = np.linspace(0, length+1e-15, nx)
    t = np.linspace(0, tmax+1e-15, nt)
    u = np.zeros([nx, nt])
    return (x, t, u)

def set_ic(x, length):
    # Set initial and boundary conditions.
    u[:, 0] = np.sin(np.pi*x/length)**2
    #boundaries are implicitly set by this initial condition
    return u

def plot_results(u, x):
    # Plot the results (initial and final conditions).
    mpl.rcParams['figure.figsize']=[15,3]
    plt.plot(x, u[:,0], 'k-', lw=2)
    plt.plot(x, u[:,-1],  'b-', lw=2)
    
    #import matplotlib.cm as cm
    #for i in range(len(u[:,0:-1:10])-1):
    #    c = cm.rainbow(i/len(u[:,0:-1:10]), 1)
    #    plt.plot(x[:], u[:,i], color=c, lw=2)
    
    plt.ylim((0,1))
    plt.xlim((0,1))
    plt.show()

def gen_matrix(nx, alpha, dt, dx):
    r  = alpha * dt / (dx*dx)
    s  = 1 - 2*r
    A = np.zeros((nx, nx), dtype=np.float128)
    i,j = np.indices(A.shape)
    A[i==j]   = s
    A[i==j-1] = r
    A[i==j+1] = r
    return A

# Load parameters from disk.
with open('fd-params.cfg', 'r') as config:
    for line in config:
        variable, value = [word.strip() for word in line.split('=')]
        exec(variable + '=' + value)

(dx, dt) = calc_params(nx, nt, alpha, length, tmax)
(x, t, u ) = create_arrays(nx, nt)
u = set_ic(x, length)
A = gen_matrix(nx, alpha, dt, dx)

# Loop through each time step.
for n in range(1, nt):
    u[:,n] = A.dot(u[:,n-1])

plot_results(u, x)

def mandelbrot(minR, maxR, minI, maxI, samples=51, iters=25):
    """
    Generate the Mandelbrot set within the boundaries of the limits
    maxR, minR, minI, maxI with samples and a maximum iteration of iters.
    """
    # Generate a mesh for the set.
    setR = np.linspace(minR, maxR, samples)
    setI = np.linspace(minI, maxI, samples)

    # Calculate the values at each point of the mesh by the escape-time
    # fractal algorithm.
    pts = np.zeros([samples, samples])
    for ii in range(1, len(setR)):
        for jj in range(1, len(setI)):
            it = 0
            x = 0.0
            y = 0.0

            xx = setR[ii]
            yy = setI[jj]

            # Look for escape---i.e., does the value settle down in a few
            # iterations or does it keep going?
            while(x * x + y * y < 4 and it < iters):
                xtmp = x * x - y * y + xx
                y = 2 * x * y + yy
                x = xtmp
                it += 1
            pts[ii, jj] = it
    return setR, setI, pts

# Plot boundaries
minR = -2.25
maxR = 0.75
minI = -1.5
maxI = 1.5

samples = 201
iters = 20

x, y, z = mandelbrot(minR, maxR, minI, maxI, samples, iters)
z = z.transpose()

mpl.rcParams['figure.figsize']=[8,8]
plt.imshow(z, interpolation='nearest')
plt.show()

def mand_mesh(params):
    """Generate a mesh for the Mandelbrot set calculation."""
    setR = np.linspace(params['r.min'], params['r.max'], params['samples'])
    setI = np.linspace(params['i.min'], params['i.max'], params['samples'])
    pts = np.zeros([params['samples'], params['samples']])
    return (setR, setI, pts)

def mand_escape(params, xx, yy):
    # Look for escape---i.e., does the value settle down in a few
    # iterations or does it keep going?
    x, y = 0.0, 0.0
    it = 0
    while(x * x + y * y < 4 and it < params['iter.max']):
        z_np1 = x * x - y * y + xx
        y = 2 * x * y + yy
        x = z_np1
        it += 1
    return it

def mand_alg(params):
    """
    Generate the Mandelbrot set within the boundaries of the limits
    r.max, r.min, i.min, i.max with samples and a maximum iteration of iter.max.
    (These are keys in params dictionary.)
    """
    # Calculate the values at each point of the mesh by the escape-time
    # fractal algorithm.
    setR, setI, pts = mand_mesh(params)
    
    for ii in range(1, params['samples']):
        for jj in range(1, params['samples']):
            xx = setR[ii]
            yy = setI[jj]
            pts[ii, jj] = mand_escape(params, xx, yy)
    return setR, setI, pts

def mand_plot(z):
    mpl.rcParams['figure.figsize']=[8,8]
    plt.imshow(z, interpolation='nearest')
    plt.show()

if __name__ == "__main__":
    params = dict()
    params['r.min'] = -2.25
    params['r.max'] =  0.75
    params['i.min'] = -1.50
    params['i.max'] =  1.50
    params['samples'] = 401
    params['iter.max']  = 20
    
    x, y, z = mand_alg(params)
    z = z.transpose()
    mand_plot(z)