# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %%
from IPython import get_ipython

# %%
import torch
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from torch.distributions import MultivariateNormal
import torch.nn as nn

get_ipython().run_line_magic('matplotlib', 'inline')

# %%
def sample_gmm(mus, sigmas, pis, N):
    all_samples = []
    all_sample_inds = []
    for ind, pi in enumerate(pis):
        comp_N = int(round(N * pi.item()))
        samples = (torch.randn((comp_N,)) * sigmas[ind]) + mus[ind]
        all_samples.append(samples)
        sample_inds = torch.full((comp_N,), ind, dtype=torch.int)
        all_sample_inds.append(sample_inds)

    return torch.cat(all_samples), torch.cat(all_sample_inds)

def plot_gmm(samples):
    if isinstance(samples, torch.Tensor):
        samples = samples.detach().numpy()
    plt.hist(samples, bins=30)
    plt.show()

# %%

if False:
    mus = torch.tensor([-20., 0., 200])
    sigmas = torch.tensor([0.5, 3.0, 0.1])
    pis = torch.tensor([0.1, 0.4, 0.5])

if True:
    mus = torch.tensor([-2., 0., 2])
    sigmas = torch.tensor([0.5, 3.0, 0.1])
    pis = torch.tensor([0.1, 0.4, 0.5])

gt_params = [mus, sigmas, pis]
N = 1000
k = 3

samples, sample_inds = sample_gmm(mus, sigmas, pis, N)
plot_gmm(samples)

# %%
import math

def normal_prob(mu, sigma, x):
    return (1.0 / (sigma * math.sqrt(2.0 * math.pi))) * torch.exp(-0.5 * ((x - mu) / sigma) ** 2)

x = torch.linspace(-10., 10, 100)
y = normal_prob(0., 1., x)
plt.plot(x, y)

# %%

def log_normal_prob(mu, sigma, x):
    return (torch.log(torch.tensor((1.0 / (sigma * math.sqrt(2.0 * math.pi))))) +
            (-0.5 * ((x - mu) / sigma) ** 2))

x = torch.linspace(-10., 10, 100)
y = log_normal_prob(0., 1., x)
plt.plot(x, y)
plt.show()

plt.plot(x, y.exp())
plt.show()

# %%

x = torch.randn((3, 3), requires_grad=True)
lse = x.exp().sum(dim=1).log().mean()
lse2 = torch.logsumexp(x, dim=1).mean()

print(lse)
print(lse2)

lse.backward()
print(x.grad)

x.grad.zero_()
print(x.grad)
lse2.backward()
print(x.grad)

# %%

class GMM(nn.Module):
    def __init__(self, k):
        super(GMM, self).__init__()
        self.k = k
        self.mus = nn.Parameter(torch.randn((k,), requires_grad=True))
        self.sigma_params = nn.Parameter(torch.full((k,), 0.0, requires_grad=True))
        self.pi_params = nn.Parameter(torch.ones((k,), requires_grad=True))

    def get_dists(self):
        pis = torch.nn.functional.softmax(self.pi_params, dim=0)
        sigmas = self.sigma_params.exp()
        return self.mus, sigmas, pis

    def forward(self, X):
        mus, sigmas, pis = self.get_dists()
        probs = torch.cat([normal_prob(mus[i], sigmas[i], X).unsqueeze(1) for i in range(k)], dim=1)
        neg_log_like = -(probs * pis).sum(dim=1).log().mean()
        # log_probs = torch.cat([log_normal_prob(mus[i], sigmas[i], X).unsqueeze(1) for i in range(k)], dim=1)
        # neg_log_like = -torch.logsumexp(log_probs + pis.log(), dim=1).mean()
        return neg_log_like

# %%

import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

model = GMM(k)
batch_sz = 100
lr = 1e-2
num_epochs = 301
log_interval = 100

sample_mu = samples.mean()
sample_std = samples.std()
norm_samples = (samples - sample_mu) / sample_std

optimizer = optim.Adam(model.parameters(), lr=lr)
train_ds = TensorDataset(norm_samples)
train_dl = DataLoader(train_ds, batch_size=batch_sz, shuffle=True)

all_mus = []
all_sigmas = []
all_pis = []

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (X,) in enumerate(train_dl):
        optimizer.zero_grad()
        loss = model(X)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    if epoch % log_interval == 0:
        print(f'epoch: {epoch} / loss: {running_loss}')
    running_loss = 0.0

    mus, sigmas, pis = model.get_dists()
    all_mus.append(mus.data.clone().unsqueeze(0) * sample_std + sample_mu)
    all_sigmas.append(sigmas.clone().unsqueeze(0) * sample_std)
    all_pis.append(pis.data.clone().unsqueeze(0))

# %%

mus, sigmas, pis = model.get_dists()
model_samples, _ = sample_gmm(mus, sigmas, pis, N)
model_samples = (model_samples * sample_std) + sample_mu
plot_gmm(model_samples)

mus = torch.cat(all_mus).detach()
sigmas = torch.cat(all_sigmas).detach()
pis = torch.cat(all_pis).detach()
params = [mus, sigmas, pis]
titles = ['mus', 'sigmas', 'pis']

fig, axs = plt.subplots(len(params), 1, figsize=(len(params), 3 * len(params)))
for pi, p in enumerate(params):
    for ki in range(k):
        axs[pi].set_ylabel(titles[pi])
        axs[pi].axhline(gt_params[pi][ki], linestyle='--', alpha=0.2)
        axs[pi].plot(torch.arange(0, len(p)), p[:, ki])

# %%
from sklearn.mixture import GaussianMixture

skgmm = GaussianMixture(n_components=3, verbose=1)
skgmm.fit(samples.unsqueeze(1))
s = skgmm.sample(300)
plot_gmm(s[0])

# %%


import numpy as np
from scipy import interpolate

def u_func(x1,x2):
    a = 4
    b = 5
    c1 = 12
    c2 = 24
    c3 = 3
    c4 = 4
    c5 = 5
    c6 = 6
    return c1*(x1**4/12 - a*x1**3/6) + c2*(x2**4/12 - b*x2**3/6) \
        + c3*x1*x2 + c4*x1 + c5*x2 + c6

def cubic_func(x1,x2):
    return x1**3 + 2*x2**3

x1 = np.array([0, 2, 4])
x2 = np.array([0, 2, 4])
y = np.zeros((len(x1),len(x2)))

for foo,i in zip(x1, range(len(x1))):
    for bar,j in zip(x2, range(len(x2))):
        y[i][j] = cubic_func(foo, bar)
        # print(y[i][j])

# for foo,i in zip(x1, range(len(x1))):
#     for bar,j in zip(x2, range(len(x2))):
#         y[i][j] = u_func(foo, bar)

# f = interpolate.interp2d(x1, x2, y, kind='cubic')

x1new = np.arange(0, 5, 1)
x2new = np.arange(0, 5, 1)
for foo in x1new:
    for bar in x2new:
        print(cubic_func(foo, bar))
# for foo,i in zip(x1new, range(len(x1new))):
#     for bar,j in zip(x2new, range(len(x2new))):
#         ynew[i][j] = u_func(foo, bar)

# print(y)
# print()
# print(ynew)
# print()
# print(f(x1new, x2new))


# find differences between 2 elements of array
# e.g. for [1, 3, 5, 7, 9] here the diff between two number is 2 so output is [2, 2, 2, 2]
# to find discrete function use diff() function and use "n" to performs this operation repeatedly

import numpy as np
r1 = np.array([1, 2, 6, 15, 36])
r2 = np.diff(r1)
r3 = np.diff(r1, n=2)
print(r2)
print(r3)

Python/Numpy/numpy 38 - ufunc differences.py


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

# # Assignment 14
#
# I have provided a Python/NumPy implementation of a [PLU Decomposition](https://johnfoster.pge.utexas.edu/numerical-methods-book/LinearAlgebra_LU.html#Python/NumPy-implementation-of-$\mathbf{PLU}$-decomposition) in the course notes.  Because this function makes use of NumPy broadcasting, it's about as fast as they can be in Python, but this efficiency is at the expense of code readability.
#
# We can make this function more readable and maintainable by using an object-oriented approach.  Your assignment is to complete the `LU` class below.  Specifically, you need to implement the `decomp` member function to perform the $\mathbf{PLU}$ decomposition and store the resulting matrices in the class attributes `P, L` and `U`, respectively.
#
# After that, implement `det` and `inverse` to compute the input matrix's determinant and inverse, respectively, with the $\mathbf{PLU}$ decomposition.  I've already implemented the forward and backward substitution methods and call them from a function called `solve`.  If the argument to `solve` is a one-dimensional NumPy array, then a single solution is returned.  However, if the input to solve is a two-dimensional NumPy array, then each *row* is interpreted as a unique right-hand side vector and a two-dimensional NumPy array is returned with each *row* being a solution vector corresponding to the row of the input argument.  **Hint**: You should be able to implement `inverse` with a single call to `solve`.
#
# I am provided a `Matrix` class definition that has been extended from the one you implemented in [assignment12](https://github.com/PGE310-Students/assignment12).  This class is instantiated as the class attribute objects `P, L` and `U` and allows for indexing operations similar to Python lists and NumPy arrays as well as the row operation functions.  Please use this class and it's member functions to implement your functions when appropriate.  See the `forward/backward_substitution` functions for help using the class.

# In[1]:

import numpy as np

class Matrix(object):

    def __init__(self, array):

        if not isinstance(array, (list, np.ndarray)):
            raise TypeError('Input matrix must be a Python list or NumPy ndarray.')
        else:
            self.mat = np.array(array, dtype=np.double)

    def __str__(self):
        return str(self.mat)

    def __call__(self):
        return self.mat

    def __getitem__(self, key):
        return self.mat[key]

    def __setitem__(self, key, value):
        self.mat[key] = value

    def row_swap(self, i, j):
        temp_row = self.mat[j].copy()
        self.mat[j] = self.mat[i]
        self.mat[i] = temp_row
        return

    def row_multiply(self, i, factor):
        self.mat[i] *= factor
        return

    def row_combine(self, i, j, factor):
        self.mat[i] -= factor * self.mat[j]
        return

# In[2]:

class LU():

    def __init__(self, A):

        self.n = A.shape[0]

        #Instantiate P, L, and U as
        #Matrix objects and store them
        #as class attributes.
        self.U = Matrix(A.copy())
        self.L = Matrix(np.eye(self.n))
        self.P = Matrix(np.eye(self.n))

        #Perform the LU decomposition at
        #class instantiation
        self.decomp()

        return

    def decomp(self):
        #Perform the PLU decomposition and
        #store the matrices P, L, and U
        #as class attributes

        #It might be useful to keep track
        #of the number of permutations that
        #will be performed on P
        self.number_of_permutations = 0

        #Loop over rows
        for i in range(self.n):

            #Permute rows if needed
            for k in range(i, self.n):
                if ~np.isclose(self.U[i, i], 0.0):
                    break
                self.U.row_swap(k, k+1)
                self.P.row_swap(k, k+1)
                self.number_of_permutations += 1

            #Eliminate entries below i with row
            #operations on U and #reverse the row
            #operations to manipulate L
            for j in range(i+1, self.n):
                factor = self.U[j, i] / self.U[i, i]
                self.L[j, i] = factor
                self.U.row_combine(j, i, factor)

    def forward_substitution(self, b):

        #Permuting b consistent w/ P
        b = np.dot(self.P(), b)

        #Allocating space for the solution vector
        y = np.zeros_like(b, dtype=np.double);

        #Here we perform the forward-substitution.
        #Initializing  with the first row.
        y[0] = b[0] / self.L[0, 0]

        #Looping over rows in reverse (from the bottom  up),
        #starting with the second to last row, because  the
        #last row solve was completed in the last step.
        for i in range(1, self.n):
            y[i] = (b[i] - np.dot(self.L[i,:i], y[:i])) / self.L[i,i]

        return y

    def back_substitution(self, y):

        #Allocating space for the solution vector
        x = np.zeros_like(y, dtype=np.double);

        #Here we perform the back-substitution.
        #Initializing with the last row.
        x[-1] = y[-1] / self.U[-1, -1]

        #Looping over rows in reverse (from the bottom up),
        #starting with the second to last row, because the
        #last row solve was completed in the last step.
        for i in range(self.n-2, -1, -1):
            x[i] = (y[i] - np.dot(self.U[i,i:], x[i:])) / self.U[i,i]

        return x

    def solve(self, b):

        #Ensure b is a NumPy array
        b = np.array(b)

        #If b is a single vector/array, add a
        #second dimension so that the function
        #call signature can be the same for a
        #single vector or multiple vectors
        if len(b.shape) == 1:
            b = b.reshape(1, -1)

        x = np.zeros_like(b, dtype=np.double)

        #Loop over all right-hand side vectors and
        #store solution
        for i in range(b.shape[0]):
            y = self.forward_substitution(b[i])
            x[i] = self.back_substitution(y)

        if x.shape[0] == 1:
            return x[0]
        else:
            return x

    def inverse(self):
        n = self.U().shape[0]
        I = np.eye(n)
        #Return the inverse solution to
        #be in true matrix form
        return self.solve(I).T

    def det(self):
        if self.number_of_permutations % 2 == 0:
            return np.diagonal(self.U()).prod()
        else:
            return -np.diagonal(self.U()).prod()


