import numpy as np
import math
class Model:
    def __init__(self,
            shape,
            activation_function,
            loss_function):

        self.shape = shape
        self.act_func = activation_function
        self.loss_func = loss_function

        self.layer = []
        for i in shape:
            self.layer.append( np.zeros(i)
                    .astype(np.longdouble) )
        self.layer = np.array(self.layer, dtype=object)

        self.weight = []
        self.bias = []
        for i in range(len(shape) - 1):
            self.weight.append( (np.random.random_sample(
                (shape[i], shape[i+1]) )*2 -1)
                .astype(np.longdouble) )

            self.bias.append((
                np.random.random_sample((shape[i+1]))*2 -1)
                .astype(np.longdouble))
        self.weight = np.array(self.weight, dtype=object)
        self.bias = np.array(self.bias, dtype=object)

    def print(self): #just a function to print the model's parameters in a better form
        for i in range(self.weight.shape[0]):
            print(self.weight[i], " ", self.bias[i], "\n")
        for i in self.layer:
            print(i)

    def forward(self, input): #(forward propagation)
        self.layer[0] = np.array(input, dtype=np.longdouble)
        self.layer[0] /= np.max(self.layer[0]) + 1e-08 #TODO
        last_index = 0
        for i in range(self.weight.shape[0]):
            last_index = i+1
            self.layer[i+1]= np.dot(self.layer[i], self.weight[i]) + self.bias[i]
            self.layer[i+1]= self.act_func[i](self.layer[i+1])
        return self.layer[last_index]

    def loss(self, input, target): #model built in loss function instead of using otim's loss function
        target = np.array(target, dtype=np.longdouble)
        return self.loss_func( self.forward(input), target, Derivative=False)


# Script to calculate slope using SciPy.

import os
import numpy as np
import scipy.ndimage
from osgeo import gdal
import ospybook as pb

in_fn = r"D:\osgeopy-data\Nepal\everest_utm.tif"
out_fn = r'D:\Temp\everest_slope_scipy2.tif'

# Write a function that calculates slope using a 3x3 window.
# This will be passed to the SciPy filter function below.
def slope(data, cell_width, cell_height):
    """Calculates slope using a 3x3 window.

    data        - 1D array containing the 9 pixel values, starting
                  in the upper left and going left to right and down
    cell_width  - pixel width in the same units as the data
    cell_height - pixel height in the same units as the data
    """
    rise = ((data[6] + (2 * data[7]) + data[8]) -
            (data[0] + (2 * data[1]) + data[2])) / \
           (8 * cell_height)
    run =  ((data[2] + (2 * data[5]) + data[8]) -
            (data[0] + (2 * data[3]) + data[6])) / \
           (8 * cell_width)
    dist = np.sqrt(np.square(rise) + np.square(run))
    return np.arctan(dist) * 180 / np.pi

# Read the data as floating point.
in_ds = gdal.Open(in_fn)
in_band = in_ds.GetRasterBand(1)
in_data = in_band.ReadAsArray().astype(np.float32)

cell_width = in_ds.GetGeoTransform()[1]
cell_height = in_ds.GetGeoTransform()[5]

# Pass your slope function, along with the pixel size, to the
# generic_filter function. This will apply your function to
# each window, and each time it will pass the extra arguments
# in to the function along with the data.
out_data = scipy.ndimage.filters.generic_filter(
    in_data, slope, size=3, mode='nearest',
    extra_arguments=(cell_width, cell_height))

# Save the output
pb.make_raster(in_ds, out_fn, out_data, gdal.GDT_Float32)
del in_ds


import os
import numpy as np
import matplotlib.pyplot as plt

def calc_rotation_matrix_2d(theta):
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)

    a11 = cos_theta
    a12 = -sin_theta
    a21 = sin_theta
    a22 = cos_theta

    return np.array([[a11, a12], [a21, a22]])

def calculate_angle_between_vectors(a, b):
    # Ref
    outer_cross = np.cross(a, b)
    dot_product = np.dot(a, b)
    theta = np.arctan2(outer_cross, dot_product)
    return theta

def calc_rotation_matrix(axis, theta):
        """
        axis: norm 1 rotation of axis vector
        theta: 0 to 2*pie rotation radian
        """
        axis = np.asarray(axis)
        axis = axis / np.linalg.norm(axis)
        n1, n2, n3 = axis

        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)

        a11 = cos_theta + n1*n1*(1 - cos_theta)
        a12 = n1*n2*(1 - cos_theta) - n3*sin_theta
        a13 = n1*n3*(1 - cos_theta) + n2*sin_theta
        a21 = n2*n1 * (1 - cos_theta) + n3*sin_theta
        a22 = cos_theta + n2*n2 * (1 - cos_theta)
        a23 = n2*n3*(1 - cos_theta) - n1*sin_theta
        a31 = n3*n1*(1 - cos_theta) - n2*sin_theta
        a32 = n3*n2*(1 - cos_theta) + n1*sin_theta
        a33 = cos_theta + n3*n3*(1 - cos_theta)

        return np.array([[a11, a12, a13], [a21, a22, a23], [a31, a32, a33]])

def rotate_vector_by_degree(axis, vector, degree):
    # this only works on 3D vector
    theta = np.deg2rad(degree)
    rot_mat = calc_rotation_matrix(axis, theta)
    translated_vector = np.matmul(rot_mat, vector)
    return translated_vector

def transform_axis_pointcloud(pointcloud, x_axis, y_axis, z_axis):

    mat = np.array([[x_axis[0], y_axis[0], z_axis[0]],
                    [x_axis[1], y_axis[1], z_axis[1]],
                    [x_axis[2], y_axis[2], z_axis[2]]])

    mat_inv = np.linalg.pinv(mat)
    pointcloud_transformed = np.matmul(mat_inv, pointcloud.T).T
    return pointcloud_transformed

def calculate_radian(vector1, vector2):
    dot = np.dot(vector1, vector2)
    norm = np.linalg.norm(vector1) * np.linalg.norm(vector2)
    cos = dot / norm
    return np.arccos(np.clip(cos, -1.0, 1.0))

def visualize_3D_pointcloud_2D(x, y, skeleton_points, figname, save_dir=None, xlim=None, ylim=None):
    fig, ax = plt.subplots()
    ax.scatter(x, y, color="black", s=1)
    if skeleton_points is not None:
        ax.scatter(skeleton_points[:, 0], skeleton_points[:, 1], color="blue", s=15)

    ax.set_aspect("equal", "datalim")
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    plt.title(figname)
    if save_dir:
        img_name = figname + ".png"
        plt.savefig(os.path.join(save_dir, img_name))
    else:
        plt.show()
    plt.close()


import matplotlib.pyplot as plt
from math import sqrt
from scipy.stats import norm
import numpy as np

def ud_function(a, b, x):
    return (x - a) / (b - a) if a <= x < b else 0 if x < a else 1

def ud_density(a, b, x):
    return 1 / (b - a) if a <= x <= b else 0

def norm_function(x, mu, sigma):
    return norm.cdf(x, mu, sqrt(sigma))

def norm_density(x, mu, sigma):
    return norm.pdf(x, mu, sqrt(sigma))

def draw_graphics(x, y_function, y_density, name):
    fig, axs = plt.subplots(2, figsize=(6, 7))

    fig.suptitle(name)
    axs[0].plot(x, y_function, color='purple')
    axs[1].plot(x, y_density, color='purple')

    axs[0].set_xlabel('x')
    axs[0].set_ylabel('F(x)')

    axs[1].set_xlabel('x')
    axs[1].set_ylabel('f(x)')

    axs[0].grid(True)
    axs[1].grid(True)
    plt.show()

def main():
    a = float(input("Input a: "))
    b = float(input("Input b: "))
    delta = b - a
    x = np.arange(a - delta / 2, b + delta / 2, 0.001)
    y_function = [ud_function(a, b, _x) for _x in x]
    y_density = [ud_density(a, b, _x) for _x in x]
    draw_graphics(x, y_function, y_density, 'Равномерное распределение')

    mu = float(input("Input mu: "))
    sigma = float(input("Input sigma: "))
    x = np.arange(-10, 10, 0.001)
    y_function = norm_function(x, mu, sigma)
    y_density = norm_density(x, mu, sigma)

    draw_graphics(x, y_function, y_density, 'Нормальное распределение')

if __name__ == '__main__':
    main()


# --------------
# Code starts here

import numpy as np
from scipy.linalg import norm,eig

# Code starts here

# Adjacency matrix
adj_mat = np.array([[0,0,0,0,0,0,1/3,0],
                   [1/2,0,1/2,1/3,0,0,0,0],
                   [1/2,0,0,0,0,0,0,0],
                   [0,1,0,0,0,0,0,0],
                  [0,0,1/2,1/3,0,0,1/3,0],
                   [0,0,0,1/3,1/3,0,0,1/2],
                   [0,0,0,0,1/3,0,0,1/2],
                   [0,0,0,0,1/3,1,1/3,0]])

# Compute eigenvalues and eigencevectrs
eigenvalues,eigenvectors=eig(adj_mat)

# Eigen vector corresponding to 1
eigen_1=abs(eigenvectors[:,0])/norm(eigenvectors[:,0],1)

# most important page
page=np.where(eigen_1==np.amax(eigen_1))[0][0]+1
print(page)

# Code ends here

# --------------
# Code starts here

# Initialize stationary vector I
init_I=np.array([1,0,0,0,0,0,0,0])

# Perform iterations for power method
for i in range(10):
    init_I=np.dot(adj_mat,init_I)/norm(init_I,1)

power_page=np.where(init_I==np.amax(init_I))[0][0]+1
print(power_page)
# Code ends here

# --------------
# Code starts here

# New Adjancency matrix
# New Adjancency matrix
new_adj_mat = np.array([[0,0,0,0,0,0,0,0],
                   [1/2,0,1/2,1/3,0,0,0,0],
                  [1/2,0,0,0,0,0,0,0],
                   [0,1,0,0,0,0,0,0],
                   [0,0,1/2,1/3,0,0,1/2,0],
                   [0,0,0,1/3,1/3,0,0,1/2],
                   [0,0,0,0,1/3,0,0,1/2],
                   [0,0,0,0,1/3,1,1/2,0]])

# Initialize stationary vector I
new_init_I=np.array([1,0,0,0,0,0,0,0])

# Perform iterations for power method
for i in range(10):
    new_init_I=np.dot(new_adj_mat,new_init_I)/norm(new_init_I,1)

print(new_init_I)

# Code ends here

# --------------
# Alpha value
alpha = 0.85

# Code starts here

# Modified adjancency matrix
G=alpha*new_adj_mat + (1-alpha)*(1/len(new_adj_mat))*np.ones(new_adj_mat.shape)
# print(G)

# Initialize stationary vector I
final_init_I=np.array([1,0,0,0,0,0,0,0])

# Perform iterations for power method
for i in range(1000):
    final_init_I=np.dot(G,final_init_I)/norm(final_init_I,1)

print(final_init_I)

# Code ends here


import numpy as np
import math
import random
from itertools import repeat

def randomKmeans(samples, k, TOL, C):
	SS_Previous = 0
	samplesLength = len(samples)
	dim = len(samples[0])

	while 1:
		P = [[] for i in repeat(None, k)]
		# find the closest cluster centroid
		for i in range(samplesLength):
			minIdx = 0
			minVal = eucliDist(samples[i], C[0], dim)
			for j in range(1,k):
				dist = eucliDist(samples[i], C[j], dim)
				if (dist < minVal):
					minIdx = j
					minVal = dist
			# assign the point to the correct cluster
			P[minIdx].append(samples[i])

		# recalculating cluster centroid
		for j in range(k):
			coords = P[j]
			if(len(coords) == 0):
				coords = random.sample(samples, 1)
			zipped = zip(*coords)
			num = len(coords)
			C[j] = [math.fsum(dList)/num for dList in zipped]

		SS_Error = 0

		# calculating total clustering error
		for idx in range(k):
			for p_idx in range(len(P[idx])):
				SS_Error += sqEucliDist(P[idx][p_idx], C[idx], dim)

		# check if no change in SSE
		delta = abs(SS_Error - SS_Previous)

		if (delta < TOL):
			break

		SS_Previous = SS_Error

	return round(SS_Error, 4), C, P

# calculating Euclidean Distance
def eucliDist(sample, center, dim):
	distance = 0
	for x in range(dim):
		distance += pow((sample[x] - center[x]), 2)
	return pow(distance,0.5)

# calculating Squared Euclidean Distance
def sqEucliDist(sample, center, dim):
	sumsq = 0
	for x in range(dim):
		sumsq += pow((sample[x] - center[x]), 2)
	return sumsq


from distanceclosure.utils import prox2dist, dist2prox
from distanceclosure.utils import dict2matrix, matrix2dict, dict2sparse
import numpy as np
from scipy.sparse import csr_matrix

P = P_true = np.array([
		[1.,.9,.1,0.],
		[.9,1.,.8,0.],
		[.1,.8,1.,.6],
		[0.,0.,.6,1.],
		], dtype=float)

D = D_true = np.array([
		[0.,.11111111,9.,np.inf],
		[.11111111,0.,0.25,np.inf],
		[9.,0.25,0.,0.66666667],
		[np.inf,np.inf,0.66666667,0.],
		], dtype=float)

#
# Test Distance Proximity Conversion
#
def test_prox2dist():
	""" Test Utils: Prox2Dist """
	assert np.isclose(dist2prox(D), P_true).all()

def test_dist2prox():
	""" Test Utils: Dist2Prox """
	assert np.isclose(prox2dist(P), D_true).all()

