# Box 8.1
import numpy as np
import pandas as pd
# import matplotlib.pyplot as plt
# import statsmodels.formula.api as sm
# import statsmodels.tsa.stattools as ts
# import statsmodels.tsa.vector_ar.vecm as vm
from scipy.optimize import minimize
# from scipy.stats import describe
from scipy.stats import pearson3

df = pd.read_csv('AUDCAD_unequal_ret.csv')

skew_, loc_, scale_ = pearson3.fit(df)
print('skew=%f loc=%f scale=%f' % (skew_, loc_, scale_))
mean, var, skew, kurt = pearson3.stats(skew_, loc_, scale_, moments='mvks')
print('mean=%f var=%f skew=%f kurt=%f' % (mean, var, skew, kurt))

ret_sim = pearson3.rvs(skew_, loc_, scale_, size=100000, random_state=0)

def g(f, R):
    return np.sum(np.log(1 + f * R), axis=0) / R.shape[0]

myf = range(1, 24)
myg = np.full(24, np.NaN)

for f in myf:
    myg[f] = g(f, ret_sim)

myg = myg[1:]
myg = pd.DataFrame(myg)
myg.plot()

minusG = lambda f: -g(f, df)
minusGsim = lambda f: -g(f, ret_sim)

# optimal leverage based on simulated returns
res = minimize(minusGsim, 0, method='Nelder-Mead')
optimalF = res.x
print('Optimal leverage=%f optimal growth rate=%f' % (optimalF, -res.fun))
# Optimal leverage=25.512625 optimal growth rate=0.005767

minR = np.min(ret_sim)
print('minR=%f' % (minR))
# minR=-0.018201

# max drawdown with optimal leverage
from calculateMaxDD import calculateMaxDD

maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF * ret_sim) - 1))
print('Max DD=%f with optimal leverage=%f' % (maxDD, optimalF))
# Max DD=-0.996312 with optimal leverage=25.512625

# max drawdown with half of optimal leverage
maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF / 2 * ret_sim) - 1))
print('Max DD=%f with half optimal leverage=%f' % (maxDD, optimalF / 2))
# Max DD=-0.900276 with half optimal leverage=12.756313

# max drawdown with 1/7 of optimal leverage
maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF / 7 * ret_sim) - 1))
print('Max DD=%f with half optimal leverage=%f' % (maxDD, optimalF / 7))
# Max DD=-0.429629 with half optimal leverage=3.644661

# max drawdown with 1/1.4 of optimal leverage for historical returns
maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF / 1.4 * df.values) - 1))
print('Max DD=%f with historical returns=%f' % (maxDD, optimalF / 1.4))
# Max DD=-0.625894 with historical returns=18.223304

D = 0.5
print('Growth rate on simulated returns using D=%3.1f of optimal leverage on full account=%f' % (
D, -minusGsim(optimalF * D)))
# Growth rate on simulated returns using D=0.5 of optimal leverage on full account=0.004317
maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF * D * ret_sim) - 1))
print('MaxDD on simulated returns using D of optimal leverage on full account=%f' % (maxDD))
# MaxDD on simulated returns using D of optimal leverage on full account=-0.900276

# CPPI
g_cppi = 0
drawdown = 0
for t in range(ret_sim.shape[0]):
    g_cppi += np.log(1 + ret_sim[t] * D * optimalF * (1 + drawdown))
    drawdown = min([0, (1 + drawdown) * (1 + ret_sim[t]) - 1])

g_cppi = g_cppi / len(ret_sim)
print('Growth rate on simulated returns using CPPI=%f' % g_cppi)
# Growth rate on simulated returns using CPPI=0.004264
print('Growth rate on historical returns using D of optimal leverage on full account=%f' % (-minusG(optimalF * D)))
# Growth rate on historical returns using D of optimal leverage on full account=0.004053
maxDD, maxDDD, i = calculateMaxDD((np.cumproduct(1 + optimalF * D * df.values) - 1))
print('MaxDD on historical returns using D of optimal leverage on full account=%f' % (maxDD))
# MaxDD on historical returns using D of optimal leverage on full account=-0.303448

# CPPI
g_cppi = 0
drawdown = 0
for t in range(df.shape[0]):
    g_cppi += np.log(1 + df.iloc[t,] * D * optimalF * (1 + drawdown))
    drawdown = np.min([0, (1 + drawdown) * (1 + df.iloc[t,]) - 1])

g_cppi = g_cppi / len(df)
print('Growth rate on historical returns CPPI=%f' % (g_cppi))

import numpy,from scipy,import statsmodels


import numpy as np

def l1_norm(estimations, true_values):
    """Calculates the L1 loss value between 2 vectors.

    The L1 loss function is defined as:
        \begin{align*}
            & L_1(\\hat{y}, y) = \\sum_{i=0}^m|y^{(i)} - \\hat{y}^{(i)}|
        \\end{align*}

    :param estimations: vector of size m containing the predicted labels
    :param true_values: vector of size m containing the true labels
    :return: the value of the L1 loss function
    """
    return np.sum(np.abs(true_values - estimations))

def l2_norm(estimations, true_values):
    """Calculates the L1 loss value between 2 vectors.

    L2 loss is defined as
        \begin{align*}
            & L_2(\\hat{y},y) = \\sum_{i=0}^m(y^{(i)} - \\hat{y}^{(i)})^2
        \\end{align*}

    :param estimations: vector of size m containing the predicted labels
    :param true_values: vector of size m containing the true labels
    :return: the value of the L2 loss function
    """
    return np.sum((true_values - estimations) ** 2)


from unittest import TestCase

import numpy as np

from SEAL.lib import compute_knot_insertion_matrix

class TestComputeKnotInsertionMatrix(TestCase):
    def test_compute_knot_insertion_matrix(self):
        """
        Given:
            p = 2
            tau = [-1, -1, -1, 0, 1, 1, 1]
            t = [-1, -1, -1, -0.5, 0, 0.5, 1, 1, 1]
        When:
            Representing the coarse linear B spline
            as a linear combination of the fine linear B splines
        Then:
            Knot Insertion Matrix A =
                [[ 1.    0.    0.    0.  ]
                [ 0.5   0.5   0.    0.  ]
                [ 0.    0.75  0.25  0.  ]
                [ 0.    0.25  0.75  0.  ]
                [ 0.    0.    0.5   0.5 ]
                [ 0.    0.    0.    1.  ]]
        """
        p = 2
        tau = [-1, -1, -1, 0, 1, 1, 1]
        t = [-1, -1, -1, -0.5, 0, 0.5, 1, 1, 1]

        expected_A = np.array([[1., 0., 0., 0.],
                               [0.5, 0.5, 0., 0.],
                               [0., 0.75, 0.25, 0.],
                               [0., 0.25, 0.75, 0.],
                               [0., 0., 0.5, 0.5],
                               [0., 0., 0., 1.]])
        computed_A = compute_knot_insertion_matrix(p, tau, t)

        for e, c in zip(expected_A, computed_A):
            np.testing.assert_array_almost_equal(e, c)


"""
For each path of simulation, the final price is:

    St = spot * exp((b - vol*vol/2) * T + vol * rand * sqrt(T))
    where rand is a random number in (0, 1)

    call option:
        each_price = max(St - strike, 0)
    put option:
        each_price = max(strike - St, 0)

The overall price is:
    sum(each_price) * exp(rate * T) / simu_num
"""

from math import exp, sqrt
from multiprocessing import Process, Queue
from random import random

from options.functions import norminv
from options.option import OptionType

# from scipy.stats import norm   # norm.ppf is about 50 times slower than norminv but no obvious accuracy improvement

class MonteCarloPricer:
    def __init__(self, simu_num=1000000, ps_num=10):
        self.simu_num = simu_num
        self.ps_num = ps_num

    def set_ps_num(self, ps_num):
        self.ps_num = ps_num

    def get_price_of_one_run(self, z):
        '''Run the simulation once and return the option price'''
        #st = self.spot * exp((self.cost_of_carry - self.vol**2 / 2) * self.expiry + self.vol * norm.ppf(random()) * sqrt(self.expiry))
        st = self.spot * exp((self.cost_of_carry - self.vol**2 / 2) * self.expiry + self.vol * norminv(random()) * sqrt(self.expiry))
        return max(z * (st - self.strike), 0)

    def _ps_slice(self, z, num, resultq):
        sum = 0
        for i in range(int(num)):
            sum += self.get_price_of_one_run(z)

        resultq.put(sum)
        resultq.close()

    def price_option(self, option):
        """
        simu_num: the number of simulation runs, usually > 100000
        ps_num: If zero, run simulation in single process mode;
                otherwise run in multiprocess mode with ps_num processes to speed up
        """
        self.spot = option.spot
        self.strike = option.strike
        self.rate = option.rate
        self.expiry = option.expiry
        self.vol = option.vol
        self.cost_of_carry = option.cost_of_carry or option.rate

        z = 1 if option.type == OptionType.CALL else -1

        sum = 0
        if self.ps_num:
            # multiprocess mode
            if (self.simu_num / self.ps_num <= 1000): # when > 1000, simu_num / ps_num ~= simu_num / ps_num +- simu_num % ps_num
                assert not self.simu_num % self.ps_num, \
                    'simu_num must be integer times of ps_num; received simu_num: {}, ps_num: {}'.format(self.simu_num, self.ps_num)
            resultq = Queue()
            processes = [Process(target=self._ps_slice, args=(z, self.simu_num / self.ps_num, resultq)) for i in range(self.ps_num)]
            for p in processes:
                p.start()
            for p in processes:
                p.join()

            while not resultq.empty():
                sum += resultq.get()

        else:
            # single process mode
            for i in range(self.simu_num):
                sum += self.get_price_of_one_run(z)

        return round(exp(- self.rate * self.expiry) * sum / self.simu_num, 4)


# encoding=utf8

import numpy as np
from collections import namedtuple
from scipy.stats import t, norm, rankdata

from mathpy._lib import _create_array
from mathpy.stats.summary import var