import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import numpy as np
import scipy.io as sio
import itertools

def find_closest_centroids_slow(X, centroids):
    return np.array([np.argmin([np.sum((x-c)**2) for c in centroids]) for x in X])

def find_closest_centroids(X, centroids):
    diff = X[np.newaxis, :, :] - centroids[:, np.newaxis, :]
    return np.argmin(np.sum(diff**2, axis=-1), axis=0)

def compute_centroids_slow(X, idx, K):
    return np.array([np.mean(X[idx.ravel() == i, :], axis=0) for i in range(K)])

def compute_centroids(X, idx, K):
    idx = np.arange(K).reshape(-1, 1) == idx.reshape(1, -1)
    return idx.dot(X)/(np.sum(idx, axis=1).reshape(-1, 1))

def run_kmeans(X, initial_centroids, max_iters, plot_progress=False):
    previous_centroids = initial_centroids
    K = len(previous_centroids)
    for _ in xrange(max_iters):
        idx = find_closest_centroids(X, previous_centroids)
        centroids = compute_centroids(X, idx, K)
        if (centroids == previous_centroids).all():
            break

        if plot_progress:
            for k in xrange(K):
                x1 = [previous_centroids[k, 0], centroids[k, 0]]
                x2 = [previous_centroids[k, 1], centroids[k, 1]]

                plt.plot(x1, x2, 'k')
                plt.plot(x1, x2, 'kx')

        previous_centroids = centroids

    if plot_progress:
        for k, color in zip(xrange(K), itertools.cycle(['r', 'b', 'g', 'k', 'm'])):
            x1, x2 = X.T
            plt.plot(x1[idx == k], x2[idx == k], color+'o')
        plt.show()

    return idx, centroids

def kmeans_init_centroids(X, K):
    temp = range(len(X))
    np.random.shuffle(temp)

    return X[temp[:K], :]

if __name__ == '__main__':

    data = sio.loadmat('ex7data2.mat')
    X = data['X']

    K = 3
    initial_centroids = np.array([[3, 3], [6, 2], [8, 5]])

    idx = find_closest_centroids(X, initial_centroids)

    print('Closest centroids for the first 3 examples: ')
    print(idx[:3])
    print('(the closest centroids should be 0, 2, 1 respectively)')

    idx, centroids = run_kmeans(X, initial_centroids, 400, True)

    img=mpimg.imread('bird_small.png')
    imgplot = plt.imshow(img)
    plt.show()

    img_size = img.shape
    X = img.reshape(img_size[0]*img_size[1], img_size[2])
    K = 16
    max_iters = 10

    initial_centroids = kmeans_init_centroids(X, K)
    idx, centroids = run_kmeans(X, initial_centroids, max_iters)

    X_recovered = centroids[idx, :]
    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
    ax1.imshow(img)
    ax2.imshow(X_recovered.reshape(img_size))

    plt.show()


import numpy as np

def fitline(x,y,dy):
    A = np.vstack((np.ones_like(x), x)).T
    C = np.diag(dy * dy)
    cov = np.linalg.inv(np.dot(A.T, np.linalg.solve(C, A)))
    b_ls, m_ls = np.dot(cov, np.dot(A.T, np.linalg.solve(C, y)))
    return (m_ls,np.sqrt(cov[1,1])),(b_ls,np.sqrt(cov[0,0]))

def fitline0(x,y,dy):
    m, dm = np.ma.average(y/x,weights=(dy/x)**(-2),returned=1)
    return m, 1/np.sqrt(dm)

def fitconstant(y, dy):
    b, db = np.ma.average(y,weights=dy**(-2),returned=1)
    return b, 1/np.sqrt(db)


import numpy as np

def getD(A):
    return [[A[i].sum() if i == j else 0 for j in range(len(A))] for i in range(len(A))]

def findInverseOfD(D):
    return [[1 / D[i][j] if D[i][j] != 0 else 0 for j in range(len(D))] for i in range(len(D))]

def getM(A):
    D = getD(A)
    Dinv = findInverseOfD(D) # np.linalg.inv(D)
    return (Dinv @ A).T

# M = (D^-1 A)^T
# M is just prob that you move from row to collumn
# M should be ^inf power
def rank(A, power=2**30):
    M = getM(A)
    M = np.linalg.matrix_power(M, power)
    #SM = getM(A)
    #IM = SM
    #M = SM
    #while power != 1:
    #    if power % 2 == 0:
    #        power = power / 2
    #        IM = IM @ IM
    #    else:
    #        power -= 1
    #        M = M @ IM
    #        IM = SM

    P = np.reshape(np.array([1/len(A) for i in range(len(A))]), (len(A), 1))
    return M @ P


import math
import numpy as np

rad = math.pi / 180.0
R = 6378137.0

def great_circle_distance(lon1, lat1, lon2, lat2):
    """
    Usage
    -----
    Compute the great circle distance, in meter, between (lon1,lat1) and (lon2,lat2)

    Parameters
    ----------
    param lat1: float, latitude of the first point
    param lon1: float, longitude of the first point
    param lat2: float, latitude of the se*cond point
    param lon2: float, longitude of the second point

    Returns
    -------x
    d: float
       Great circle distance between (lon1,lat1) and (lon2,lat2)
    """

    dlat = rad * (lat2 - lat1)
    dlon = rad * (lon2 - lon1)
    a = (math.sin(dlat / 2.0) * math.sin(dlat / 2.0) +
         math.cos(rad * lat1) * math.cos(rad * lat2) *
         math.sin(dlon / 2.0) * math.sin(dlon / 2.0))
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    d = R * c
    return d

def great_circle_distance_traj(lons1, lats1, lons2, lats2, l1, l2):
    """
    Usage
    -----
    Compute pairwise great circle distance, in meter, between longitude/latitudes coordinates.

    Parameters
    ----------
    param lats1: float, latitudes of the firs trajectories
    param lons1: float, longitude of the trajectories
    param lats2: float, latitudes of the se*cond trajectories
    param lons2: float, longitudess of the second trajectories
    param l1 : int, length of the first trajectories
    param l2 : int, length of the second trajectories

    Returns
    -------x
    d: float
       Great circle distance between (lon1,lat1) and (lon2,lat2)
    """

    mdist = np.empty((l1, l2), dtype=float)
    for i in range(l1):
        for j in range(l2):
            mdist[i, j] = great_circle_distance(lons1[i], lats1[i], lons2[j], lats2[j])
    return mdist

def initial_bearing(lon1, lat1, lon2, lat2):
    """
    Usage
    -----
    Bearing between (lon1,lat1) and (lon2,lat2), in degree.

    Parameters
    ----------
    param lat1: float, latitude of the first point
    param lon1: float, longitude of the first point
    param lat2: float, latitude of the second point
    param lon2: float, longitude of the second point

    Returns
    -------
    brng: float
           Bearing between (lon1,lat1) and (lon2,lat2), in degree.

    """
    dlon = rad * (lon2 - lon1)
    y = math.sin(dlon) * math.cos(rad * (lat2))
    x = math.cos(rad * (lat1)) * math.sin(rad * (lat2)) - math.sin(rad * (lat1)) * math.cos(rad * (lat2)) * math.cos(
        dlon)
    ibrng = math.atan2(y, x)

    return ibrng

def cross_track_distance(lon1, lat1, lon2, lat2, lon3, lat3, d13):
    """
    Usage
    -----
    Angular cross-track-distance of a point (lon3, lat3) from a great-circle path between (lon1, lat1) and (lon2, lat2)
    The sign of this distance tells which side of the path the third point is on.

    Parameters :
    ----------
    param lat1: float, latitude of the first point
    param lon1: float, longitude of the first point
    param lat2: float, latitude of the second point
    param lon2: float, longitude of the second point
     param lat3: float, latitude of the third point
    param lon3: float, longitude of the third point

    Usage
    -----
    crt: float
         the (angular) cross_track_distance

    """

    theta13 = initial_bearing(lon1, lat1, lon3, lat3)  # bearing from start point to third point
    theta12 = initial_bearing(lon1, lat1, lon2, lat2)  # bearing from start point to end point

    crt = math.asin(math.sin(d13 / R) * math.sin(theta13 - theta12)) * R

    return crt

def along_track_distance(crt, d13):
    """
    Usage
    -----
    The along-track distance from the start point (lon1, lat1) to the closest point on the the path
    to the third point (lon3, lat3).

    Parameters
    ----------
    param crt : float, cross_track_distance
    param d13 : float, along_track_distance

    Returns
    -------
    alt: float
         The along-track distance
    """

    alt = math.acos(math.cos(d13 / R) / math.cos(crt / R)) * R
    return alt

def point_to_path(lon1, lat1, lon2, lat2, lon3, lat3, d13, d23, d12):
    """
    Usage
    -----
    The point-to-path distance between point (lon3, lat3) and path delimited by (lon1, lat1) and (lon2, lat2).
    The point-to-path distance is the cross_track distance between the great circle path if the projection of
    the third point lies on the path. If it is not on the path, return the minimum of the
    great_circle_distance between the first and the third or the second and the third point.

    Parameters
    ----------
    param lat1: float, latitude of the first point
    param lon1: float, longitude of the first point
    param lat2: float, latitude of the second point
    param lon2: float, longitude of the second point
    param lat3: float, latitude of the third point
    param lon3: float, longitude of the third point
    param d13 : float, great circle distance between (lon1, lat1) and (lon3, lat3)
    param d23 : float, great circle distance between (lon2, lat2) and (lon3, lat3)
    param d12 : float, great circle distance between (lon1, lat1) and (lon2, lat2)

    Returns
    -------

    ptp : float
          The point-to-path distance between point (lon3, lat3) and path delimited by (lon1, lat1) and (lon2, lat2)

    """
    crt = cross_track_distance(lon1, lat1, lon2, lat2, lon3, lat3, d13)
    d1p = along_track_distance(crt, d13)
    d2p = along_track_distance(crt, d23)
    if (d1p > d12) or (d2p > d12):
        ptp = np.min((d13, d23))
    else:
        ptp = np.abs(crt)
    return ptp


import numpy as np

def gauss(x, mu, sigma):
    """
    Given x, mu and sigma, this function returns the normalized gausian distribution value at x.
    """
    return np.exp(-0.5 * ((x - mu) / sigma) ** 2) / sigma / np.sqrt(2 * np.pi)

def get_ihhomogeneous_broadening_parameters(fwhm, no_of_points):
    """
    Given the full width half maximum (FWHM) of your inhomogenous broadening and number of points, this function returns
    the frequencies of the broadening, the probability profile of each frequency and the constant frequency step that
    was used.
    """
    db_fwhm = fwhm * 2 * np.pi
    db_sig = db_fwhm / (2 * np.sqrt(2 * np.log(2)))
    db_x = np.linspace(-5 * db_sig, 5 * db_sig, no_of_points)
    db_dx = db_x[1] - db_x[0]
    db_y = gauss(db_x, 0, db_sig)
    prb_prof = db_y*db_dx

    return db_x, prb_prof, db_dx


import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
import argparse
import copy
from numpy.core.fromnumeric import repeat
from knn import KNN

has_converged = False

def KMeans(data, k):
    """
    Takes in unlabeled data and a number of means to compute.
    Step 1: Choose k data points at random to create a mean at.
    Step 2: For each data point, compute which mean is closest and assign it to that mean.
    Step 3: Recompute the mean based on the data points assigned to it.
    Repeat step 2 and 3 until convergence.
    Yields frames of the animation
    """
    random_indices = np.random.choice(len(data), size=k, replace=False)
    previous_means = []
    means = data[random_indices, :]

    def converged(curr, prev):
        if len(curr) != len(prev):
            return False
        for a, b in zip(curr, prev):
            if not np.array_equal(np.array(a), np.array(b)):
                return False
        global has_converged
        has_converged = True
        return True

    while not converged(means, previous_means):
        previous_means = copy.deepcopy(means)
        points_segregated = [[] for _ in range(k)]
        for point in data:
            min_dist = np.inf
            for mean_idx, mean in enumerate(means):
                dist = np.linalg.norm(mean - point)
                if dist < min_dist:
                    min_dist = dist
                    assigned_mean = mean_idx
            points_segregated[assigned_mean].append(point)
        yield means, points_segregated

        for i, points in enumerate(points_segregated):
            means[i] = sum(points) / len(points)
        yield means, points_segregated

parser = argparse.ArgumentParser()
parser.add_argument("-k", type=int, help="The K in Kmeans (the number of groups)", default=3)
args = parser.parse_args()

# I can run any implemented algorithm from here
# Let's say the total dataset comes from 3 normally distributed sources
data_source_0 = np.random.multivariate_normal((-0.5, -0.5), np.identity(2)*0.2, 100)
data_source_1 = np.random.multivariate_normal((1.5, 1.), np.identity(2)*0.2, 100)
data_source_2 = np.random.multivariate_normal((0.5, 1.5), np.identity(2)*0.2, 100)
unlabeled_data = np.append(data_source_0, data_source_1, 0)
unlabeled_data = np.append(unlabeled_data, data_source_2, 0)

# Just data fig
fig, ax = plt.subplots()
color_array = ["red", "green", "blue", "cyan", "magenta", "yellow", "black", "grey"]
data_scats = [ax.scatter([], [], s=10) for i in range(args.k)]
mean_scat = ax.scatter([], [], marker="*", s=250)

