# numpy_scipy.py
'''
Demonstration of numpy and scipy functionalities.
'''

#
# Numpy arrays (easy).
#

# Create a 1d array from a list.
a = np.array([1, 2, 3])
# Create a 2d array from a list.
b = np.array([[1, 2], [3, 4]])
# Access the array's elements.
a[0]
b[0, 0]
# Create array with specific data type.
a = np.array([1, 2, 3], dtype=np.float16)

# Create a default array.
a = np.zeros(10)
b = np.zeros([3, 3])
c = np.ones([2, 3])
d = np.eye(3)
e = np.random.random([10, 10])
f = np.arange(3, 20, 3)
g = np.linspace(0, 1, 11)

# Create a meshgrid from two or more 1d arrays.
x = np.linspace(0, 1, 6)
y = np.linspace(0, 10, 6)
yy, xx = np.meshgrid(x, y)
xx, yy = np.meshgrid(x, y, indexing='ij')
z = np.linspace(0, 100, 6)
yy, xx, zz = np.meshgrid(x, y, z)
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')

# Arrays can be accessed similarly to lists including slicing.
b[-1, -2]
a[1:-3] = 0.3
a[1::2]
a[-1::-2]

# Access some information of the arrays.
e.shape
a.dtype
e.max()
e.min()
e.ndim
e.size

# Array operations are for the most part element wise.
x = np.ones(5)
y = np.zeros(5)
z = x + y

# A few useful function to manipulate arrays.
e.swapaxes(0, 1)
e.reshape(5, 20)
e.astype(np.float16)

# Like strings and lists, numpy arrays are mutable object.
a = np.array([1, 2, 3])
b = a
b[0] = 5
print(a)
# Works with slicing.
a = np.linspace(0, 1, 11)
b = a[::2]
b[0] = 5
b[1] = 6
print(a)

# To copy, one needs to use 'deep copy'.
a = np.array([1, 2, 3])
b = a.copy()
b[0] = 5
print(a)


#
# Numpy functions and constants.
#

# Numpy functions mostly act on numpy arrays, which make computation much faster.
a = np.random(5)
np.sin(a)
np.cos(a)
np.arctan2(x, y)
np.max()
np.min()
a = np.random.random(4)
b = np.random.random(4)
np.dot(a, b)
np.median(a)
np.mean(a)
np.isnan()
np.isinf()
np.real(3 + 2j)
np.imag(3 + 2j)
np.histogram(np.random.random(100))
np.sum(a)

# Some numpy constants.
np.NaN
np.Inf
np.e
np.pi
np.float128
np.int16


#
# Numpy arrays (advanced).
#

# Array masks (Boolean array indexing).
a = np.linspace(0, 1, 11)
a > 0.5
a[a > 0.5]
age = np.array([18, 21, 19, 23, 21, 20, 19.5])
marks = np.array([1, 5, 2, 4, 1, 3, 4])
np.mean(marks[age > 20])
# Better: create a Boolean mask array.
mask = age > 20
np.mean(marks[mask])

# Find value in array.
np.where(a > 0.5)

# Array shape can be changed in various ways.

# Flatten arrays. Return the flatten arrays.
a = np.eye(3)
b = a.flatten()
c = a.ravel()
# Flatten returns a copy, ravel a reference.
b[0] = 5
print(a)
c[0] = 5
print(a)

# Stack arrays
a = np.ones(3)
b = np.array([1, 2, 3])
np.hstack([a, b])
np.vstack([a, b])

# Different dimensions.
c = np.eye(3)
np.vstack([c, b])
b = b.reshape([3, 1])
np.hstack([c, b])

# Broadcasting.
a = np.linspace(0, 1, 3)
e = np.eye(3)
x = a + 5
y = e + 5
x = a + e

# Create higher dimensional arrays.
a = np.array([[1, 2, 3], [10, 11, 12]])
b = np.array([[1, 2, 3, 14], [10, 11, 12, 13]])
x = a.reshape(2, 3, 1)
y = b.reshape(2, 1, 4)
z = x + y

# Numpy can broadcast arrays with compatible dimensions.
# It starts comparing the trailing dimension and works its way forward.
# Compatible are dimensions that are equal or 1.
x = np.linspace(0, 1, 11)
y = np.reshape(np.linspace(0, 1, 11), [11, 1])
z = x + y

a = np.array([[0, 1, 2], [10, 11, 12], [20, 21, 22]])
a = a.reshape([3, 3, 1])
b = np.array([11, 12, 13])
c = a + b


#
# Numpy matrices.
#

# Matrix creation and manipulation.
M = np.matrix(np.eye(3))
M[0, 1] = 3
N = M.copy()
N[1, 0] = 4

# Matrix opertations.
M + N
M*N
M.T
M.I
np.linalg.eig(M)
M.trace()


# Numpy vs. Scipy.
# Both, numpy and scipy contain linear algebra modules.
# However, there are more functionalities in scipy.


#
# Numpy IO
#

# Write and read clear text data files.
t = np.linspace(0, 10, 11)
x = np.random.random(11)
y = np.random.random(11)
data = np.vstack([t, x, y]).transpose()
np.savetxt('time_series.dat', data, delimiter=',', header='t    x    y')
my_data = np.loadtxt('time_series.dat', delimiter=',')

# Write and read binary files.
np.save('data', data)
my_data = np.load('data.npy')


#
# Least square finding.
#

from scipy.optimize import leastsq

def my_fit_function(x, p):
    return p[2]*x**2 + p[1]*x + p[0]

def residuals(p, x, y):
    return (y - my_fit_function(x, p))**2

# Initial data.
n = 201
x = np.linspace(-2, 2, n)
y = x**2 + 0.5*x - 0.1 + (np.random.random(n)-0.5)/5
# Initial parameter guess.
p0 = np.array([1, 1, 1])
p = leastsq(residuals, p0, args=(x, y))
q = leastsq(residuals, p0, args=(x, y), full_output=True)


#
# FFT
#

# Can be also done using the numpy fft.
from scipy.fftpack import fft

x = np.linspace(-3*np.pi, 3*np.pi, 1000)
y = 5*np.exp(1j*x) + 3*np.exp(1j*x*3) + 0.5*np.exp(-1j*x*7)
f = fft(y)/len(y)

# n-dimensional fft
from scipy.fftpack import fftn


#
# Solve ode numerically.
#

import scipy

def fp(f, t):
    return np.array([10*(f[1]-f[0]), -f[1] + f[0]*(28-f[2]), f[0]*f[1] - 8./3*f[2]])

f0 = np.array([1, 1, 1])
t = np.linspace(0, 200, 10000)
f = scipy.integrate.odeint(fp, f0, t)

plt.plot(f[:, 0], f[:, 1])