def ttest(y1, y2=None, mu=None, var_equal=False):
    r"""
    Performs one and two-sample t-tests.

    Parameters
    ----------
    y1
        First sample to test
    y2
        Second sample. Optional
    mu
        Optional, sets the mean for comparison in the one sample t-test. Default 0.
    var_equal
        Optional, default False. If False, Welch's t-test for unequal variance and
        sample sizes is used. If True, equal variance between samples is assumed
        and Student's t-test is used.

    Returns
    -------
    namedtuple
        Namedtuple containing following values:
        t-value
        degrees of freedom
        p-value
        confidence intervals
        sample means

    Notes
    -----
    Welch's t-test is an adaption of Student's t test and is more performant when the
    sample variances and size are unequal. The test still depends on the assumption of
    the underlying population distributions being normally distributed.

    Welch's t test is defined as:

    .. math::

        t = \frac{\bar{X_1} - \bar{X_2}}{\sqrt{\frac{s_{1}^{2}}{N_1} + \frac{s_{2}^{2}}{N_2}}}

    where:

    :math:`\bar{X}` is the sample mean, :math:`s^2` is the sample variance, :math:`n` is the sample size

    If the :code:`var_equal` argument is True, Student's t-test is used, which assumes the two samples
    have equal variance. The t statistic is computed as:

    .. math::

        t = \frac{\bar{X}_1 - \bar{X}_2}{s_p \sqrt{\frac{1}{n_1} + \frac{1}{n_2}}

    where:

    .. math::

        s_p = \sqrt{\frac{(n_1 - 1)s^2_{X_1} + (n_2 - 1)s^2_{X_2}}{n_1 + n_2 - 2}

    References
    ----------
    Rencher, A. C., & Christensen, W. F. (2012). Methods of multivariate analysis (3rd Edition).

    Student's t-test. (2017, June 20). In Wikipedia, The Free Encyclopedia.
        From https://en.wikipedia.org/w/index.php?title=Student%27s_t-test&oldid=786562367

    """
    y1 = _create_array(y1)[0]

    n1 = len(y1)
    s1 = var(y1)
    ybar1 = np.mean(y1)

    if y2 is not None:
        y2 = _create_array(y2)[0]
        n2 = len(y2)
        s2 = var(y2)
        ybar2 = np.mean(y2)

        if var_equal is False:
            tval = float((ybar1 - ybar2) / np.sqrt(s1 / n1 + s2 / n2))
        else:
            sp = np.sqrt(((n1 - 1.) * s1 + (n2 - 1.) * s2) / (n1 + n2 - 2.))
            tval = float((ybar1 - ybar2) / (sp * np.sqrt(1. / n1 + 1. / n2)))

    else:
        ybar2, n2, s2 = 0.0, 1.0, 0.0
        if mu is None:
            mu = 0.0

        tval = float((ybar1 - mu) / np.sqrt(s1 / n1))

    dof = degrees_of_freedom(y1, y2)
    pvalue = _student_t_pvalue(np.absolute(tval), dof)
    intervals = _t_conf_int((ybar1, n1, s1), dof=dof, y=(ybar2, n2, s2))

    if y2 is not None:
        tTestResult = namedtuple('tTestResult', ['tvalue', 'dof', 'pvalue', 'confint', 'x_mean', 'y_mean'])

        tt = tTestResult(tvalue=tval, dof=dof, pvalue=pvalue, confint=intervals, x_mean=ybar1, y_mean=ybar2)

    else:
        tTestResult = namedtuple('tTestResult', ['tvalue', 'dof', 'pvalue', 'confint', 'x_mean'])
        tt = tTestResult(tvalue=tval, dof=dof, pvalue=pvalue, confint=intervals, x_mean=ybar1)

    return tt

def mann_whitney(y1, y2, continuity=True):
    r"""
    Performs the nonparametric Mann-Whitney U test of two independent sample groups.

    Parameters
    ----------
    y1
        One-dimensional array-like (Pandas Series or DataFrame, Numpy array, list, or dictionary)
        designating first sample
    y2
        One-dimensional array-like (Pandas Series or DataFrame, Numpy array, list, or dictionary)
        designating second sample to compare to first
    continuity
        Boolean, optional. If True, apply the continuity correction of :math:`\frac{1}{2}` to the
        mean rank.

    Returns
    -------
    namedtuple
        Namedtuple of following entries that contain resulting Mann-Whitney test statistics.
        Mann-Whitney U Test Statistic: The U Statistic of the Mann-Whitney test
        Mean Rank: The mean rank of U statistic
        Sigma: the standard deviation of U
        z-value: The standardized value of U
        p-value: p-value of U statistic compared to critical value

    Notes
    -----
    The Mann-Whitney U test is a nonparametric hypothesis test that tests the null hypothesis that
    there is an equally likely chance that a randomly selected observation from one sample will be
    less than or greater than a randomly selected observation from a second sample. Nonparametric
    methods are so named since they do not rely on the assumption of normality of the data.

    The test statistic in the Mann-Whitney setting is denoted as :math:`U` and is the minimum of
    the summed ranks of the two samples. The null hypothesis is rejected if :math:`U \leq U_0`,
    where :math:`U_0` is found in a table for small sample sizes. For larger sample sizes,
    :math:`U` is approximately normally distributed.

    The test is nonparametric in the sense it uses the ranks of the values rather than the values
    themselves. Therefore, the values are ordered then ranked from 1 (smallest value) to the largest
    value. Ranks of tied values get the mean of the ranks the values would have received. For example,
    for a set of data points :math:`\{4, 7, 7, 8\}` the ranks are :math:`\{1, 2.5, 2.5, 4\}`. The
    :math:`2.5` rank comes from :math:`2 + 3 = 5 / 2`. The ranks are then added for the values for
    both samples. The sum of the ranks for each sample are typically denoted by :math:`R_k` where
    :math:`k` is a sample indicator.

    :math:`U` for the two samples in the test, is given by:

    References
    ----------
    Mann–Whitney U test. (2017, June 20). In Wikipedia, The Free Encyclopedia.
        From https://en.wikipedia.org/w/index.php?title=Mann%E2%80%93Whitney_U_test&oldid=786593885

    """
    n1, n2 = len(y1), len(y2)

    ranks = np.concatenate((y1, y2))

    ranks = rankdata(ranks, 'average')

    ranks = ranks[:n1]

    n = n1 + n2

    u1 = n1 * n2 + (n1 * (n1 + 1)) / 2. - np.sum(ranks)
    u2 = n1 * n2 - u1

    u = np.minimum(u1, u2)
    mu = (n1 * n2) / 2. + (0.5 * continuity)

    rankcounts = np.unique(ranks, return_counts=True)[1]

    sigma = np.sqrt(((n1 * n2) * (n + 1)) / 12. * (1 - np.sum(rankcounts ** 3 - rankcounts) / float(n ** 3 - n)))
    z = (np.absolute(u - mu)) / sigma
    p = 1-norm.cdf(z)

    MannWhitneyResult = namedtuple('MannWhitneyResult', ['u', 'meanrank', 'sigma', 'zvalue', 'pvalue'])

    mwr = MannWhitneyResult(u=u, meanrank=mu, sigma=sigma, zvalue=z, pvalue=p)

    return mwr

def degrees_of_freedom(y1, y2=None, var_equal=False):
    r"""
    Computes the degrees of freedom of one or two samples.

    Parameters
    ----------
    y1
        First sample to test
    y2
        Second sample. Optional.
    var_equal
        Optional, default False. If False, Welch's t-test for unequal variance and
        sample sizes is used. If True, equal variance between samples is assumed
        and Student's t-test is used.

    Returns
    -------
    float
        the degrees of freedom

    Notes
    -----
    When Welch's t test is used, the Welch-Satterthwaite equation for approximating the degrees
    of freedom should be used and is defined as:

    .. math::

        \large v \approx \frac{\left(\frac{s_{1}^2}{N_1} +
        \frac{s_{2}^2}{N_2}\right)^2}{\frac{\left(\frac{s_1^2}{N_1^{2}}\right)^2}{v_1} +
        \frac{\left(\frac{s_2^2}{N_2^{2}}\right)^2}{v_2}}

    If the two samples are assumed to have equal variance, the degrees of freedoms become simply:

    .. math::

        v = n_1 + n_2 - 2

    In the case of one sample, the degrees of freedom are:

    .. math::

        v = n - 1

    References
    ----------
    Rencher, A. C., & Christensen, W. F. (2012). Methods of multivariate analysis (3rd Edition).

    Welch's t-test. (2017, June 16). In Wikipedia, The Free Encyclopedia.
        From https://en.wikipedia.org/w/index.php?title=Welch%27s_t-test&oldid=785961228

    """
    y1 = _create_array(y1)[0]
    n1 = len(y1)
    s1 = var(y1)
    v1 = n1 - 1

    if y2 is not None:
        y2 = _create_array(y2)[0]
        n2 = len(y2)
        s2 = var(y2)
        v2 = n2 - 1

        if var_equal is False:
            v = np.power((s1 / n1 + s2 / n2), 2) / (np.power((s1 / n1), 2) / v1 + np.power((s2 / n2), 2) / v2)
        else:
            v = n1 + n2 - 2

    else:
        v = v1

    return float(v)

def _t_conf_int(x, dof, y=None):

    xbar, xn, xvar = x[0], x[1], x[2]

    if y is not None:
        ybar, yn, yvar = y[0], y[1], y[2]

        low_interval = (xbar - ybar) + t.ppf(0.025, dof) * np.sqrt(xvar / xn + yvar / yn)
        high_interval = (xbar - ybar) - t.ppf(0.025, dof) * np.sqrt(xvar / xn + yvar / yn)
    else:
        low_interval = xbar + 1.96 * np.sqrt((xbar * (1 - xbar)) / xn)
        high_interval = xbar - 1.96 * np.sqrt((xbar * (1 - xbar)) / xn)

    return float(low_interval), float(high_interval)

def _student_t_pvalue(n, dof, test='two-tail'):
    p = (1. - t.cdf(n, dof))

    if test == 'two-tail':
        p *= 2.

    return p


import numpy as np
import scipy.ndimage as spim
from porespy.tools import get_tqdm
from porespy import settings
from loguru import logger
tqdm = get_tqdm()

def random_cantor_dust(shape, n, p=2, f=0.8):
    r"""
    Generates an image of random cantor dust

    Parameters
    ----------
    shape : array_like
        The shape of the final image.  If not evenly divisible by $p**n$
        it will be increased to the nearest size that is.
    n : int
        The number of times to iteratively divide the image.
    p : int (default = 2)
        The number of divisions to make on each iteration.
    f : float (default = 0.8)
        The fraction of the set to keep on each iteration.

    Returns
    -------
    dust : ndarray
        A boolean image of a random Cantor dust

    """
    # Parse the given shape and adjust if necessary
    shape = np.array(shape)
    trim = np.mod(shape, (p**n))
    if np.any(trim > 0):
        shape = shape - trim + p**n
        logger.warning(f"Requested shape being changed to {shape}")
    im = np.ones(shape, dtype=bool)
    divs = []
    if isinstance(n, int):
        for i in range(1, n):
            divs.append(p**i)
    else:
        for i in n:
            divs.append(p**i)
    for i in tqdm(divs, **settings.tqdm):
        sh = (np.array(im.shape)/i).astype(int)
        mask = np.random.rand(*sh) < f
        mask = spim.zoom(mask, zoom=i, order=0)
        im = im*mask
    return im