def init():
    ax.set_xlim(min(unlabeled_data[:, 0]) - 0.05, max(unlabeled_data[:, 0]) + 0.05)
    ax.set_ylim(min(unlabeled_data[:, 1]) - 0.05, max(unlabeled_data[:, 1]) + 0.05)
    return mean_scat, *data_scats

def animate(frame, generator):
    try:
        means, point_sets = next(generator)
    except StopIteration:
        return mean_scat, *data_scats

    mean_scat.set_offsets(means)
    mean_scat.set(facecolor=color_array, edgecolor="black")
    for i, (data_scat, point_set) in enumerate(zip(data_scats, point_sets)):
        data_scat.set_offsets(point_set)
        data_scat.set(color=color_array[i])
    return mean_scat, *data_scats

def frame_generator():
    i = 0
    global has_converged
    while not has_converged:
        yield i
        i += 1

generator = KMeans(unlabeled_data, k=args.k)
anim = FuncAnimation(fig, animate, frames=frame_generator, interval=100, fargs=(generator,), init_func=init, blit=True, repeat=True)
writergif = matplotlib.animation.PillowWriter(fps=10)

anim.save(f"animations/kmeans_anim_k{args.k}.gif", writer=writergif)


# https://dsp.stackexchange.com/q/76560/50076 ################################
# -*- coding: utf-8 -*-
import numpy as np
from numpy.fft import ifft, ifftshift
from kymatio.scattering1d.filter_bank import morlet_1d
from ssqueezepy.visuals import plot as _plot, plotscat

def plot(*args, **kw):
    if 'title' in kw:
        kw['title'] = (kw['title'], {'fontsize': 18})
    mx = re_im_max(args[0]) * 1.05
    kw['ylims'] = (-mx, mx)
    _plot(*args, **kw)

def re_im_max(x):
    return max(np.abs(x.real).max(), np.abs(x.imag).max())

def re_im_parts(x, y):
    A =  x.real * y.real
    B = -x.imag * y.imag
    C =  x.real * y.imag
    D =  x.imag * y.real
    return A, B, C, D

def l2(x):
    return np.sqrt(np.sum(np.abs(x)**2))

#%%# Morlet visuals ##########################################################
N = 256
pf = morlet_1d(N, xi=.025, sigma=.1/20)
pt = ifftshift(ifft(pf))
#%%
kw = dict(complex=1, show=1, ticks=(1, 0))
plot(pt, **kw, title="psi: analytic Morlet")

#%%
A, B, C, D = re_im_parts(pt, pt)
plot(A + 1j*B, **kw, title="(psi * psi).real")
plot(C + 1j*D, **kw, title="(psi * psi).imag")
plot(pt * pt,  **kw,
     title="(psi * psi).real + (psi * psi).imag = psi * psi")

#%%
pta = np.conj(pt)
plot(pta, **kw, title="psia: anti-analytic Morlet")

#%%
plot(pt.imag + 1j*pta.imag, **kw, title="psia = psi.real - psi.imag")

#%%
A, B, C, D = re_im_parts(pt, pta)
plot(A + 1j*B, **kw, title="(psi * psia).real")
plot(C + 1j*D, **kw, title="(psi * psia).imag")
plot(pt * pta, **kw,
     title="(psi * psia).real + (psi * psia).imag = psi * psia")

#%%# Random sequence visuals #################################################
def viz_seq(x, show=1, complex=1, **kw):
    mx = re_im_max(x) * 1.1
    plotscat(x, complex=complex, ylims=(-mx, mx), show=show, **kw,
             hlines=(0, {'color': 'tab:red', 'linewidth': 1}),
             ticks=(1, 0), auto_xlims=0)