def test_dist2prox_prox2dist():
	""" Test Utils: Prox2Dist & Dist2Prox """
	assert np.isclose(dist2prox(prox2dist(P)) , P).all()
#
# Test Data Conversion
#
def test_matrix2dict():
	""" Test Utils: matrix 2 dict """
	m = [[0, 1, 3], [1, 0, 2], [3, 2, 0]]
	d = matrix2dict(m)
	assert (d == {0: {0: 0, 1: 1, 2:3}, 1: {0: 1, 1: 0, 2:2}, 2: {0: 3, 1:2, 2:0}} )

def test_dict2matrix():
	""" test Utils: dict 2 matrix """
	d = {0: {0: 0, 1: 1, 2:3}, 1: {0: 1, 1: 0, 2:2}, 2: {0: 3, 1:2, 2:0}}
	m = dict2matrix(d)
	assert (m == np.array([[0, 1, 3], [1, 0, 2], [3, 2, 0]]) ).all()

def test_dict2sparse():
	""" Test Utils: dict 2 sparse """
	d = {0: {0: 0, 1: 1, 2:3}, 1: {0: 1, 1: 0, 2:2}, 2: {0: 3, 1:2, 2:0}}
	s = dict2sparse(d)
	t = csr_matrix(np.array([[0, 1, 3], [1, 0, 2], [3, 2, 0]]))
	assert np.allclose(s.A, t.A)


import superimport

import pyprobml_utils as pml
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import scipy
#import scipy.special as sc
from scipy import stats
from sklearn.neighbors import KernelDensity
import math

S = np.array([[3.1653, -0.0262], [-0.0262, 0.6477]])
dof = 3
Sigma = S
nr = 3
nc = 3
nsamples = 9
Xs1 = np.linspace(0.1, 200, 2000)
Xs2 = np.linspace(0.1, 40, 400)
Xs3 = np.linspace(0.1, 10, 100)
Xsa = {0:Xs1, 1:Xs2}
Xsb = {0:Xs2, 1:Xs3}

def wishart_sample(dof, sigma, nsamples):
  d = np.size(sigma, 0)
  C = np.linalg.cholesky(Sigma)
  S = np.zeros((d, d, nsamples))
  for i in range(nsamples):
      t = np.random.randn(dof, d)
      Z = np.matmul(t, C)
      S[:, :, i] = np.matmul(Z.T, Z)
  return S

def gamma_log_prob(a, b, X):
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.gamma.html
    # a=shape, b=rate=1/scale
  logZ = scipy.special.gammaln(a) - np.multiply(a, np.log(b))
  logp = np.multiply((a-1), np.log(X)) - np.multiply(b, X) - logZ
  return logp

def get_cov_ellipse(cov, centre, nstd):
    eigvals, eigvecs = np.linalg.eigh(cov)

    order = eigvals.argsort()[::-1]
    eigvals, eigvecs = eigvals[order], eigvecs[:, order]

    vx, vy = eigvecs[:,0][0], eigvecs[:,0][1]
    theta = np.arctan2(vy, vx)
    theta = math.degrees(theta)
    width, height = 2 * nstd * np.sqrt(eigvals)

    return width, height, theta

np.random.seed(4)
M = dof*Sigma
R = np.corrcoef(M)
S = wishart_sample(dof, Sigma, nsamples)

# Plots of some samples from Wishart distribution:
fig, ax = plt.subplots(nr, nc, figsize=(10,10))
mu = np.array([0, 0])
j = 0
for r in range(nr):
    for c in range(nc):
      w, h, theta = get_cov_ellipse(S[:, :, j], mu, nstd=3)
      j = j + 1
      ell = mpl.patches.Ellipse(xy=mu, width=w, height=h, angle = theta, ec='black', fc='none') #facecolor='none'
      ax[r, c].add_patch(ell)
      ax[r, c].set_aspect('equal')
      ax[r, c].set_xlim(-25, 25)
      ax[r, c].set_ylim(-25, 25)
      ax[r, c].autoscale()
      ax[r, c].plot(mu[0], mu[1], marker='x')
plt.tight_layout()
fig.suptitle("Wi(dof=%i, S), E=[9.5, -0.1; -0.1, 1.9], rho=-0.018" %dof)
xmin = -17
xmax = 17
ymin = -10
ymax = 10
custom_xlim = (xmin, xmax)
custom_ylim = (ymin, ymax)
plt.setp(ax, xlim=custom_xlim, ylim=custom_ylim)
pml.savefig('wishart_samples.pdf')
plt.show()

marg1a = dof / 2
marg1b = 1/(2*Sigma[0, 0])
marg2a = dof / 2
marg2b = 1/(2*Sigma[1, 1])

#Plots of marginals
logp = gamma_log_prob(marg1a, marg1b, Xs1)
expo = np.exp(logp)
plt.figure()
plt.plot(Xs1, expo)
plt.title("Marginal sigma1_squared ")
pml.savefig('wishart_sigma1.pdf')
plt.show()

logp = gamma_log_prob(marg2a, marg2b, Xs2)
expo = np.exp(logp)
plt.figure()
plt.plot(Xs2, expo)
plt.title("Marginal sigma2_squared ")
pml.savefig('wishart_sigma2.pdf')
plt.show()

# Plot of correlation coefficient
n = 1000
Rs = wishart_sample(dof, Sigma, nsamples)
for s in range(nsamples):
        Rs[:, :, s] = np.corrcoef(Rs[:, :, s])    #cov2cor(Rs(:, :, s));
data = np.squeeze(Rs[0, 1, :])
kde = KernelDensity(bandwidth=1.0, kernel='gaussian')
kde.fit(data.reshape(-1, 1))
x = np.linspace(data.min()-2, data.max()+2, 100)
logprob = kde.score_samples(x[:, None])
plt.figure()
plt.title('Rho')
plt.plot(x, np.exp(logprob))
pml.savefig('wishart_rho.pdf')
plt.show()


import numpy as np

def integer_sequences(L, S, nondecr=False, m=None, M=None):
    """
    Generate sequences of non-negative integers.

    Parameters:
    L:         the length of the sequences
    S:         the sum of the integers in each sequence

    Optional parameters:
    nondecr:   (boolean) return only non-decreasing sequences (default: False)
    m:         tuple of length L; gives lower bounds for coefficients of list (default: None)
    M:         tuple of length L; gives upper bounds for coefficients of list (default: None)
    """

    # If M and m are not given, use the following defaults.
    if M is None:
        M = (S,)*L
    if m is None:
        m = (0,)*L

    # If length=0 and sum=0 then yield the empty tuple.
    # Otherwise, yield nothing.
    if L==0:
        if S==0:
            yield tuple()
    # If length=1 and sum lies within given boundaries, yield 1-tuple.
    elif L==1:
        if m[0]<=S<=M[0]:
            yield (S,)
    # If length>1, then loop through possible values for first coefficient,
    # and recursively generate (L-1)-tuples that give the remaining coefficients.
    elif L>1:
        for first in range(m[0], min(S,M[0])+1):
            m_next = m[1:] if nondecr==False else (first,)+m[2:]
            for tail in integer_sequences(L=L-1, S=S-first, nondecr=nondecr, m=m_next, M=M[1:]):
                yield (first,)+tail

def matrices_from_degree_sequence(deg_seq=tuple()):
    """
    Generate all matrices that can occur as adjacency matrices
    of r-marked graphs whose degree_sequence() equals deg_seq.

    These are symmetric (r×r)-matrices with non-negative integer coefficients
    with even coefficients on the diagonal, whose row sum equals deg_seq.
    """
    L = len(deg_seq)

    M = np.zeros((L,L), dtype=int)
    if L==0:
        yield M
    elif L>0:
        # Range over possible values for top left entries:
        for top_left in range(0,deg_seq[0]+1,2):
            M[0,0] = top_left
            # Range over sequences that make up the rest of the first row:
            for top_row_rest in integer_sequences(L=L-1, S=deg_seq[0]-top_left, M=deg_seq[1:]):
                M[0,1:] = M[1:,0] = top_row_rest
                # Compute the row sum of the remaining bottom right square
                row_sum_rem = tuple(deg_seq[i] - M[0,i] for i in range(1,L))

                # Loop over all possible bottom right squares:
                for BRS in matrices_from_degree_sequence(deg_seq=row_sum_rem):
                    M[1:,1:]=BRS
                    yield M.copy()

def increasing_injections(r,s,m=1):
    """
    Generates all increasing tuples of length r
    with values in {1, ..., s}.

    Optional argument:
    m: minimum value of first entry (default=1, used for recursion)
    """
    if r==0:
        yield tuple()
    elif r==1:
        for i in range(m,s+1):
            yield (i,)
    else:
        for first in range(m,s-r+2):
            for tail in increasing_injections(r-1,s,m=first+1):
                yield (first,)+tail


import numpy as np
import pandas as pd
import scipy.stats
def drawdown(return_series: pd.Series):
    """
    Takes a Time Series of asset returns and computes and returns a dataframe that contains
    1) Wealth Idex
    2) Previous Peaks
    3) Percentage of drawdowns
    """
    wealth_index = 1000*(1+return_series).cumprod()
    previous_peaks = wealth_index.cummax()
    drawdowns = (wealth_index-previous_peaks)/previous_peaks
    return pd.DataFrame(dict(Wealth = wealth_index, Peaks = previous_peaks, Drawdowns = drawdowns))

def get_ffme_returns():
    """
    Load the Farma-French Dataset for the returns of the top and bottom deciles by Market Cap
    """
    me_m = pd.read_csv('Portfolios_Formed_on_ME_monthly_EW.csv', header=0, index_col=0, na_values = -99.99)
    rets = me_m[['Lo 10', 'Hi 10']]
    rets.columns = ['SmallCap', 'LargeCap']
    rets = rets/100
    rets.index = pd.to_datetime(rets.index, format="%Y%m").to_period('M')
    return rets

def get_hfi_returns():
    """
    Load the Edhec Hedge Fund Dataset for the returns of the top and bottom deciles by Market Cap
    """
    hfi = pd.read_csv('edhec-hedgefundindices.csv', header=0, index_col=0, parse_dates=True)
    hfi = hfi/100
    hfi.index = hfi.index.to_period('M')
    return hfi

def semideviation(r):
    """
    Returns the semideviation aka negative semideviation of r
    r must be a series or a dataframe
    """
    is_negative = r<0  #produced the mask
    return r[is_negative].std(ddof=0)

"""
semideviation can also be done in a more efficient way
def semideviation(r):
    return r[r<0].std(ddof=0)
"""

def skewness(r):
    """
    Alternative method to scipy.stats.skew()
    Computes the skewness of the supplied Series or Dataframe
    Retruns a float or a series
    """
    demeaned_r = r - r.mean()
    #use the population standard deviation so set dof=0
    sigma_r = r.std(ddof=0)
    exp = (demeaned_r**3).mean()
    return exp/(sigma_r**3)

def kurtosis(r):
    """
    Alternative method to scipy.stats.kurtosis()
    Computes the kurtosis of the supplied Series or Dataframe
    Retruns a float or a series
    """
    demeaned_r = r - r.mean()
    #use the population standard deviation so set dof=0
    sigma_r = r.std(ddof=0)
    exp = (demeaned_r**4).mean()
    return exp/(sigma_r**4)

def is_normal(r, level = 0.01):
    """
    Applies the Jarque-Bera test to determine if a series is normal or not
    Test is applied at the 1% level by default
    Returns true if the hypothesis of normality is accepted otherwise returns False
    """
    statistic, p_value = scipy.stats.jarque_bera(r)
    return p_value>level

def var_historic(r, level=5):
    """
    Returns the historic value at risk at a specified level
    i.e returns the number such that the 'level' percent of the returns
    fall below that number, and the (100-level) percent are above
    """
    if isinstance(r, pd.DataFrame):
        return r.aggregate(var_historic, level=level) #calls this function for every column if the series is a dataframe
    elif isinstance(r, pd.Series):
        return -np.percentile(r, level)
    else:
        raise TypeError('Expected r to be a series or a dataframe')

from scipy.stats import norm
def var_gaussian(r, level=5, modified=False):
    """
    Returns the parametric Gaussian VaR of a series or dataframe
    """
    #compute the z score
    z = norm.ppf(level/100)
    if modified:
        s = skewness(r)
        k= kurtosis(r)
        z = (z+
        (z**2 - 1)*s/6 +
        (z**3 - 3*z)*(k-3)/24 -
        (2*z**3 - 5*z)*(s**2)/36)

    return -(r.mean() + z*r.std(ddof=0))

def cvar_historic(r, level=5):
    """
    Computes the conditional VaR of series or dataframe
    """
    if isinstance(r, pd.Series):
        is_beyond = r <= -var_historic(r, level=level)
        return -r[is_beyond].mean()
    elif isinstance(r, pd.DataFrame):
        return r.aggregate(cvar_historic, level=level)
    else:
        raise TypeError('Expected r to be series or dataframe')


# coding: utf-8

# ### DEMOPT04
# # Maximization of banana function by various methods

# $$f(x,y)=-100*(y-x^2)^2-(1-x)^2$$
# starting at \[0,1\].
#

# In[1]:

import numpy as np
from compecon import OP
np.set_printoptions(4, linewidth=120, suppress=True)
import matplotlib.pyplot as plt

from compecon.demos import demo

# In[2]:

''' Set up the problem '''
x0 = [1, 0]
banana = OP(lambda x: -100 * (x[1] - x[0] ** 2)**2 - (1 - x[0]) ** 2,
            x0, maxit=250, print=True, all_x=True)

# In[3]:

x = banana.qnewton()
J = banana.jacobian(x)
E = np.linalg.eig(banana.hessian(x))[0]
print('x = ', x, '\nJ = ', J, '\nE = ', E)