def sierpinski_foam(dmin, n, ndim=2, max_size=1e9):
    r"""
    Generates an image of a Sierpinski carpet or foam

    Parameters
    ----------
    dmin : int
        The size of the smallest square in the final image
    n : int
        The number of times to iteratively tile the image
    ndim : int
        The number of dimensions of the desired image, can be 2 or 3. The
        default value is 2.

    Returns
    -------
    foam : ndarray
        A boolean image of a Sierpinski gasket or foam

    """
    def _insert_cubes(im, n):
        if n > 0:
            n -= 1
            shape = np.asarray(np.shape(im))
            im = np.tile(im, (3, 3, 3))
            im[shape[0]:2*shape[0], shape[1]:2*shape[1], shape[2]:2*shape[2]] = 0
            if im.size < max_size:
                im = _insert_cubes(im, n)
        return im

    def _insert_squares(im, n):
        if n > 0:
            n -= 1
            shape = np.asarray(np.shape(im))
            im = np.tile(im, (3, 3))
            im[shape[0]:2*shape[0], shape[1]:2*shape[1]] = 0
            if im.size < max_size:
                im = _insert_squares(im, n)
        return im

    im = np.ones([dmin]*ndim, dtype=int)
    if ndim == 2:
        im = _insert_squares(im, n)
    elif ndim == 3:
        im = _insert_cubes(im, n)
    return im


import numpy as np

def compute_pi(n_samples=1000000):
    """Estimate the value of pi by sampling points inside a unit square and
    checking how many of them fall inside the enclosed circle
    :param n_samples: number of samples to use
    :return: estimated value of pi
    """
    n, n_inner = 0, 0

    for i in range(n_samples):
        x = np.random.random() - 0.5
        y = np.random.random() - 0.5
        if is_inside_circle(x, y):
            n_inner += 1
        n += 1

    pi = 4 * n_inner / n
    return pi

def is_inside_circle(x, y):
    """Check if a given point falls inside the enclosed circle
    :param x: x coordinate of the point
    :param y: y coordinate of the point
    :return: a boolean True/False
    """
    r = 0.5  # radius of the circle
    if x ** 2 + y ** 2 < r ** 2:
        return True
    else:
        return False

if __name__ == '__main__':
    pi = compute_pi()
    print('Approximate value of pi:', pi)
    print('Exact value of pi:', np.pi)


#!/usr/bin/env python3
"""
Computational demo of the exponential map for the quaternion representation of SO3.
Quaternions here are stored as arrays [w, i, j, k]. (NOT the ROS TF convention).

"""
import numpy as np
from matplotlib import pyplot
from mpl_toolkits.mplot3d import Axes3D

npl = np.linalg
PI = np.pi

################################################## FUNCTIONS

def exp(v):
    """
    Returns the unit-quaternion corresponding to a rotation-vector v.

    """
    mag = npl.norm(v)
    if np.isclose(mag, 0): return np.array([1, 0, 0, 0])
    angle_by_2 = np.mod(mag, 2*PI) / 2
    axis = np.divide(v, mag)
    return np.insert(np.sin(angle_by_2)*axis, 0, np.cos(angle_by_2))

def log(q):
    """
    Returns a rotation-vector corresponding to a unit-quaternion q.

    """
    vnorm = npl.norm(q[1:])
    if np.isclose(vnorm, 0): return np.zeros(3)
    if q[0] < 0: q = np.negative(q)  # optional convention
    qnorm = npl.norm(q)
    return 2*np.arccos(q[0]/qnorm)*(q[1:]/vnorm)

def compose(ql, qr):
    """
    Returns the product of two quaternions as ql*qr.

    """
    return np.array([-ql[1]*qr[1] - ql[2]*qr[2] - ql[3]*qr[3] + ql[0]*qr[0],
                      ql[1]*qr[0] + ql[2]*qr[3] - ql[3]*qr[2] + ql[0]*qr[1],
                     -ql[1]*qr[3] + ql[2]*qr[0] + ql[3]*qr[1] + ql[0]*qr[2],
                      ql[1]*qr[2] - ql[2]*qr[1] + ql[3]*qr[0] + ql[0]*qr[3]])

def invert(q):
    """
    Returns the inverse of a unit-quaternion q.

    """
    return np.array([q[0], -q[1], -q[2], -q[3]])

def act(q, v):
    """
    Returns the transformation (rotation) of a vector v by a unit-quaternion q.

    """
    return compose(compose(q, np.insert(v, 0, 0)), invert(q))[1:]

################################################## TESTS

# Test exp
assert np.allclose(exp([0, 0, PI/2]),
                   [np.sqrt(2)/2, 0, 0, np.sqrt(2)/2])  # 90deg rotation about z axis
assert np.allclose(exp([PI, 0, 0]),
                   [0, 1, 0, 0])  # 180deg rotation about x axis

# Test log
random_quat = np.random.sample(4)
random_quat /= npl.norm(random_quat)
assert np.allclose(exp(log(random_quat)), random_quat)  # exp and log are inverses

################################################## DEMO

# Arbitrary mean
q0 = [1, 0, 0, 0]

# Arbitrary covariance
P = [[0.01,   0,   0],
     [   0, 0.1,   0],
     [   0,   0, 0.5]]  # rad^2

# Generate samples
samples = []
for i in range(3000):

    # Sample a random tangent-vector to SO3
    v = np.random.multivariate_normal(np.zeros(3), P)  # Gaussian is arbitrary choice

    # Perturb mean by this tangent
    sample = compose(exp(v), q0)
    samples.append(sample)

    # ^ In the above, if q0 converts B-coordinates to A-coordinates, then composing exp(v)
    # on the left means the noise is in A-coordinates, while composing exp(v) on the right
    # would mean that the noise is in B-coordinates.

# Transform some point by the SO3 samples to demonstrate the distribution
point0 = [1, 0, 0]
points_x, points_y, points_z = zip(*[act(sample, point0) for sample in samples])

# Visualize distribution
fig = pyplot.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(points_x, points_y, points_z, c='b', alpha=0.075)
ax.set_title("Samples", fontsize=22)
ax.set_xlabel("x", fontsize=22)
ax.set_ylabel("y", fontsize=22)
ax.set_zlabel("z", fontsize=22)
pyplot.show()


"""Utils"""
import math
import numpy as np

def make_topic_array_from_tuple_list(weight_tuples, num_topics, scale_factor=1.0):
    topic_array = [0 for i in range(num_topics)]
    total = float(sum([weight_tuple[1] for weight_tuple in weight_tuples])) # For normalization
    for weight_tuple in weight_tuples:
        topic_index = weight_tuple[0]
        topic_weight = weight_tuple[1]
        topic_array[topic_index] = (topic_weight) * scale_factor  # For normalization + scaling
    return topic_array

def make_tuple_list_from_topic_array(topic_array):
	tuple_list = []
	for i in range(len(topic_array)):
		tuple_list.append((i,topic_array[i]))
	return tuple_list

def gaussian(x, mean, variance):
    a = math.sqrt(2*math.pi*variance)
    b = mean
    c = variance
    dist_squared = np.sum((x - b)**2)
    return a*math.exp(-1*dist_squared/(2*c*c))

def print_median_std_from_clusters(clusters):
    counts = []
    num_points = sum([len(cluster) for cluster in clusters])
    mean = num_points/len(clusters)
    SD_sum = 0.0
    for i in range(len(clusters)):
        cluster = clusters[i]
        count = len(cluster)
        counts.append(count)
        v = (mean - count)**2
        SD_sum += v
    ave_var = SD_sum/len(clusters)
    SD = math.sqrt(ave_var)
    med = median(counts)
    print "Standard deviation", SD
    print "median", med

def median(mylist):
    sorts = sorted(mylist)
    length = len(sorts)
    if not length % 2:
        return (sorts[length / 2] + sorts[length / 2 - 1]) / 2.0
    return sorts[length / 2]


from __future__ import annotations

import numpy as np
import torch
from torch import Tensor

# TODO: base `non_dominated_set_2d` and `fast_non_dominated_sort` on Tensor
def non_dominated_set_2d(y, minimize=True):
    """
    Argument
    --------
    y : numpy 2d array,
        where the each solution occupies a row
    """
    y = np.asarray(y)
    N, _ = y.shape

    if isinstance(minimize, bool):
        minimize = [minimize]

    minimize = np.asarray(minimize).ravel()
    assert len(minimize) == 1 or minimize.shape == (N,)
    y *= (np.asarray([-1] * N) ** minimize).reshape(-1, 1)

    _ = np.argsort(y[:, 0])[::-1]
    y2 = y[_, 1]
    ND = []
    for i in range(N):
        v = y2[i]
        if not any(v <= y2[ND]) or len(ND) == 0:
            ND.append(i)
    return _[ND]

def fast_non_dominated_sort(fitness):
    fronts = []
    dominated_set = []
    mu = fitness.shape[1]
    n_domination = np.zeros(mu)

    for i in range(mu):
        p = fitness[:, i]
        p_dominated_set = []
        n_p = 0

        for j in range(mu):
            q = fitness[:, j]
            if i != j:
                # TODO: verify this part
                # check the strict domination
                # allow for duplication points on the same front
                if all(p <= q) and not all(p == q):
                    p_dominated_set.append(j)
                elif all(p >= q) and not all(p == q):
                    n_p += 1

        dominated_set.append(p_dominated_set)
        n_domination[i] = n_p

    # create the first front
    fronts.append(np.nonzero(n_domination == 0)[0].tolist())
    n_domination[n_domination == 0] = -1

    i = 0
    while True:
        for p in fronts[i]:
            p_dominated_set = dominated_set[p]
            n_domination[p_dominated_set] -= 1

        _front = np.nonzero(n_domination == 0)[0].tolist()
        n_domination[n_domination == 0] = -1

        if len(_front) == 0:
            break
        fronts.append(_front)
        i += 1

    return fronts

def is_non_dominated(Y: Tensor, deduplicate: bool = True) -> Tensor:
    r"""Computes the non-dominated front.

    Note: this assumes maximization.

    Args:
        Y: A `(batch_shape) x n x m`-dim tensor of outcomes.
        deduplicate: A boolean indicating whether to only return
            unique points on the pareto frontier.

    Returns:
        A `(batch_shape) x n`-dim boolean tensor indicating whether
        each point is non-dominated.
    """
    if Y.shape[-2] == 0:
        return torch.zeros(Y.shape[:-1], dtype=torch.bool, device=Y.device)
    Y1 = Y.unsqueeze(-3)
    Y2 = Y.unsqueeze(-2)
    dominates = (Y1 >= Y2).all(dim=-1) & (Y1 > Y2).any(dim=-1)
    nd_mask = ~(dominates.any(dim=-1))
    if deduplicate:
        # remove duplicates
        # find index of first occurrence  of each unique element
        indices = (Y1 == Y2).all(dim=-1).long().argmax(dim=-1)
        keep = torch.zeros_like(nd_mask)
        keep.scatter_(dim=-1, index=indices, value=1.0)
        return nd_mask & keep
    return nd_mask


import numpy as np
import scipy

def chToDir(newpath):
    import os
    os.chdir(newpath)

def getRandomMatrix(m, k, mean=0, std=1, seed=-1):
    '''
    Generate randomly(Gaussian) a matrix; the shape is m by k.
    '''
    if seed > 0:
        np.random.seed(seed)
    return np.random.normal(mean, std, m*k).reshape((m,k))