def zero_sum_seq(N):
    x = np.random.randn(N) + 1j*np.random.randn(N)
    x[N//2:] = 0

    slc = x[:N//2][::-1]
    x[N//2:] = slc.real - 1j * slc.imag
    x -= x.mean()

    x.real *= (l2(x.imag) / l2(x.real))
    return x

#%%
np.random.seed(0)
N = 12
x = zero_sum_seq(N)
y = zero_sum_seq(N)

viz_seq(x,     title="x | sum = {:.2f}".format(x.sum()))
viz_seq(y,     title="y | sum = {:.2f}".format(y.sum()))
viz_seq(x * x, title="x * x | sum = {:.2f}".format((x*x).sum()))
viz_seq(y * y, title="y * y | sum = {:.2f}".format((y*y).sum()))
viz_seq(x * y, title="x * y | sum = {:.2f}".format((x * y).sum()))

SignalProcessing/Q76560 - Analyticity - Symmetries/main.py


from math import *
import numpy as np
from scipy.integrate import quad

class Distances():

    def __init__(self,z0=0,z1=.3,z2=1.5,Om=.25,Ol=.75,Or=8.4e-5,H0=70):

        self.Om = Om
        self.Or = Or
        self.Ol = Ol
        self.z0 = z0
        self.z1 = z1
        self.z2 = z2
        self.H0 = H0

    def comoving_distance(self,z1,z2):
        invE = lambda z: 1./np.sqrt(self.Or*(1+z)**4+self.Om*(1+z)**3+self.Ol)
        res = quad(invE,z1,z2)[0]*3e5/self.H0
        return res

    def angular_diameter_distance(self,z1,z2):
        invE = lambda z: 1/np.sqrt(self.Or*(1+z)**4+self.Om*(1+z)**3+self.Ol)
        res = quad(invE,z1,z2)[0]*3e5/self.H0
        return res/(1+z2)

    def DA(self,OR=0.):
        DH = 3e5/self.H0
        Dm1 = self.comoving_distance(self.z0,self.z1)
        Dm2 = self.comoving_distance(self.z0,self.z2)
        return (Dm2*np.sqrt(1+OR*Dm1**2/DH**2)-Dm1*np.sqrt(1+OR*Dm2**2/DH**2))/(1+self.z2)

    def get_Sigmacr(self):
        G = 6.67e-11*3.24e-23*1.989e30*1e-6  #Mpc Msun^-1 (km/s)^2
        c = 3e5 #km/s
        Da1 = self.angular_diameter_distance(self.z0,self.z1)
        Da2 = self.angular_diameter_distance(self.z0,self.z2)
        Da12 = self.angular_diameter_distance(self.z1,self.z2)
        return c**2/(4*pi*G) * Da2/(Da1*Da12)

if __name__=="__main__":

    d = Distances()
    Dm1 =  d.comoving_distance(0,.3)
    Dm2 = d.comoving_distance(0,1.5)
    Da1 = d.angular_diameter_distance(0,.3)
    Da2 = d.angular_diameter_distance(0,1.5)
    Da12 = d.angular_diameter_distance(.3,1.5)
    Sigma_cr = d.get_Sigmacr()  #Msun Mpc^-2
    print("The angular diameter distances are found to be: ","{:.2f},{:.2f},{:.2f}".format(Da1,Da2,Da12),"Mpc.")
    print("The critical mass density is found to be: ","%.3e" % Sigma_cr,"Msun/Mpc^2.")


from scipy.stats import norm
import numpy as np

def get_wiener_increment(dt, t):
    """
    Get the Wiener (white-noise) increment dW for timestep dt

    dW is sampled from Gaussian with mean of 0, variance of dt

    :param dt: the time increment
    :return: a Wiener increment dW
    """
    return norm.rvs(scale=dt)

def get_coherent_wiener(mean, dt):
    return norm.rvs(loc=mean, scale=dt)

def make_coherent_wiener_func(mean):
    """
    Return a "coherent" (non-zero mean) wiener process with given mean
    :param mean: the centre of the wiener process
    :return: a Wiener process with offset mean, taking one argument dt and returning the Wiener increment
    """
    return lambda dt, t: get_coherent_wiener(mean, dt)

def get_cosine_increment(frequency, amplitude, dt, t):
    # u(t) = A sin(2 pi f t)
    # du(t) = 2 pi f A cos(2 pi f t) dt
    return 2 * np.pi * frequency * amplitude * np.cos(2 * np.pi * frequency * t) * dt

def make_cosine_input(frequency, amplitude=1):
    return lambda dt, t: get_cosine_increment(frequency, amplitude, dt, t)


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

#calculate distance function
def cal_distance(test,cal):
    distance_list=[]#list of distance
    for i in range(len(cal)):
        distance=[]
        distance.append(cal[i, 4])#append corresponding label
        d=test-cal[i,:]#test sample value-training sample value
        #calculate distance
        distance.append(math.sqrt(d[0,0]**2+d[0,1]**2+d[0,2]**2+d[0,3]**2))
        distance_list.append(distance)
    #print(distance_list)
    return distance_list

if __name__=='__main__':
    iris_data=pd.read_csv("iris.data",header=None)
    labels_codes=pd.Categorical(iris_data[4]).codes
    for i in range(150):
        iris_data.loc[i,4]=labels_codes[i]
    datalist=iris_data.values.tolist()
    random.seed(17)
    random.shuffle(datalist)
    data_set=np.mat(datalist)
    data_set=np.vsplit(data_set,5)

    sub_set1=data_set[0].copy()
    sub_set2=data_set[1].copy()
    sub_set3=data_set[2].copy()
    sub_set4=data_set[3].copy()
    sub_set5=data_set[4].copy()

    average_acc=[]
    k_list=[]
    for K in range(1,120):
        if K%3!=0:
            k_list.append(K)
            accuracy=[]
            for i in range(5):
                test_data = data_set[4 - i].copy()
                if i == 0:
                    cal_data = np.vstack((sub_set1, sub_set2, sub_set3, sub_set4))
                elif i == 1:
                    cal_data = np.vstack((sub_set1, sub_set2, sub_set3, sub_set5))
                elif i == 2:
                    cal_data = np.vstack((sub_set1, sub_set2, sub_set4, sub_set5))
                elif i == 3:
                    cal_data = np.vstack((sub_set1, sub_set3, sub_set4, sub_set5))
                else:
                    cal_data = np.vstack((sub_set2, sub_set3, sub_set4, sub_set5))
                right = 0
                for x in range(len(test_data)):
                    d_set = np.mat(cal_distance(test_data[x], cal_data))#list of distance
                    d_set = (d_set[np.lexsort(d_set.T)])[0, :, :]#sort distance list
                    p_wk = [0, 0, 0]#be used to record P(wk|x)
                    #calculate P(wk|x) for each test sample
                    for y in range(K):
                        if d_set[y, 0] == 0:
                            p_wk[0] = p_wk[0] + 1 / K
                        elif d_set[y, 0] == 1:
                            p_wk[1] = p_wk[1] + 1 / K
                        else:
                            p_wk[2] = p_wk[2] + 1 / K
                    #print(p_wk)
                    #calculate accuracy
                    if p_wk.index(max(p_wk)) == test_data[x, 4]:
                        right = right + 1
                accuracy.append(right / len(test_data))
            accuracy=np.mat(accuracy)
            print(accuracy)
            average_acc.append(np.mean(accuracy))
    plt.scatter(k_list,average_acc)
    plt.title('KNN')
    plt.xlabel('hyperparameter K')
    plt.ylabel('average accuracy')
    plt.show()
    k_max_list=[]
    for z in range(len(average_acc)):
        if average_acc[z]==max(average_acc):
            k_max_list.append(k_list[z])
    print("highest average accuracy:",round(max(average_acc),3))
    print("corresponding K:",k_max_list)


# -*- coding: utf-8 -*-
"""
Created on Mon Dec 25 09:17:02 2017

@author: philt
"""

import numpy as np
import matplotlib.pyplot as plt

# load data
X = []
Y = []

# Can use cv reader  - Try later
for line in open('data_1d.csv'):
    x, y = line.split(',')
    X.append(float(x))
    Y.append(float(y))

# Turn the input into numpy arrays
X = np.array(X)
Y = np.array(Y)

plt.scatter(X, Y)
plt.show()

# Apply eqns for linear regression  a,b
denominator = X.dot(X) - X.mean()*X.sum()

a = (X.dot(Y) - Y.mean() * X.sum()) / denominator
b = (Y.mean() * X.dot(X) - X.mean() * X.dot(Y)) / denominator

# Predicted Y
Yhat = a*X + b
plt.scatter(X,Y)
plt.plot(X, Yhat)
plt.show()

# Calculate R-Squared
d1 = Y - Yhat
d2 = Y - Y.mean()
r2 = 1 - d1.dot(d1) / d2.dot(d2)
print("The R2 is :  ", r2)


import numpy as np

def sigmoid(x):
	return 1.0 / (1 + np.exp(-x))

def y_hat(weights, bias, x):
	return np.dot(weights, x) + bias

def cost(y, output):
	return -(y*np.log(output) - (1-y)*np.log(1-output))

def gradient_descent(x, y, weights, bias, learnrate):
	y_h = y_hat(weights, bias, x)
	weights += learnrate * (y-y_h) * x
	bias += learnrate * (y-y_h)
	return weights, bias

Deep Learning with PyTorch/Part 0: Intro to ML/gradient_descent.py

Deep Learning with PyTorch/Lesson 1/gradient_descent.py


import numpy as np

def mohr_c(stress_x, stress_y, shear):
    """

    mohr_c() outputs two values for the circle, center, and radius

    output:
        C, R
        C is x-value of the circle center
        R is the radius of the circle
    """
    C = (stress_x + stress_y) / 2
    R = ((((stress_x - stress_y) / 2) ** 2) + shear ** 2) ** 0.5
    return C, R

def c_array(C, R, n=100):
    """

    c_array() outputs two NumPy arrays of x and y values that creates the circle

    output:
        x, y
        x is a NumPy array with n-values that contains the x-values to build the circle
        y is a NumPy array with n-values that contains the y-value to build the circle
    """
    t = np.linspace(0, 2 * np.pi, n + 1)
    x = R * np.cos(t) + C
    y = R * np.sin(t)
    return x, y

def X_Y(stress_x, stress_y, shear):
    """build the arrays that describe the line X-Y between the known points"""
    X = np.array([stress_x, stress_y])
    Y = np.array([-shear, shear])
    return X, Y


import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

def sigmoid_prime(z):
    return sigmoid(z) * (1 - sigmoid(z))


"""
Name: Pham Tuan Anh
Class: K63K2
MSSV: 18020116

You should understand the code you write.
"""

import cv2
import numpy as np
import argparse

def conv_sum(a, b):
    s = 0
    for i in range(len(a)):
        for j in range(len(a[0])):
            s += a[i][j]*b[i][j]
    return s

def my_convolution(I, g, mode='valid', boundary='zero_padding'):
    h, w = len(g), len(g[0])
    H, W = I.shape[0], I.shape[1]

    if mode == 'valid':
        output_h = H - (h - 1)
        output_w = W - (w - 1)
    else:
        output_h = H
        output_w = W

    output = [[0 for _ in range(output_w)] for _ in range(output_h)]
    for i in range(output_h):
        for j in range(output_w):
            output[i][j] = conv_sum(I[i:i + h, j:j + w], g)

    return output

# Calculate gradient magnitude and gradient direction
def gradient(img, s):
    img = np.array(img)

    Mx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    My = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

    smoothed_img = cv2.GaussianBlur(img, (s, s), 0)
    img_x = my_convolution(smoothed_img, Mx)
    img_y = my_convolution(smoothed_img, My)

    gradient_magnitude = np.sqrt(np.square(img_x) + np.square(img_y))
    gradient_direction = np.arctan2(img_y, img_x)

    return (gradient_magnitude, gradient_direction)

# Non-maximum suppression
def non_maximum_suppression(gradient_magnitude, gradient_direction):
    (h, w) = gradient_magnitude.shape[:2]
    out_img = np.zeros((h, w), dtype=np.int32)

    theta = gradient_direction * 180 / np.pi
    theta[theta < 0] += 180
    pi = 180
    for i in range(1, h - 1):
        for j in range(1, w - 1):
            direction = theta[i, j]

            r_pixel = 255
            l_pixel = 255

            # angle 0
            if (0 <= direction < pi / 8) or (7 * pi / 8 <= direction <= pi):
                r_pixel = gradient_magnitude[i, j+1]
                l_pixel = gradient_magnitude[i, j-1]

            # angle 45
            elif (pi / 8 <= direction < 3 * pi / 8):
                r_pixel = gradient_magnitude[i+1, j-1]
                l_pixel = gradient_magnitude[i-1, j+1]

            # angle 90
            elif (3 * pi / 8 <= direction < 5 * pi / 8):
                r_pixel = gradient_magnitude[i+1, j]
                l_pixel = gradient_magnitude[i-1, j]

                # angle 135
            elif (5 * pi / 8 <= direction < 7 * pi / 8):
                r_pixel = gradient_magnitude[i-1, j-1]
                l_pixel = gradient_magnitude[i+1, j+1]

            if (gradient_magnitude[i, j] >= r_pixel) and (gradient_magnitude[i, j] >= l_pixel):
                out_img[i, j] = gradient_magnitude[i, j]
            else:
                out_img[i, j] = 0

    return out_img

# Hysteresis threshold:
def threshold(non_max_img, lowThresholdRatio=0.05, highThresholdRatio=0.1):
    highThreshold = non_max_img.max() * highThresholdRatio
    lowThreshold = lowThresholdRatio * highThreshold

    (h, v) = non_max_img.shape[:2]
    out_img = np.zeros((h, v), dtype=np.int32)

    weak = np.int32(25)
    strong = np.int(255)

    strong_i, strong_j = np.where(non_max_img >= highThreshold)
    weak_i, weak_j = np.where((lowThreshold <= non_max_img) & (non_max_img <= highThreshold))
    irrelevant_i, irrelevant_j = np.where(non_max_img < lowThreshold)

    out_img[strong_i, strong_j] = strong
    out_img[weak_i, weak_j] = weak

    return (out_img, weak, strong)

# Determine which pixels are part of real edges
def hysteresis(img, weak, strong=255):
    (h, w) = img.shape[:2]
    for i in range(1, h-1):
        for j in range(1, w-1):
            if (img[i, j] == weak):
                if (img[i+1, j-1] == strong) or (img[i+1, j] == strong) or (img[i+1, j+1] == strong) or (img[i, j-1] == strong) or (img[i, j+1] == strong) or (img[i-1, j-1] == strong) or (img[i-1, j] == strong) or (img[i-1, j+1] == strong):
                    img[i, j] = strong
                else:
                    img[i, j] = 0

    return img

# Canny algorithm implementation
def canny(input_file, output_file, s):
    img = cv2.imread(input_file, cv2.IMREAD_GRAYSCALE)

    # Apply Gaussian blur
    smoothed_img = cv2.GaussianBlur(img, (s, s), 0)
    gradient_magnitude, gradient_direction = gradient(smoothed_img, s)
    output = non_maximum_suppression(gradient_magnitude, gradient_direction)
    cv2.imshow("Output 1", np.array(output).astype(np.uint8))
    output, weak, strong = threshold(output)
    cv2.imshow("Output 2", np.array(output).astype(np.uint8))
    output = hysteresis(output, weak, strong)
    cv2.imwrite(output_file, np.array(output).astype(np.uint8))
    cv2.imshow("Output 3", np.array(output).astype(np.uint8))
    cv2.waitKey(0)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", "-i", type=str, help="Path to input image")
    parser.add_argument("--output_file", "-o", type=str, help="Path to output image")
    parser.add_argument("--size", "-s", type=int, default=3, help="Size of gaussian filter")

    args = parser.parse_args()
    canny(args.input_file, args.output_file, args.size)


import numpy as np

def q_from_v(env, V, s, gamma=0.99):
    q = np.zeros(env.nA)
    for a in range(env.nA):
        for prob, s_next, reward in env.MDP[s][a]:
            q[a] += prob * (reward + gamma * V[s_next])
    return q

def policy_evaluation(env, policy, gamma=0.99, theta=1e-8):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            q = q_from_v(env, V, s, gamma)
            v = np.sum(policy[s] * q)
            delta = max(delta, np.abs(V[s] - v))
            V[s] = v
        if delta < theta:
            break
    return V

def policy_improvement(env, V, gamma=0.99):
    policy = np.zeros([env.nS, env.nA]) / env.nA
    for s in range(env.nS):
        q = q_from_v(env, V, s, gamma)
        a = np.argmax(q)
        policy[s] = np.eye(env.nA)[a]
    return policy

def policy_iteration(env, gamma=0.99, theta=1e-8):
    # 1. initialization
    V = np.zeros(env.nS)
    policy = np.ones([env.nS, env.nA]) / env.nA
    while True:
        # 2. policy evaluation
        V = policy_evaluation(env, policy, gamma, theta)
        # 3. policy improvement
        new_policy = policy_improvement(env, V, gamma)
        if np.all(policy == new_policy):
            # if policy-stable
            break
        policy = new_policy.copy()
    return policy, V

def value_iteration(env, gamma=0.99, theta=1e-8):
    V = np.zeros(env.nS)
    policy = np.ones([env.nS, env.nA]) / env.nA
    while True:
        delta = 0
        for s in range(env.nS):
            q = q_from_v(env, V, s, gamma)
            v = np.max(q)
            delta = max(delta, np.abs(V[s] - v))
            V[s] = v
        if delta < theta:
            policy = policy_improvement(env, V, gamma)
            break
    return policy, V


# This feature can be used to reduce the memory size consumed by the feature layer of the top MLP.
# Suppose we have n sparse features, each sparse features is represented by an embedding of size d,
# then, we can represent the sparse embeddings by a matrix X = (n, d). The dot product between sparse
# features is X(X^T), which is a symmetric matrix of (n, n) and will be fed into the top MLP.
# Actually We only need the upper or lower traingles to eliminate duplication. If n is large,
# such as, n = 1000, then the number of dot features fed into the MLP will be n^2/2 = 50,000.
# Considering the layer size 4096, the weight parameters will be a matrix (n^2/2, 4096), which
# may consume a large amount of precious memory resources.

# To reduce the number of dot features, we introduce a parameter called arch-projec-size (k) to compress
# the embeddings. We introduce a parameter matrix Y = (n, k) to compute the weighted sum of the
# dot features. The compressed embeddings is represented by (X^T)Y. Then, we compute the compressed dot
# features by X(X^T)Y = (n, k). Therefore, we can reduce the dot features fed into MLP from n*n/2
# to n*k.

import sys
import torch
import torch.nn as nn
import numpy as np

"""
Compute the projected dot features
T: (batch_size, n, d), batched raw embeddings
x: dense features
proj_layer: the projection layer created by create_proj
"""
def project(T, x, proj_layer):

  TT = torch.transpose(T, 1, 2)
  # TS = torch.reshape(TT, (-1, TT.size(2)))
  # TC = proj_layer(TS)
  # TR = torch.reshape(TC, (-1, T.shape[2], k))
  TR = proj_layer(TT)
  Z  = torch.bmm(T, TR)
  Zflat = Z.view((T.shape[0], -1))
  R = torch.cat([x] + [Zflat], dim=1)

  return R

"""
Create the project layer
n: number of sparse features
m: projection size
"""
def create_proj(n, m):
  # build MLP layer by layer
  layers = nn.ModuleList()
  # construct fully connected operator
  LL = nn.Linear(int(n), int(m), bias=True)

  # initialize the weights
  # with torch.no_grad():
  # custom Xavier input, output or two-sided fill
  mean = 0.0  # std_dev = np.sqrt(variance)
  std_dev = np.sqrt(2 / (m + n))  # np.sqrt(1 / m) # np.sqrt(1 / n)
  W = np.random.normal(mean, std_dev, size=(m, n)).astype(np.float32)
  std_dev = np.sqrt(1 / m)  # np.sqrt(2 / (m + 1))
  bt = np.random.normal(mean, std_dev, size=m).astype(np.float32)
  # approach 1
  LL.weight.data = torch.tensor(W, requires_grad=True)
  LL.bias.data = torch.tensor(bt, requires_grad=True)
  # approach 2: constant value ?
  layers.append(LL)

  return torch.nn.Sequential(*layers)


from liegroups import SE3, SO3
from scipy.io import loadmat
from math import cos, sin, sqrt
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from scipy.linalg import block_diag

D = np.array(
    [
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1],
        [0, 0, 0],
    ]
)

def skew(r):
    return np.array(
        [
            [0, -r[2], r[1]],
            [r[2], 0, -r[0]],
            [-r[1], r[0], 0],
        ]
    )

def so3_exp(psi):

    angle = np.linalg.norm(psi)

    axis = psi / angle
    s = np.sin(angle)
    c = np.cos(angle)

    return c * np.eye(3) + (1 - c) * np.outer(axis, axis) - s * skew(axis)

def so3_log(C):

    cos_angle = 0.5 * np.trace(C) - 0.5

    # Clip cos(angle) to its proper domain to avoid NaNs from rounding errors
    cos_angle = np.clip(cos_angle, -1.0, 1.0)
    angle = np.arccos(cos_angle)

    if np.isclose(angle, 0.0):
        tmp = C - np.eye(3)

    else:
        tmp = (angle / (2 * np.sin(angle))) * (C - C.T)

    return np.array([tmp[2, 1], tmp[0, 2], tmp[1, 0]])

def so3_jacob(psi):
    angle = np.linalg.norm(psi)

    if np.isclose(angle, 0.0):
        return np.eye(3) + 0.5 * skew(psi)

    axis = psi / angle
    s = np.sin(angle)
    c = np.cos(angle)

    return (
        (s / angle) * np.eye(3)
        + (1 - s / angle) * np.outer(axis, axis)
        + (1 - c) / angle * skew(axis)
    )

def se3_exp(d, psi):

    # Compute jacobian
    J = so3_jacob(psi)

    C = so3_exp(psi)
    return np.block([[C, (J @ d).reshape(3, 1)], [0, 0, 0, 1]])

def se3_log(T):
    psi = so3_log(T[0:3, 0:3])

    J = so3_jacob(psi)
    r = np.linalg.inv(J) @ T[0:3, 3]

    return np.hstack((r, psi))

class BatchEstimator:
    """
    Batch Estimator
    """

    def __init__(self, dataset):
        """ """

        # Load all input data
        dataset = loadmat(dataset)

        # Extract scalar values that are stored in 2D arrays
        for key in dataset:
            if key.startswith("__"):
                continue
            if dataset[key].size == 1:
                dataset[key] = dataset[key][0][0]
            setattr(self, key, dataset[key])

        self.rho_v_c_v = self.rho_v_c_v.reshape(
            -1,
        )
        self.v_var = (
            self.v_var.reshape(
                -1,
            )
            * 3
        )
        self.w_var = self.w_var.reshape(
            -1,
        )
        self.y_var = self.y_var.reshape(
            -1,
        )

        # Time steps
        self.time_steps = [
            self.t[:, k] - self.t[:, k - 1] for k in range(1, len(self.t[0]))
        ]

        # Number of landmark points in the map
        self.num_points = 20

        self.T_c_v = np.block(
            [[self.C_c_v, -self.C_c_v @ self.rho_v_c_v.reshape(3, 1)], [0, 0, 0, 1]]
        )

    def ground_truth_k(self, k):
        C_vk_i = so3_exp(self.theta_vk_i[:, k])
        # C_vk_i = SO3.exp(self.theta_vk_i[:, k]).as_matrix()
        T = SE3.from_matrix(
            np.block(
                [[C_vk_i, -C_vk_i @ self.r_i_vk_i[:, k].reshape(3, 1)], [0, 0, 0, 1]]
            ),
            normalize=True,
        )
        return T

    def ground_truth(self, k1=1215, k2=1714):
        """
        return the ground truth poses.
        """
        # theta_vk_i is in axis-angle representation
        # https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Rotating_a_vector
        return [self.ground_truth_k(k) for k in range(k1, k2 + 1)]

    def forward(self, T_prev, psi, d):
        """Forward Dead Reckoning step."""
        dT = SE3.from_matrix(se3_exp(d, -psi)).inv()
        return dT.dot(T_prev)

    def dead_reckoning(self, k1=1215, k2=1714, init=None):
        """
        Estimate position using only dead reckoning
        from the input velocities
        """
        states = []

        # Set initial state using ground-truth values
        if init:
            states.append(init)
        else:
            states.append(self.ground_truth_k(k1))

        # For each time step, computes the difference in
        # rotation and position due to the linear and angular
        # velocities.
        for k in range(k1, k2):

            # Rotation = Angular velocity * time period
            psi = self.w_vk_vk_i[:, k] * self.time_steps[k]

            # Distance = Linear velocity * time period
            d = self.v_vk_vk_i[:, k] * self.time_steps[k]

            states.append(self.forward(states[-1], psi, d))

        return states

    # Given a point p in the camera frame, project it
    # on the images such that pixels are returned
    def g(self, p):
        """Given a point and pose project into camera."""

        return np.array(
            [
                self.fu * p[0] / p[2] + self.cu,
                self.fv * p[1] / p[2] + self.cv,
                self.fu * (p[0] - self.b) / p[2] + self.cu,
                self.fv * p[1] / p[2] + self.cv,
            ]
        )

    # Given the state {r,C} of the vehicle, project the
    # i-th point expressed wrt the inertial frame
    # into the camera frame and return its [x,y,z] coords.
    def p_i(self, T, i):
        """Project point into camera frame."""
        return D.T @ self.T_c_v @ T.as_matrix() @ np.block([self.rho_i_pj_i[:, i], 1])

    # Jacobian of g() wrt to point p
    def dg_dp(self, p):
        """Camera jacobian."""
        x = p[0]
        y = p[1]
        z = p[2]
        z_squared = z ** 2
        dg_dp = np.array(
            [
                [self.fu / z, 0, -self.fu * x / z_squared],
                [0, self.fv / z, -self.fv * y / z_squared],
                [self.fu / z, 0, -self.fu * (x - self.b) / z_squared],
                [0, self.fv / z, -self.fv * y / z_squared],
            ]
        )
        return dg_dp

    def dg_dp_numerical(self, P):
        """Camera jacobian."""
        d = 0.00001
        dp_dx = self.g(P + np.array([d, 0, 0])) - self.g(P)
        dp_dy = self.g(P + np.array([0, d, 0])) - self.g(P)
        dp_dz = self.g(P + np.array([0, 0, d])) - self.g(P)

        return np.block(
            [dp_dx.reshape(4, 1) / d, dp_dy.reshape(4, 1) / d, dp_dz.reshape(4, 1) / d]
        )

    def dp_dx(self, T, p):
        """
        C: Orientation of the vehicle wrt inertial frame
        r: Position of the vehicle wrt inertial frame, expressed in inertial frame
        p: Position of the observed point wrt inertial frame
        """
        return D.T @ self.T_c_v @ SE3.odot(T.as_matrix() @ np.block([p, 1]))

    # Jacobian of the transformation that brings a point on the floor
    # to a position in the camera frame. This transformation is p_i()
    def dp_dx_numerical(self, T, p):
        """
        C: Orientation of the vehicle wrt inertial frame
        r: Position of the vehicle wrt inertial frame, expressed in inertial frame
        p: Position of the observed point wrt inertial frame
        """
        p = np.block([p, 1])
        T = T.as_matrix()
        d = 0.00001
        dp_de1 = D.T @ (
            SE3.exp(np.array([d, 0, 0, 0, 0, 0])).as_matrix() @ T @ p - T @ p
        )
        dp_de2 = D.T @ (
            SE3.exp(np.array([0, d, 0, 0, 0, 0])).as_matrix() @ T @ p - T @ p
        )
        dp_de3 = D.T @ (
            SE3.exp(np.array([0, 0, d, 0, 0, 0])).as_matrix() @ T @ p - T @ p
        )
        dp_de4 = D.T @ (
            SE3.exp(np.array([0, 0, 0, d, 0, 0])).as_matrix() @ T @ p - T @ p
        )
        dp_de5 = D.T @ (
            SE3.exp(np.array([0, 0, 0, 0, d, 0])).as_matrix() @ T @ p - T @ p
        )
        dp_de6 = D.T @ (
            SE3.exp(np.array([0, 0, 0, 0, 0, d])).as_matrix() @ T @ p - T @ p
        )

        return np.block(
            [
                dp_de1.reshape(3, 1) / d,
                dp_de2.reshape(3, 1) / d,
                dp_de3.reshape(3, 1) / d,
                dp_de4.reshape(3, 1) / d,
                dp_de5.reshape(3, 1) / d,
                dp_de6.reshape(3, 1) / d,
            ]
        )

    # Jacobian of g() wrt to the perturbations given
    # current pose {C,r} and observed point p
    def dg_dx(self, T, p):
        # Perform chain-rule
        return self.dg_dp(p) @ self.dp_dx(T, p)

    # Jacobian of the motion model wrt the perturbation
    # If you have the previous pose and the current pose,
    # this can be computed as Ad( T_curr inv(T_prev) ).
    # The provided poses needs to be transformation matrices.
    def motion_jacobian(self, T_prev, T_curr):
        F = SE3.adjoint(T_curr.dot(T_prev.inv()))
        return F

    # Return the error between two poses expressed as
    # transformation matrices. The error is a 6x1 vector.
    def pose_error(self, pose_1, pose_2):
        error = SE3.log(pose_1.dot(pose_2.inv()))
        return error

    # Return the error, in pixels between an observation
    # and the predicted pixel position (4x1) of a landmark
    def measurement_error(self, obs_1, obs_2):
        return obs_1 - obs_2

    def measurement_to_point(self, p):
        ul = p[0]
        vl = p[1]
        ur = p[2]
        vr = p[3]
        z = self.fu * self.b / (ul - ur)
        x = (ul - self.cu) * z / self.fu
        y = (vl - self.cv) * z / self.fv
        return np.array([x, y, z, 1])

    def verify_projections(self, k1=1215, k2=1714):

        ground_truth = self.ground_truth(k1=k1, k2=k2)

        for k in range(k1, k2):
            T = ground_truth[k - k1]
            C = T.rot.as_matrix()
            r = T.trans
            for j in range(self.num_points):
                if np.all(self.y_k_j[:, k, j] != np.array([-1, -1, -1, -1])):
                    P = self.p_i(T, j)

                    # Error terms
                    error = self.y_k_j[:, k, j] - self.g(P)
                    print(error)

    def gauss_newton(self, k1=1215, k2=1714, max_iterations=8, init=None):
        """
        Estimate position using only dead reckoning
        from the input velocities
        """

        # Initalize the operating trajectory by using the
        # ground-truth for the first pose and then dead-reckoning
        # from there until the end of the trajectory.
        dead_reckoning = self.dead_reckoning(k1=k1, k2=k2, init=init)
        T0 = dead_reckoning[0]
        T_op = dead_reckoning

        # Number of points in batch
        batch_size = (k2 - k1) + 1

        # Holds the perturbations used on each iteration
        delta_x = np.zeros((batch_size, 6))

        # Width of H jacobian matrix
        width = 6 * batch_size

        # Iterate refinement until convergence
        for iteration in range(max_iterations):

            # Initialize lists of arrays
            ev = []  # Motion errors throughout trajectory
            ey = []  # Measurement errors for observations
            Gs = []  # Lower part of the H jacobian matrix
            Fs = []  # Upper part of the H jacobian matrix

            # Variance P0 of the initial pose
            # TODO: Why 0.2 ?
            sigma_init = 1  # 0.2
            # Wv_inv is the inverse of the first part of the W matrix
            # with diagonal entries related to process noise.
            Wv_inv = 6 * [sigma_init]
            # Wy_inv is the inverse of the second part of the W matrix
            # with diagonal entries related to measurement noise.
            Wy_inv = []

            # Initialize the motion error list
            ev = list(self.pose_error(T0, T_op[0]).flatten())

            # Each "row" of Fs consists in a 6 x 6K array
            # The first one has only np.eye(6) at the front of the row
            F = np.zeros((6, width))
            F[:, 0:6] = np.eye(6)
            Fs.append(F)

            # This loop generates the upper part of the H jacobian matrix
            # and the upper part of the error matrix.
            for k in range(k1, k2):
                # Relative k index starting at 0 when k = k1
                k_rel = k - k1

                # Current pose T_curr and previous pose T_prev
                T_curr = T_op[k_rel + 1]
                T_prev = T_op[k_rel]

                # Angular perturbation = Angular velocity * time period
                psi = self.w_vk_vk_i[:, k] * self.time_steps[k]

                # Linear perturbation = Linear velocity * time period
                d = self.v_vk_vk_i[:, k] * self.time_steps[k]

                # Given the previous pose, use the motion model to predict
                # the current pose.
                T_pred = self.forward(T_prev, psi, d)
                # Compare predicted pose with current pose
                error = self.pose_error(T_pred, T_curr)
                # Append the error to the list
                ev += list(error.flatten())

                # Since the Q matrix is diagonal, the inverse of this matrix
                # is obtained by taking the reciprocal of the diagonal elements.
                # The variance of the linear and angular process noise must be multiplied
                # by the squared period of the timestep since the variance is expressed
                # in [m/s] and [rad/s] but we want [m] and [rad].
                qr_inv = list(1 / (self.v_var * (self.time_steps[k] ** 2)))
                qc_inv = list(1 / (self.w_var * (self.time_steps[k] ** 2)))
                q_inv = qr_inv + qc_inv
                # Append the inverse of the Q matrix to the list
                Wv_inv += q_inv

                # Motion jacobian of the previous timestep F_{k-1}
                F = self.motion_jacobian(T_prev, T_curr)

                # Most of the elements of the row are zeros
                F_row = np.zeros((6, width))
                # Horizontal offset where to place the -1*F matrix
                offset = 6 * k_rel
                F_row[:, offset : offset + 6] = -1 * F
                # Horizontal offset where to place the identity matrix
                offset = offset + 6
                F_row[:, offset : offset + 6] = np.eye(6)

                # Append the produced row to the list
                Fs.append(F_row)

            # This loop generates the lower part of the H jacobian matrix
            # and the lower part of the error matrix.
            for k in range(k1, k2 + 1):
                # Produced matrices for the k-th timestep
                Gk = []  # Measurement jacobian
                eyk = []  # Observation errors
                Rk_inv = []  # Inverted measurement covariance matrix

                # Relative k index starting at 0 when k = k1
                k_rel = k - k1

                # For convenience
                T = T_op[k_rel]
                C = T.rot.as_matrix()
                r = T.trans

                # Iterate through all of the possible landmarks
                for j in range(self.num_points):
                    # Pixels values (4x1) of the j-th landmark
                    # observed at the k-th timestep.
                    landmark = self.y_k_j[:, k, j]

                    # If the landmark is invalid, each pixel will be set to -1
                    if landmark[0] > 0:

                        # Project the j-th point in camera frame
                        P_cam = self.p_i(T, j)
                        # Project a point in space to a point in images
                        P_image = self.g(P_cam)

                        # Compute the measurement error and append to list
                        meas_error = self.measurement_error(landmark, P_image)
                        error = list(meas_error.flatten())
                        eyk += error

                        # Append the inverse of R
                        Rk_inv += list(1 / self.y_var)

                        # Jacobian of g() wrt perturbations given pose and landmark
                        # G = self.dg_dx(T, self.rho_i_pj_i[:,j])
                        G = self.dg_dp(P_cam) @ self.dp_dx(T, self.rho_i_pj_i[:, j])
                        Gk.append(G)

                # Add the optimization variables to the problem
                # if you were able to observe more than a certain
                # number of landmarks.
                if len(Gk) > 0:
                    ey += eyk
                    Wy_inv += Rk_inv

                    # Each G jacobian is a 4x6 matrix
                    G_row = np.zeros((4 * len(Gk), width))
                    # Offset each jacobian below and right of the previous one
                    for i, Gki in enumerate(Gk):
                        vert_offset = 4 * i  # Vertical offset
                        horz_offset = 6 * k_rel  # Horizontal offset
                        G_row[
                            vert_offset : vert_offset + 4, horz_offset : horz_offset + 6
                        ] = Gki
                    Gs.append(G_row)

            # Finally, the H matrix can be produced from its upper and lower parts
            H = np.vstack(tuple(Fs + Gs))
            # Build the inverse covariance matrix
            W_inv = np.diag(tuple(Wv_inv + Wy_inv))
            # Build the error matrix
            e = np.array(ev + ey)

            # For Ax=b and solving for x
            A = H.T @ W_inv @ H
            b = H.T @ W_inv @ e

            # We solve for the optimal perturbations \delta_x
            # for the whole trajectory. Each perturbation is a
            # 6x1 vector.
            x_opt = np.linalg.solve(A, b).flatten()

            # Reshape the very 6*N long vector in a Nx6 matrix where
            # N is the number of timesteps in the trajectory.
            delta_x = np.reshape(x_opt, (len(T_op), 6))

            # To update the operating point of the optimization we
            # need to unstack all optimal perturbations and then
            # use T = exp(x^)T to update the operating point T
            # which represents our best guess in terms of pose
            norms = []
            for k in range(len(T_op)):

                T_op[k] = SE3.exp(delta_x[k]).dot(T_op[k])

                l2_norm = np.linalg.norm(delta_x[k])
                norms = norms + [l2_norm]
            print("Average L2 norm of perturbations: {}".format(np.mean(norms)))

            if np.mean(norms) < 1e-2:
                break

        # A is the inverse covariance matrix
        return T_op, np.linalg.inv(A).diagonal().reshape(batch_size, 6)

def plot3d(states):
    """
    Plot the 3D trajectory
    """

    fig = plt.figure()
    ax = plt.axes(projection="3d")

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    rs = [state.trans for state in states]
    x = [r[0] for r in rs]
    y = [r[1] for r in rs]
    z = [r[2] for r in rs]

    ax.plot3D(x, y, z)
    plt.show()

def plot3d_compare(data_dict, scatter_data_dict):
    """
    Plot the 3D trajectory
    """

    fig = plt.figure()
    ax = plt.axes(projection="3d")

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    for data in data_dict:

        rs = [state.trans for state in data_dict[data]]
        x = [r[0] for r in rs]
        y = [r[1] for r in rs]
        z = [r[2] for r in rs]

        ax.plot3D(x, y, z, label=data)

    for data in scatter_data_dict:

        x = [r[0] for r in scatter_data_dict[data]]
        y = [r[1] for r in scatter_data_dict[data]]
        z = [r[2] for r in scatter_data_dict[data]]
        ax.scatter(x, y, z, label=data)

    plt.legend()
    plt.show()

def plot_errors_batch(estimates, cov, ground_truth, prefix):

    trans_errors = [
        T_star.trans - T.trans for T_star, T in zip(estimates, ground_truth)
    ]

    rot_errros = [
        SO3.vee(np.eye(3) - T_star.rot.dot(T.rot.inv()).as_matrix())
        for T_star, T in zip(estimates, ground_truth)
    ]

    # Plot translation errors
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True, figsize=(8, 10))

    ax1.set_title("Translation Error: x")
    ax1.set_ylabel("Error (m)")
    ax1.plot([e[0] for e in trans_errors])
    ax1.plot([3 * sqrt(cov[i][0]) for i in range(len(trans_errors))], ":r")
    ax1.plot([-3 * sqrt(cov[i][0]) for i in range(len(trans_errors))], ":r")

    ax2.set_title("Translation Error: y")
    ax2.set_ylabel("Error (m)")
    ax2.plot([e[1] for e in trans_errors])
    ax2.plot([3 * sqrt(cov[i][1]) for i in range(len(trans_errors))], ":r")
    ax2.plot([-3 * sqrt(cov[i][1]) for i in range(len(trans_errors))], ":r")

    ax3.set_title("Translation Error: z")
    ax3.set_ylabel("Error (m)")
    ax3.set_xlabel("tk")
    ax3.plot([e[2] for e in trans_errors])
    ax3.plot([3 * sqrt(cov[i][2]) for i in range(len(trans_errors))], ":r")
    ax3.plot([-3 * sqrt(cov[i][2]) for i in range(len(trans_errors))], ":r")

    fig.savefig(prefix + "_trans_errors.pdf")

    # Plot rotational errors
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True, figsize=(8, 10))

    ax1.set_title("Rotational Error: x")
    ax1.set_ylabel("Error (rad)")
    ax1.plot([e[0] for e in rot_errros])
    ax1.plot([3 * sqrt(cov[i][3]) for i in range(len(trans_errors))], ":r")
    ax1.plot([-3 * sqrt(cov[i][3]) for i in range(len(trans_errors))], ":r")

    ax2.set_title("Rotational Error: y")
    ax2.set_ylabel("Error (rad)")
    ax2.plot([e[1] for e in rot_errros])
    ax2.plot([3 * sqrt(cov[i][4]) for i in range(len(trans_errors))], ":r")
    ax2.plot([-3 * sqrt(cov[i][4]) for i in range(len(trans_errors))], ":r")

    ax3.set_title("Rotational Error: z")
    ax3.set_xlabel("tk")
    ax3.set_ylabel("Error (rad)")
    ax3.plot([e[2] for e in rot_errros])
    ax3.plot([3 * sqrt(cov[i][5]) for i in range(len(trans_errors))], ":r")
    ax3.plot([-3 * sqrt(cov[i][5]) for i in range(len(trans_errors))], ":r")

    fig.savefig(prefix + "_rot_errors.pdf")