# In[4]:

''' Plots options '''
steps_options = {'marker': 'o',
                 'color': (0.2, 0.2, .81),
                 'linewidth': 1.0,
                 'markersize': 3,
                 'markerfacecolor': 'white',
                 'markeredgecolor': 'red'}

contour_options = {'levels': -np.exp(np.arange(7,0.25,-0.5)),
                   'colors': '0.25',
                   'linewidths': 0.5}

# In[5]:

''' Data for coutours '''
n = 40  # number of grid points for plot per dimension
xmin = [-0.7, -0.2]
xmax = [ 1.2,  1.2]

X0, X1 = np.meshgrid(*[np.linspace(a, b, n) for a, b in zip(xmin, xmax)])
Y = banana.f([X0.flatten(), X1.flatten()])
Y.shape = (n, n)

fig = plt.figure(figsize=[12,4])

for it, method in enumerate(banana.search_methods.keys()):
    ''' Solve problem with given method '''
    print('\n\nMaximization with method ' + method.upper())
    x = banana.qnewton(SearchMeth=method)
    print('x =', x)

    ''' Plot the result '''
    demo.subplot(1, 3, it + 1, method.upper(),'x','y')

    plt.contour(X0, X1, Y, **contour_options)
    plt.plot(*banana.x_sequence, **steps_options)
    plt.plot(1, 1, 'r*', markersize=15)
    plt.title(method.upper() + " search")
    plt.xlabel('x', verticalalignment='top')
    plt.ylabel('y', verticalalignment= 'bottom')
    plt.axis((xmin[0], xmax[0], xmin[1], xmax[1]))

plt.show()

# ## Using Scipy
#
# As of this version of CompEcon, the Nelder Mead method has not been implemented. However, we can still use it with the help of the **scipy.optimize.minimize** function. To this end, we must rewrite the banana function (change its sign) so that we switch from our original maximization problem to one of minimization.
#
#
#

# In[6]:

from scipy.optimize import minimize

# In[7]:

x0 = [1, 0]
def banana2(x):
    return 100 * (x[1] - x[0] ** 2)**2 + (1 - x[0]) ** 2

# In[8]:

res = minimize(banana2, x0, method='Nelder-Mead')
print(res)

demo.savefig([fig])


"""

Path Planning with Bezier curve.

author: Atsushi Sakai(@Atsushi_twi)

"""

import scipy.special
import numpy as np
import matplotlib.pyplot as plt

show_animation = True

def calc_4points_bezier_path(sx, sy, syaw, ex, ey, eyaw, offset):
    """
    Compute control points and path given start and end position.

    :param sx: (float) x-coordinate of the starting point
    :param sy: (float) y-coordinate of the starting point
    :param syaw: (float) yaw angle at start
    :param ex: (float) x-coordinate of the ending point
    :param ey: (float) y-coordinate of the ending point
    :param eyaw: (float) yaw angle at the end
    :param offset: (float)
    :return: (numpy array, numpy array)
    """
    dist = np.sqrt((sx - ex) ** 2 + (sy - ey) ** 2) / offset
    control_points = np.array(
        [[sx, sy],
         [sx + dist * np.cos(syaw), sy + dist * np.sin(syaw)],
         [ex - dist * np.cos(eyaw), ey - dist * np.sin(eyaw)],
         [ex, ey]])

    path = calc_bezier_path(control_points, n_points=100)

    return path, control_points

def calc_bezier_path(control_points, n_points=100):
    """
    Compute bezier path (trajectory) given control points.

    :param control_points: (numpy array)
    :param n_points: (int) number of points in the trajectory
    :return: (numpy array)
    """
    traj = []
    for t in np.linspace(0, 1, n_points):
        traj.append(bezier(t, control_points))

    return np.array(traj)

def bernstein_poly(n, i, t):
    """
    Bernstein polynom.

    :param n: (int) polynom degree
    :param i: (int)
    :param t: (float)
    :return: (float)
    """
    return scipy.special.comb(n, i) * t ** i * (1 - t) ** (n - i)

def bezier(t, control_points):
    """
    Return one point on the bezier curve.

    :param t: (float) number in [0, 1]
    :param control_points: (numpy array)
    :return: (numpy array) Coordinates of the point
    """
    n = len(control_points) - 1
    return np.sum([bernstein_poly(n, i, t) * control_points[i] for i in range(n + 1)], axis=0)

def bezier_derivatives_control_points(control_points, n_derivatives):
    """
    Compute control points of the successive derivatives of a given bezier curve.

    A derivative of a bezier curve is a bezier curve.
    See https://pomax.github.io/bezierinfo/#derivatives
    for detailed explanations

    :param control_points: (numpy array)
    :param n_derivatives: (int)
    e.g., n_derivatives=2 -> compute control points for first and second derivatives
    :return: ([numpy array])
    """
    w = {0: control_points}
    for i in range(n_derivatives):
        n = len(w[i])
        w[i + 1] = np.array([(n - 1) * (w[i][j + 1] - w[i][j])
                             for j in range(n - 1)])
    return w

def curvature(dx, dy, ddx, ddy):
    """
    Compute curvature at one point given first and second derivatives.

    :param dx: (float) First derivative along x axis
    :param dy: (float)
    :param ddx: (float) Second derivative along x axis
    :param ddy: (float)
    :return: (float)
    """
    return (dx * ddy - dy * ddx) / (dx ** 2 + dy ** 2) ** (3 / 2)

def plot_arrow(x, y, yaw, length=1.0, width=0.5, fc="r", ec="k"):  # pragma: no cover
    """Plot arrow."""
    if not isinstance(x, float):
        for (ix, iy, iyaw) in zip(x, y, yaw):
            plot_arrow(ix, iy, iyaw)
    else:
        plt.arrow(x, y, length * np.cos(yaw), length * np.sin(yaw),
                  fc=fc, ec=ec, head_width=width, head_length=width)
        plt.plot(x, y)

def main():
    """Plot an example bezier curve."""
    start_x = 10.0  # [m]
    start_y = 1.0  # [m]
    start_yaw = np.radians(180.0)  # [rad]

    end_x = -0.0  # [m]
    end_y = -3.0  # [m]
    end_yaw = np.radians(-45.0)  # [rad]
    offset = 3.0

    path, control_points = calc_4points_bezier_path(
        start_x, start_y, start_yaw, end_x, end_y, end_yaw, offset)

    # Note: alternatively, instead of specifying start and end position
    # you can directly define n control points and compute the path:
    # control_points = np.array([[5., 1.], [-2.78, 1.], [-11.5, -4.5], [-6., -8.]])
    # path = calc_bezier_path(control_points, n_points=100)

    # Display the tangent, normal and radius of cruvature at a given point
    t = 0.86  # Number in [0, 1]
    x_target, y_target = bezier(t, control_points)
    derivatives_cp = bezier_derivatives_control_points(control_points, 2)
    point = bezier(t, control_points)
    dt = bezier(t, derivatives_cp[1])
    ddt = bezier(t, derivatives_cp[2])
    # Radius of curvature
    radius = 1 / curvature(dt[0], dt[1], ddt[0], ddt[1])
    # Normalize derivative
    dt /= np.linalg.norm(dt, 2)
    tangent = np.array([point, point + dt])
    normal = np.array([point, point + [- dt[1], dt[0]]])
    curvature_center = point + np.array([- dt[1], dt[0]]) * radius
    circle = plt.Circle(tuple(curvature_center), radius,
                        color=(0, 0.8, 0.8), fill=False, linewidth=1)

    assert path.T[0][0] == start_x, "path is invalid"
    assert path.T[1][0] == start_y, "path is invalid"
    assert path.T[0][-1] == end_x, "path is invalid"
    assert path.T[1][-1] == end_y, "path is invalid"

    if show_animation:  # pragma: no cover
        fig, ax = plt.subplots()
        ax.plot(path.T[0], path.T[1], label="Bezier Path")
        ax.plot(control_points.T[0], control_points.T[1],
                '--o', label="Control Points")
        ax.plot(x_target, y_target)
        ax.plot(tangent[:, 0], tangent[:, 1], label="Tangent")
        ax.plot(normal[:, 0], normal[:, 1], label="Normal")
        ax.add_artist(circle)
        plot_arrow(start_x, start_y, start_yaw)
        plot_arrow(end_x, end_y, end_yaw)
        ax.legend()
        ax.axis("equal")
        ax.grid(True)
        plt.show()

def main2():
    """Show the effect of the offset."""
    start_x = 10.0  # [m]
    start_y = 1.0  # [m]
    start_yaw = np.radians(180.0)  # [rad]

    end_x = -0.0  # [m]
    end_y = -3.0  # [m]
    end_yaw = np.radians(-45.0)  # [rad]

    for offset in np.arange(1.0, 5.0, 1.0):
        path, control_points = calc_4points_bezier_path(
            start_x, start_y, start_yaw, end_x, end_y, end_yaw, offset)
        assert path.T[0][0] == start_x, "path is invalid"
        assert path.T[1][0] == start_y, "path is invalid"
        assert path.T[0][-1] == end_x, "path is invalid"
        assert path.T[1][-1] == end_y, "path is invalid"

        if show_animation:  # pragma: no cover
            plt.plot(path.T[0], path.T[1], label="Offset=" + str(offset))

    if show_animation:  # pragma: no cover
        plot_arrow(start_x, start_y, start_yaw)
        plot_arrow(end_x, end_y, end_yaw)
        plt.legend()
        plt.axis("equal")
        plt.grid(True)
        plt.show()

if __name__ == '__main__':
    main()
    #  main2()


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# simple moving average
def sma(df, periods = 20):
    df['SMA'] = df['Close'].rolling(window = periods).mean()
    return df

# exponential moving average
def ema(df, periods = 20):
    k =2/(periods + 1)
    df.loc[df.index[0], 'EMA'] = df.loc[df.index[0], 'Close']
    for index in df[:-1].index:
        df.loc[index+1, 'EMA'] = k*df.loc[index+1, 'Close']+(1-k)*df.loc[index, 'EMA']
    return df


# rayleigh distribution used in signal processing.
# it has two parameters.
# scale - (standard deviation) decides how flat the distribution will be (default 1.0).
# size - shape of returned array.
from numpy import random
import matplotlib.pyplot as plt
import seaborn as sns

arr1 = random.rayleigh(scale=2, size=10)
print(arr1)
sns.distplot(arr1, hist=False)
plt.show()
# at unit standard deviation of rayleigh and df=2 in chi square represents same distribution

Python/Numpy/numpy 28 - rayleigh distribution.py


import numpy as np

def normalize(x: np.ndarray) -> np.ndarray:
    assert x.ndim == 1, 'x must be a vector (ndim: 1)'
    return x / np.linalg.norm(x)

def look_at(
    eye,
    target,
    up,
) -> np.ndarray:
    """Returns transformation matrix with eye, at and up.
    Parameters
    ----------
    eye: (3,) float
        Camera position.
    target: (3,) float
        Camera look_at position.
    up: (3,) float
        Vector that defines y-axis of camera (z-axis is vector from eye to at).
    Returns
    -------
    T_cam2world: (4, 4) float (if return_homography is True)
        Homography transformation matrix from camera to world.
        Points are transformed like below:
            # x: camera coordinate, y: world coordinate
            y = trimesh.transforms.transform_points(x, T_cam2world)
            x = trimesh.transforms.transform_points(
                y, np.linalg.inv(T_cam2world)
            )
    """
    eye = np.asarray(eye, dtype=float)

    if target is None:
        target = np.array([0, 0, 0], dtype=float)
    else:
        target = np.asarray(target, dtype=float)

    if up is None:
        up = np.array([0, 0, -1], dtype=float)
    else:
        up = np.asarray(up, dtype=float)

    assert eye.shape == (3,), 'eye must be (3,) float'
    assert target.shape == (3,), 'target must be (3,) float'
    assert up.shape == (3,), 'up must be (3,) float'

    # create new axes
    z_axis = normalize(target - eye)
    x_axis = normalize(np.cross(up, z_axis))
    y_axis = normalize(np.cross(z_axis, x_axis))

    # create rotation matrix: [bs, 3, 3]
    R = np.vstack((x_axis, y_axis, z_axis))
    t = eye

    T_cam2world = np.zeros([4,4])
    T_cam2world[:3,:3] = R.T
    T_cam2world[:3, 3] = t
    T_cam2world[3,3] = 1.
    return T_cam2world


"""
Given an integer n, your task is to count how many strings of length n can be formed under the following rules:
    Each character is a lower case vowel ('a', 'e', 'i', 'o', 'u')
    Each vowel 'a' may only be followed by an 'e'.
    Each vowel 'e' may only be followed by an 'a' or an 'i'.
    Each vowel 'i' may not be followed by another 'i'.
    Each vowel 'o' may only be followed by an 'i' or a 'u'.
    Each vowel 'u' may only be followed by an 'a'.
Since the answer may be too large, return it modulo 10^9 + 7.
Example:
        Input - n = 1
        Output - 5
        Explanation - All possible strings are: "a", "e", "i" , "o" and "u".

        Input - n = 2
        Output - 10

Constraints:
        Time complexity - O(N) [ o(5N) ]
        Space complexity - O(N) [ o(5N) ]
        1 <= n <= 2 * 10^4

Hints:
      Use dynamic programming.
      Let dp[i][j] be the number of strings of length i that ends with the j-th vowel.
      Deduce the recurrence from the given relations between vowels.
"""