def getOrthogonalMatrix(m, k, mean=0, std=1, seed = -1):
    '''
    Generate randomly(Gaussian) a matrix; the shape is m by k.
    And then QR s.t. the columns of produced matrix are orthogonal.
    Q.T@Q=I
    '''
    if seed > 0:
        np.random.seed(seed)
    H = getRandomMatrix(m, k, mean, std, seed)
    Q, R = np.linalg.qr(H)
    return Q

def getOrthogonalMatrixUniform(m, k, seed = -1):
    '''
    Generate randomly(Uniform) a matrix;
    The columns of produced matrix are orthogonal
    Q.T@Q=I
    '''
    if seed > 0:
        np.random.seed(seed)

    H = np.random.rand(m,k)
    Q, R = np.linalg.qr(H)
    return Q

def UncenterEigen(At):
    '''
    Uncenter PCA: all of eigen values and eigenvectors of At@At.T
    '''
    s, v =scipy.linalg.eigh(At.dot(At.T))

    sorted_idxes = np.argsort(-s)
    s = s[sorted_idxes]
    v = v[:, sorted_idxes]

    return s, v # sorted

def centered_PCA(X):
    mean = np.mean(X, axis=1)
    mean = mean.reshape((mean.shape[0], 1))
    X_red = X - mean
    e, V = scipy.linalg.eigh(X_red@X_red.T)

    sorted_idxes = np.argsort(-e)
    e = e[sorted_idxes]
    V = V[:, sorted_idxes]

    return (mean, e, V) # sorted

def EVD(B,k=-1):
    '''
    The full EVD of B or R.
    @input: B or R
    B = AA.T
    R = (1/n)AA.T
    '''
    e, V = scipy.linalg.eigh(B)

    sorted_idxes = np.argsort(-e)
    e = e[sorted_idxes]
    V = V[:, sorted_idxes]

    return (e, V)

def EVDSparse(B,k):
    '''
    The full EVD of B or R.
    @input: B or R
    B = AA.T
    R = (1/n)AA.T
    '''
    e, V = scipy.sparse.linalg.eigsh(B,k)

    sorted_idxes = np.argsort(-e)
    e = e[sorted_idxes]
    V = V[:, sorted_idxes]

    return (e, V)

def EVDnotPSD(B):
    '''
    The full EVD of B or R.
    @input: B or R
    B = AA.T
    R = (1/n)AA.T
    B and R are not PSD.
    '''
    e, V = scipy.linalg.eig(B)
    # print('D:', e)
    # print('B-B.T:\n', B-B.T)
    # print('B:\n', B)

    sorted_idxes = np.argsort(-e)
    e = e[sorted_idxes]
    V = V[:, sorted_idxes]

    return (e.real, V.real)

def isPSD(B):
    '''B is a sqaure matrix'''
    e, V = scipy.linalg.eig(B)
    e = e.real
    for i in range(len(e)):
        if e[i] <= 0:
            print(e[i])
            return False
    return True
def randomSign():

    v = np.random.randint(0, 2)
    if v == 0:
        return 1
    if v ==1:
        return -1

def randomInt(low, high, size):
    '''
    return a list
    with all values are integers, no repeating,
    and with len equals to size.
    '''

    ret = []
    while(len(ret) < size):
        v = np.random.randint(low, high)
        if v not in ret:
            ret.append(v)
    return ret

def normalizeP(P):
    '''
    P is a list
    '''

    Z = sum(P)
    if Z != 0:
        P = P/Z
    return P

def normalizeDict(P):
    if sum(P.values()) == 0:
        factor = 0
    else:
        factor=1.0/sum(P.values())
    normalised_P = {k: v*factor for k, v in P.items()}
    return normalised_P

def argmaxDict(P):
    import operator
    return max(P.items(), key=operator.itemgetter(1))[0]

def sortDictbyKey(d):
    import collections
    return collections.OrderedDict(sorted(d.items()))

def printred(args):
    CRED = '\033[91m'
    CEND = '\033[0m'
    print(CRED+args+CEND)

def trueUncenterEVD(At, k):

    m,n = At.shape
    C = (At@At.T)/n
    s, v =scipy.linalg.eigh(C)

    sorted_idxes = np.argsort(-s)
    s = s[sorted_idxes]
    v = v[:, sorted_idxes]
    tD = s[0:k]
    tV = v[:,0:k]

    return tD,tV

def trueCenterEVD(At, k):
    m,n = At.shape
    mean = np.mean(At, axis=1)
    mean = mean.reshape((mean.shape[0], 1))
    At_c = At - mean
    C = (At_c@At_c.T)/n
    s, v =scipy.linalg.eigh(C)

    sorted_idxes = np.argsort(-s)
    s = s[sorted_idxes]
    v = v[:, sorted_idxes]
    tD = s[0:k]
    tV = v[:,0:k]

    return tD,tV

def coverToList(data, method='integer'):
    '''
    Convert string list to float or integer list
    input should be something like [30, 31, 32] or [30.1, 31.3, 32.0]
    method: float, or integer(default)
    '''
    data = data.split("[")[1]
    data = data.split("]")[0]
    data = data.split(",")
    data = list(data)
    ret =[]
    for i, item in enumerate(data):
        item = item.split(' ')
        item = item[len(item)-1]
        if method == 'float':
            ret.append(float(item))
        if method == 'integer':
            ret.append(int(item))
    return ret


import numpy as np
import random

# 8 steps
next_step = [(2,1), (1,2), (2,-1), (1,-2), (-2,1), (-1,2), (-2,-1), (-1,-2)]

class KnightTour:
    def __init__(self, size):
        self.size = size
        self.chessboard = np.zeros((self.size, self.size))

    # randon initional position
    def initiate_pos(self):
        return (random.randint(0,self.size -1), random.randint(0,self.size -1))

    # check the history that it is a avalivable step which is never visiting.
    def check_his(self, new_x, new_y):
        return self.chessboard[new_x][new_y] == 0

    # check next step not out of the chessboard
    def check_pos(self, new_x, new_y):
        return new_x >= 0 and new_x < self.size and new_y >= 0 and new_y < self.size

    # next step
    def move(self, x, y):
        next_pos = []
        for step in range(8):
            new_x = x + next_step[step][0]
            new_y = y + next_step[step][1]
            if  self.check_pos(new_x, new_y) and self.check_his(new_x, new_y):
                next_pos.append([new_x, new_y])
        return next_pos

    # find all steps
    def travel(self):
        x, y = self.initiate_pos()
        self.chessboard[x][y] = 1

        total_steps = self.size**2
        for counter in range(2, total_steps+1):
            pos = self.move(x, y)
            if len(pos) == 0:
                return False
            else :
                first_step = pos[0]

                # move to the bound position first
                for j in pos:
                    if len(self.move(j[0], j[1])) <= len(self.move(first_step[0], first_step[1])):
                        first_step = j

                x, y = first_step
                self.chessboard[x][y] = counter
        return True

if __name__ == '__main__':

    board_size = int(input("Please enter the n of the n*n chessboard :"))
    test = KnightTour(board_size)
    result = test.travel()

    while not result:
        test = KnightTour(board_size)
        result = test.travel()

    print('Success!')
    print('Pathway :')
    print(test.chessboard)


# -*- coding: utf-8 -*-
# Copyright (c) 2020, Matgenix SRL, All rights reserved.
# Distributed open source for academic and non-profit users.
# Contact Matgenix for commercial usage.
# See LICENSE file for details.

"""Example usage of pysisso for a regression using sklearn interface."""

import matplotlib.pyplot as plt
import numpy as np

from pysisso.sklearn import SISSORegressor

# Define general parameters
TITLE = "f(x) = 0.5*x^3 + 0.5*x^2 - 4.0*x - 4.0"
NPOINTS = 100  # Number of data points
SIGMA = 0.5  # Randomness in the data points
PLOT_FIGURES = True  # whether to interactively plot the figures with matplotlib
SAVE_FIGURES = False  # whether to save the matplotlib figures to a file
CLEAN_RUN_DIR = True  # whether to remove the SISSO_dir after the fit

# Set the random seed to always keep the same figure
np.random.seed(42)

# Define the function:
# f(x) = 0.5*x^3 + 0.5*x^2 - 4.0*x - 4.0 (roots = [-2.0, -1.0, 2.0])
def fun(xx, const=1.0):
    return 0.5 * xx ** 3 + 0.5 * xx ** 2 - 4.0 * 0.5 * xx - 4.0 * const

# Define the data set
X = np.random.uniform(-2.5, 2.5, NPOINTS)
y = fun(X) + np.random.normal(0.0, scale=SIGMA, size=NPOINTS)

# Plot true function and data
xlin = np.linspace(-3, 3, 1000)
ylin = fun(xlin)
fig, subplot = plt.subplots()
subplot.plot(xlin, ylin, "-", color="C0", label="True function")
subplot.plot(X, y, "o", color="C1", label="Data")
subplot.set_xlabel("x")
subplot.set_ylabel("f(x)")
subplot.set_title(TITLE)
subplot.legend()
if SAVE_FIGURES:
    fig.savefig("true_data.pdf")
if PLOT_FIGURES:
    plt.show()

# Define the regressor, fit the data and predict
sisso_regressor = SISSORegressor(
    rung=1,
    opset="(+)(*)(^2)(^3)(^-1)(cos)(sin)",
    desc_dim=3,
    clean_run_dir=CLEAN_RUN_DIR,
)
X = X.reshape(-1, 1)  # only one feature, X is initially defined as 1D, sklearn needs 2D
sisso_regressor.fit(X, y)
ylin_pred = sisso_regressor.predict(xlin)

# Plot the true and predicted functions, together with the data
fig, subplot = plt.subplots()
subplot.plot(xlin, ylin, "-", color="C0", label="True function")
subplot.plot(X, y, "o", color="C1", label="Data")
subplot.plot(xlin, ylin_pred, "-", color="C2", label="Predicted function")
subplot.set_xlabel("x")
subplot.set_ylabel("f(x)")
subplot.set_title(TITLE)
subplot.legend()
if SAVE_FIGURES:
    fig.savefig("true_data_pred.pdf")
if PLOT_FIGURES:
    plt.show()


#!/usr/bin/env python
# Demonstrate criterion of Chauvenet to exclude poor data

from numpy.random import normal
from scipy.special import erf, erfc
import numpy
from kapteyn import kmpfit
from matplotlib.pyplot import figure, show, rc