# Question 5
estimator = BatchEstimator("data.mat")

# # Part 1:
# k1 = 1215
# k2 = 1714
# ground_truth = estimator.ground_truth(k1=k1, k2=k2)
# gauss_newton, cov = estimator.gauss_newton(k1=k1, k2=k2)
# plot_errors_batch(gauss_newton, cov, ground_truth, "part1")

# # Part 2
# window_size = 50
# poses = []
# covariences = []
# for k in range(k1, k2):
#     if k == k1:
#         init = estimator.ground_truth_k(k1)
#     else:
#         init = poses[-1]

#     gauss_newton, cov = estimator.gauss_newton(k1=k, k2=k + window_size, init=init)

#     poses.append(gauss_newton[0])
#     covariences.append(cov[0])

# covariences = np.array(covariences)
# plot_errors_batch(poses, covariences, ground_truth, "part2")

# window_size = 10
# poses = []
# covariences = []
# for k in range(k1, k2):
#     if k == k1:
#         init = estimator.ground_truth_k(k1)
#     else:
#         init = poses[-1]

#     gauss_newton, cov = estimator.gauss_newton(k1=k, k2=k + window_size, init=init)

#     poses.append(gauss_newton[0])
#     covariences.append(cov[0])

# covariences = np.array(covariences)
# plot_errors_batch(poses, covariences, ground_truth, "part3")