import numpy as np
class Solution:
    def countVowelPermutation(self, n: int) -> int:
        dp = np.zeros((n, 5), dtype="Int64")
        rules = {0: [1], 1: [0, 2], 2: [0, 1, 3, 4], 3: [2, 4], 4: [0]}

        for i in range(n):
            for j in range(5):
                if i == 0:
                    dp[i][j] = 1
                    continue
                for k in rules[j]:
                    dp[i][k] = (dp[i - 1][j]%1000000007 + dp[i][k]%1000000007)%1000000007

        return sum(dp[n - 1])%1000000007


# Question 1 Lab 04
# AB Satyaprakash 180123062

# imports ----------------------------------------------------------------------
from math import sqrt
import matplotlib.pyplot as plt
import numpy as np

# ------------------------------------------------------------------------------
# (a) Plotting the Markowitz efficient frontier and
# (b) Tabulating Return, Risk and Weights for 10 different values on the efficient frontier
# Given Mean vector and Covariance matrix
m = np.array([0.1, 0.2, 0.15])
C = np.array([[0.005, -0.010, 0.004],
              [-0.010, 0.040, -0.002],
              [0.004, -0.002, 0.023]])
u = np.array([1, 1, 1])

m_trans = m.transpose()
C_trans = C.transpose()
u_trans = u.transpose()
C_inv = np.linalg.inv(C)

ret, sig = np.zeros(100), np.zeros(100)
print(' Return - Risk   -   W1    -    W2    -    W3\n')

# Finding and tabulating the values
for i in range(1, 101):
    ret[i-1] = i*0.005
    y = np.array([[ret[i-1]], [1]])
    M = np.array([[m.dot(C_inv).dot(m_trans), u.dot(C_inv).dot(m_trans)],
                  [m.dot(C_inv).dot(u_trans), u.dot(C_inv).dot(u_trans)]])
    lam = 2*(np.linalg.inv(M).dot(y))
    w = (lam[0][0]*(m.dot(C_inv))+lam[1][0]*(u.dot(C_inv)))/2
    sig[i-1] = sqrt(w.dot(C).dot(w.transpose()))
    if i % 10 == 0:
        print('  {}   {}    {}    {}     {}  '.format(
            '%.2f' % ret[i-1], '%.4f' % sig[i-1], '%.4f' % w[0], '%.4f' % w[1], '%.4f' % w[2]))

# Plotting the efficient frontier
plt.plot(sig, ret, color='red')
plt.grid()
plt.title('Markowitz Efficient Frontier for the given data')
plt.xlabel('Value of Risk (σ)')
plt.ylabel('Value of Return (μ)')
plt.show()

# (c) risk = 15%, calculate the maximum and minimum return and corresponding portfolios
A = u.dot(C_inv).dot(u_trans)
B = u.dot(C_inv).dot(m_trans)
V = m.dot(C_inv).dot(m_trans)
d = A*V - B*B
risk = 0.15

# Maximum return
ret_max = (2*B+sqrt(4*B*B-4*A*(V-d*risk*risk)))/(2*A)
y = np.array([[ret_max], [1]])
M = np.array([[m.dot(C_inv).dot(m_trans), u.dot(C_inv).dot(m_trans)],
              [m.dot(C_inv).dot(u_trans), u.dot(C_inv).dot(u_trans)]])
lam = 2*(np.linalg.inv(M).dot(y))
w = (lam[0][0]*(m.dot(C_inv))+lam[1][0]*(u.dot(C_inv)))/2
print('\n\nMaximum Return Portfolio for 15% risk --')
print('Return = {}% and Weights : {},  {},  {}'.format('%.4f' %
                                                       (ret_max*100), '%.4f' % w[0], '%.4f' % w[1], '%.4f' % w[2]))

# Minimum return
ret_min = (2*B-sqrt(4*B*B-4*A*(V-d*risk*risk)))/(2*A)
y = np.array([[ret_min], [1]])
M = np.array([[m.dot(C_inv).dot(m_trans), u.dot(C_inv).dot(m_trans)],
              [m.dot(C_inv).dot(u_trans), u.dot(C_inv).dot(u_trans)]])
lam = 2*(np.linalg.inv(M).dot(y))
w = (lam[0][0]*(m.dot(C_inv))+lam[1][0]*(u.dot(C_inv)))/2
print('\n\nMinimum Return Portfolio for 15% risk --')
print('Return = {}% and Weights : {},  {},  {}'.format('%.4f' %
                                                       (ret_min*100), '%.4f' % w[0], '%.4f' % w[1], '%.4f' % w[2]))

# (d) return = 18%, calculate the minimum risk portfolio
y = np.array([[0.18], [1]])  # using return = 18%
M = np.array([[m.dot(C_inv).dot(m_trans), u.dot(C_inv).dot(m_trans)],
              [m.dot(C_inv).dot(u_trans), u.dot(C_inv).dot(u_trans)]])
lam = 2*(np.linalg.inv(M).dot(y))
w = (lam[0][0]*(m.dot(C_inv))+lam[1][0]*(u.dot(C_inv)))/2
risk = sqrt(w.dot(C).dot(w.transpose()))
print('\n\nMinimum Risk Portfolio for 18% return --')
print('Risk = {}% and Weights : {},  {},  {}'.format('%.4f' %
                                                     (risk*100), '%.4f' % w[0], '%.4f' % w[1], '%.4f' % w[2]))

# (e) Given µrf = 10%, compute the market portfolio and plot the CML
wm = ((m-0.1*u).dot(C_inv))/((m-0.1*u).dot(C_inv).dot(u_trans))
um = wm.dot(m_trans)
sigm = sqrt(wm.dot(C).dot(wm.transpose()))

print('\nFor 10% risk-free return --')
print('Weights : {},  {},  {}'.format('%.4f' % wm[0], '%.4f' % wm[1], '%.4f' % wm[2]))
print('Return on market portfolio = {}%'.format('%.4f' % (um*100)))
print('Risk on market portfolio = {}%\n'.format('%.4f' % (sigm*100)))

x = np.arange(0, 1.1, 0.1)
y = 0.1 + ((um-0.1)/sigm)*x
plt.scatter(sigm, um, marker=(5, 2), color='blue', label='Market Portfolio')
plt.scatter([0], [0.1], marker='o', color='green', label='Zero Risk Portfolio')
plt.plot(x, y, color='green', label='CML')
plt.plot(sig, ret, color='red', label='Markowitz Efficient Frontier')
plt.title('Capital Market Line and Markowitz Efficient Frontier')
plt.xlabel('Value of Risk (σ)')
plt.ylabel('Value of Return (μ)')
plt.legend()
plt.grid()
plt.show()

plt.plot(x, y, color='green')
plt.title('Capital Market Line')
plt.xlabel('Value of Risk (σ)')
plt.ylabel('Value of Return (μ)')
plt.grid()
plt.show()

# (f) Find 2 portfolios (having both risky and risk free assets) with risk = 10% and 25%.
y = 0.1+(um-0.1)/sigm*0.1
rf = (y-um)/(0.1-um)
risked = (1-rf)*wm

print('\nPortfolio for 10% risk --')
print('Risk-Free Asset : {}'.format('%.4f' % rf))
print('Risky Assets Weights : {},  {},  {}'.format('%.4f' %
                                                   risked[0], '%.4f' % risked[1], '%.4f' % risked[2]))

y = 0.1+(um-0.1)/sigm*0.25
rf = (y-um)/(0.1-um)
risked = (1-rf)*wm

print('\nPortfolio for 25% risk --')
print('Risk-Free Asset : {}'.format('%.4f' % rf))
print('Risky Assets Weights : {},  {},  {}'.format('%.4f' %
                                                   risked[0], '%.4f' % risked[1], '%.4f' % risked[2]))

# Question 1 ends --------------------------------------------------------------

Semester 6/MA 374 (Financial Engg. Lab)/Lab 4/180123062_AB_q1.py


# Python code for Mandelbrot Fractal

# Import necessary libraries
from PIL import Image, ImageEnhance
from PIL import Image
import numpy as np
import colorsys
import time
import os

# setting the width and height of the output image
# 460 480
# 720 576
# 720 480
# 1280 720
# 1440 1080
# 1920 1080
# 2048 1080
# 3840 2160
# 4096 2160

WIDTH = 1280
HEIGHT = 720

# a function to return a tuple of colors
# as integer value of rgb

# def rgb(i):
#     red = int(i / 255.0)
#     green = int(1.0+i/10.0)
#     blue = int(0.5*i/100.0)
#     color = (red, green, blue)
#     return color

def rgb(i):
    color = 255 * np.array(colorsys.hsv_to_rgb(i / 255.0, i/64, i/128.0))
    return tuple(color.astype(int))

# function defining a mandelbrot
def mandelbrot(x, y, dimension):
    c = complex(x, y)
    z = 0
    for i in range(1, 100):
        if abs(z) > 2:
            return rgb(i)
        # z =complex(abs(z.real),abs(z.imag))**2+c #Burning ship
        # z = z.conjugate()**2+c
        z = complex(z.real,abs(z.imag))**2+c
        # z = z**dimension+ c #Mandelbrot
    return (0, 0, 0)

# Main function

def evaluateMandelbrot(dimension):
    # get the path to the current python file directory
    path = os.path.dirname(os.path.realpath(__file__))

    # creating the new image in RGB mode
    img = Image.new('RGB', (WIDTH, HEIGHT))
    pixels = img.load()

    start = time.time()
    for x in range(img.size[0]):

        # displaying the progress as percentage
        print("{:.2f} %".format(x / WIDTH * 100.0))

        for y in range(img.size[1]):
            pixels[x, y] = mandelbrot((x - (0.65 * WIDTH)) / (0.2*WIDTH),
                                      (y - (0.5*HEIGHT))/(0.3*HEIGHT), dimension)

    #image brightness enhancer
    enhancer = ImageEnhance.Contrast(img)

    factor = 10 #gives original image
    img = enhancer.enhance(factor)

    # to display the created fractal after
    # completing the given number of iterations
    img.save(path + "/mandelbrot.png")

    print("Done!")

if __name__ == '__main__':
    evaluateMandelbrot(2)


# coordinates.py
# Jacob Hummel
import numpy as np

#===============================================================================
def cartesian_to_spherical(coords):
    coords['r_sph'] = np.sqrt(np.square(coords.x) + np.square(coords.y) + np.square(coords.z))
    coords['theta_sph'] = np.arccos(coords.z/coords.r_sph)
    coords['phi_sph'] = np.arctan2(coords.y,coords.x)

def spherical_to_cartesian(r, theta, phi):
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)
    return x,y,z

def cartesian_to_cylindrical(coords):
    coords['r_cyl'] = np.sqrt(np.square(coords.x) + np.square(coords.y))
    coords['theta_cyl'] = np.arctan2(coords.y,coords.x)
    #return coords.drop(['x', 'y', 'z'], axis=1)

def cylindrical_to_cartesian(r, theta, z):
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    return x,y,z

def define_unit_vectors(pos):
    e1 = np.zeros_like(pos)
    e2 = np.zeros_like(pos)
    e3 = np.zeros_like(pos)
    e1[...,0] = 1.
    e2[...,1] = 1.
    e3[...,2] = 1.
    return e1, e2, e3

def cartesian_to_spherical_unit_vectors(coords):
    unit_x, unit_y, unit_z = define_unit_vectors(coords)
    r,theta,phi = cartesian_to_spherical(coords[...,0],coords[...,1],coords[...,2])
    r = r[...,np.newaxis]
    theta = theta[...,np.newaxis]
    phi = phi[..., np.newaxis]
    unit_r = coords / r
    unit_phi = unit_y * np.cos(phi) - unit_x * np.sin(phi)
    unit_theta = (unit_x * np.cos(theta) * np.cos(phi)
                  + unit_y * np.cos(theta) * np.sin(phi)
                  - unit_z * np.sin(theta))
    return unit_r, unit_theta, unit_phi

def spherical_to_cartesian_unit_vectors(sph):
    unit_r, unit_theta, unit_phi = define_unit_vectors(sph)
    r = sph[...,0, np.newaxis]
    theta = sph[...,1, np.newaxis]
    phi = sph[...,2, np.newaxis]
    unit_x = (unit_r * np.sin(theta) * np.cos(phi)
              + unit_theta * np.cos(theta) * np.cos(phi)
              - unit_phi * np.sin(phi))
    unit_y = (unit_r * np.sin(theta) * np.sin(phi)
              + unit_theta * np.cos(theta) * np.sin(phi)
              + unit_phi * np.cos(phi))
    unit_z = unit_r * np.cos(theta) - unit_theta * np.sin(theta)
    return unit_x, unit_y, unit_z

def cartesian_to_spherical_velocities(coords, uvw):
    unit_r, unit_theta, unit_phi = cartesian_to_spherical_unit_vectors(coords)
    # np.einsum('ij,ij->i',...) is einstein notation for dot product.
    vr = np.einsum('ij,ij->i', uvw, unit_r)
    vtheta = np.einsum('ij,ij->i', uvw, unit_theta)
    vphi = np.einsum('ij,ij->i', uvw, unit_phi)
    return vr, vtheta, vphi

def spherical_to_cartesian_velocities(sph, vsph):
    unit_x, unit_y, unit_z = spherical_to_cartesian_unit_vectors(sph)
    # np.einsum('ij,ij->i',...) is einstein notation for dot product.
    u = np.einsum('ij,ij->i', vsph, unit_x)
    v = np.einsum('ij,ij->i', vsph, unit_y)
    w = np.einsum('ij,ij->i', vsph, unit_z)
    return u, v, w

