# Loading modules
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: Week 05 - Numpy II and Scipy
Today's Agenda

Numpy II
Scipy

Numpy II
Last time in Week 05, we covered Numpy and Matplotlib. This time we will be focusing on more advanced concepts of Numpy.
End of explanation
"""


x = np.array([1,2,3,5,6,7,8,10],dtype=float)
x

y = np.arange(10)
y

z = np.linspace(0,100,50)
z

h = np.random.randn(100)
h

"""
Explanation: Review
As a review, let's explore some of the concepts that were introduced last time, in Numpy I.
Create 1D-arrays
We introduced how to create a 1D-array
End of explanation
"""


print('Min X: {0:.3f} \t Max X: {1:.3f}'.format(np.min(x), np.max(x)) )

"""
Explanation: Handling arrays
These are just a few of the different ways to create numpy arrays.
You can also use functions like np.max() and np.min() to get the maximum and minimum values, respectively.
End of explanation
"""


zz = x**2 + 3*x**3
zz

"""
Explanation: Apply mathematical functions
End of explanation
"""


zz_idx = np.where((zz>= 100)&(zz <= 500))[0]
print('zz_idx: {0}'.format(zz_idx))

zz[zz_idx]

"""
Explanation: Conditionals
Find the indices of the elements in an array that meet some criteria.
In this example, we're finding all the elements that are within 100 and 500 in array "zz".
End of explanation
"""


h1 = np.random.randint(10, 50, 50)
h1

"""
Explanation: Manipulating Arrays
There are a lot of things we can do to a numpy array.
End of explanation
"""


np.size(h1)

h1.shape

A = np.array([[1,2,3,4,5],
              [6,7,8,9,10],
              [12,13,14,16,17],
              [13,45,67,89,90] ])
A

np.shape(A)

"""
Explanation: We can get the overall size and shape of the array.
We cane use the functions numpy.size and numpy.shape to get the total number of elements in an array and the shape of the array, respectively.
End of explanation
"""


A_t = np.transpose(A)
A_t

"""
Explanation: You can also transpose array A.
End of explanation
"""


np.arange(0,10,1)

np.arange(0,20,5)

np.arange(-40,21,10)

"""
Explanation: Why are Numpy arrays better than lists:
- Python lists are very general.
- Lists do not support matrix and dot multiplications, etc.
- Numpy arrays are memory efficient.
- Numpy arrays are statically typed and homogeneous.
- They are fast at mathematical functions.
- They can be used in compiled languages, e.g. C and Fortran.
Array-generating functions
For large arrays it is inpractical to initialize the data manually, using normal Python lists. Instead, we can use many of the Numpy functions to generate arrays of different forms.
numpy.arange
We use this one to create a sequence of ordered elements
End of explanation
"""


B = np.linspace(0,50)
B

B = np.linspace(0,100, 20)
B

"""
Explanation: linspace and logspace
We use these functions to created ordered lists, separated by intervals in real- and log-space.
End of explanation
"""


B = np.logspace(0,3,25)
B

"""
Explanation: Array of 25 elements from $10^{0}$ to $10^{3}$, with base of 10.
End of explanation
"""


B = np.logspace(0,10,11, base=np.e)
B

"""
Explanation: Creating an array of 11 elements from $e^{0}$ to $e^{10}$, with the base == numpy.e
End of explanation
"""


from numpy import random

# Uniform random numbers in [0,1]
random.rand(5,5)

# 20 Random integers from 10 to 30
random.randint(10,30,20)

"""
Explanation: Random Data
End of explanation
"""


np.zeros(20)

"""
Explanation: Arrays of zeros and ones.
End of explanation
"""


nelem = 10
C = np.ones(10)
C

for ii in range(C.size):
    C[ii] = random.rand()
C

"""
Explanation: You can use these to populate other arrays
End of explanation
"""


np.diag(random.randint(10,20,5))

"""
Explanation: Diagonals
You can also construct an array with another array as the diagonal
End of explanation
"""


M = random.rand(10,5)
M

"""
Explanation: Indexing
You can choose which values to select.
Normally, you select the rows first, and then the cols of a numpy.ndarray.
End of explanation
"""


M[1,:]

"""
Explanation: Selecting the 1st row
End of explanation
"""


M[:,1]

"""
Explanation: The 2nd column
End of explanation
"""


M[1:3, 2:4]

"""
Explanation: Select a range of columns and rows
End of explanation
"""


A = random.rand(3,3)
np.fill_diagonal(A, np.nan)
A

B = np.arange(0,9).reshape((3,3))
B

"""
Explanation: You can easily use this to create a mask, for when you are cleaning your data.
End of explanation
"""


A_mask = np.isfinite(A)
A_mask

B[A_mask]

"""
Explanation: Appying the mask from $A \to B$
End of explanation
"""


# Creating my bin edges
bins = np.arange(0,13)
bins

# Generating Data
data = 10*random.rand(100)
data

"""
Explanation: Binning you data
This is probably one of the best functions of Numpy.
You can use this to bin you data, and calculate means, standard deviations, etc.
numpy.digitize
End of explanation
"""


# Defining statistical function to use
stat_func = np.nanmean
# Binning the data
data_bins = np.digitize(data, bins)
data_bins

"""
Explanation: Now I want to bin my data and calculate the mean for each bin
End of explanation
"""


failval = -10
bins_stat = np.array([stat_func(data[data_bins == ii]) \
                      if len(data[data_bins == ii]) > 0 \
                      else failval \
                      for ii in range(1,len(bins))])
bins_stat = np.asarray(bins_stat)
bins_stat

"""
Explanation: Calculating the mean for each of the bins
End of explanation
"""


import math

def myceil(x, base=10):
    """
    Returns the upper-bound integer of 'x' in base 'base'.

    Parameters
    ----------
    x: float
        number to be approximated to closest number to 'base'

    base: float
        base used to calculate the closest 'largest' number

    Returns
    -------
    n_high: float
        Closest float number to 'x', i.e. upper-bound float.

    Example
    -------
    >>>> myceil(12,10)
      20
    >>>>
    >>>> myceil(12.05, 0.1)
     12.10000 
    """
    n_high = float(base*math.ceil(float(x)/base))

    return n_high

def myfloor(x, base=10):
    """
    Returns the lower-bound integer of 'x' in base 'base'

    Parameters
    ----------
    x: float
        number to be approximated to closest number of 'base'

    base: float
        base used to calculate the closest 'smallest' number

    Returns
    -------
    n_low: float
        Closest float number to 'x', i.e. lower-bound float.

    Example
    -------
    >>>> myfloor(12, 5)
    >>>> 10
    """
    n_low = float(base*math.floor(float(x)/base))

    return n_low

def Bins_array_create(arr, base=10):
    """
    Generates array between [arr.min(), arr.max()] in steps of `base`.

    Parameters
    ----------
    arr: array_like, Shape (N,...), One-dimensional
        Array of numerical elements

    base: float, optional (default=10)
        Interval between bins

    Returns
    -------
    bins_arr: array_like
        Array of bin edges for given arr

    """
    base = float(base)
    arr  = np.array(arr)
    assert(arr.ndim==1)
    arr_min  = myfloor(arr.min(), base=base)
    arr_max  = myceil( arr.max(), base=base)
    bins_arr = np.arange(arr_min, arr_max+0.5*base, base)

    return bins_arr

def Mean_std_calc_one_array(x1, y1, arr_len=0, statfunc=np.nanmean,
                           failval=np.nan, error='std',
                           base=10.):
    """
    Calculates statistics of two arrays, e.g. scatter,
    error in `statfunc`, etc.
    
    Parameters
    ----------
    x1: array-like, shape (N,)
        array of x-values
    
    y1: array-like, shape (N,)
        array of y-values
    
    arr_len: int, optional (default = 0)
        minimum number of elements in the bin
    
    statfunc: numpy function, optional (default = numpy.nanmean)
        statistical function used to evaluate the bins
    
    failval: int or float, optional (default = numpy.nan)
        Number to use to replace when the number of elements in the 
        bin is smaller than `arr_len`
    
    error: string, optional (default = 'std')
        type of error to evaluate
        Options:
            - 'std': Evaluates the standard deviation of the bin
            - 'stat': Evaluates the error in the mean/median of each bin
            - 'none': Does not calculate the error in `y1`
    base: float
        Value of bin width in units of that of `x1`
    
    Returns
    --------
    x1_stat: array-like, shape (N,)
        `stat_func` of each bin in `base` spacings for x1

    y1_stat: array-like, shape (N,)
        `stat_func` of each bin in `base` spacings for y1
    """
    x1 = np.asarray(x1)
    y1 = np.asarray(y1)
    assert((x1.ndim==1) & (y1.ndim==1))
    assert((x1.size >0) & (y1.size>0))
    n_elem = len(x1)
    ## Computing Bins
    x1_bins  = Bins_array_create(x1, base=base)
    x1_digit = np.digitize(x1, x1_bins)
    ## Computing Statistics in bins
    x1_stat = np.array([statfunc(x1[x1_digit==ii]) 
                                 if len(x1[x1_digit==ii])>arr_len 
                                 else failval
                                 for ii in range(1,x1_bins.size)])
    y1_stat = np.array([statfunc(y1[x1_digit==ii])
                                 if len(y1[x1_digit==ii])>arr_len 
                                 else failval
                                 for ii in range(1,x1_bins.size)])
    ## Computing error in the data
    if error=='std':
        stat_err = np.nanstd
        y1_err = np.array([stat_err(y1[x1_digit==ii]) 
                                     if len(y1[x1_digit==ii])>arr_len 
                                     else failval
                                     for ii in range(1,x1_bins.size)])
    if error!='none':
        y1_err = np.array([stat_err(y1[x1_digit==ii])/np.sqrt(len(y1[x1_digit==ii]))
                          if len(y1[x1_digit==ii])>arr_len
                          else failval
                          for ii in range(1,x1_bins.size)])
        if (stat_func==np.median) or (stat_func==np.nanmedian):
            y1_err *= 1.253
    else:
        y1_err = np.zeros(y1.stat.size)
    
    return x1_stat, y1_stat, y1_err

"""
Explanation: You can put all of this into a function that estimates errors and more...
End of explanation
"""


import numpy as np

# Defining arrays
x_arr = np.arange(100)
y_arr = 50*np.random.randn(x_arr.size)

# Computing mean and error in the mean for `x_arr` and `y_arr`
x_stat, y_stat, y_err = Mean_std_calc_one_array(x_arr, y_arr,
                                               statfunc=np.nanmean,
                                               failval=np.nan,
                                               base=10)
x_stat2, y_stat2, y_err2 = Mean_std_calc_one_array(x_arr, y_arr,
                                               statfunc=np.nanmedian,
                                               failval=np.nan,
                                               base=10)

plt.style.use('seaborn-notebook')
plt.clf()
plt.close()
fig = plt.figure(figsize=(12,8))
ax  = fig.add_subplot(111,facecolor='white')
ax.plot(x_arr, y_arr, 'ro', label='Data')
ax.errorbar(x_stat, y_stat, yerr=y_err, color='blue', marker='o',
           linestyle='--',label='Mean')
ax.errorbar(x_stat2, y_stat2, yerr=y_err2, color='green', marker='o',
           linestyle='--',label='Median')
ax.set_xlabel('X axis', fontsize=20)
ax.set_ylabel('Y axis', fontsize=20)
ax.set_title('Data and the Binned Data', fontsize=24)
plt.legend(fontsize=20)
plt.show()

"""
Explanation: Example of using these function:
End of explanation
"""


A = np.array([[n+m*10 for n in range(5)] for m in range(5)])
A

n, m = A.shape

B = A.reshape((1,n*m))
B

A_f = A.flatten()
A_f

C = random.rand(A.size)
C

C.shape

# Stacking the two arrays
D = np.column_stack((A_f,C))
D

# Selecting from 3rd to 11th row
D[2:10]

"""
Explanation: With this function, it is really easy to apply statistics on binned data, as well as to estimate errors on the data.
Reshaping, resizing and stacking arrays
One can always modify the shape of a numpy.ndarray, as well as append it to a pre-existing array.
End of explanation
"""


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

np.concatenate((a,b))

np.concatenate((a,b.T), axis=1)

"""
Explanation: np.concatenate
You can also concadenate different arrays
End of explanation
"""


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

# `B` is now referring to the same array data as `A`
B = A

"""
Explanation: Copy and "Deep Copy"
Sometimes it is important to create new copies of arrays and other objects. For this reason, one uses numpy.copy to create new copies of arrays
End of explanation
"""


B[0,0] = 10
B

A

"""
Explanation: If we make any changes to B, A will also be affected by this change.
End of explanation
"""


B = np.copy(A)

# Modifying `B`
B[0,0] = -5
B

A

"""
Explanation: To get a completely independent, new object, you would use:
End of explanation
"""


import scipy as sc

"""
Explanation: The array A was not affected by this changed. This is important when you're constantly re-defining new arrays
Scipy - Library of Scientific Algorithms for Python
SciPy provides a large number of higher-level scientif algorithms.
It includes:
- Special Functions
- Integration
- Optimization
- Interpolation
- Fourier Transforms
- Signal Processing
- Linear Algebra
- Statistics
- Multi-dimensional image processing
End of explanation
"""


from scipy.interpolate import interp1d

def f(x):
    return np.sin(x)

n = np.arange(0, 10)  
x = np.linspace(0, 9, 100)

y_meas = f(n) + 0.1 * np.random.randn(len(n)) # simulate measurement with noise
y_real = f(x)

linear_interpolation = interp1d(n, y_meas)
y_interp1 = linear_interpolation(x)

cubic_interpolation = interp1d(n, y_meas, kind='cubic')
y_interp2 = cubic_interpolation(x)

fig, ax = plt.subplots(figsize=(15,6))
ax.set_facecolor('white')
ax.plot(n, y_meas, 'bs', label='noisy data')
ax.plot(x, y_real, 'k', lw=2, label='true function')
ax.plot(x, y_interp1, 'r', label='linear interp')
ax.plot(x, y_interp2, 'g', label='cubic interp')
ax.legend(loc=3, prop={'size':20});
ax.tick_params(axis='both', which='major', labelsize=20)
ax.tick_params(axis='both', which='minor', labelsize=15)

"""
Explanation: Interpolation
You can use Scipy to interpolate your data.
You would use the interp1d function to interpolate your function.
End of explanation
"""


Lbox = 250.
Npts = 1000
# Creating cartesian coordinates
x = np.random.uniform(0, Lbox, Npts)
y = np.random.uniform(0, Lbox, Npts)
z = np.random.uniform(0, Lbox, Npts)
sample1 = np.vstack([x, y, z]).T
sample1

sample1.shape

"""
Explanation: KD-Trees
You can also use SciPy to calculate KD-Trees for a set of points
End of explanation
"""


from scipy.spatial import cKDTree

# Initializing KDTree
KD_obj = cKDTree(sample1)

N_neighbours = cKDTree.count_neighbors(KD_obj, KD_obj, 50) - \
                cKDTree.count_neighbors(KD_obj, KD_obj, 30)
print("Number of Neighbours: {0}".format(N_neighbours))

"""
Explanation: Let's say we want to know how many points are within distances of 30 and 50 from other points. To know this, you construct a KD-Tree
End of explanation
"""


k_nearest = 4
dist_k, dist_k_idx = cKDTree.query(KD_obj, sample1, k_nearest)
dist_k

"""
Explanation: Let's say you want to get the distances to the Nth-nearest neighbor.
End of explanation
"""


dist_k_idx

"""
Explanation: You can also get the indices
End of explanation
"""


pairs = KD_obj.query_ball_tree(KD_obj, 30)
pairs[0:10]

"""
Explanation: The first columns corresponds to itself.
You can also find pairs that are separated by at most a distance r
End of explanation
"""