ground_truth = estimator.ground_truth()
dead_reckoning = estimator.dead_reckoning()
points = [estimator.rho_i_pj_i[:, i] for i in range(20)]
# gauss_newton = estimator.gauss_newton()

# estimator.verify_projections()

# Test point cloud allignment
# point_cloud_allignment = estimator.point_cloud_allignment(num_iterations=10)
gauss_newton, _ = estimator.gauss_newton(max_iterations=8)
plot3d_compare(
    {
        "Ground Truth": [T.inv() for T in ground_truth],
        "Dead Reckoning": [T.inv() for T in dead_reckoning],
        "Gauss Newton": [T.inv() for T in gauss_newton],
    },
    {
        "Points": points,
        "Intertial Frame": [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]],
        # "Point Cloud Allignment": [point_cloud_allignment[i].inv().trans for i in range(len(point_cloud_allignment))],
    },
)

# # Test Gauss Newton
# dead_reckoning = estimator.dead_reckoning()
# gauss_newton = estimator.gauss_newton(num_iterations=8)
# plot3d_compare(
#     {
#         "Dead Reckoning": dead_reckoning,
#         "Ground Truth": ground_truth,
#         "Gauss Newton": gauss_newton,
#     },
#     {
#         "Points": points
#     }
# )


import matplotlib.pyplot as plt
from statistics import mean
import numpy as np
import scipy as sp
from math import sqrt
import pandas as pd
from pprint import pprint
import statsmodels.api as sm
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.compat import lzip

#loads df from file
df = pd.read_csv("vehicle.csv", index_col='Year', thousands=",")

# This calculates each one independently.
# However, it was still good practice, I suppose.

# Defines variables
df['y'] = df['Vehicle']
df['x1'] = df['GDP']
df['x2'] = df['Population']

#Means of each variable
mean_y = mean(df['y'])
mean_x1 = mean(df['x1'])
mean_x2 = mean(df['x2'])

# Variance from mean for each variable
df['var_y'] = df['y'] - mean_y
df['var_x1'] = df['x1'] - mean_x1
df['var_x2'] = df['x2'] - mean_x2

#products of variance
df['product_of_var_x1'] = df['var_y'] * df['var_x1']
df['product_of_var_x2'] = df['var_y'] * df['var_x2']

#Predicted slope (beta)
beta_x1 = sum(df['product_of_var_x1']) / sum([n * n for n in df['var_x1']])
beta_x2 = sum(df['product_of_var_x2']) / sum([n * n for n in df['var_x2']])

#intercept (alpha, or sometimes beta zero)
alpha_x1 = mean_y - (beta_x1 * mean_x1)
alpha_x2 = mean_y - (beta_x2 * mean_x2)

#print(df.drop(['Vehicle', 'GDP', 'Population'], axis =1))

with open("prelimdata.csv", "w") as file:
    df.drop(['Vehicle', 'GDP', 'Population'], axis =1).to_csv(file)

df['predicted_y1'] = alpha_x1 + (beta_x1 * df['x1'])
df['predicted_y2'] = alpha_x2 + (beta_x2 * df['x2'])

plt.plot(df['x1'], df['predicted_y1'], color='red')
plt.plot(df['x1'], df['y'], marker='o', linewidth=0, markersize = 1.6, color='black')
plt.title("Vehicles vs GDP")
plt.xlabel("GDP")
plt.ylabel("Vehicles")
# plt.savefig("gdp_explore.png")
plt.show()

plt.plot(df['x2'], df['predicted_y2'], color='red')
plt.plot(df['x2'], df['y'], marker='o', linewidth=0, markersize = 1.6, color='black')
plt.title("Vehicles vs Population")
plt.xlabel("Population")
plt.ylabel("Vehicles")
# plt.savefig("pop_explore.png")
plt.show()

# This uses some Python to complete the actual thing.
# I wanted to do this all manually to better understand the material,
# But every guide wanted me to use some library or another.

X = df[['GDP', "Population"]]
y = df['Vehicle']
X = sm.add_constant(X)
model = sm.OLS(y, X).fit()
predictions = model.predict(X)

print(model.summary())
print(model.fvalue)
print(model.f_pvalue)

plt.plot(df['GDP'], predictions, color='teal')
plt.title("Vehicles vs GDP")
plt.xlabel("GDP")
plt.plot(df['GDP'], y, marker='o', linewidth=0, markersize = 1.6, color='black')
# plt.savefig("stats_gdp.png")
plt.show()

plt.plot(df['Population'], predictions, color='teal')
plt.title("Vehicles vs Population")
plt.xlabel("Population")
plt.plot(df['Population'], y, marker='o', linewidth=0, markersize = 1.6, color='black')
# plt.savefig("stats_pop.png")
plt.show()

vif_df = pd.DataFrame()
vif_df["VIF Factor"] = [variance_inflation_factor(X.values, i) for i in range(X.shape[1])]
vif_df["features"] = X.columns
print(vif_df)

name = ['Lagrange multiplier statistic', 'p-value',
        'f-value', 'f p-value']
test = sm.stats.het_breuschpagan(model.resid, model.model.exog)
test = lzip(name, test)

pprint(test)

import numpy,import scipy,import statsmodels,from statsmodels


import random
import cmath
import numpy as np
import pandas as pd
from qlazy import QState, DensOp

MIN_DOUBLE = 0.000001

def classical_joint_entropy(A,B):

    code_num_A = max(A) + 1
    code_num_B = max(B) + 1

    prob = np.zeros((code_num_A,code_num_B))
    for i in range(len(A)):
        prob[A[i]][B[i]] += 1
    prob = prob / sum(map(sum, prob))

    ent = 0.0
    for i in range(code_num_A):
        for j in range(code_num_B):
            if abs(prob[i][j]) < MIN_DOUBLE:
                ent -= 0.0
            else:
                ent -= prob[i][j] * np.log2(prob[i][j])

    return ent

def classical_entropy(A):

    code_num = max(A) + 1

    prob = np.zeros(code_num)
    for a in A:
        prob[a] += 1.0
    prob = prob / sum(prob)

    ent = 0.0
    for p in prob:
        if abs(p) < MIN_DOUBLE:
            ent -= 0.0
        else:
            ent -= p * np.log2(p)

    return ent