def chauvenet(x, y, mean=None, stdv=None):
   #-----------------------------------------------------------
   # Input:  NumPy arrays x, y that represent measured data
   #         A single value of a mean can be entered or a
   #         sequence of means with the same length as
   #         the arrays x and y. In the latter case, the
   #         mean could be a model with best-fit parameters.
   # Output: It returns a boolean array as filter.
   #         The False values correspond to the array elements
   #         that should be excluded
   #
   # First standardize the distances to the mean value
   # d = abs(y-mean)/stdv so that this distance is in terms
   # of the standard deviation.
   # Then the  CDF of the normal distr. is given by
   # phi = 1/2+1/2*erf(d/sqrt(2))
   # Note that we want the CDF from -inf to -d and from d to +inf.
   # Note also erf(-d) = -erf(d).
   # Then the threshold probability = 1-erf(d/sqrt(2))
   # Note, the complementary error function erfc(d) = 1-erf(d)
   # So the threshold probability pt = erfc(d/sqrt(2))
   # If d becomes bigger, this probability becomes smaller.
   # If this probability (to obtain a deviation from the mean)
   # becomes smaller than 1/(2N) than we reject the data point
   # as valid. In this function we return an array with booleans
   # to set the accepted values.
   #
   # use of filter:
   # xf = x[filter]; yf = y[filter]
   # xr = x[~filter]; yr = y[~filter]
   # xf, yf are cleaned versions of x and y and with the valid entries
   # xr, yr are the rejected values from array x and y
   #-----------------------------------------------------------
   if mean is None:
      mean = y.mean()           # Mean of incoming array y
   if stdv is None:
      stdv = y.std()            # Its standard deviation
   N = len(y)                   # Lenght of incoming arrays
   criterion = 1.0/(2*N)        # Chauvenet's criterion
   d = abs(y-mean)/stdv         # Distance of a value to mean in stdv's
   d /= 2.0**0.5                # The left and right tail threshold values
   prob = erfc(d)               # Area normal dist.
   filter = prob >= criterion   # The 'accept' filter array with booleans
   return filter                # Use boolean array outside this function

def residuals(p, data):
   a, b = p
   x, y = data
   return (y-a-b*x)

def prob( Xlo, Xhi ):
   sq2 = numpy.sqrt(2.0)
   # Left tail and right tail area
   return 1.0 - 0.5 * (erf(Xhi/sq2) - erf(Xlo/sq2))

# Artificial data
N = 100
a0 = 2; b0 = 3
x = numpy.linspace(0.0, 2.0, N)
y = a0 + b0*x + normal(0.0, 1.0, N)  # Mean 0, sigma 1

# Add outlier(s)
y[40] = 14
y[7] = 12
y[-5] = 0.2

# Prepare plot
fig = figure()
rc('legend', fontsize=8)
frame = fig.add_subplot(1,1,1)
frame.plot(x, y, 'go', label="data")
frame.set_xlabel("x")
frame.set_ylabel("y")
frame.set_title("Exclude poor data with criterion of Chauvenet")

fitter = kmpfit.Fitter(residuals=residuals, data=(x,y))
fitter.fit(params0=(a0,b0))

print("======== Fit results all data included ==========")
print("Params:                      ", fitter.params)
print("Uncertainties:               ", fitter.xerror)
print("Errors assuming red.chi^2=1: ", fitter.stderr)
print("Iterations:                  ", fitter.niter)
print("Function ev:                 ", fitter.nfev)
print("dof:                         ", fitter.dof)
print("chi^2, rchi2:                ", fitter.chi2_min, fitter.rchi2_min)
print("Status:                      ", fitter.status)

#mean = y.mean() # Use this for data with slope 0
a, b = fitter.params
mean = a + b*x
stdv = (y-mean).std()
criterion = 1.0/(2*N)
for xf, yf, m in zip(x, y, mean):
   d = abs(yf-m)/stdv
   P = prob(-d, d)
   if P < criterion:
      print("Data y=%f has distance %fsigma to mean (%f), prob.=%f" %(yf,d,m,P))
      print("Based on Chauvenet's criterion (prob=%g), we reject data point y=%g"\
      %(criterion,yf))
      frame.plot((xf,), (yf,), 'rx', ms=20)

# The function chauvenet() does it the NumPy way
filter = chauvenet(x, y)
xf = x[filter]; yf = y[filter]
xe = x[~filter]; ye = y[~filter]
print("Excluded by function chauvenet() are:")
print(list(zip(xe, ye)))
fitter = kmpfit.Fitter(residuals=residuals, data=(xf,yf))
fitter.fit(params0=(a0,b0))

print("\n======== Fit results with filtered data ==========")
print("Params:                      ", fitter.params)

frame.set_ylim(0, 1.1*y.max())
frame.plot(x, a+b*x, 'g', label="Fit unfilterd data")
a, b = fitter.params
frame.plot(x, a+b*x, 'b', label="Fit filterd data")
leg = frame.legend(loc=2)
show()


import numpy as np

def l2_regularization(W, reg_strength):
    """
    Computes L2 regularization loss on weights and its gradient

    Arguments:
      W, np array - weights
      reg_strength - float value

    Returns:
      loss, single value - l2 regularization loss
      gradient, np.array same shape as W - gradient of weight by l2 loss
    """
    # TODO: Copy from the previous assignment
    loss, grad = reg_strength*np.sum(W**2), 2*reg_strength*W
    return loss, grad

def softmax(predictions):
    '''
    Computes probabilities from scores
    Arguments:
      predictions, np array, shape is either (N) or (batch_size, N) -
        classifier output
    Returns:
      probs, np array of the same shape as predictions -
        probability for every class, 0..1
    '''
    # TODO implement softmax
    # Your final implementation shouldn't have any loops
    # Done below
    if len(predictions.shape) == 1:
        # batch_size of 1
        probs = predictions - np.max(predictions)
        probs = np.exp(probs) / np.sum(np.exp(probs))
    elif len(predictions.shape) > 1:
        probs = predictions - np.max(predictions, axis=1)[:, np.newaxis]
        probs = np.exp(probs) / np.sum(np.exp(probs), axis=1)[:, np.newaxis]
    return probs

def cross_entropy_loss(probs, target_index):
    '''
    Computes cross-entropy loss
    Arguments:
      probs, np array, shape is either (N) or (batch_size, N) -
        probabilities for every class
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)
    Returns:
      loss: single value
    '''
    # TODO implement cross-entropy
    # Your final implementation shouldn't have any loops
    # Done below
    if len(probs.shape) == 1:
        # batch_size of 1
        row_index = None
    elif len(probs.shape) > 1:
        N = probs.shape[0]
        row_index = range(N)
    H = -np.mean(np.log(probs[row_index, target_index]))
    return H

def softmax_with_cross_entropy(preds, target_index):
    """
    Computes softmax and cross-entropy loss for model predictions,
    including the gradient

    Arguments:
      predictions, np array, shape is either (N) or (batch_size, N) -
        classifier output
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)

    Returns:
      loss, single value - cross-entropy loss
      dprediction, np array same shape as predictions - gradient of predictions by loss value
    """
    # TODO: Copy from the previous assignment
    if len(preds.shape) == 1:
        # batch_size of 1
        N = 1
        row_index = None
    elif len(preds.shape) > 1:
        N = preds.shape[0]
        row_index = range(N)
    d_preds = softmax(preds)
    loss = cross_entropy_loss(d_preds, target_index)
    d_preds[row_index, target_index] -= 1
    d_preds /= N
    return loss, d_preds

class Param:
    """
    Trainable parameter of the model
    Captures both parameter value and the gradient
    """

    def __init__(self, value):
        self.value = value
        self.grad = np.zeros_like(value)

class ReLULayer:
    def __init__(self):
        pass

    def forward(self, X):
        # TODO: Implement forward pass
        # Hint: you'll need to save some information about X
        # to use it later in the backward pass
        self.cache = X
        return np.maximum(X, 0)

    def backward(self, d_out):
        """
        Backward pass

        Arguments:
        d_out, np array (batch_size, num_features) - gradient
           of loss function with respect to output

        Returns:
        d_result: np array (batch_size, num_features) - gradient
          with respect to input
        """
        # TODO: Implement backward pass
        # Your final implementation shouldn't have any loops
        X = self.cache
        d_result = np.multiply(d_out, X > 0)
        return d_result

    def params(self):
        # ReLU Doesn't have any parameters
        return {}

class FullyConnectedLayer:
    def __init__(self, n_input, n_output):
        self.W = Param(0.001 * np.random.randn(n_input, n_output))
        self.B = Param(0.001 * np.random.randn(1, n_output))
        self.X = None

    def forward(self, X):
        # TODO: Implement forward pass
        # Your final implementation shouldn't have any loops
        self.X = X
        W, B = self.params()['W'].value, self.params()['B'].value
        out = np.dot(X, W) + B
        return out

    def backward(self, d_out):
        """
        Backward pass
        Computes gradient with respect to input and
        accumulates gradients within self.W and self.B

        Arguments:
        d_out, np array (batch_size, n_output) - gradient
           of loss function with respect to output

        Returns:
        d_result: np array (batch_size, n_input) - gradient
          with respect to input
        """
        # TODO: Implement backward pass
        # Compute both gradient with respect to input
        # and gradients with respect to W and B
        # Add gradients of W and B to their `grad` attribute

        # It should be pretty similar to linear classifier from
        # the previous assignment
        W = self.params()['W'].value
        X = self.X

        # Gradient wrt to weights
        self.params()['W'].grad = np.dot(X.T, d_out)

        # Gradient wrt to bias
        self.params()['B'].grad = np.sum(d_out, axis=0).reshape(1, -1)

        # Gradient wrt to input
        d_input = np.dot(d_out, W.T)

        return d_input

    def params(self):
        return {'W': self.W, 'B': self.B}