def rotation_matrix(axis, angle):
    if axis == 'x':
        rot = ((1., 0., 0.),
               (0., np.cos(angle), -np.sin(angle)),
               (0., np.sin(angle), np.cos(angle)))
    elif axis == 'y':
        rot = ((np.cos(angle), 0.,np.sin(angle)),
               (0., 1., 0.),
               (-np.sin(angle), 0., np.cos(angle)))
    elif axis == 'z':
        rot = ((np.cos(angle), -np.sin(angle), 0.),
               (np.sin(angle), np.cos(angle), 0.),
               (0., 0.,1.))
    else:
        u = axis[0]
        v = axis[1]
        w = axis[2]
        u2= u*u
        v2= v*v
        w2= w*w
        cos = np.cos(angle)
        sin = np.sin(angle)
        rot = ((u2 + (1-u2)*cos, u*v*(1-cos) - w*sin, u*w*(1-cos) + v*sin),
               (u*v*(1-cos) + w*sin, v2 + (1-v2)*cos, v*w*(1-cos) - u*sin),
               (u*w*(1-cos) - v*sin, v*w*(1-cos) + u*sin, w2 + (1-w2)*cos))
    rot = np.asarray(rot)
    return rot

def rotate(coords, axis, angle, verbose=False):
    rot = rotation_matrix(axis,angle)
    if verbose:
        print "Rotating about the {}-axis by {:6.3f} radians.".format(axis,angle)
        print "Rotation Matrix:"
        print rot
    coords[coords.columns] = np.dot(coords,rot)
    return coords


# -*- coding: utf-8 -*-
# Chapter 11 Working with Unlabeled Data - Clustering Analysis

# Grouping objects by similarity using k-means
from sklearn.datasets import make_blobs
X, y = make_blobs(n_samples=150,
                               n_features=2,
                               centers=3,
                               cluster_std=0.5,
                               shuffle=True,
                               random_state=0)

import matplotlib.pyplot as plt
plt.scatter(X[:, 0], X[:,1], c='white', marker='o', s=50)
plt.grid()
plt.show()

# Apply k-means to our sample dataset using the KMeans class from scikit-learn's cluster module:
from sklearn.cluster import KMeans
km = KMeans(n_clusters=3,
                        init='random',
                        n_init=10,
                        max_iter=300,
                        tol=1e-04,
                        random_state=0)
y_km = km.fit_predict(X)

plt.scatter(X[y_km==0,0],
                X[y_km==0,1],
                s=50,
                c='lightgreen',
                marker='s',
                label='cluster 1')
plt.scatter(X[y_km==1,0],
                X[y_km==1,1],
                s=50,
                c='orange',
                marker='o',
                label='cluster 2')
plt.scatter(X[y_km==2,0],
                X[y_km==2,1],
                s=50,
                c='lightblue',
                marker='v',
                label='cluster 3')
plt.scatter(km.cluster_centers_[:,0],
                km.cluster_centers_[:,1],
                s=250,
                marker='*',
                c='red',
                label='centroids')
plt.legend()
plt.grid()
plt.show()
# Using the elbow method to find the optimal number of clusters
print 'Distortion: %.2f' % km.inertia_
distortions = []
for i in range(1, 11):
    km = KMeans(n_clusters=i,
                            init='k-means++',
                            n_init=10,
                            max_iter=300,
                            random_state=0)
    km.fit(X)
    distortions.append(km.inertia_)
plt.plot(range(1, 11), distortions, marker='o')
plt.xlabel('Number of clusters')
plt.ylabel('Distortion')
plt.show()

# Quantifying the quality of clustering via silhouette plots
km = KMeans(n_clusters=3,
                        init='k-means++',
                        n_init=10,
                        max_iter=300,
                        tol=1e-04,
                        random_state=0)
y_km = km.fit_predict(X)

import numpy as np
from matplotlib import cm
from sklearn.metrics import silhouette_samples
cluster_labels = np.unique(y_km)
n_clusters = cluster_labels.shape[0]
silhouette_vals = silhouette_samples(X,
                                                            y_km,
                                                            metric='euclidean')
y_ax_lower, y_ax_upper = 0, 0
yticks = []
for i, c in enumerate(cluster_labels):
    c_silhouette_vals = silhouette_vals[y_km == c]
    c_silhouette_vals.sort()
    y_ax_upper += len(c_silhouette_vals)
    color = cm.jet(i / n_clusters)
    plt.barh(range(y_ax_lower, y_ax_upper),
                 c_silhouette_vals,
                 height=1.0,
                 edgecolor='none',
                 color=color)
    yticks.append((y_ax_lower + y_ax_upper) / 2)
    y_ax_lower += len(c_silhouette_vals)
silhouette_avg = np.mean(silhouette_vals)
plt.axvline(silhouette_avg,
                 color="red",
                 linestyle="--")
plt.yticks(yticks, cluster_labels + 1)
plt.ylabel('Cluster')
plt.xlabel('Silhouette coefficient')
plt.show()

'''
To see how a silhouette plot looks for a relatively bad clustering, let's seed the
k-means algorithm with two centroids only:'''
km = KMeans(n_clusters=2,
                        init='k-means++',
                        n_init=10,
                        max_iter=300,
                        tol=1e-04,
                        random_state=0)
y_km = km.fit_predict(X)
plt.scatter(X[y_km==0,0],
                X[y_km==0,1],
                s=50, c='lightgreen',
                marker='s',
                label='cluster 1')
plt.scatter(X[y_km==1,0],
X[y_km==1,1],
s=50,
c='orange',
marker='o',
label='cluster 2')
plt.scatter(km.cluster_centers_[:,0],
                km.cluster_centers_[:,1],
                s=250,
                marker='*',
                c='red',
                label='centroids')
plt.legend()
plt.grid()
plt.show()

'''
Next we create the silhouette plot to evaluate the results. Please keep in mind that
we typically do not have the luxury of visualizing datasets in two-dimensional
scatterplots in real-world problems, since we typically work with data in higher
dimensions:
'''
cluster_labels = np.unique(y_km)
n_clusters = cluster_labels.shape[0]
silhouette_vals = silhouette_samples(X,
                                                            y_km,
                                                            metric='euclidean')
y_ax_lower, y_ax_upper = 0, 0
yticks = []
for i, c in enumerate(cluster_labels):
    c_silhouette_vals = silhouette_vals[y_km == c]
    c_silhouette_vals.sort()
    y_ax_upper += len(c_silhouette_vals)
    color = cm.jet(i / n_clusters)
    plt.barh(range(y_ax_lower, y_ax_upper),
                 c_silhouette_vals,
                 height=1.0,
                 edgecolor='none',
                 color=color)
    yticks.append((y_ax_lower + y_ax_upper) / 2)
    y_ax_lower += len(c_silhouette_vals)
silhouette_avg = np.mean(silhouette_vals)
plt.axvline(silhouette_avg, color="red", linestyle="--")
plt.yticks(yticks, cluster_labels + 1)
plt.ylabel('Cluster')
plt.xlabel('Silhouette coefficient')
plt.show()

# Organizing clusters as a hierarchical tree
import pandas as pd
import numpy as np
np.random.seed(123)
variables = ['X', 'Y', 'Z']
labels = ['ID_0', 'ID_1', 'ID_2', 'ID_3', 'ID_4']
X = np.random.random_sample([5, 3]) * 10
df = pd.DataFrame(X, columns=variables, index=labels)
df

# Performing hierarchical clustering on a distance matrix
from scipy.spatial.distance import pdist, squareform
row_dist = pd.DataFrame(squareform(pdist(df, metric='euclidean')), columns=labels, index=labels)
row_dist

from scipy.cluster.hierarchy import linkage
help(linkage)

#Incorrect approach
from scipy.cluster.hierarchy import linkage
row_clusters = linkage(row_dist, method='complete', metric='euclidean')
# Correct approach
row_clusters = linkage(pdist(df, metric='euclidean'), method='complete')
# Correct approach
row_clusters = linkage(df.values, method='complete', metric='euclidean')

pd.DataFrame(row_clusters,
                         columns=['row label 1',
                                          'row label 2',
                                          'distance',
                                          'no. of items in clust.'],
                          index=['cluster %d' %(i+1) for i in
                                      range(row_clusters.shape[0])])

from scipy.cluster.hierarchy import dendrogram
# make dendrogram black (part 1/2)
# from scipy.cluster.hierarchy import set_link_color_palette
# set_link_color_palette(['black'])
row_dendr = dendrogram(row_clusters, labels=labels,
                                          # make dendrogram black (part 2/2)
                                          # color_threshold=npl.inf
                                          )
plt.tight_layout()
plt.ylabel('Euclidean distance')
plt.show()

# Attaching dedrograms to a heat map
'''1. We create a new figure object and define the x axis position, y axis
position, width, and height of the dendrogram via the add_axes attribute.
Furthermore, we rotate the dendrogram 90 degrees counter-clockwise.
The code is as follows:'''
fig = plt.figure(figsize=(8,8), facecolor='white')
axd = fig.add_axes([0.09,0.1,0.2,0.6])
row_dendr = dendrogram(row_clusters, orientation='right')
# note: for matplotlib >= v1.5.1, please use orientation=‘left’
'''2. Next we reorder the data in our initial DataFrame according to the clustering
labels that can be accessed from the dendrogram object, which is essentially a
Python dictionary, via the leaves key. The code is as follows:'''
df_rowclust = df.ix[row_dendr['leaves'][::-1]]

'''3. Now we construct the heat map from the reordered DataFrame and position
it right next to the dendrogram:'''
axm = fig.add_axes([0.23,0.1,0.6,0.6])
cax = axm.matshow(df_rowclust,
                                  interpolation='nearest', cmap='hot_r')
'''4. Finally we will modify the aesthetics of the heat map by removing the axis
ticks and hiding the axis spines. Also, we will add a color bar and assign
the feature and sample names to the x and y axis tick labels, respectively.
The code is as follows:'''
axd.set_xticks([])
axd.set_yticks([])
for i in axd.spines.values():
    i.set_visible(False)
fig.colorbar(cax)
axm.set_xticklabels([''] + list(df_rowclust.columns))
axm.set_yticklabels([''] + list(df_rowclust.index))
plt.show()

# Applying agglomerative clustering via scikit-learn
from sklearn.cluster import AgglomerativeClustering
ac = AgglomerativeClustering(n_clusters=2, affinity='euclidean', linkage='complete')
labels = ac.fit_predict(X)
print 'Cluster labels: %s' % labels

# Locating regions of high density via DBSCAN
'''let's create a new dataset of half-moon-shaped structures to compare k-means clustering,
hierarchical clustering, and DBSCAN:'''
from sklearn.datasets import make_moons
X, y = make_moons(n_samples=200, noise=0.05, random_state=0)
plt.scatter(X[:, 0], X[:, 1])
plt.show()
'''We will start by using the k-means algorithm and complete linkage clustering to see
whether one of those previously discussed clustering algorithms can successfully
identify the half-moon shapes as separate clusters. The code is as follows:'''
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(8,3))
km = KMeans(n_clusters=2,
                        random_state=0)
y_km = km.fit_predict(X)
ax1.scatter(X[y_km==0,0],
                  X[y_km==0,1],
                  c='lightblue',
                  marker='o',
                  s=40,
                  label='cluster 1')
ax1.scatter(X[y_km==1,0],
                   X[y_km==1,1],
                   c='red',
                   marker='s',
                   s=40,
                   label='cluster 2')
ax1.set_title('K-means clustering')
ac = AgglomerativeClustering(n_clusters=2, affinity='euclidean', linkage='complete')
y_ac = ac.fit_predict(X)
ax2.scatter(X[y_ac==0, 0],
                   X[y_ac==0, 1],
                   c='lightblue',
                   marker='o',
                   s=40,
                   label='cluster 1')
ax2.scatter(X[y_ac==0, 0],
                   X[y_ac==0, 1],
                   c='red',
                   marker='s',
                   s=40,
                   label='cluster 2')
ax2.set_title('Agglomerative clustering')
plt.legend()
plt.show()

'''Finally, let's try the DBSCAN algorithm on this dataset to see if it can find the two
half-moon-shaped clusters using a density-based approach:'''
from sklearn.cluster import DBSCAN
db = DBSCAN(eps=0.2, min_samples=5, metric='euclidean')
y_db = db.fit_predict(X)
plt.scatter(X[y_db==0, 0], X[y_db==0, 1], c='lightblue', marker='o', s=40, label='cluster 1')
plt.scatter(X[y_db==1, 0], X[y_db==1, 1], c='red', marker='s', s=40, label='cluster 2')
plt.legend()
plt.show()

self_practice/Chapter 11 Working with Unlabeled Data - Clustering Analysis.py


import numpy as np
from numpy import linalg
from abc import abstractmethod
import pandas as pd
import math

pd.options.display.float_format = '{:,.10f}'.format
np.set_printoptions(suppress=True, precision=10)

TOR = pow(10.0, -5)

class NewtonMethod(object):

    def __init__(self):
        return

    @abstractmethod
    def f(self, x):
        return NotImplementedError('Implement f()!')

    @abstractmethod
    def jacobian(self, x):
        return NotImplementedError('Implement jacobian()!')

    @abstractmethod
    def run(self, x):
        return NotImplementedError('Implement run()!')