def classical_mutual_information(A,B):

    ent_A = classical_entropy(A)
    ent_B = classical_entropy(B)
    ent_AB = classical_joint_entropy(A,B)

    return ent_A + ent_B - ent_AB

def holevo_quantity(X,de):

    samp_num = len(X)
    code_num = len(de)

    prob = np.zeros(code_num)
    for x in X:
        prob[x] += 1.0
    prob = prob / sum(prob)

    de_total = DensOp.mix(densop=de, prob=prob)

    holevo = de_total.entropy()
    for i in range(code_num):
        holevo -= prob[i]*de[i].entropy()

    return holevo

def transmit(X,de,povm):

    samp_num = len(X)
    dim_X = len(de)
    dim_Y = len(povm)

    Y = np.array([0]*samp_num)

    prob_list = [None]*len(X)
    for i in range(samp_num):
        prob_list[i] = de[X[i]].probability(povm=povm)
        r = random.random()
        p = 0.0
        mes = dim_Y - 1
        for k in range(dim_Y-1):
            p += prob_list[i][k]
            if r < p:
                mes = k
                break
        Y[i] = mes

    return Y

def make_densop(basis):

    qs = [QState(vector=b) for b in basis]
    de = [DensOp(qstate=[q], prob=[1.0]) for q in qs]

    return de

def make_povm(vec):

    return [np.outer(v,v.conjugate())/2.0 for v in vec]

def random_sample(code_num,samp_num):

    return np.array([random.randint(0,code_num-1) for _ in range(samp_num)])

if __name__ == '__main__':

    SQRT1      = cmath.sqrt(1/3)
    SQRT2      = cmath.sqrt(2/3)
    EXT2       = cmath.exp(2*cmath.pi*1j/3)
    EXT4       = cmath.exp(4*cmath.pi*1j/3)

    basis      = [np.array([1.0, 0.0]),
                  np.array([SQRT1, SQRT2]),
                  np.array([SQRT1, SQRT2*EXT2]),
                  np.array([SQRT1, SQRT2*EXT4])]

    basis_orth = [np.array([0.0, 1.0]),
                  np.array([SQRT2, -SQRT1]),
                  np.array([SQRT2, -SQRT1*EXT2]),
                  np.array([SQRT2, -SQRT1*EXT4])]

    de = make_densop(basis)

    code_num = 4
    samp_num = 100
    trial = 100

    povm_name  = ['#1','#2']
    povm_basis = [basis, basis_orth]

    for b in povm_basis:

        povm = make_povm(b)

        mutual = []
        holevo = []
        for _ in range(trial):

            X = random_sample(code_num,samp_num)
            Y = transmit(X, de, povm)

            mutual.append(classical_mutual_information(X,Y))
            holevo.append(holevo_quantity(X,de))

        df = pd.DataFrame({'holevo quantity':holevo,'mutual information':mutual})
        holevo_mean = df['holevo quantity'].mean()
        holevo_std  = df['holevo quantity'].std()
        mutual_mean = df['mutual information'].mean()
        mutual_std  = df['mutual information'].std()

        print("== povm: {:} ==".format(povm_name.pop(0)))
        print("[holevo quantity]")
        print("- mean = {0:.4f} (std = {1:.4f})".format(holevo_mean, holevo_std))
        print("[mutual information]")
        print("- mean = {0:.4f} (std = {1:.4f})".format(mutual_mean, mutual_std))
        print()


import random
import typing

import numpy as np

class ARMASim:

    def __init__(
            self,
            _phi_arr: typing.Sequence[float] = (),
            _theta_arr: typing.Sequence[float] = (),
            _mu: float = 0.0, _sigma: float = 1.0,
    ):
        # AR part
        self.phi_arr = np.array(_phi_arr)
        assert np.all(np.abs(self.phi_arr) < 1)
        self.inv_phi_arr = self.phi_arr[::-1]
        self.phi_num = len(self.phi_arr)
        # MA part
        self.theta_arr = np.array(_theta_arr)
        assert np.all(np.abs(self.theta_arr) < 1)
        self.inv_theta_arr = self.theta_arr[::-1]
        self.theta_num = len(self.theta_arr)

        self.mu = _mu
        self.sigma = _sigma
        assert self.sigma > 0

        self.ret = [random.gauss(self.mu, self.sigma) for _ in range(self.phi_num)]
        self.latent = [random.gauss(0, self.sigma) for _ in range(self.theta_num)]

    def sample(self) -> float:
        new_value = self.mu
        if self.phi_num > 0:  # AR part
            new_value += self.inv_phi_arr.dot(
                self.ret[-self.phi_num:]
            )
        if self.theta_num > 0:  # MA part
            new_value += self.inv_theta_arr.dot(
                self.latent[-self.theta_num:]
            )

        new_info = random.gauss(0, self.sigma)
        new_value += new_info
        self.ret.append(new_value)
        self.latent.append(new_info)

        return new_value

    def sample_n(self, _num: int) -> np.ndarray:
        for _ in range(_num):
            self.sample()
        return np.array(self.ret[-_num:])


# -*- coding: utf-8 -*-
"""
Created on Thu Jan  6 09:36:16 2022

@author: OMEN
"""

import numpy as np
import matplotlib.pyplot as plt
fx= lambda x:x**3+4*x**2-10
a=1
b=2
tolerancia= 0.0001
n=1
#proceso
tramo=abs(b-a)
while not(tramo<=tolerancia):
    fa=fx(a)
    fb=fx(b)
    c=b-fb*(a-b)/(fa-fb)
    fc=fx(c)

    cambio=np.sign(fa)*np.sign(fc)
    print('{0:2d} {1:18.8f} {2:12.10f}'.format(n,c,tramo))
    n=n+1
    if cambio>0:
        tramo=abs(c-a)
        a=c
    else:
        tramo= abs(b-c)
        b=c
raiz = c

x=np.linspace(1,2)
ox=0*x
plt.figure()
plt.plot(x,fx(x))
plt.plot(x,ox,'k--')


import numpy as np
import matplotlib.pyplot as plt

# y = w1 * x1 + w2 * x2 + b

n = 30
epoch = 5
np.random.seed(0)
# x = np.random.uniform(0, 5, n)
x = np.random.rand(n,2)
x.shape

#x = np.reshape( np.arange(60), (30,2))
#x.shape

d = np.random.uniform(0, 5, n)
d.shape

#d = 2 * x[:,0]

w = np.random.rand(1,2)
b = np.random.uniform(0.001, 0.002, 1)

alpha = 0.9

#fig1 = plt.figure()
#ax1 = fig1.add_subplot(111)
#ax1.plot(x, d, 'ro')

epdist = np.zeros((epoch,2))
wdata = np.zeros((epoch,2))
bdata = np.zeros((epoch,1))

for j in range(0,epoch):
        print(j)
        for i in range(0,n):
                y = x[i,:].dot(w.T) + b
                e = d[i] - y
                dw = alpha * e * x[i,:] # (1,2)
                db = alpha * e
                w += dw
                b += db
                # print(w, b)
        wdata[j,:] = w
        bdata[j,:] = b
        d1 = x.dot(w.T) + b
        dist = d1 - y
        dist1 = np.mean(abs(dist)) # MAE
        epdist[j,0] = j ; epdist[j,1] = dist1 # Matrix for epoch # and MAE

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(epdist[:,0], epdist[:,1])
ax.set_xlabel('Number of Epoch')
ax.set_ylabel('Mean Absolute Error')
ax.set_title('Trend of MAE')

# Regression visualization
import seaborn as sns; sns.set(color_codes=True)
ax = sns.regplot(x=x[:,0], y=d)

# Result visualizaion
y_pred = np.zeros((epoch, n))
for j in range(0, epoch):
    y_pred[j] = x.dot(wdata[j, :].T) + bdata[j]

import pandas as pd

num = np.repeat(np.arange(epoch), n)
x1 = np.tile(x[:,0],epoch)
y1 = np.concatenate(y_pred)
df = pd.DataFrame({'epoch':num ,'x':x1, 'y':y1})

sns.lmplot("x", "y", data=df, hue='epoch', fit_reg=True)

'''
#fig1 = plt.figure()
#ax1 = fig1.add_subplot(111)
for j in range(0, epoch):
    y_pred[j] = x.dot(wdata[j,:].T) + bdata[j]

    #plt.plot(x, d, 'ro')
    #plt.plot(x[:,0], y_pred[j])

    fit = np.polyfit(x[:, 0], y_pred[j], 1)
    fit_fn = np.poly1d(fit)

    ax1.plot(x[:,0], y_pred[j], 'yo', x, fit_fn(x), '--k')

#colormap = plt.cm.gist_ncar
#colors = [colormap(i) for i in np.linspace(0, 1,len(ax1.lines))]

fig1 = plt.figure()
ax1 = fig1.add_subplot(111)
ax1.set_color_cycle(['red', 'green', 'blue', 'yellow'])
ax1.plot(x, d, 'ro')

for j in range(0, epoch):
    y_pred[j] = x.dot(wdata[j, :].T) + bdata[j]
    # plt.plot(x, d, 'ro')
    # plt.plot(x[:,0], y_pred[j])

    fit = np.polyfit(x[:, 0], y_pred[j], 1)
    fit_fn = np.poly1d(fit)

    ax1.plot(x[:,0], y_pred[j])
'''


# Do not import any additional 3rd party external libraries as they will not
# be available to AutoLab and are not needed (or allowed)

import numpy as np
import os

# The following Criterion class will be used again as the basis for a number
# of loss functions (which are in the form of classes so that they can be
# exchanged easily (it's how PyTorch and other ML libraries do it))

class Criterion(object):
    """
    Interface for loss functions.
    """

    # Nothing needs done to this class, it's used by the following Criterion classes

    def __init__(self):
        self.logits = None
        self.labels = None
        self.loss = None

    def __call__(self, x, y):
        return self.forward(x, y)

    def forward(self, x, y):
        raise NotImplemented

    def derivative(self):
        raise NotImplemented

class SoftmaxCrossEntropy(Criterion):
    """
    Softmax loss
    """

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

    def softmax(x):
        x = x - x.max(axis=1, keepdims=True)
        exp_scores = np.exp(x)
        p = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)
        return p

    def forward(self, x, y):
        """
        Argument:
            x (np.array): (batch size, 10)
            y (np.array): (batch size, 10)
        Return:
            out (np.array): (batch size, )
        """
        self.logits = x
        self.labels = y

        m = y.shape[0]
        # get unnormalized probabilities
        p = SoftmaxCrossEntropy.softmax(x)

        # alternate way of computing loss
        #log_likelihood = np.dot(-y, (np.log(p)).T)
        #loss = log_likelihood.diagonal()

        loss = np.sum(np.multiply(-y, np.log(p)), axis=1)
        return loss

    def derivative(self):
        """
        Return:
            out (np.array): (batch size, 10)
        """
        x = self.logits
        y = self.labels
        p = SoftmaxCrossEntropy.softmax(x)

        return (p - y)


import numpy as np
import matplotlib.pyplot as plt

def signal(t, freq_1=4.0, freq_2=7.0):
    return np.sin(freq_1*2*np.pi*t) + np.sin(freq_2*2*np.pi*t)

state = np.random.RandomState(12345)
sample_size = 2**7
sample_t = np.linspace(0, 4, sample_size, dtype=np.float64)
sample_y = signal(sample_t) + state.standard_normal(sample_size)
sample_d = 4. / (sample_size - 1) # Spacing for linspace array
true_signal = signal(sample_t)

from numpy import fft

fig1, ax1 = plt.subplots()
ax1.plot(sample_t, sample_y, "k.", label="Noisy signal")
ax1.plot(sample_t, true_signal, "k--", label="True signal")

ax1.set_title("Sample signal with noise")
ax1.set_xlabel("Time")
ax1.set_ylabel("Amplitude")
ax1.legend()

spectrum = fft.fft(sample_y)