class ConvolutionalLayer:
    def __init__(self, in_channels, out_channels,
                 filter_size, padding):
        '''
        Initializes the layer

        Arguments:
        in_channels, int - number of input channels
        out_channels, int - number of output channels
        filter_size, int - size of the conv filter
        padding, int - number of 'pixels' to pad on each side
        '''

        self.filter_size = filter_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.W = Param(
            np.random.randn(filter_size, filter_size,
                            in_channels, out_channels)
        )

        self.B = Param(np.zeros(out_channels))

        self.padding = padding

    def forward(self, X):
        batch_size, height, width, channels = X.shape

        # We need to save X to perform backward pass
        self.X = X
        pad_width, filter_size = self.padding, self.filter_size

        # Padding X:
        if self.padding:
            npad = ((0, 0), (pad_width, pad_width), (pad_width, pad_width), (0, 0))
            X = np.pad(X, npad, 'constant', constant_values=0)
        stride = 1
        out_height = int((height - filter_size + 2*pad_width) / stride + 1)
        out_width = int((width - filter_size + 2*pad_width) / stride + 1)
        # TODO: Implement forward pass
        # Hint: setup variables that hold the result
        # and one x/y location at a time in the loop below
        # It's ok to use loops for going over width and height
        # but try to avoid having any other loops
        conv_layer_out = np.zeros((batch_size, out_height, out_width, self.out_channels))
        W_reshaped = self.W.value.reshape(-1, self.out_channels)
        for y in range(out_height):
            for x in range(out_width):
                # TODO: Implement forward pass for specific location
                Xf = X[:, y: y + filter_size, x: x + filter_size, :]
                out = (Xf.reshape(batch_size, -1)).dot(W_reshaped)
                conv_layer_out[:, y, x, :] = out + self.B.value
        return conv_layer_out

    def backward(self, d_out):
        # Hint: Forward pass was reduced to matrix multiply
        # You already know how to backprop through that
        # when you implemented FullyConnectedLayer
        # Just do it the same number of times and accumulate gradients
        X = self.X
        pad_width, filter_size = self.padding, self.filter_size
        if self.padding:
            npad = ((0, 0), (pad_width, pad_width), (pad_width, pad_width), (0, 0))
            X = np.pad(X, npad, 'constant', constant_values=0)

        batch_size, height, width, channels = X.shape
        _, out_height, out_width, out_channels = d_out.shape

        # TODO: Implement backward pass
        # Same as forward, setup variables of the right shape that
        # aggregate input gradient and fill them for every location
        # of the output
        Wgrad = np.zeros(self.W.value.shape)
        d_input = np.zeros(X.shape)
        W_reshaped = self.W.value.reshape(-1, self.out_channels)
        # Try to avoid having any other loops here too
        for y in range(out_height):
            for x in range(out_width):
                # TODO: Implement backward pass for specific location
                # Aggregate gradients for both the input and
                # the parameters (W and B)

                # Gradient wrt weights
                Xf = X[:, y: y + self.filter_size, x: x + self.filter_size, :]
                out = Xf.reshape(batch_size, -1).T.dot(d_out[:, y, x, :])
                Wgrad += out.reshape(self.filter_size, self.filter_size, channels, out_channels)

                # Gradient wrt input
                gradinput = d_out[:, y, x, :].dot(W_reshaped.T).reshape(batch_size, self.filter_size, self.filter_size, channels)
                d_input[:, y: y + self.filter_size, x: x + self.filter_size, :] += gradinput

        # Cut off the padding
        if self.padding:
            d_input = d_input[:, pad_width:-pad_width, pad_width:-pad_width, :]

        # Gradient wrt to weights
        self.params()['W'].grad = Wgrad

        # Gradient wrt to bias
        self.params()['B'].grad = np.sum(d_out.reshape(-1, out_channels), axis=0)

        return d_input

    def params(self):
        return {'W': self.W, 'B': self.B}

class MaxPoolingLayer:
    def __init__(self, pool_size, stride):
        '''
        Initializes the max pool

        Arguments:
        pool_size, int - area to pool
        stride, int - step size between pooling windows
        '''
        self.pool_size = pool_size
        self.stride = stride
        self.X = None

    def forward(self, X):
        self.X = X
        batch_size, height, width, channels = X.shape
        pool_size, stride = self.pool_size, self.stride
        out_height = int((height - pool_size)/stride+1)
        out_width = int((width - pool_size)/stride+1)
        # TODO: Implement maxpool forward pass
        # Hint: Similarly to Conv layer, loop on
        # output x/y dimension
        maxpoolout = np.zeros((batch_size, out_height, out_width, channels))
        self.mask = {}
        for y in range(out_height):
            for x in range(out_width):
                X_F = X[:, y*stride: y*stride + pool_size, x*stride: x*stride + pool_size, :]
                maxpoolout[:, y, x, :] = np.amax(X_F, axis=(1, 2)).reshape(batch_size, channels)

                mask = np.zeros_like(X_F)
                X_F_rshpd = X_F.reshape(batch_size, pool_size*pool_size, channels)
                idx = X_F_rshpd.argmax(1)
                ax1, ax2 = np.indices((batch_size, channels))
                mask.reshape(mask.shape[0], mask.shape[1] * mask.shape[2], mask.shape[3])[ax1, idx, ax2] = 1
                self.mask[(y, x)] = mask
        return maxpoolout

    def backward(self, d_out):
        # TODO: Implement maxpool backward pass
        pool_size, stride = self.pool_size, self.stride
        batch_size, height, width, channels = self.X.shape
        _, out_height, out_width, out_channels = d_out.shape
        d_input = np.zeros((batch_size, height, width, channels))
        for y in range(out_height):
            start_h = y*stride
            end_h = start_h + pool_size
            for x in range(out_width):
                start_w = x*stride
                end_w = start_w + pool_size
                mask = self.mask[(y, x)]
                d_input[:, start_h: end_h, start_w: end_w, :] += \
                    d_out[:, y, x, :].reshape(batch_size, 1, 1, channels)*mask
        return d_input

    def params(self):
        return {}

class Flattener:
    def __init__(self):
        self.X_shape = None

    def forward(self, X):
        batch_size, height, width, channels = X.shape
        self.X_shape = X.shape

        # TODO: Implement forward pass
        # Layer should return array with dimensions
        # [batch_size, hight*width*channels]
        return X.reshape(batch_size, -1)

    def backward(self, d_out):
        # TODO: Implement backward pass
        return d_out.reshape(self.X_shape)

    def params(self):
        # No params!
        return {}


"""
    A Monte Carlo approach to solving the cube problem

    author: Thomas Stainer
    date: 26/09/2019 @ 21:34
"""

import math
import numpy as np
from scipy import stats

showplot = True
NITERS = 10000

# make a graph of a cube with 8 vertices
# we could indeed use a library such as networkx
# but we keep it simple and vanilla for
# educational purposes
CUBEGRAPH = {
    'nodes': [
        1, 2, 3, 4, 5, 6, 7, 8
    ],
    'edges': [
        (1, 2),
        (1, 4),
        (1, 6),
        (2, 3),
        (2, 7),
        (3, 4),
        (3, 8),
        (4, 5),
        (5, 6),
        (5, 8),
        (6, 7),
        (7, 8)
    ]
}

def getadjacent(graph, vertex):
    """
        vertex is the node indexnumber i.e 1, 2, 3, ... or 8
        and returns all adjacent nodes

        For example:
            getadjacent(1) -> [2, 4, 6]
            getadjacent(2) -> [1, 3, 7]
            getadjacent(3) -> [2, 4, 8]
            ....
    """
    nodes = []
    for n1, n2 in graph['edges']:
        if n1 == vertex:
            nodes.append(n2)
        if n2 == vertex:
            nodes.append(n1)
    return sorted(nodes)

def randomstep(graph, start, steps=1):
    """
        Take a random step(s) from a starting node.
        Return the node ending on.

        Since all edges are equal all weights are 1

        This function assumes weight =1, to extend
        we would need to define weights in edge data
    """
    if start < min(graph['nodes']) or start > max(graph['nodes']):
        raise RuntimeError("Invalid node number: {}".format(start))

    if steps == 0:
        return start

    # get all adjacent nodes
    anodes = getadjacent(graph, start)

    # generate a random number drawn from a uniform distribution
    rmin, rmax = 0, 1
    prob = np.random.uniform(rmin, rmax, 1)[0]

    # equal weights
    # this remains unchanged every call, should be moved
    # out but if weights change then code is needed here
    weights = [(rmax-rmin)/len(anodes)]*len(anodes)
    r = rmin
    for w, n in zip(weights, anodes):
        r += w
        if r > prob:
            # use recursion for this problem
            return randomstep(graph, n, steps=steps-1)

    return None

def getnumberofsteps(graph, start, end):
    nextnode = start
    count = 0
    while nextnode != end or (count == 0 and nextnode == end):
        nextnode = randomstep(graph, nextnode, steps=1)
        count += 1
    return count

def estimateaveragewalk(graph, start, end, niter=1000, op=np.mean):
    steps = []
    for _ in range(niter):
        steps.append(getnumberofsteps(graph, start, end))
    return op(steps)

for node in CUBEGRAPH['nodes']:
    # 1 -> X
    # all others are same due to symmetry
    averagesteps = estimateaveragewalk(CUBEGRAPH, 1, node, niter=NITERS)
    mostcommon = estimateaveragewalk(CUBEGRAPH, 1, node, niter=NITERS, op=stats.mode)
    print("1 -> {} = {:.3f} with dominant path of {} step(s)".format(node, averagesteps, mostcommon.mode[0]))

if showplot:
    import matplotlib.pyplot as plt
    from scipy.optimize import curve_fit
    from scipy.stats import norm
    from tqdm import tqdm

    # check distribution of final node for 8 steps
    nodes = []
    for _ in range(10000):
        nodes.append(randomstep(CUBEGRAPH, 1, steps=8))

    fig = plt.figure()
    plt.hist(nodes, bins=range(min(CUBEGRAPH['nodes']), max(CUBEGRAPH['nodes'])),
        facecolor='r', edgecolor='black', linewidth=1.2, alpha=0.4)

    # check distribution of path lengths for 1 -> 1
    steps = []
    for _ in range(100000):
        steps.append(getnumberofsteps(CUBEGRAPH, 1, 1))

    fig = plt.figure()
    n, bins, patches = plt.hist(steps, bins=range(0, 100), facecolor='r',
        edgecolor='black', linewidth=1.2, alpha=0.4)
    total = 0
    for i,j in enumerate(n):
        total += i*j/100000

    print(total)

    plt.xlabel("number of steps 1 -> 1", fontsize=16)
    plt.ylabel("count", fontsize=16)
    plt.yscale('log', nonposy='clip')
    # print(n)

    # run walk 1 -> 1 and get average many times to estimate uncertainty
    iters = 100
    values = []
    for _ in tqdm(range(iters)):
        values.append(estimateaveragewalk(CUBEGRAPH, 1, 1, niter=1000))

    fig = plt.figure()
    n, bins, patches = plt.hist(values, density=True, bins=20, facecolor='k',
        edgecolor='black', linewidth=2, alpha=0.2)

    # normalise values
    mu, std = norm.fit(values)

    xmin, xmax = plt.xlim()
    x = np.linspace(xmin, xmax, 100)
    p = norm.pdf(x, mu, std)
    plt.plot(x, p, 'k', linewidth=2)
    plt.xlabel("number of steps 1 -> 1", fontsize=16)
    plt.ylabel("density", fontsize=16)
    plt.yscale('linear')
    plt.title("Fit results: mu = {:.2f},  std = {:.2f}".format(mu, std))

    plt.show()


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

# # Mine management model
#
# **Randall Romero Aguilar, PhD**
#
# This demo is based on the original Matlab demo accompanying the  Computational Economics and Finance 2001 textbook by Mario Miranda and Paul Fackler.
#
# Original (Matlab) CompEcon file: **demddp03.m**
#
# Running this file requires the Python version of CompEcon. This can be installed with pip by running
#
#     !pip install compecon --upgrade
#
# Last updated: 2021-Oct-01
#

# ## About
#
# A mine operator must decide how much ore to extract from a mine that will be shut down and abandoned after $T$ years of operation. The price of extracted ore is $p$ dollars per ton, and the total cost of extracting $x$ tons of ore in any year, given that the mine contains $s$ tons at the beginning of the year, is $c(s, x)$ dollars. The mine currently contains $\bar{s}$ tons of ore. Assuming the amount of ore extracted in any year must be an integer number of tons, what extraction schedule maximizes profits?

# ## Initial tasks
#