class Newton(NewtonMethod):

    def __init__(self):
        super(NewtonMethod, self).__init__()

    def f(self, x):
        sol = np.zeros(len(x))
        sol[0] = 3 * x[0] - math.cos(x[1] * x[2]) - 1.0 / 2.0
        sol[1] = pow(x[0], 2) - 81 * pow(x[1] + 0.1, 2) + math.sin(x[2]) + 1.06
        sol[2] = math.exp(-x[0] * x[1]) + 20 * x[2] + (10 * math.pi - 3.0) / 3.0
        return sol

    def jacobian(self, x):
        jac = np.zeros(shape=(3, 3))
        jac[0][0] = 3.0
        jac[0][1] = x[2] * math.sin(x[1] * x[2])
        jac[0][2] = x[1] * math.sin(x[1] * x[2])
        jac[1][0] = 2 * x[0]
        jac[1][1] = -162 * (x[1] + 0.1)
        jac[1][2] = math.cos(x[2])
        jac[2][0] = -x[1] * math.exp(-x[0] * x[1])
        jac[2][1] = -x[0] * math.exp(-x[0] * x[1])
        jac[2][2] = 20
        return jac

    def run(self, x):
        """
        given x_0 in R^3 as a starting point.

        :param x: x_0 as described
        :return: the minimizer x* of f
        """
        df = pd.DataFrame(columns=['x' + str(i + 1) for i in range(len(x))] + ['residual', 'actual-residual'])

        row = len(df)
        df.loc[row] = [xe for xe in x] + [np.nan, np.nan]

        while True:
            jac = self.jacobian(x)
            f = -self.f(x)
            y = linalg.solve(jac, f)
            nx = x + y
            residual = linalg.norm(x - nx, np.inf)
            x = nx

            row = len(df)
            df.loc[row] = [nxe for nxe in nx] + [residual, np.nan]
            if residual < TOR:
                break

        for i in range(len(df)):
            xk = np.array([df.loc[i][j] for j in range(len(x))])
            df.loc[i][4] = linalg.norm(x - xk, np.inf)

        print(df)

def main():
    x0 = np.array([0.1, 0.1, -0.1])
    Newton().run(x0)
    return

if __name__ == '__main__':
    main()


import numpy as np
np.random.seed(42)
import matplotlib.pyplot as plt

def plot_hist(x, num_bins):
    y = np.histogram(x, bins=num_bins)
    plt.vlines(y[1][:-1], ymin=0, ymax=y[0])

def plot_scatter(x):
    fig, ax = plt.subplots(1, 1)
    plt.scatter(x[:, 0], x[:, 1], s=5)
    ax.axis('tight')

a = np.arange(10)
print(a)

x = np.random.randn(1_000_000)
bins = 1000
plot_hist(x, bins)

x = np.random.normal(loc=5.0, scale=0.2, size=1_000_000)


import sys
import numpy as np

A = [[31, -13, 0, 0, 0, -10, 0, 0, 0],
    [-13, 35, -9, 0, -11, 0, 0, 0, 0],
    [0, -9, 31, -10, 0, 0, 0, 0, 0],
    [0, 0, -10, 79, -30, 0, 0, 0, -9],
    [0, 0, 0, -30, 57, -7, 0, -5, 0],
    [0, 0, 0, 0, -7, 47, -30, 0, 0],
    [0, 0, 0, 0, 0, -30, 41, 0, 0],
    [0, 0, 0, 0, -5, 0, 0, 27, -2],
    [0, 0, 0, -9, 0, 0, 0, -2, 29]]

b = [-15, 27, -23, 0, -20, 12, -7, 7, 10]

def find_shape(A, b):
    """
    获取 n
    """

    # n行 m列
    n1, m1 = A.shape
    n2, = b.shape

    # print(n1, m1)
    # print(n2)

    # 判断矩阵形状
    if n1 != m1 or n1 != n2:
        print('Error martix shape!')
        sys.exit()
    else:
        N = n1

    return N

def MGauss(A, b):
    """
    列主元 Gauss 消去 消元
    """

    # 判断矩阵形状
    N = find_shape(A, b)

    # 列主元高斯消元
    for k in range(0, N):
        p = k
        maxabs = abs(A[k, k])
        # 找列最大值
        for i in range(k+1, N):
            if abs(A[i, k]) > maxabs:
                p = i
                maxabs = abs(A[i, k])
        # print('maxabs', maxabs)
        # 最大值为 0
        if maxabs == 0:
            print('Singular')
            sys.exit()
        # 最大值不在对角线, 则交换两行
        if p != k:
            A[[p,k],:] = A[[k,p],:]
            b[[p,k]] = b[[k,p]]
        # print('exchange r{0} and r{1}:\r\n'.format(k, p), A)
        # 消元，将对角线以下变为 0
        for i in range(k+1, N):
            m_ik = A[i, k] / A[k, k]
            for j in range(0, N):
                A[i, j] -= A[k, j] * m_ik
            b[i] -= b[k] * m_ik
        # print('After Elimination:\r\n', np.concatenate((A,np.asarray([b]).T), axis = 1))

    if A[N-1, N-1] == 0:
        print('Singular')
        sys.exit()

    return A, b

def bring_back(A, b):
    """
    列主元 Gauss 消去 回带
    """

    # 判断矩阵形状
    N = find_shape(A, b)

    # 回带
    X = np.zeros(N)
    X[N-1] = b[N-1] / A[N-1, N-1]
    for i in range(0, N-1):
        k = N-2-i
        sigma = sum(A[k, j]*X[j] for j in range(k+1, N))
        # print('k:{0}, sigma:{1}'.format(k, sigma))
        X[k] = (b[k] - sigma) / A[k, k]

    return X

# 列主元高斯消去求解
def MGauss_Caculate(A, b):
    A_np = np.asarray(A, dtype = float)
    b_np = np.asarray(b, dtype = float)

    # 列主元、回带
    A_G, b_G = MGauss(A_np, b_np)
    x_G = bring_back(A_G, b_G)

    # print('A_G:b_G\r\n', np.concatenate((A_G,np.asarray([b_G]).T), axis = 1))

    return x_G

def main():
    """
    main
    """

    x_G = MGauss_Caculate(A, b)
    print('x_G', x_G)

if __name__ == "__main__":
    main()


import numpy as np

def L2(v, *args):
    """Returns the L2 norm of a vector v

    INPUTS
    =======
    v: list
       Vector of which the L2 norm to be calculated
    *arg: list, optional
       Weight vector

    RETURNS
    ========
    L2_norm: float
       float unless len(v) != len(args[0])
       in which case a ValueError exception is raised

    EXAMPLES
    =========
    >>> L2([3, 4])
    5.0
    >>> L2([3, 4], [1, 1])
    5.0
    """
    s = 0.0 # Initialize sum
    if len(args) == 0: # No weight vector
        for vi in v:
            s += vi * vi
    else: # Weight vector present
        w = args[0] # Get the weight vector
        if (len(w) != len(v)): # Check lengths of lists
            raise ValueError("Length of list of weights must match length of target list.")
        for i, vi in enumerate(v):
            s += w[i] * w[i] * vi * vi
    return np.sqrt(s)


'''
Project A: An optical ray tracer
Aim: A module to model an optical system using the principle of geometric optics with different types of optical surfaces,
and subsequently optimising the curvature of a singlet lens. This module contains codes for plotting all the graphs.

Written by Son-Gyo Jung; Group F; Lab Group A9; 04.02.16
'''

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from math import *
import scipy.optimize as spo
import genpolarSJ as gp

#Normalisation function and Snell's law are placed at the beginning for convenience.
def normalise(x):
    """
    Normalises a given 3D vector.

    Parameter:
        x - a 3D vector that is an array.
    """
    if x is not None and len(x) == 3:
        return x/np.sqrt(np.dot(x,x))

    else:
        raise Exception("Three coordinates are required in the array.")
        return None

def Snell(k1, s_normal, n1, n2): #note: k1 must be a array
    '''
    Refracts the incident rays using Snell's law.

    Parameters:
        k1 - the direction of the incident ray.
        s_normal - vector normal to the surface.
        n1,n2 - the refractive indices either side of the surface.
    '''

    n_12 = n1/n2
    k1_hat = normalise(k1)
    s_normal_hat = normalise(s_normal)
    cos_angle1 = fabs(np.dot(k1_hat, s_normal_hat))
    angle1 = acos(cos_angle1)

    if sin(angle1) >= n2/n1:
        print "Total internal reflection."

        return None

    else:
        k2 = n_12*k1_hat + (n_12*cos_angle1 - np.sqrt(1 - n_12**2 * (1-cos_angle1)**2)) * s_normal_hat #vector form of Snell's law
        k2_hat = normalise(k2)

        return k2_hat

class ray:
    """
    Class defining the positions and directions of optical rays.

    Parameters:
        p - the position of the ray.
        k - the direction of the ray.
    """

    def __init__(self, p=[0.0, 0.0, 0.0], k=[0.0, 0.0, 1.0]):
        self.__positions = list()
        self.__directions = list()

        if p is None and len(p) != 3:
            raise Exception("p requires exactly 3 arguments.")

        else:
            self.__positions.append(np.array(p))

        if k is None and len(k) != 3:
            raise Exception("k requires exactly 3 arguments.")

        else:
            self.__directions.append(normalise(np.array(k))) #append normalised k as an array

    def p(self):
        '''
        Returns the current position of the ray.
        '''
        return self.__positions[-1]

    def k(self):
        '''
        Returns the current direction of the ray.
        '''
        return self.__directions[-1]

    def vertices(self):
        '''
        Returns the entire points through which the ray propagated.
        '''
        return self.__positions

    def klist(self):
        '''
        Returns the list of directions of the ray.
        '''
        return self.__directions

    def append(self, new_p=None, new_k=None):
        '''
        Add new position and direction of the ray

        Parameters:
            new_p - the new position of the ray.
            new_k - the new direction of the ray.
        '''

        if new_p is not None and len(new_p) == 3:
            self.__positions.append(np.array(new_p))

        else:
            raise Exception("P should only take 3 arguments.")

        if new_k is not None and len(new_k) == 3:
            self.__directions.append(np.array(new_k))

        else:
            raise Exception("K should only take 3 arguments.")

class OpticalElement: #provided in BlackBoard
    '''
    A base class for optical elements being employed; it propagates a ray through the optical element
    '''

    def propagate_ray(self, ray):

        raise NotImplementedError()

class SphericalRefraction(OpticalElement):
    """
    Initialises the interface of refraction with either a spherical lens or a plane lens, where the latter corresponds to the curvature of zero.
    It returns the location of the intercepts and the new direction calculated by calling the Snell's law function.

    Parameters:
        z0 - the intercept of the surface with the z-axis.
        curv - the curvature of the surface.
        n1,n2 - the refractive indices either side of the surface.
        apr - aperture radius which is the maximum extent of the surface from the optical axis.
    """

    def __init__(self, z0=[0, 0, 100], curv=0.03, n1=1.0, n2=1.5168, apr=1000):
        self.__z0 = z0
        self.__curv = curv
        self.__n1 = n1
        self.__n2 = n2
        self.__apr = apr

    def intercept(self, ray):
        '''
        Determines the point of interception of the ray with the given lens.
        '''

        if self.__curv != 0: #when the lens is spherical
            R = 1.0/self.__curv
            self.__centre = np.array([0., 0., R + self.__z0[2]])
            r = ray.p() - self.__centre
            discriminant = np.inner(r, normalise(ray.k()))**2 - (np.inner(r, r) - R**2) #using the equation provided in BlackBoard.

            if discriminant < 0:
                print "No real solution as the ray is complex."
                return None #No real solutions so terminate

            Lplus = -np.inner(r, normalise(ray.k())) + np.sqrt(discriminant) #Possible distance to the intercept
            Lminus = -np.inner(r, normalise(ray.k())) - np.sqrt(discriminant) #Possible distance to the intercept
            interceptplus = ray.p() + Lplus * normalise(ray.k()) #coordinate of the possible intercept
            interceptminus = ray.p() +Lminus * normalise(ray.k()) #coordinate of the possible intercept

            if min(self.__z0[2],self.__centre[2]) <= interceptminus[2] <= max(self.__z0[2],self.__centre[2]):
                return interceptminus

            if min(self.__z0[2],self.__centre[2]) <= interceptplus[2] <= max(self.__z0[2],self.__centre[2]):
                return interceptplus

        else: #when the lens is planar
            number_of_khatz = (self.__z0[2] - ray.p()[2])/normalise(ray.k())[2] #See lab book page 112
            intercept_in_plane = ray.p() + number_of_khatz * normalise(ray.k())
            return intercept_in_plane

    def propagate_ray(self,ray):
        '''
        Determines the new position and direction of the ray after propagating through the lens.
        '''
        intersection = self.intercept(ray)
        if intersection is None:
            print "No intersections with the lens."
            return None

        else: #determining the direction of the surface normal vector (using the convention taught in the vector calculus course)
            if self.__curv == 0:
                new_s_normal_hat = np.array([0., 0., -1.])

            elif self.__curv > 0:
                new_s_normal_hat = normalise(intersection - self.__centre)

            elif self.__curv < 0:
                new_s_normal_hat = normalise(self.__centre - intersection)

            else:
                raise Exception("Curvature is not defined.")

            newdir = Snell(normalise(ray.k()), new_s_normal_hat, self.__n1, self.__n2) #calling Snell's law
            newpos = intersection

            ray.append(newpos, newdir)

class OutputPlane(SphericalRefraction): #take SphericalRefraction class as an argument to inherit the methods
    '''
    Terminates the propagation of the ray by creating a plane with a normal vector parallel to the z-axis.
    Note: the output plane is modelled using a planar lens that has infinite aperture radius.
    '''

    def __init__(self, output_z):
        SphericalRefraction.__init__(self,[0,0,output_z], 0.00, 1., 1., float('inf')) #design output plane using a plane lens with infinite aperture radius