freq = fft.fftfreq(sample_size, sample_d)
pos_freq_i = np.arange(1, sample_size//2, dtype=int)

psd = np.abs(spectrum[pos_freq_i])**2 + np.abs(spectrum[-pos_freq_i])**2

fig2, ax2 = plt.subplots()
ax2.plot(freq[pos_freq_i], psd)
ax2.set_title("PSD of the noisy signal")
ax2.set_xlabel("Frequency")
ax2.set_ylabel("Density")

filtered = pos_freq_i[psd > 1e4]

new_spec = np.zeros_like(spectrum)
new_spec[filtered] = spectrum[filtered]
new_spec[-filtered] = spectrum[-filtered]

new_sample = np.real(fft.ifft(new_spec))

fig3, ax3 = plt.subplots()
ax3.plot(sample_t, true_signal, color="#8c8c8c", linewidth=1.5, label="True signal")
ax3.plot(sample_t, new_sample, "k--", label="Filtered signal")
ax3.legend()
ax3.set_title("Plot comparing filtered signal and true signal")
ax3.set_xlabel("Time")
ax3.set_ylabel("Amplitude")

plt.show()


# coding: utf-8

# # SIR version 01.01

# ## El modelo epidémico de SIR

# Una descripción matemática simple de la propagación de una enfermedad en una población es el llamado modelo SIR, que divide la población (fija) de N individuos en tres "compartimentos" que pueden variar en función del tiempo, t:
#
# - S( t) son aquellos susceptibles pero aún no infectados con la enfermedad;
# - I( t) es la cantidad de individuos infecciosos;
# - R( t) son aquellos individuos que se han recuperado de la enfermedad y ahora tienen inmunidad contra ella.
#
# El modelo SIR describe el cambio en la población de cada uno de estos compartimentos en términos de dos parámetros, $\beta$ y $\gamma$.
#
# $\beta$ describe la tasa de contacto efectiva de la enfermedad: un individuo infectado entra en contacto con otros $\beta \cdot N$ individuos por unidad de tiempo (de los cuales la fracción que es susceptible a contraer la enfermedad es $S/N$) $\gamma$ es la tasa de recuperación media: es decir, $\frac{1}{\gamma}$ es el período de tiempo promedio durante el cual un individuo infectado puede transmitirlo.
# Además hemos incluido el caso de decesos, los que representamos con el factor $\mu$.

# Las ecuaciones diferenciales que describen este modelo fueron derivadas primero por Kermack y McKendrick [Proc.R.Soc.A,115,772(1927)]:
#
# $$\frac{dS}{dt}=−\beta \cdot \frac{S \cdot I}{N}$$
#
# $$ \frac{dI}{dt} =\beta \cdot \frac{S\cdot I}{N}−\gamma \cdot I - \mu \cdot I$$
#
# $$\frac{dR}{dt} = \gamma \cdot I $$
#
# $$\frac{dr}{dt} = \mu \cdot I $$
#

# El siguiente código de Python integra estas ecuaciones para una enfermedad caracterizada por los parámetros β = 0.2, 1/γ = 10 [dias] en una población de N = 1000 (quizás 'gripe en una escuela).
#
# El modelo se inicia con una sola persona infectada el día 0: I( 0 ) = 1.
#
# Las curvas trazadas de S(t), I(t) y R(t) están diseñados para ver los resultados obtenidos.

# In[37]:

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

# In[10]:

# Una cuadrícula de puntos de tiempo (en días).
t = np.linspace(0, 160, 160)

# In[15]:

# Las ecuaciones diferenciales del modelo SIR.
def deriv(y, t, N, beta, gamma, mu):
    S, I, R, r = y
    dSdt = -beta * S * I / N
    dIdt = beta * S * I / N - gamma * I - mu * I
    dRdt = gamma * I
    drdt = mu * I
    return dSdt, dIdt, dRdt, drdt

# In[16]:

# Población total, N.
N = 10000
# Número inicial de individuos infectados, recuperados y fallecidos, I0, R0 y r0.
I0, R0, r0 = 1, 0, 0
# Todos los demás, S0, son susceptibles a la infección inicialmente.
S0 = N - I0 - R0 - r0

# In[36]:

# beta: tasa media de contacto,
# gamma (en 1/dias): tasa media de recuperación.
# mu: tasa media de decesos    # 0.28
beta, gamma, mu = 0.335, 1./14, 0.006

# In[ ]:

# Vector de condiciones iniciales
y0 = S0, I0, R0, r0

# In[18]:

# Integra las ecuaciones SIR sobre la cuadrícula de tiempo, t.
ret = odeint(deriv, y0, t, args=(N, beta, gamma, mu))
S, I, R, r = ret.T

# In[35]:

# Grafico los datos en tres curvas separadas para S(t), I(t) y R(t)
with plt.style.context('seaborn'):
    plt.plot(t, S, 'b', alpha=0.5, lw=2, label='Susceptibles')
    plt.plot(t, I, 'r', alpha=0.5, lw=2, label='Infectados')
    plt.plot(t, R, 'g', alpha=0.5, lw=2, label='Recuperados con inmunidad')
    plt.plot(t, r, 'k', alpha=0.5, lw=2, label='Decesos')
    plt.xlabel('Tiempo [días]')
    plt.ylabel('Número ('+str(N)+'%)')
    plt.ylim(-N/20,N*1.1)
    plt.xlim(0, 120)
    plt.grid(b=True, which='major', c='w', lw=2, ls='-')
    plt.legend(loc= 'upper left', prop={'size': 16})
    plt.show()


# from the jupyter notebook provided by the class

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader, Dataset, TensorDataset
import torchvision.datasets as ds
import pylab as plt

def load_mnist(datadir='./data_cache'):
    train_ds = ds.MNIST(root=datadir, train=True,
                           download=True, transform=None)
    test_ds = ds.MNIST(root=datadir, train=False,
                          download=True, transform=None)

    def to_xy(dataset):
        X = np.array(dataset.data) / 255.0  # [0, 1]
        Y = np.array(dataset.targets)
        return X, Y

    X_tr, Y_tr = to_xy(train_ds)
    X_te, Y_te = to_xy(test_ds)
    return X_tr, Y_tr, X_te, Y_te

X_tr, Y_tr, X_te, Y_te = load_mnist()

i = np.random.choice(len(X_tr))
plt.imshow(X_tr[i], cmap='gray')
plt.title(f'digit: {Y_tr[i]}')

print('original X_tr:', X_tr.shape)

# select 500 random examples
n = 500
I = np.random.choice(len(X_tr), n, replace=False)
X = X_tr[I]
Y = (Y_tr[I] % 2) * 2.0 - 1 # odd/even --> +1/-1
X = X.reshape(-1,  28*28) # flatten

print('reshaped X:', X.shape)
print('reshaped Y:', Y.shape)

# problem 3 part 2

def dLdbeta(X,Y,beta):
    if len(X.shape)==1:
        return np.multiply(np.dot(X.T,X), beta) - np.multiply(X.T,Y)
    else:
        return np.matmul(np.matmul(X.T,X), beta) - np.matmul(X.T,Y)

# this sgd uses all samples with minibatches of size 1. may want to add minibatches to this for better tradeoff between convergence and computational efficiency
def sgd(X, Y, learning_rate=0.01, n_epochs=50):

    beta = np.random.randn(784, )  # start with random beta

    epoch = 0

    while epoch < n_epochs:

        inds = np.arange(len(X))
        np.random.shuffle(inds) # shuffle the inds
        for i in inds:
            beta = beta-learning_rate*dLdbeta(X[i],Y[i],beta)
        epoch += 1
        learning_rate = learning_rate / 1.02 # gradually reducing the learning rate

    return beta

def gd(X, Y, learning_rate=0.0001, n_epochs=50):

    beta = np.random.randn(784, )  # start with random beta

    epoch = 0

    while epoch < n_epochs:

        beta = beta-learning_rate*dLdbeta(X,Y,beta)
        epoch += 1
        learning_rate = learning_rate / 1.02 # gradually reducing the learning rate

    return beta

# running SGD
print("SGD resutls")
beta_sgd = sgd(X,Y)

# double check:
train_error = np.linalg.norm(np.matmul(X,beta_sgd)-Y)**2/Y.shape[0]
print(f"train error: {train_error}")
print(f"train acc: {100*np.sum(np.round(np.matmul(X,beta_sgd))==Y)/Y.shape[0]}%")

# get the 'test error'
test_error = np.linalg.norm(np.matmul(X_te.reshape(-1,  28*28),beta_sgd)-Y_te)**2/Y_te.shape[0]
print(f"test error: {test_error}")
print(f"test acc: {100*np.sum(np.round(np.matmul(X_te.reshape(-1,  28*28),beta_sgd))==Y_te)/Y_te.shape[0]}%")

# running GD
print("GD resutls")
beta_gd = gd(X,Y)

# double check:
train_error = np.linalg.norm(np.matmul(X,beta_gd)-Y)**2/Y.shape[0]
print(f"train error: {train_error}")
print(f"train acc: {100*np.sum(np.round(np.matmul(X,beta_gd))==Y)/Y.shape[0]}%")

# get the 'test error'
test_error = np.linalg.norm(np.matmul(X_te.reshape(-1,  28*28),beta_gd)-Y_te)**2/Y_te.shape[0]
print(f"test error: {test_error}")
print(f"test acc: {100*np.sum(np.round(np.matmul(X_te.reshape(-1,  28*28),beta_gd))==Y_te)/Y_te.shape[0]}%")


import numpy
import matplotlib.pyplot as plt
#
"""
Fixed Parameters
================

SB_SIGMA - Stefan-Boltzman constant in W/(m**2 K**4)
PI       - delicious old pi
"""
SB_SIGMA = 5.7E-8
PI = 3.141592653589793

"""
Sun's Parameters
----------------
F_SOL - Temperature of the sun (K)
R_SOL - Radius of the sun (m)
"""
T_SOL = 5800
R_SOL = 6.96000E8

"""
Earth's Parameters
------------------
R - distance of Earth from sun (m)
"""
R = 1.5E11

"""
C_AT - heat capacity: atmosphere only
C_OS - heat capacity: ocean surface 50m (real ca. 70m)
C_OU - heat capacity: upper ocean layer 250m (real ca. 360m)
C_OD - heat capacity: ocean 3900m, 70% of Earth surface

Note: Heat capacity is given in J/(K*m^2) to simplify calculations.

The specific heat capacity of air (Cp) is ~1000 J/(kg*K) or, given the density
of air (ro) of 1 kg/m^3, the volumetric heat capacity Cv of ~1000 J/(K*m^3).
[ Cv=Cp*ro=1000 J/(K*m^3) ].

In ro*Cp*dT/dt=dQ/dx, Q is heat flow in W/m^2 or J/(s*m^2), which is given by
f_net. To get dT, we need to divide heat flow by a unit of length. In this exercise,
volumetric heat capacity (c=ro*cp) is already multiplied by 10000 (m), i.e. the
average thickness of the troposphere, and therefore skip the division of the heat
flow by length.

Example:
C_AT= ro*Cp*H where H is the height of the troposphere (~10km).
"""
C_AT = 1.0E7
C_OS = 2.0E8
C_OU = 1.0E9
C_OD = 1.6E10

"""
ALBEDO_0  - albedo mean constant
TRANSM_0  - transmissivity mean constant
"""
ALBEDO_0 = 0.3
TRANSM_0 = 0.64262

"""
Tasks for You!
==============

1) Calculate the sun's output F_SOL and radiative flux PHI_SOL
"""
F_SOL   = SB_SIGMA * T_SOL**4
PHI_SOL = F_SOL * 4 * PI * R_SOL**2
print("F_SOL:", F_SOL)
print("PHI_SOL:", PHI_SOL)

"""
2) Calculate F_SW_IN_0, the mean incoming shorwave radiation per area on Earth in W/m^2, i.e. the solar constant.
"""
E_0 = PHI_SOL / (4 * PI * R**2)  # solar constant
F_SW_IN_0 = E_0 / 4
print("E_0:", E_0)
print("F_SW_IN_0:", F_SW_IN_0)

"""
3) Calculate the predicted temperature on Earth T_EARTH assuming no atmosphere
"""
T_EARTH = ( ( E_0 * (1 - ALBEDO_0) ) / (4 * SB_SIGMA) )**(1/4)
print("T_EARTH:", T_EARTH)

"""
You'll notice that the temperature of task 3 deviates somewhat from the actual temperature on Earth.
What we need is to consider greenhouse gas related transmissivity.

4) Let temperature evolve over time!
Initialise temperature and let it evolve over time until it reaches equilibrium.
Make sure to include albedo and transmissivity in your calculations.

HINT: To calculate temperature change for a specific time step based on net incoming radiation,
remember dT (change in temperature) = dQ (change in heat/energy) / C (heat capacity), incoming
radiation is measured in W/m^2 (J/(s*m^2)) and heat capacity is given in J/(K*m^2) here.

Note: When calculating the radiation emitted by Earth, a multiplication by the factor of 0.95 is
necessary to take into consideration the emissive propertiesof the Earth's surface.

4-a: Initialise temperature at 300K and let it evolve using a heat capacity of the atmosphere only.
4-b: Change the initial temperature, run your code again and see what happens. Does this make sense?
4-c: Pick a different heat capacity and see what happens. Does this make sense?
"""

# Simulation Parameters
i_dt  = 10                                # time step size in days
nt    = 50                                # number of time steps
dt    = 60*60*24*float(i_dt)              # time step size in SI units (s)
ck    = C_AT                              # selection of heat capacity
#
# Initialisation
temp = [0.0]*nt                           # initialise temperature array
temp[0]=300                               # initial temperature
#
# run simulation
for i in range(0,nt-1):
    #
    # radiative fluxes
    f_sw_in  = F_SW_IN_0
    f_sw_out = ALBEDO_0 * f_sw_in
    f_lw_out = 0.95 * SB_SIGMA * temp[i]**4
    f_net    = f_sw_in - f_sw_out - f_lw_out * TRANSM_0
    #
    # temperature for next time step
    temp[i+1] = temp[i] + dt*(1/ck) * f_net
#    print(temp[i])
#
# plotting everything
x = [0.0]*nt
y = temp
#
for i in range (1,nt):
    x[i] = x[i-1]+i_dt
#
# create plot
plt.figure(0)                          # number your figures
plt.plot(x, y)                         # create a simple x-y plot
plt.title('temperature evolution')     # set a title
plt.xlabel('days')                     # label the x axis
plt.ylabel('temperature (K)')          # label the y axis
#plt.axis([0, nt, 280, 300])            # give x and y bounds
plt.show()