# In[1]:

import numpy as np
import pandas as pd
from compecon import DDPmodel, getindex

# ## Model Parameters

# Assuming a market price $p=1$, initial stock of ore $\bar{s}=100$, and annual discount factor $\delta = 0.9$

# In[2]:

price = 1
sbar  = 100
delta = 0.9

# ### State Space

# This is a finite horizon, deterministic model with time $t$ measured in years. The state
# variable $s \in \{0, 1, 2, \dots, \bar{s}\}$ is the amount of ore remaining in the mine at the beginning of the year, measured in tons.

# In[3]:

S = np.arange(sbar + 1)      # vector of states
n = S.size                   # number of states

# ### Action Space

# The action variable $x \in \{0, 1, 2, \dots, s\}$ is the amount of ore extracted over the year, measured in tons.

# In[4]:

X = np.arange(sbar + 1)      # vector of actions
m = X.size                   # number of actions

# ### Reward Function

# The reward function is $f(s, x) = px − c(s, x)$. The cost of extraction is $c(s, x) = \frac{x^2}{1+s}$.
#
# Here, the reward is set to negative infinity if the extraction level exceeds the available stock in order to preclude the choice of an infeasible
# action:

# In[5]:

f = np.full((m, n), -np.inf)
for c, s in enumerate(S):
    for r, x in enumerate(X):
        if x <= s:
            f[r, c] = price * x - (x ** 2) / (1 + s)

# ### State Transition Function

# The state transition function is $g(s, x) = s − x$
#
# Here, the routine `getindex` is used to find the index of the following period’s state.

# In[6]:

g = np.empty_like(f)
for r, x in enumerate(X):
    snext = S - x
    g[r] = getindex(snext, S)

# ## Model Structure

# The value of the mine, given that it contains $s$ tons of ore at the beginning of year $t$, satisfies the Bellman equation
#
# \begin{equation} V_t(s) = max_{x\in\{0,1,\dots,s\}} \left\{px−c(s, x) + \delta V_{t+1}(s−x)\right\} \end{equation}
#
# subject to the terminal condition $V_{T+1}(s) = 0$
#
# To solve and simulate this model, use the CompEcon class ```DDPmodel```.

# In[7]:

model = DDPmodel(f, g, delta)
model.solve()

# By default, the `solve()` method uses Newton's algorithm. This and other default settings can be changed when solving the model. For example,
# ```python
# model.solve(algorithm='funcit', print=True)
# ```
# solves the model by function iteration, printing a summary of each iteration to screen.
#
# In either case, `solve()` updates the model itself, storing the $n$ vector of values `.value`, the $n$ vector of optimal actions `.policy`, and the $n\times n$ controlled state `.transition` probability.
#

# In[8]:

solution = pd.DataFrame({
    'Stock': S,
    'Extraction': X[model.policy],
    'Value': model.value}).set_index('Stock')

# ## Analysis

# ### Simulate Model
# The path followed by the stock level over time is computed by the `simulate()` method. Here, the simulation assumes an initial stock level of 100 and 15 years in duration.

# In[9]:

sinit = S.max()
nyrs = 15
t = np.arange(nyrs + 1)
spath, xpath = model.simulate(sinit, nyrs)

simul = pd.DataFrame({
    'Year': t,
    'Stock': S[spath],
    'Extraction': X[xpath]}).set_index('Year')
simul

# ### Plot Optimal Policy

# In[10]:

ax = solution['Extraction'].plot(title='Optimal Extraction Policy')
ax.set(ylabel='Extraction');

# ### Plot Value Function
# The value of the firm is very nearly proportional to the stock level.

# In[11]:

ax = solution['Value'].plot(title='Optimal Value Function')
ax.set(ylabel='Value');

# ### Plot State and Extraction Path
# As seen in this figure, the content of the mine is optimally exhausted in 15 years.

# In[12]:

ax = simul.plot(subplots=True, layout=[1,2], legend=None)
for i, lab in enumerate(simul.columns):
    ax[0, i].set(ylabel=lab);

_build/jupyter_execute/notebooks/ddp/01 Mine management model.py


import numpy as np
from numpy.testing import assert_allclose
from pytest import raises

from fastats.linear_algebra import matrix_minor

def test_matrix_minor():

    A = np.array([[3, 13, 14, 10, 12],
                  [8, 15, 4, 16, 5],
                  [6, 11, 7, 9, 17]])  # <- eliminate this row (idx = 2)
    #                  \
    #                   eliminate this column (idx = 1)

    output = matrix_minor(A, 2, 1)
    expected = np.array([[3, 14, 10, 12],
                         [8,  4, 16, 5]])

    assert_allclose(output, expected)

def test_matrix_minor_2x2():

    A = np.array([[3, 13],
                  [8, 15]])

    output = matrix_minor(A, 1, 1)
    expected = np.array([[3]])
    assert_allclose(output, expected)

def test_matrix_minor_perimeter():

    A = np.array([[3, 13, 14, 10, 12],
                  [8, 15, 4, 16, 5],
                  [6, 11, 7, 9, 17]])

    output = matrix_minor(A, 2, 4)
    expected = np.array([[3, 13, 14, 10],
                         [8, 15, 4, 16]])

    assert_allclose(output, expected)

def test_matrix_minor_interior():

    A = np.array([[3, 13, 14, 10, 12],
                  [8, 15, 4, 16, 5],
                  [6, 11, 7, 9, 17]])

    output = matrix_minor(A, 1, 3)
    expected = np.array([[3, 13, 14, 12],
                         [6, 11, 7, 17]])

    assert_allclose(output, expected)

def test_matrix_minor_raises_if_idx_out_of_bounds():

    A = np.array([[3, 13, 14, 10, 12],
                  [8, 15, 4, 16, 5],
                  [6, 11, 7, 9, 17]])

    with raises(AssertionError):
        _ = matrix_minor(A, 0, 5)  # 5 exceeds index of final column

    with raises(AssertionError):
        _ = matrix_minor(A, 3, 0)  # 3 exceeds index of final row

def test_matrix_minor_raises_if_array_not_at_least_2x2():

    A = np.array([3, 13, 14, 10, 12]).reshape(1, 5)

    with raises(AssertionError):
        _ = matrix_minor(A, 0, 2)

    A = np.array([3, 13, 14, 10, 12]).reshape(5, 1)

    with raises(AssertionError):
        _ = matrix_minor(A, 2, 0)

if __name__ == '__main__':
    import pytest
    pytest.main([__file__])


from IMLearn.learners import UnivariateGaussian, MultivariateGaussian
import numpy as np
# import plotly.graph_objects as go
# import plotly.io as pio
from matplotlib import pyplot as plt

# pio.templates.default = "simple_white"

def test_univariate_gaussian():
    # Question 1 - Draw samples and print fitted model
    mu = 10
    sigma = 1
    univariategaussian = UnivariateGaussian().fit(np.random.normal(loc=mu, scale=sigma, size=1000))
    print("({},{})".format(univariategaussian.mu_, univariategaussian.var_))

    # Question 2 - Empirically showing sample mean is consistent
    mu = 10
    sigma = 1
    x = np.arange(10, 1000, 10)
    mu_array = []
    sigma_array = []
    for i in x:
        univariategaussian = UnivariateGaussian().fit(np.random.normal(loc=mu, scale=sigma, size=i))
        mu_array.append(np.abs(univariategaussian.mu_ - mu))
    plt.plot(x, mu_array)
    plt.legend(["mu"])
    plt.xlabel("Sample Size")
    plt.ylabel("Absolute distance from real value")
    plt.title("The different errors according to different sample size.")
    plt.show()

    # Question 3 - Plotting Empirical PDF of fitted model
    univariategaussian = UnivariateGaussian().fit(np.random.normal(loc=mu, scale=sigma, size=1000))
    t = np.linspace(6, 14, 1000)
    y = univariategaussian.pdf(t)
    plt.scatter(t, y)
    plt.legend(["PDF"])
    plt.xlabel("Sample value")
    plt.ylabel("Density of probability")
    plt.title("PDF function for mu=10, sigma=1")  # TODO: answer the Q3
    plt.show()

def test_multivariate_gaussian():
    # Question 4 - Draw samples and print fitted model
    mu = np.array([0, 0, 4, 0]).T
    cov = np.array([[1, 0.2, 0, 0.5], [0.2, 2, 0, 0], [0, 0, 1, 0], [0.5, 0, 0, 1]])
    samples = np.random.multivariate_normal(mu, cov, 1000)
    mvg = MultivariateGaussian().fit(samples)
    print()
    print("expectations:")
    print(mvg.mu_)
    print("cov matrix:")
    print(mvg.cov_)

    # Question 5 - Likelihood evaluation
    f3 = f1 = np.linspace(-10, 10, 200)
    y = np.zeros(shape=(200, 200))
    b = True
    max = 0
    maxi = (0, 0)
    samples = np.random.multivariate_normal(mu, cov, 1000)
    for i in range(200):
        for j in range(200):
            mu = np.array([f1[i], 0, f3[j], 0]).T
            y[i][j] = MultivariateGaussian.log_likelihood(mu, cov, samples)
            if b:
                max = y[i][j]
                b = False
            if max < y[i][j]:
                max = y[i][j]
                maxi = (i, j)

    plt.imshow(y, extent=[-10, 10, -10, 10])
    plt.title("log-likelihood as a function of f1 and f3")
    plt.colorbar()
    plt.xlabel("feature 3")
    plt.ylabel("feature 1")
    plt.show()

    # Question 6 - Maximum likelihood
    print(maxi)

if __name__ == '__main__':
    np.random.seed(0)
    test_univariate_gaussian()
    test_multivariate_gaussian()


import math
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
startTime = datetime.now()

def angular_momentum(posArray, velArray, massArray):
    '''calculates the total angular momentum for a system of objects
    '''
    numObjects=len(massArray)
    numDimensions=len(posArray[0])
    momentum=[[0.0]*numDimensions for j in range(numObjects)]
    angMomen=[[0.0]*numDimensions for j in range(numObjects)]
    angMomenTot=[0.0]*numDimensions

    for j in range(numObjects):
        #take cross product of postion and velocity
        angMomen[j]=np.cross(posArray[j],velArray[j])
        #multiply each component by mass
        for i in range(numDimensions):
            angMomen[j][i]=massArray[j]*angMomen[j][i]
        #print angMomen[j]
    #sum up all angular momenta
    for j in range(numObjects):
        for i in range(numDimensions):
            angMomenTot[i]=angMomen[j][i]+angMomenTot[i]
    return angMomenTot

def total_energy(posArray, velArray, massArray):
    '''calculates the total energy (gravitational and kinetic) for a system of objects
    '''
    numObjects=len(massArray)
    numDimensions=len(posArray[0])
    G=math.pow(2.0*math.pi,2.0)
    gravPotential=0.0
    kineticEnergy=0.0
    speedArray=np.linalg.norm(velArray,axis=1)

    for j in range(numObjects):
        for k in range(j+1,numObjects):
            r = pythag_pos_diff(posArray[j],posArray[k])
            gravPotential+=(-G*massArray[j]*massArray[k])/r
    for j in range(numObjects):
        kineticEnergy+=(1.0/2.0)*massArray[j]*(speedArray[j]**2.0)

    energyTot=gravPotential+kineticEnergy
    return energyTot