class raybundle(list):
    '''
    Generates a bundle of collimated rays using the genpolar module.
    Note: the parameters can be adjusted to vary the beam radius, distribution and the number of rays.

    Parameters:
        n - Number of rings (this includes the central ray).
        rmax - Parameter which determines the maximum radius. (Note: the array created starts with zero radius by default.)
        m - The step-size of number of rays in each subsequent ring.
    '''

    def __init__(self, n=6, rmax=6., m=6,k=[0., 0., 1.]):
        for r, t in gp.rtuniform(n, rmax, m):
            self.append(ray([r*cos(t), r*sin(t), 0], normalise(k)))

    def propagate_ray(self, lens_element):
        '''
        Propagates the given rays through a single optical element at a time.
        '''
        for r in self:
            lens_element.propagate_ray(r)

    def RMS(self):
        '''
        Caculates the Root-Mean-Square of the spot radius at the paraxial focus for the bundle of rays;
        that is the RMS deviation of the ray positions from the optical axis.
        '''

        rxy_squared = list()

        for r in self:
            rxy_squared.append(r.p()[0]**2 + r.p()[1]**2)

        RMS = np.sqrt(np.mean(rxy_squared))*1e-3

        print 'RMS is %sm' % (RMS)

        return RMS

    def plot2D(self):
        '''
        Plots the paths of the rays in two-dimensions.
        '''
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for r in self:
            x, y, z = zip(*r.vertices())
            ax.plot(z, x, 'b')
        ax.set_title("Ray propagation in z-x plane")
        ax.set_xlabel("z (mm)")
        ax.set_ylabel("x (mm)")
        plt.grid()
        plt.show()

    def plot3D(self):
        '''
        Plots the paths of the rays in three-dimensions.
        '''
        fig = plt.figure()
        ax = fig.add_subplot(111, projection = '3d')
        for r in self:
            x, y, z = zip(*r.vertices())
            ax.plot(z, x, y) #use c = np.random.rand(3,1) to make it pretty
        ax.set_title("Ray propagation in 3D")
        ax.set_xlabel("z (mm)")
        ax.set_ylabel("y (mm)")
        ax.set_zlabel("x (mm)")
        plt.grid()
        plt.show()

    def plot2DOutput(self):
        '''
        Plots the rays that terminate in the output plane as a scatter graph in the x-y plane.
        '''
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for r in self:
            x, y, z = zip(*r.vertices())
            ax.plot( x[-1],y[-1], 'bo')
        ax.set_title("Output in the x-y Plane")
        ax.set_xlabel("x (mm)")
        ax.set_ylabel("y (mm)")
        plt.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        plt.grid()
        plt.show()

#Extension tasks and additional calculators
def ParaxialFocus(fristlens, secondlens): #to calcuate the paraxial focus
    '''
    Accurately estimates the location of the paraxial focus in the z-direction by utilising a paraxial ray.

    Parameters:
        fristlens - the front/frist optical surface defined using the SphericalRefraction class.
        secondlens - the back/second optical surface defined using the SphericalRefraction class.
    '''

    paraxial_ray = ray([0.00001, 0., 0.], [0., 0., 1.,]) #defines the paraxial ray
    fristlens.propagate_ray(paraxial_ray)

    if secondlens is not None:
        secondlens.propagate_ray(paraxial_ray)

    paraxial_focus = paraxial_ray.p()[2] - paraxial_ray.p()[0] * paraxial_ray.k()[2]/paraxial_ray.k()[0] # using eq. x = m(z-f), see lab book p.132

    print 'Paraxial focus is at %smm' % (paraxial_focus)

    return paraxial_focus

def DiffractionLimit(fristlens, secondlens, FirstLensPos, rmax, n, image_dis, optimisation): #to calculate the diffraction limit
    '''
    Accurately estimates the diffraction limit by employing the ParaxialFocus function.

    Parameters:
        fristlens, secondlens - first and second optical surface defined using the SphericalRefraction class.
        rmax, n - defined using parameters of raybundle class.
    '''

    wavelength = 588e-9
    diameter = np.arange (0.0, rmax, rmax/n)[-1] * 2e-3

    if optimisation is True:
        f_length = image_dis*1e-3 - FirstLensPos*1e-3 - 2.5*1e-3 #for optimising biconvex lens with 5mm thickness

    elif  secondlens is not None:
        f_length = ParaxialFocus(fristlens, secondlens)*1e-3 - FirstLensPos*1e-3 - 2.5*1e-3 #for a lens of 5mm thickness

    else:
        f_length = ParaxialFocus(fristlens, secondlens)*1e-3 - FirstLensPos*1e-3 #for a single refracting surface

    D_limit = wavelength*f_length/diameter

    print 'Diffraction limit is %sm' % (D_limit)

    return D_limit

def OptimiseLenses(lensposition,image_dis, n, rmax, m, k): #to optimise a singlet lens into its best form for a given image distance.
    '''
    Optimises the curvatures of both lenses by minimising the RMS of spot radius at a given image distance;
    Note: both the positions of the lenses and the output plane is required.

    Parameters:
        lensposition - essentially the z-coordinate of z0.
        image_dis - the position of the image (i.e. the position of the output plane).
        n, rmax, m, k - defined using parameters of raybundle class.
    '''

    def singlet(curvatures):
        '''
        Parameter:
            curvatures - an optimizer which minimises the RMS spot radius using the downhill simplex algorithm.
            [Args: singlet - the objective function to be minimized; trial_solution - the initial guesses.]
        '''

        first_lens = SphericalRefraction(z0=[0., 0., lensposition], curv=curvatures[0], n1=1.0, n2=1.5168, apr=float('inf'))
        second_lens = SphericalRefraction(z0=[0., 0., lensposition+5], curv=curvatures[1], n1=1.5168, n2=1.0, apr=float('inf'))
        opp = OutputPlane(image_dis) #output plane positioned at the give image distance

        b = raybundle(n, rmax, m, k) #substitute with new values using ExecuteRayTracerSJ.py
        b.propagate_ray(first_lens)
        b.propagate_ray(second_lens)
        b.propagate_ray(opp)

        rxy_squared = list()

        for r in b:
            rxy_squared.append(r.p()[0]**2 + r.p()[1]**2)
        RMS = np.sqrt(np.mean(rxy_squared))*1e-3
        print 'RMS is %sm' % (RMS)

        return RMS

    trial_solutions = [0.05, -0.05]
    curvatures = spo.fmin(singlet, trial_solutions)

    print 'The best form curvatures for a given image distance of %s mm are: first lens curvature = %s/mm, second lens curvature = %s/mm.' % (image_dis, curvatures[0], curvatures[1])

def RMSvsDiffractionLimit(lensposition, image_dis, n, rmax, rmin, steps, m, k, optcurv_1, optcurv_2): #function to plot RMS spot radius and diffraction limit against increasing beam radius
    rmaxlist = np.linspace(rmin, rmax, steps)
    RMSlist = list()
    DLlist = list()

    first_lens = SphericalRefraction(z0=[0., 0., lensposition], curv=optcurv_1, n1=1.0, n2=1.5168, apr=float('inf'))
    second_lens = SphericalRefraction(z0=[0., 0., lensposition+5], curv=optcurv_2, n1=1.5168, n2=1.0, apr=float('inf'))
    opp = OutputPlane(image_dis)

    for rmax in rmaxlist:
        b = raybundle(n, rmax, m, k) #substitute with new values using ExecuteRayTracerSJ.py
        b.propagate_ray(first_lens)
        b.propagate_ray(second_lens)
        b.propagate_ray(opp)
        RMSlist.append(b.RMS()) #append all the RMS values

        DLlist.append(DiffractionLimit(first_lens, second_lens, lensposition, rmax, n, image_dis, optimisation=True)) #append the diffraction limits calculated using the DiffractionLimit function

    plt.figure()
    line_up, = plt.plot(rmaxlist,RMSlist,'g--', label='RMS Spot Radius')
    line_down, = plt.plot(rmaxlist,DLlist,'r-', label='Diffraction Limit')

    plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
    plt.grid()
    plt.xlabel("Beam radius (mm)")
    plt.ylabel('Limiting factor (m)')
    plt.title('Diffraction Limit and RMS Spot Size with Beam Radius \n for Optimised Biconvex lens', fontsize=13)
    plt.autoscale(tight=True)
    plt.legend(handles=[line_up, line_down])
    plt.show()

def RMSPlanoConvex(lensposition, n, rmax, rmin, steps, m, k, optcurv_1, optcurv_2): #to plot the behaviour of RMS spot radius with increasing beam size when using the planocovex lens
    rmaxlist = np.linspace(rmin, rmax, steps)
    RMSlist = list()
    RMSlist2 = list() #both orientations considered

    first_lens = SphericalRefraction(z0=[0., 0., lensposition], curv=optcurv_1, n1=1.0, n2=1.5168, apr=float('inf'))
    second_lens = SphericalRefraction(z0=[0., 0., lensposition+5], curv=optcurv_2, n1=1.5168, n2=1.0, apr=float('inf'))
    third_lens = SphericalRefraction(z0=[0., 0., lensposition], curv=optcurv_2, n1=1.0, n2=1.5168, apr=float('inf'))
    fourth_lens = SphericalRefraction(z0=[0., 0., lensposition+5], curv=-optcurv_1, n1=1.5168, n2=1.0, apr=float('inf'))
    opp1 = OutputPlane(ParaxialFocus(first_lens,  second_lens))
    opp2 = OutputPlane(ParaxialFocus(third_lens, fourth_lens))

    for rmax in rmaxlist:
        b = raybundle(n, rmax, m, k) #substitute with new values using ExecuteRayTracerSJ.py
        b.propagate_ray(first_lens)
        b.propagate_ray(second_lens)
        b.propagate_ray(opp1)
        RMSlist.append(b.RMS()) #list of all the RMS values

    for rmax in rmaxlist:
        b2 = raybundle(n, rmax, m, k) #substitute with new values using ExecuteRayTracerSJ.py
        b2.propagate_ray(third_lens)
        b2.propagate_ray(fourth_lens)
        b2.propagate_ray(opp2)
        RMSlist2.append(b2.RMS()) #list of all the RMS values

    plt.figure()
    line_up, = plt.plot(rmaxlist,RMSlist, 'g-', label='Convex-Plano')
    line_down, = plt.plot(rmaxlist,RMSlist2, 'b--', label='Plano-Convex')
    plt.legend(handles=[line_up, line_down], loc=2)
    plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
    plt.grid()
    plt.xlabel("Beam radius (mm)")
    plt.ylabel('RMS Spot Size (m)')
    plt.title('RMS spot size with beam radius\n(planoconvex lenses)')
    plt.autoscale(tight=True)
    plt.show()


#!/usr/bin/env python
"""
Simple vector classes. (Python3)

Why not use numpy instead? A couple of reasons:
1. performance, working with small numpy arrays introduces a lot of overhead
2. .x and .y class properties
3. clear implementation for readability and porting to other languages.

 Copyright (c) 2021 AIGRO B.V. - Jev Kuznetsov
 License: BSD

Special thanks for these information sources and references:

* [vector.py](https://gist.github.com/mostley/3819375)
* [StackOverflow](https://stackoverflow.com/a/43542669)
* [vector.py](https://github.com/betados/vector_2d/blob/develop/vector_2d/vector.py)

"""

from math import atan2, hypot, sin, cos
import numpy as np

def rotation_matrix(theta):
    """ create 2-d rotation matrix """
    return np.array([[cos(theta), sin(theta)], [-sin(theta), cos(theta)]])

class Vector:
    """ 2d vector, compatible with numpy """

    def __init__(self, *args):
        """ create from x,y or xy """
        try:
            self.x, self.y = float(args[0]), float(args[1])
        except (IndexError, TypeError):
            self.x = float(args[0][0])
            self.y = float(args[0][1])

    def __repr__(self):
        return f"(x={self.x:.3f}, y={self.y:.3f})"

    def __add__(self, other):
        return Vector(self.x + other.x, self.y + other.y)

    def __sub__(self, other):
        return Vector(self.x - other.x, self.y - other.y)

    def __eq__(self, other):
        return (self.x == other.x) & (self.y == other.y)

    def cross(self, other):
        return (self.x * other.y) - (self.y * other.x)

    def dot(self, other):
        return self.x * other.x + self.y * other.y

    def inner(self, other):
        return self.dot(other)

    def rotate(self, angle):
        a = np.dot(self, rotation_matrix(angle))
        return Vector(a[0], a[1])

    def round(self, precision=5):
        r = np.round(self, precision)
        return Vector(r)

    def angle(self, other):
        """ angle between vector, *relative* to other, range (-pi..pi)  """
        return -atan2(self.cross(other), self.inner(other))

    def __mul__(self, scalar):
        """ scalar multiplification """
        return Vector(self.x * scalar, self.y * scalar)

    def __rmul__(self, scalar):
        return self.__mul__(scalar)

    def __truediv__(self, scalar):  # note __div__ was removed in Python3
        return Vector(self.x / scalar, self.y / scalar)

    def __neg__(self):
        return Vector(-self.x, -self.y)

    def __abs__(self):
        return hypot(self.x, self.y)

    def distance(self, other):
        return hypot(self.x-other.x, self.y-other.y)

    @classmethod
    def from_polar(cls, r, theta):
        return Vector(r * cos(theta), r * sin(theta))

    def __getitem__(self, item):
        """ make subscriptable """
        return [self.x, self.y][item]

    def __array__(self, dtype=None) -> np.array:
        """ make this class compatable with numpy operations """
        if dtype:
            return np.array([self.x, self.y], dtype=dtype)
        else:
            return np.array([self.x, self.y])

def point_on_line(a: Vector, b: Vector, p: Vector):
    """
    project point p to line defined by a and b
    """
    ap = p - a
    ab = b - a
    return a + ap.dot(ab) / ab.dot(ab) * ab


#!/usr/bin/python
#  -*- coding:utf-8 -*-

import numpy as np
from scipy import stats
import math
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