def pythag_pos_diff(objPosArray1,objPosArray2):

    #calculates the pythagorian position difference between two n-dimensional position arrays
    pythagSum=0
    for l in range(len(objPosArray1)):
        pythagSum+=math.pow(objPosArray1[l]-objPosArray2[l],2)
    return math.pow(pythagSum,1.0/2.0)

def nbody_sim_leapfrog(posArray,velArray, massArray,stepSize,stepCount):
    '''Returns an array of computed position steps for the inputed objects
    '''
    numObjects=len(massArray)
    numDimensions=len(posArray[0])
    angStart=angular_momentum(posArray, velArray, massArray)
    energyStart=total_energy(posArray, velArray, massArray)
    #0. calculate first velocity step with forward euler
    #1. update position
    #2. compute acceleration
    #3. update velocity
    G=math.pow(2.0*math.pi,2.0)
    historicPosArray=[[[]for j in range(numDimensions)] for j in range(numObjects)]
    historicVelArray=[[[]for j in range(numDimensions)] for j in range(numObjects)]
    historicAngMomenDrift=[]
    historicEngDrift=[]

    #start with forward Euler
    accelArray=[[0]*numDimensions for j in range(numObjects)]
    for j in range(numObjects):
        for l in range(numDimensions):
            for k in range(j+1,numObjects):
                #print posArray[j],posArray[k]
                newAccel=G*massArray[k]*(posArray[k][l]-posArray[j][l]) / math.pow(pythag_pos_diff(posArray[j],posArray[k]),3.0)
                oppositeNewAccel=-newAccel*massArray[j]/massArray[k]
                accelArray[j][l]+=newAccel
                accelArray[k][l]+=oppositeNewAccel
    #print accelArray[1]
    #update velocity
    for j in range(numObjects):
        for l in range(numDimensions):
            velArray[j][l]+=(1.0/2.0)*stepSize*accelArray[j][l]
            historicVelArray[j][l].append(velArray[j][l])

    #simulation loop using leap frog
    for i in range(stepCount):
        #update positon
        for j in range(numObjects):
            for l in range(numDimensions):
                historicPosArray[j][l].append(posArray[j][l])
                posArray[j][l]+=stepSize*velArray[j][l]
        #compute acceleration
        accelArray=[[0]*numDimensions for j in range(numObjects)]
        for j in range(numObjects):
            for l in range(numDimensions):
                for k in range(j+1,numObjects):
                    #print posArray[j],posArray[k]
                    newAccel=G*massArray[k]*(posArray[k][l]-posArray[j][l]) / math.pow(pythag_pos_diff(posArray[j],posArray[k]),3.0)
                    oppositeNewAccel=-newAccel*massArray[j]/massArray[k]
                    accelArray[j][l]+=newAccel
                    accelArray[k][l]+=oppositeNewAccel
        #print accelArray[1]
        #update velocity
        for j in range(numObjects):
            for l in range(numDimensions):
                velArray[j][l]+=stepSize*accelArray[j][l]
                historicVelArray[j][l].append(velArray[j][l])

        #compute drift
        ang=angular_momentum(posArray, velArray, massArray)
        energy=total_energy(posArray, velArray, massArray)
        angMomenDrift=np.subtract(ang,angStart)
        historicAngMomenDrift.append((angMomenDrift[0]**2 + angMomenDrift[1]**2 + angMomenDrift[2]**2)**(1.0/2.0))
        historicEngDrift.append((energy-energyStart))

    print 'Execution Duration: '+str(datetime.now() - startTime)
    speedup=((3.154E+07)*stepSize*stepCount)/((datetime.now() - startTime).total_seconds())
    print 'Speedup: '+str(speedup)
    print 'total drift'
    print np.subtract(ang,angStart), (energy-energyStart)
    return historicPosArray, historicAngMomenDrift, historicEngDrift

def plot_planets(historicPosArray,numObjects, numDimensions):

    #for j in range(numDimensions-1):
    for i in range(numObjects):
        plt.plot(historicPosArray[i][0],historicPosArray[i][1],'.-')
    plt.gca().set_aspect('equal', adjustable='box')

numObjects=10
numDimensions=3
#Quantiy      Units
#time      |   Year
#distance  |   AU
#velocity  |   AU/Year
#mass      |   Solar Mass

#Number       Object
#0         |   Sun
#1         |   Mercury
#2         |   Venus
#3         |   Earth + Moon
#4         |   Mars
#5         |   Jupiter
#6         |   Saturn
#7         |   Uranus
#8         |   Neptune

#generate velocity and position vectors
#the arrays are defined as ____Array[jth object][dimensions x,y,z: 0,1,2]
#velocity array is half a timestep ahead of the position and acceleration arrays
posArray=[[0.0]*numDimensions for j in range(numObjects)]
velArray=[[0.0]*numDimensions for j in range(numObjects)]

#generate massArray[jth object]
massArray=[0.5, 1.661E-07, 2.4484E-06, 3.00346E-06, 3.2279E-07, 0.0009543, 0.0002857, 4.36576E-05, 5.15E-05, 0.5]

#positions and velocitys of solar system bodies in baryocentric coords
#all starting values at A.D. 2016-Feb-29 00:00:00.0000

# x_sun = -M_jup*x_jup
# T = ((x_jup - x_sun) / (M_jup+M_sun))^(1/2)
# v_jup=2*pi*x_jup / T
# v_sun = -M_jup*v_jup
jupOrbit=0.05
sunOrbit=-jupOrbit
period= ((jupOrbit-sunOrbit)**3)*31.59
jupSpeed=2.0*math.pi*jupOrbit / (period)
sunSpeed=-jupSpeed
print period
print jupSpeed
#SUN 1
posArray[0]=[3.769360061446780E-03+sunOrbit,  1.812773374712511E-03, -1.624014729877611E-04]
velArray[0]=[2.116907618849916E-07*365,  sunSpeed + 6.980633760687217E-06*365, -1.203681274902981E-08*365]

#SUN 2
posArray[-1]=[3.769360061446780E-03+jupOrbit,  1.812773374712511E-03, -1.624014729877611E-04]
velArray[-1]=[2.116907618849916E-07*365,  jupSpeed + 6.980633760687217E-06*365, -1.203681274902981E-08*365]

#MERCURY
posArray[1]=[6.384933295823990E-02, -4.515747976027313E-01, -4.272075514855849E-02]
velArray[1]=[2.224830448309151E-02*365,  5.147225564402190E-03*365, -1.621110695570192E-03*365]

#VENUS
posArray[2]=[1.167830458441351E-01, -7.165914711359284E-01, -1.653385613745591E-02]
velArray[2]=[1.984550800517364E-02*365,  3.079327873621045E-03*365, -1.103086914976536E-03*365]

#EARTH
posArray[3]=[-9.247648258319329E-01,  3.469717738912940E-01, -1.730700077332298E-04]
velArray[3]=[-6.279271560859209E-03*365, -1.617967468331688E-02*365,  2.694319332791631E-07*365]

#MARS
posArray[4]=[-1.510114581167425E+00, -5.693088681586485E-01,  2.502401748406479E-02]
velArray[4]=[5.461536539056124E-03*365, -1.189036297311848E-02*365, -3.833644710566484E-04*365]

#JUPITER
posArray[5]=[-5.291527066741596E+00,  1.182309344528905E+00,  1.134242057706169E-01]
velArray[5]=[-1.733768918092524E-03*365, -7.008096404329502E-03*365,  6.792695786000429E-05*365]

#SATURN
posArray[6]=[-3.421132666518896E+00, -9.407865807544542E+00,  2.997343479755205E-01]
velArray[6]=[4.937187945581996E-03*365, -1.922883900630874E-03*365, -1.633640215930448E-04*365]

#URANUS
posArray[7]=[1.879301283705500E+01,  6.762792046709805E+00, -2.183506612863105E-01]
velArray[7]=[-1.360455775595364E-03*365,  3.517436778076412E-03*365,  3.073947954095653E-05*365]

#NEPTUNE
posArray[8]=[ 2.801943446698039E+01, -1.060190333542911E+01, -4.274100493913166E-01]
velArray[8]=[1.090032425561670E-03*365,  2.955023769315658E-03*365, -8.560834225044430E-05*365]

print '\n\n--------------------------------------------------\nLEAPFROG REAL DATA\n--------------------------------------------------'
#print 'Start Postions',posArray
#print 'Start Velocities',velArray
#print 'Masses',massArray

stepSize=10*period/(365.25)#*100000000*0.928514127163)
stepCount=100000
print 'Step Size: ' +str(stepSize*365.25) +' days'
print 'Simulation Duration: '+str(stepSize*stepCount) +' years OR '+str(stepSize*stepCount*365.25)+' days'

historicPosArray,historicAngMomenDrift, historicEngDrift = nbody_sim_leapfrog(posArray,velArray,massArray,stepSize,stepCount)

plot_planets(historicPosArray, numObjects, numDimensions)
plt.title('Binary Sun with Real Solar System Data From 2/29/2016\n'+str(stepCount)+' steps of '+str(stepSize*365.25)+' days',fontsize=15)
plt.xlabel('AU',fontsize=20)
plt.ylabel('AU',fontsize=20)
plt.gca().set_aspect('equal', adjustable='box')
plt.show()

#plt.plot(historicPosArray[0][0],historicPosArray[0][1],'.--')
#plt.plot(historicPosArray[-1][0],historicPosArray[-1][1],'.--')
##plt.plot(historicPosArray[1][0],historicPosArray[1][1],'.--')
##plt.plot(historicPosArray[2][0],historicPosArray[2][1],'.--')
#plt.title('Binary Sun with Mercury and Venus Zoom\n'+str(stepCount)+' steps of '+str(stepSize*365.25)+' days',fontsize=15)
#plt.xlabel('AU',fontsize=20)
#plt.ylabel('AU',fontsize=20)
#plt.gca().set_aspect('equal', adjustable='box')
#plt.show()
#
#plt.plot(historicAngMomenDrift)
#plt.title('Leapfrog Method\nSun and Jupiter\n'+str(stepCount)+' steps of '+str(stepSize*365.25)+' days',fontsize=15)
#plt.xlabel('Steps',fontsize=20)
#plt.ylabel('Total Angular Momentum Drift',fontsize=15)
#plt.show()
#
#plt.plot(historicEngDrift)
#plt.title('Leapfrog Method\nSun and Jupiter\n'+str(stepCount)+' steps of '+str(stepSize*365.25)+' days',fontsize=15)
#plt.xlabel('Steps',fontsize=20)
#plt.ylabel('Total Energy Drift',fontsize=15)
#plt.show()