def calc_statistics(x):
    n = x.shape[0]  # 样本个数

    # 手动计算
    m = 0
    m2 = 0
    m3 = 0
    m4 = 0
    for t in x:
        m += t
        m2 += t * t
        m3 += t ** 3
        m4 += t ** 4
    m /= n
    m2 /= n
    m3 /= n
    m4 /= n

    mu = m
    sigma = np.sqrt(m2 - mu * mu)
    skew = (m3 - 3 * mu * m2 + 2 * mu ** 3) / sigma ** 3
    kurtosis = (m4 - 4 * mu * m3 + 6 * mu * mu * m2 - 4 * mu ** 3 * mu + mu ** 4) / sigma ** 4 - 3
    print('手动计算均值、标准差、偏度、峰度：', mu, sigma, skew, kurtosis)

    # 使用系统函数验证
    mu = np.mean(x, axis=0)
    sigma = np.std(x, axis=0)
    skew = stats.skew(x)
    kurtosis = stats.kurtosis(x)
    return mu, sigma, skew, kurtosis

if __name__ == '__main__':
    d = np.random.randn(100000)
    print(d)
    mu, sigma, skew, kurtosis = calc_statistics(d)
    print('函数库计算均值、标准差、偏度、峰度：', mu, sigma, skew, kurtosis)
    # 一维直方图
    mpl.rcParams[u'font.sans-serif'] = 'SimHei'
    mpl.rcParams[u'axes.unicode_minus'] = False
    y1, x1, dummy = plt.hist(d, bins=50, normed=True, color='g', alpha=0.75)
    t = np.arange(x1.min(), x1.max(), 0.05)
    y = np.exp(-t ** 2 / 2) / math.sqrt(2 * math.pi)
    plt.plot(t, y, 'r-', lw=2)
    plt.title(u'高斯分布，样本个数：%d' % d.shape[0])
    plt.grid(True)
    plt.show()

    d = np.random.randn(100000, 2)
    mu, sigma, skew, kurtosis = calc_statistics(d)
    print('函数库计算均值、标准差、偏度、峰度：', mu, sigma, skew, kurtosis)
    # 二维图像
    N = 30
    density, edges = np.histogramdd(d, bins=[N, N])
    print('样本总数：', np.sum(density))
    density /= density.max()
    x = y = np.arange(N)
    t = np.meshgrid(x, y)
    fig = plt.figure(facecolor='w')
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(t[0], t[1], density, c='r', s=15 * density, marker='o', depthshade=True)
    ax.plot_surface(t[0], t[1], density, cmap=cm.Accent, rstride=2, cstride=2, alpha=0.9, lw=0.75)
    ax.set_xlabel(u'X')
    ax.set_ylabel(u'Y')
    ax.set_zlabel(u'Z')
    plt.title(u'二元高斯分布，样本个数：%d' % d.shape[0], fontsize=20)
    plt.tight_layout(0.1)
    plt.show()


import numpy as np

def DFT(M):
    M = np.array(M)
    try:
        N, _ = M.shape
    except:
        N, = M.shape
    FM = np.array(range(N))
    FM = FM.reshape(-1, 1).dot(FM.reshape(1, -1))
    FM = np.exp(FM * (-2 * np.pi * 1j / N))
    print('Discrete Fourier Transform Matrix:\n', FM)
    return FM.dot(M)

def IDFT(M):
    M = np.array(M)
    try:
        N, _ = M.shape
    except:
        N, = M.shape
    IFM = np.array(range(N))
    IFM = IFM.reshape(-1, 1).dot(IFM.reshape(1, -1))
    IFM = np.exp(IFM * (2 * np.pi * 1j / N))
    IFM = IFM /  N
    print('Inverse Discrete Fourier Transform Matrix:\n', IFM)
    return IFM.dot(M)


import numpy as np
import pandas as pd
import os

def new_alt_values(x): # this function converts original latitude values into the new altitude values
    lat_deltas = x - min_lat
    y_1 = lat_deltas * np.sin(rot_ang_rad)
    conv_alt = y_1 * alt_conv_factor
    output_alt = conv_alt + des_min_alt
    return output_alt

def new_lat_values(x): # this function converts original latitude values into the new latitude values
    lat_deltas = x - min_lat
    x_1 = lat_deltas * np.cos(rot_ang_rad)
    output_lat = x_1 + min_lat
    return output_lat

file_name = input('Enter name of file: ')
rot_ang_deg = input('Please enter desired rotation angle from 0 - 180 degrees: ')
rot_ang_rad = float(rot_ang_deg) * 0.0174533
csv_df = pd.read_csv(file_name, delimiter = ',')
min_lat = csv_df['latitude'][csv_df['latitude'].idxmin()]

if csv_df.columns[2] == 'altitude(m)' : # this runs if the original .CSV file was generated with metric units
    des_min_alt = input('Please enter desired minimum altitude in meters: ')
    des_min_alt = float(des_min_alt)
    alt_conv_factor = 110947.2
    alt_type = 'm'
    csv_df['altitude(m)'] = csv_df['latitude'].apply(new_alt_values)
    print('Minimum Altitude:', des_min_alt, 'meters')

elif csv_df.columns[2] == 'altitude(ft)' : # this runs if the original .CSV file was generated with imperial units
    des_min_alt = input('Please enter desired minimum altitude in feet: ')
    des_min_alt = float(des_min_alt)
    alt_conv_factor = 364000.0
    alt_type = 'ft'
    csv_df['altitude(ft)'] = csv_df['latitude'].apply(new_alt_values)
    print('Minimum Altitude:', des_min_alt, 'feet')

else :
    print('Error with input file!')

csv_df['latitude'] = csv_df['latitude'].apply(new_lat_values)

output_file_name = file_name[0:-4] + '_' + str(des_min_alt) + '_' + alt_type + '_' + rot_ang_deg + '_degrees.csv' # this sets the name of the output file to the minimum altitude and rotation angle
output_file_path = os.path.dirname(os.path.abspath(__file__)) # this is the same path as this running script
csv_df.to_csv(output_file_path + '/' + output_file_name, index = False)

print('Exported' , output_file_name, 'to', output_file_path)


# Spatial Filter
#
# Applies a Lowpass, Highpass, Band Reject or Bandpass circular symmetric filter
# by convolution
#
import sys,os
import numpy as np
import scipy.misc as sm
import scipy.special as ss
#
# Import the module with the I/O scaffolding of the External Attribute
#
sys.path.insert(0, os.path.join(sys.path[0], '..'))
import extattrib as xa

#
# The attribute parameters
#
xa.params = {
	'Inputs': ['Input'],
	'StepOut' : {'Value': [9,9], 'Minimum': [9,9], 'Hidden': False, 'Same': True},
	'Par_0' : {'Name': 'Norm. Spatial Frequency', 'Value': 0.5},
	'Select' : {'Name': 'Type', 'Values': ['Low Pass', 'High Pass', 'Band Pass', 'Band Reject'], 'Selection': 0},
	'Help'  : 'http://waynegm.github.io/OpendTect-Plugin-Docs/external_attributes/Spatial_Filter_Circular.html'
}
#
# Define the compute function
#
def doCompute():
#
#	Compute the filter kernel
#
	nil = xa.SI['nrinl']
	nxl = xa.SI['nrcrl']
	centre_trace_x = nil//2
	centre_trace_y = nxl//2
	freq = xa.params['Par_0']['Value']
	type = xa.params['Select']['Selection']
	kernelFunc = lpKernel if type==0 else hpKernel if type==1 else bpKernel if type==2 else brKernel
	kernel = np.zeros((nil,nil,1))
	N = nil//2
	if (N%2 == 0):
		N=N-1
		kernel[1:2*N+2,1:2*N+2,1] = kernelFunc(N, freq)
	else:
		kernel = kernelFunc(N, freq)
#
#	This is the trace processing loop
#
	while True:
		xa.doInput()
#
#	Get the input
#
		indata = xa.Input['Input']
#
#	Apply the kernel
#
		outdata = np.sum(kernel * indata, axis=(0,1))
#------------------------------------------------------------------------------------
#
		xa.Output = outdata
		xa.doOutput()

def lpKernel(N, freq):
#
# Lowpass filter kernel generator
#
# N is the filter half-size, must be odd
# freq is the normalised cutoff frequency
#
# Returns the filter kernel of size (2N+1,2N+1)
#
    num = 2*N + 1
    result = np.zeros((num,num,1))
    for m in range(N+1):
        i = m+N
        im = -m+N
        for n in range(N+1):
            j = n+N
            jm = -n+N
            if (m==0 and n==0):
                result[i,j] = np.pi*freq**2/4
            else:
                pandq = m%2 + n%2
                mandn = np.sqrt(m*m+n*n)
                val = (sm.factorial2(N)**4 * np.pi**pandq * freq * ss.jn(1,np.pi * freq * mandn)) / ( 2**(pandq+1) * sm.factorial2(N+m) * sm.factorial2(N-m) * sm.factorial2(N+n) *sm.factorial2(N-n) *mandn )
                result[i,j] = val
                result[j,i] = val
                result[im,j] = val
                result[i,jm] = val
                result[im,jm] = val

    return result/np.sum(result)

def hpKernel(N, freq):
#
# Highpass circular symmetric filter kernel generator
#
# N is the filter half-size, must be odd
# freq is the normalised cutoff frequency
#
# Returns the filter kernel of size (2N+1,2N+1)
#
    result = lpKernel(N, freq)
    for m in range(-N,N+1):
        i = m+N
        for n in range(-N,N+1):
            j = n+N
            if (m==0 and n==0):
                result[i,j] = 1-result[i,j]
            else:
                result[i,j] = -result[i,j]
    return result

def brKernel(N, freq, halfwidth=0.1):
#
# Band Reject circular symmetric filter kernel generator
#
# N is the filter half-size, must be odd
# freq is the normalised centre frequency of the reject band
# halfwidth controls the aperture of the reject band to freq +/- halfwidth
#
# Returns the filter kernel of size (2N+1,2N+1)
#
    kernel_lp = lpKernel(N, freq-halfwidth)
    kernel_hp = hpKernel(N, freq+halfwidth)
    result = kernel_lp + kernel_hp
    return result

def bpKernel(N, freq , halfwidth=0.1):
#
# Bandpass circular symmetric filter kernel generator
#
# N is the filter half-size, must be odd
# freq is the normalised centre frequency of the pass band
# halfwidth controls the aperture of the pass band to freq +/- halfwidth
#
# Returns the filter kernel of size (2N+1,2N+1)
#
    result = brKernel(N, freq, halfwidth)
    for m in range(-N,N+1):
        i = m+N
        for n in range(-N,N+1):
            j = n+N
            if (m==0 and n==0):
                result[i,j] = 1-result[i,j]
            else:
                result[i,j] = -result[i,j]
    return result
#
# Assign the compute function to the attribute
#
xa.doCompute = doCompute
#
# Do it
#
xa.run(sys.argv[1:])


# Bengisu Ayan - 2236974
# Ceren Gürsoy - 2237485

from math import sqrt
import numpy as np
import cv2

# fourier transform the image
def fourier_transform(image):
    image_f_trans = np.fft.fft2(image)
    image_f_shift = np.fft.fftshift(image_f_trans)  # flip the corners

    return image_f_shift

# butterworth filter
def butter_worth_filter(size_x, size_y, r, n):
    b_filter = np.zeros([size_x, size_y])
    for u in range(size_x):
        for v in range(size_y):
            val = pow((u - int(size_x / 2)), 2) + pow((v - int(size_y / 2)), 2)
            if val != 0:  # to prevent zero division
                x = pow(r / sqrt(val), 2 * n)
                b_filter[u][v] = 1 / (1 + x)

    return b_filter

# apply butterworth filter to image
def apply_filter(f_shifted_image, r, n):
    b_filter = butter_worth_filter(f_shifted_image.shape[0], f_shifted_image.shape[1], r, n)
    filtered_image = np.multiply(f_shifted_image, b_filter)

    return filtered_image

# inverse transform the image
def inverse_fourier_transform(filtered_image):
    image_inv_f_shift = np.fft.ifftshift(filtered_image)  # flip the corners back
    image_inv_f_trans = np.fft.ifft2(image_inv_f_shift)

    return image_inv_f_trans

B1 = cv2.imread('B1.jpg', 0)  # read image as grayscale
B2 = cv2.imread('B2.jpg', 0)  # read image as grayscale
B3 = cv2.imread('B3.jpg', 0)  # read image as grayscale

b1_f_shift = fourier_transform(B1)  # fourier transform the image
b1_filtered = apply_filter(b1_f_shift, 80, 2)  # apply butterworth filter to image with r=80 and n=2
b1_inv_f_shift = inverse_fourier_transform(b1_filtered)  # inverse transform the image
b1_abs = np.abs(b1_inv_f_shift)  # take absolute in order to save
cv2.imwrite('B1_output.png', b1_abs)  # save image

b2_f_shift = fourier_transform(B2)  # fourier transform the image
b2_filtered = apply_filter(b2_f_shift, 80, 2)  # apply butterworth filter to image with r=80 and n=2
b2_inv_f_shift = inverse_fourier_transform(b2_filtered)  # inverse transform the image
b2_abs = np.abs(b2_inv_f_shift)  # take absolute in order to save
cv2.imwrite('B2_output.png', b2_abs)  # save image

b3_f_shift = fourier_transform(B3)  # fourier transform the image
b3_filtered = apply_filter(b3_f_shift, 50, 2)  # apply butterworth filter to image with r=50 and n=2
b3_inv_f_shift = inverse_fourier_transform(b3_filtered)  # inverse transform the image
b3_abs = np.abs(b3_inv_f_shift)  # take absolute in order to save
cv2.imwrite('B3_output.png', b3_abs)  # save image