import numpy as np
A = np.arange(2, 14).reshape((3, 4))
B = np.arange(14, 2, -1).reshape((3, 4))  # -1代表步长
print(A)
print(np.argmin(A))  # 计算最小值索引
print(np.argmax(A))  # 计算最大值索引
print(np.mean(A))  # 计算平均值
print(A.mean())
print(np.average(A))  # 计算平均值
print(np.median(A))  # 计算中位数
print(np.cumsum(A))  # 数组个数仍然一样，但是第二个数是前两数之和，第三数是前三数之和以此类推
#  [ 2  5  9 14 20 27 35 44 54 65 77 90]
print(np.diff(A))  # 后一个数减去前一个数的差
"""[[1 1 1]
 [1 1 1]
 [1 1 1]]"""
print(np.nonzero(A))
"""(array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=int64),
array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3], dtype=int64)) 前面一个数组代表横坐标后一个数组代表纵坐标"""
print(np.sort(A))
"""[[ 2  3  4  5]
 [ 6  7  8  9]
 [10 11 12 13]]"""
print(A.T.dot(A))  # 矩阵转置
print(np.clip(A, 5, 9))  # 将小于5的数划归到5，大于9的数划归到9，范围之中的数保持不变
print(np.mean(A, axis=0))  # axis表示维数，0表示列，1表示行进行计算mean代表平均数值


import numpy as np

def wmean(x, w=None):
    """ weighted mean """
    if w is None:
        return np.mean(x)
    else:
        return np.sum(x * w) / np.sum(w)

def wstd(x, w=None):
    """
    Weighted standard deviation

    Parameters
    ----------
    x : array
        data
    w : array
        weights

    Returns
    -------
    wstd

    Ref
    ---
    https://www.itl.nist.gov/div898/software/dataplot/refman2/ch2/weightsd.pdf

    """
    if w is None:
        w = np.ones_like(x)
    if len(x) < 1:
        return 0.
    else:
        _n_nonzero = np.sum(w>0)
        if _n_nonzero > 1:
            _wmean = wmean(x, w)
            return np.sqrt(np.sum(w*(x-_wmean)**2) / np.sum(w) /(_n_nonzero-1)*_n_nonzero)
        else:
            return 0.


from numpy import pi

def comp_inductance(self):
    """Compute the end winding inductance based on permeance coefficient stored in EndWindingRect object
    from "Design of Rotating Electrical Machines", J. Pyrhonen, second edition

    Parameters
    ----------
    self: EndWindingRect
        A EndWindingRect object

    Returns
    -------
    Lew : float
        end winding inductance [H].
    """

    if self.parent is not None:
        winding = self.parent
    else:
        return 0

    p = winding.p
    Ntcoil = winding.Ntcoil

    mu0 = 4 * pi * 1e-7

    if winding.Lewout is not None:
        l_ew = winding.Lewout
    else:
        l_ew = 0

    # Calculate width of end winding
    w_ew = 2 * self.comp_length()

    # Calculate Eq(4.100)
    l_lambda = 2 * l_ew * self.lambda_length + w_ew * self.lambda_width

    # Calculate end-winding inductance from Eq (4.98) p.261
    Lw = 2 / p * Ntcoil ** 2 * mu0 * l_lambda

    return Lw


#!/usr/bin/python
"""
This is a mixture of two gaussians example

We mode the proportion of one distribution with a random variable.

For the sample error it would be nice to use:

precision = Gamma('precision', alpha=0.1, beta=0.1

as it is the conjugate prior of the precision.

However, this produced odd results.  In order to get this example to
work properly I had to modify the Uniform range for the precision parameter
until the results were reasonable.  Oddly enough increasing this range
will decrease DIC (good), but the visible fit is worse.

"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import pymc as mc

## example from elem. stat. learning book
book_mu1 = 4.62
book_mu2 = 1.06
book_var1 = 0.87
book_var2 = 0.77
book_pi = 0.546

x1 = np.array([-0.39,0.12,0.94,1.67,1.76,2.44,3.72,4.28,4.92,5.53])
x2 = np.array([ 0.06,0.48,1.01,1.68,1.80,3.25,4.12,4.60,5.28,6.22])
x  = np.hstack([x1,x2])

## the mixing parameter
pi = mc.Uniform('pi',lower=0,upper=1)
delta = mc.Bernoulli( "delta",p=pi,size=x.size) # produces 1 with proportion p.

## other model parameters
sigma = mc.Uniform("sigma", 0.0,0.5)
var1 = mc.Uniform('var1',0,1)
var2 = mc.Uniform('var2',0,1)
mode1 = mc.Normal("mode1",mu=0,tau=var1**-2)
mode2 = mc.Normal("mode2",mu=0,tau=var2**-2)

@mc.deterministic(plot=False)
def mix_model(delta=delta,mode1=mode1,mode2=mode2):
    return (1.0-delta)*mode1 + delta*mode2

obs = mc.Normal("obs",mu=mix_model,tau=sigma**-2,value=x,observed=True)
model = mc.Model({"mode1":mode1,"mode2":mode2,"sigma":sigma,"pi":pi,
                  "var1":var1,"var2":var2,"obs":obs})

M = mc.MCMC(model)
M.sample(100000,50000)
mc.Matplot.plot(M)

## relabel the estimates if necessary
if M.trace('mode1')[:].mean() > M.trace('mode2')[:].mean():
    est_mu1  = M.trace('mode1')[:].mean()
    est_var1 = M.trace('var1')[:].mean()
    est_mu2  = M.trace('mode2')[:].mean()
    est_var2 = M.trace('var2')[:].mean()
    est_pi   = M.trace('pi')[:].mean()
else:
    est_mu1  = M.trace('mode2')[:].mean()
    est_var1 = M.trace('var2')[:].mean()
    est_mu2  = M.trace('mode1')[:].mean()
    est_var2 = M.trace('var1')[:].mean()
    est_pi   = 1.0 - M.trace('pi')[:].mean()

print 'book,model'
print 'mu1', book_mu1,est_mu1
print 'mu2', book_mu2,est_mu2
print 'var1', book_var1,est_var1
print 'var2', book_var2,est_var2
print 'pi', book_pi, est_pi
print 'DIC',M.dic

fig = plt.figure(1)
ax = fig.add_subplot(111)
n,bins,patches = ax.hist(x,15,normed=1,facecolor='gray',alpha=0.75)

p1 = mlab.normpdf(bins, book_mu1, np.sqrt(book_var1))
p2 = mlab.normpdf(bins, book_mu2, np.sqrt(book_var2))
l1 = plt.plot(bins, p1, 'r--', linewidth=1)
l2 = plt.plot(bins, p2, 'r--', linewidth=1)

## add a 'best fit' line (results from here)
p3 = mlab.normpdf(bins, est_mu1, np.sqrt(est_var1))
p4 = mlab.normpdf(bins, est_mu2, np.sqrt(est_var2))
l3 = plt.plot(bins, p3, 'k-', linewidth=1)
l4 = plt.plot(bins, p4, 'k-', linewidth=1)

plt.xlabel('x')
plt.ylabel('freq')
plt.ylim([0,0.8])

ax.legend((l1[0], l3[0]), ('Book Estimate', 'My Estimate') )
#plt.savefig('../TwoComponentGauss.png')
plt.show()


"""
This Python3 script creates a Pythagoras tree in SVG or PNG format

Requires numpy and matplotlib to create pngs.

Get started with : python Pythagoras_tree.py --help
"""
import argparse
from math import atan2, ceil, pi, sqrt
from typing import List

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm, patches

def pythagoras_tree(ratio: float = 1.0, nb_levels: int = 12):
    """
    Compute Pythagoras_tree
    The Pythagoras Tree is a plane fractal constructed from squares.
    It is named after Pythagoras because each triple of touching squares
    encloses a right triangle, in a configuration traditionally used to
    depict the Pythagorean theorem.
    http://en.wikipedia.org/wiki/Pythagoras_tree

          All these arguments are optional: the function can run with
          argument.
    Output :
          - Matrix M: Pythagoras tree is stored in a matrix M.
            This matrix has 5 columns.
            Each row corresponds to the coordinate of each square of the tree
            The two first columns give the bottom-left position of each
            square. The third column corresponds to the orientation angle of
            each square. The fourth column gives the size of each square. The
            fifth column specifies the level of recursion of each square.
            The first row corresponds to the root of the tree. It is always
            the same
            M[0,:] = [0 -1 0 1 1];
            The leaf located at row i will give 2 leaves located at 2*i and
            2*i+1.
    """
    # pylint: disable=too-many-locals
    # Check inputs
    if ratio <= 0:
        raise Exception("Length of ratio has to be greater than zero")
    if int(nb_levels) != float(nb_levels):
        raise Exception("The number of level has to be integer")

    # Compute constants
    c_d = sqrt(1.0 + ratio ** 2)
    # Normalized length 1
    c_1 = 1.0 / c_d
    # Normalized length 2
    c_2 = ratio / c_d
    # Translation pattern
    tr_pat = np.array(
        [[0.0, 1.0 / (1.0 + ratio ** 2)], [1.0, 1.0 + ratio / (1.0 + ratio ** 2)]]
    )
    # Defines the first rotation angle
    alpha1 = atan2(ratio, 1.0)
    # Defines the second rotation angle
    alpha2 = alpha1 - pi / 2.0
    # Number of elements (square)
    nb_elements = 2 ** (nb_levels + 1) - 1
    # Matrice containing the tree
    pt_arrray = np.zeros((nb_elements, 5))
    # Initialization of the tree
    pt_arrray[0, :] = [0.0, -1.0, 0.0, 1.0, 1.0]

    # Compute the level of each square contained in the resulting matrix
    offset = 0
    for i in range(nb_levels + 1):
        tmp = 2 ** i
        pt_arrray[offset : offset + tmp, 4] = i
        offset += tmp

    def mat_rot(angle_rad: float) -> np.ndarray:
        c_a = np.cos(angle_rad)
        s_a = np.sin(angle_rad)
        return np.array([[c_a, -s_a], [s_a, c_a]])

    # Compute the position and size of each square wrt its parent
    for i in range(1, nb_elements, 2):
        j = (i + 1) // 2 - 1
        t_m = pt_arrray[j, 3] * mat_rot(pt_arrray[j, 2]) @ tr_pat
        t_x = t_m[0, :] + pt_arrray[j, 0]
        t_y = t_m[1, :] + pt_arrray[j, 1]
        theta1 = (pt_arrray[j, 2] + alpha1) % (2.0 * pi)
        theta2 = (pt_arrray[j, 2] + alpha2) % (2.0 * pi)
        pt_arrray[i, 0:4] = [t_x[0], t_y[0], theta1, pt_arrray[j, 3] * c_1]
        pt_arrray[i + 1, 0:4] = [t_x[1], t_y[1], theta2, pt_arrray[j, 3] * c_2]
    return pt_arrray

def _svg_write_metadata(write):
    write("\t<title>Pythagoras tree</title>\n")
    write("\t<metadata>\n")
    write("\t\t<rdf:RDF>\n")
    write("\t\t\t<cc:Work\n")
    write('\t\t\t\trdf:about="">\n')
    write("\t\t\t\t<dc:format>image/svg+xml</dc:format>\n")
    write("\t\t\t\t<dc:type\n")
    write("\t\t\t\t\trdf:resource=" + '"http://purl.org/dc/dcmitype/StillImage"/>\n')
    write("\t\t\t\t<dc:title>Pythagoras tree</dc:title>\n")
    write("\t\t\t\t<dc:creator>\n")
    write("\t\t\t\t\t<cc:Agent>\n")
    write("\t\t\t\t\t\t<dc:title>Guillaume Jacquenot</dc:title>\n")
    write("\t\t\t\t\t</cc:Agent>\n")
    write("\t\t\t\t</dc:creator>\n")
    write("\t\t\t\t<cc:license\n")
    write(
        "\t\t\t\t\t\trdf:resource="
        + '"http://creativecommons.org/licenses/by-nc-sa/3.0/"/>\n'
    )
    write("\t\t\t</cc:Work>\n")
    write("\t\t\t<cc:License\n")
    write(
        "\t\t\t\trdf:about=" + '"http://creativecommons.org/licenses/by-nc-sa/3.0/">\n'
    )
    write("\t\t\t\t<cc:permits\n")
    write(
        "\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#Reproduction"/>\n'
    )
    write("\t\t\t\t<cc:permits\n")
    write(
        "\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#Reproduction"/>\n'
    )
    write("\t\t\t\t<cc:permits\n")
    write(
        "\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#Distribution"/>\n'
    )
    write("\t\t\t\t<cc:requires\n")
    write("\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#Notice"/>\n')
    write("\t\t\t\t<cc:requires\n")
    write("\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#Attribution"/>\n')
    write("\t\t\t\t<cc:prohibits\n")
    write(
        "\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#CommercialUse"/>\n'
    )
    write("\t\t\t\t<cc:permits\n")
    write(
        "\t\t\t\t\trdf:resource="
        + '"http://creativecommons.org/ns#DerivativeWorks"/>\n'
    )
    write("\t\t\t\t<cc:requires\n")
    write("\t\t\t\t\trdf:resource=" + '"http://creativecommons.org/ns#ShareAlike"/>\n')
    write("\t\t\t</cc:License>\n")
    write("\t\t</rdf:RDF>\n")
    write("\t</metadata>\n")

def pythagor_tree_write2svg(
    pt_array: np.ndarray,
    ratio: float = 0.8,
    nb_levels: int = 12,
    colormap_name: str = "summer",
):
    """
    A svg file giving a vectorial display of the tree. The name of
    file is generated from the parameter m, n, colormap. The file is
    stored in the current folder.
    """
    # pylint: disable=too-many-locals
    display_metadata = True

    nb_elements = pt_array.shape[0]
    length_offset = pt_array[nb_elements - 1, 3] + 0.1
    min_x = np.min(pt_array[:, 0] - sqrt(2) * pt_array[:, 3]) - length_offset
    max_x = np.max(pt_array[:, 0] + sqrt(2) * pt_array[:, 3]) + length_offset
    min_y = np.min(pt_array[:, 1]) - length_offset
    max_y = np.max(pt_array[:, 1] + sqrt(2) * pt_array[:, 3]) + length_offset

    # Compute the color of tree
    colormap = cm.get_cmap(colormap_name)
    nb_pixels = 100
    fig_w = ceil(nb_pixels * (max_x - min_x))
    fig_h = ceil(nb_pixels * (max_y - min_y))
    filename = (
        "Pythagoras_tree_"
        + str(ratio).replace(".", "_")
        + "__"
        + str(nb_levels)
        + "__"
        + colormap_name
        + ".svg"
    )
    with open(filename, "wt", encoding="utf-8") as fid:
        write = fid.write
        write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n')
        if not display_metadata:
            write('<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"\n')
            write('  "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">\n')
        write('<svg width="{0}" height="{1}" version="1.1"\n'.format(fig_w, fig_h))
        # w(['<svg width="12cm" height="4cm" version="1.1"\n']); % Wfig,

        # w(['<svg width="15cm" height="10cm" '...
        #              'viewBox="0 0 %d %d" version="1.1"\n'],...
        #              Wfig,Hfig);
        if display_metadata:
            write('\txmlns:dc="http://purl.org/dc/elements/1.1/"\n')
            write('\txmlns:cc="http://creativecommons.org/ns#"\n')
            write('\txmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"\n')
        write('\txmlns:svg="http://www.w3.org/2000/svg"\n')
        write('\txmlns="http://www.w3.org/2000/svg"\n')
        write('\txmlns:xlink="http://www.w3.org/1999/xlink">\n')

        if display_metadata:
            _svg_write_metadata(write)
        write("\t<defs>\n")
        write('\t\t<rect width="{0}" height="{0}" \n'.format(nb_pixels))
        write('\t\t\tx="0" y="0"\n')
        write('\t\t\tstyle="fill-opacity:1;stroke:#00d900;stroke-opacity:1"\n')
        write('\t\t\tid="squa"\n')
        write("\t\t/>\n")
        write("\t</defs>\n")
        write(
            '\t<g transform="translate({0} {1}) rotate(180) " >\n'.format(
                round(nb_pixels * max_x), round(nb_pixels * max_y)
            )
        )
        for i in range(nb_levels + 1):
            write(
                '\t\t<g style="fill:{0};" >\n'.format(
                    matplotlib.colors.rgb2hex(
                        colormap(1.0 - float(i) / (nb_levels + 1))
                    )
                )
            )
            offset = 2 ** i - 1
            for j in range(2 ** i):
                k = j + offset
                write(
                    (
                        '\t\t\t<use xlink:href="#squa" '
                        + 'transform="translate({0:+010.5f} {1:+010.5f}) '
                        + 'rotate({2:3.1f}) scale({3:8.6f})" />\n'
                    ).format(
                        nb_pixels * pt_array[k, 0],
                        nb_pixels * pt_array[k, 1],
                        pt_array[k, 2] * 180 / pi,
                        pt_array[k, 3],
                    )
                )
            write("\t\t</g>\n")
        write("\t</g>\n")
        write("</svg>\n")

def pythagor_tree_plot(
    pt_array: np.ndarray, colormap_name: str = "summer", output_filename: str = "lm.png"
):
    """Plot a Pythagoras tree for a PNG format"""
    colormap = cm.get_cmap(colormap_name)
    fig, axis = plt.subplots()
    for i in range(pt_array.shape[0]):
        c_x = pt_array[i, 0]
        c_y = pt_array[i, 1]
        theta = pt_array[i, 2] * 180.0 / pi
        s_i = pt_array[i, 3]
        rect = patches.Rectangle(
            [c_x, c_y],
            s_i,
            s_i,
            angle=theta,
            ec="none",
            color=colormap(1.0 - i / (pt_array[-1, 4] + 1)),
        )
        axis.add_patch(rect)
    plt.xlim([-4, 4])
    plt.ylim([-1.5, 3.5])
    # plt.gca().relim()
    plt.gca().set_aspect("equal", adjustable="box")
    plt.axis("off")
    fig.savefig(output_filename, bbox_inches="tight", dpi=300)

def get_list_of_colormaps() -> List[str]:
    """Return the list of avalaible colors"""
    return [
        "autumn",
        "bone",
        "colorcube",
        "cool",
        "copper",
        "flag",
        "gray",
        "hot",
        "hsv",
        "jet",
        "lines",
        "pink",
        "prism",
        "spring",
        "summer",
        "white",
        "winter",
    ]

def is_colormap(cmap: str) -> bool:
    """This function returns True if 'cmap' is a valid colormap"""
    return cmap in get_list_of_colormaps()

def main():
    """Main entrypoint"""
    parser = argparse.ArgumentParser(
        description="This script creates a SVG image of a Pythagoras tree, a"
        + " plane fractal constructed from squares."
    )
    parser.add_argument(
        "-r",
        "--ratio",
        type=float,
        default=1.0,
        help="r ( r > 0 ) is the relative length of one of the side "
        + "right-angled triangle. "
        + "The second side of the right-angle is taken to be one. "
        + "To have a symmetric tree, r has to be 1.",
    )
    parser.add_argument(
        "-l",
        "--level",
        type=int,
        default=10,
        help="n is the level of recursion. The number of elements of "
        + "tree is equal to 2**(n+1)-1. A reasonable number for n "
        + "is 10.",
    )
    parser.add_argument(
        "-p",
        "--plot",
        action="store_true",
        help="Option used to display the tree as a PNG image with matplotlib",
    )
    parser.add_argument(
        "-c",
        "--colormap",
        type=str,
        default="summer",
        help="Matplotlib colormap used to generate color of the different "
        + "levels of the tree. "
        + "Possible values are {0}".format(" ".join(get_list_of_colormaps())),
    )
    args = parser.parse_args()

    # Create a matrix containing the informations representing the tree
    # Each row represents a single square
    pt_array = pythagoras_tree(ratio=args.ratio, nb_levels=args.level)

    if args.plot:
        # Display the tree
        filename = (
            "Pythagoras_tree_"
            + str(args.ratio).replace(".", "_")
            + "__"
            + str(args.level)
            + "__"
            + args.colormap
            + ".png"
        )
        pythagor_tree_plot(
            pt_array, colormap_name=args.colormap, output_filename=filename
        )

    # Write results to an SVG file
    pythagor_tree_write2svg(pt_array, args.ratio, args.level, args.colormap)

if __name__ == "__main__":
    main()


'''

Machine Learning - Normal Data Distribution
Normal Data Distribution

In the previous chapter we learned how to create a completely random array, of a given size, and between two given values.

In this chapter we will learn how to create an array where the values are concentrated around a given value.

In probability theory this kind of data distribution is known as the normal data distribution, or the Gaussian data distribution, after the mathematician Carl Friedrich Gauss who came up with the formula of this data distributio
'''

import matplotlib.pyplot as plt
import numpy as np

speed=np.random.normal(5.0,1.0,100000)
plt.hist(speed,100)
print(speed)
plt.show()

'''

Note: A normal distribution graph is also known as the bell curve because of it's characteristic shape of a bell.
Histogram Explained

We use the array from the numpy.random.normal() method, with 100000 values,  to draw a histogram with 100 bars.

We specify that the mean value is 5.0, and the standard deviation is 1.0.

Meaning that the values should be concentrated around 5.0, and rarely further away than 1.0 from the mean.

And as you can see from the histogram, most values are between 4.0 and 6.0, with a top at approximately 5.0.
'''

'''

Machine Learning - Scatter Plot
Scatter Plot

A scatter plot is a diagram where each value in the data set is represented by a dot.

The Matplotlib module has a method for drawing scatter plots, it needs two arrays of the same length, one for the values of the x-axis, and one for the values of the y-axis:

x = [5,7,8,7,2,17,2,9,4,11,12,9,6]

y = [99,86,87,88,111,86,103,87,94,78,77,85,86]

The x array represents the age of each car.

The y array represents the speed of each car.
'''


"""
Approximation of functions by linear combination of basis functions in
function spaces and the least squares method (or the Galerkin method).
2D version.
"""
import sympy as sym
import numpy as np

def least_squares(f, psi, Omega, symbolic=True, print_latex=False):
    """
    Given a function f(x,y) on a rectangular domain
    Omega=[[xmin,xmax],[ymin,ymax]],
    return the best approximation to f(x,y) in the space V
    spanned by the functions in the list psi.
    """
    N = len(psi) - 1
    A = sym.zeros(N+1, N+1)
    b = sym.zeros(N+1, 1)
    x, y = sym.symbols('x y')
    print('...evaluating matrix...')
    for i in range(N+1):
        for j in range(i, N+1):
            print(('(%d,%d)' % (i, j)))

            integrand = psi[i]*psi[j]
            if symbolic:
                I = sym.integrate(integrand,
                                 (x, Omega[0][0], Omega[0][1]),
                                 (y, Omega[1][0], Omega[1][1]))
            if not symbolic or isinstance(I, sym.Integral):
                # Could not integrate symbolically, use numerical int.
                print(('numerical integration of', integrand))
                integrand = sym.lambdify([x,y], integrand)
                I = sym.mpmath.quad(integrand,
                                   [Omega[0][0], Omega[0][1]],
                                   [Omega[1][0], Omega[1][1]])
            A[i,j] = A[j,i] = I
        integrand = psi[i]*f
        if symbolic:
            I = sym.integrate(integrand,
                             (x, Omega[0][0], Omega[0][1]),
                             (y, Omega[1][0], Omega[1][1]))
        if not symbolic or isinstance(I, sym.Integral):
            # Could not integrate symbolically, use numerical int.
            print(('numerical integration of', integrand))
            integrand = sym.lambdify([x,y], integrand)
            I = sym.mpmath.quad(integrand,
                               [Omega[0][0], Omega[0][1]],
                               [Omega[1][0], Omega[1][1]])
        b[i,0] = I
    print()
    print(('A:\n', A, '\nb:\n', b))
    if symbolic:
        c = A.LUsolve(b)  # symbolic solve
        # c is a sympy Matrix object, numbers are in c[i,0]
        c = [c[i,0] for i in range(c.shape[0])]
    else:
        c = sym.mpmath.lu_solve(A, b)  # numerical solve
    print(('coeff:', c))

    u = sum(c[i]*psi[i] for i in range(len(psi)))
    print(('approximation:', u))
    print(('f:', sym.expand(f)))
    if print_latex:
        print((sym.latex(A, mode='plain')))
        print((sym.latex(b, mode='plain')))
        print((sym.latex(c, mode='plain')))
    return u

def comparison_plot(f, u, Omega, plotfile='tmp', title=''):
    """Compare f(x,y) and u(x,y) for x,y in Omega in a plot."""
    x, y = sym.symbols('x y')

    f = sym.lambdify([x,y], f, modules="numpy")
    u = sym.lambdify([x,y], u, modules="numpy")
    # When doing symbolics, Omega can easily contain symbolic expressions,
    # assume .evalf() will work in that case to obtain numerical
    # expressions, which then must be converted to float before calling
    # linspace below
    for r in range(2):
        for s in range(2):
            if not isinstance(Omega[r][s], (int,float)):
                Omega[r][s] = float(Omega[r][s].evalf())

    resolution = 41  # no of points in plot
    xcoor = np.linspace(Omega[0][0], Omega[0][1], resolution)
    ycoor = np.linspace(Omega[1][0], Omega[1][1], resolution)
    xv, yv = np.meshgrid(xcoor, ycoor)
    # Vectorized functions expressions does not work with
    # lambdify'ed functions without the modules="numpy"
    exact  = f(xv, yv)
    approx = u(xv, yv)
    error = exact - approx

    import matplotlib.pyplot as plt
    plt.figure()
    contours = plt.contour(xv, yv, error, 8)     # 8 contour lines
    plt.clabel(contours, inline=1, fontsize=10)  # labels
    if title: plt.title(title)
    if plotfile:
        plt.savefig('%s_error_c.pdf' % plotfile)
        plt.savefig('%s_error_c.png' % plotfile)
    fig = plt.figure()
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    #ax = fig.add_subplot(111) #, projection='3d')
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(xv, yv, error, rstride=1, cstride=1,
                           cmap=cm.coolwarm, linewidth=0,
                           antialiased=False)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    if title: plt.title(title)
    if plotfile:
        plt.savefig('%s_error_s.pdf' % plotfile)
        plt.savefig('%s_error_s.png' % plotfile)
    plt.show()

if __name__ == '__main__':
    print('Module file not meant for execution.')


from app.utils.methods import BaseMethod
from sympy import Symbol
from sympy.parsing.sympy_parser import parse_expr

class Bisection(BaseMethod):

    def __init__(self, x0, x1, tolerance, iterations, function, **kwargs):
        self.xl = x0
        self.xh = x1
        self.tol = tolerance
        self.iter = iterations
        self.function = function
        self.array = []

    def run(self):
        f = parse_expr(self.function)
        x = Symbol("x")
        fxl = f.subs(x, self.xl)
        fxh = f.subs(x, self.xh)

        if fxl == 0:
            return self.success_response(f"{self.xl} is a root")
        elif fxh == 0:
            return self.success_response(f"{self.xh} is a root")
        elif fxl * fxh > 0:
            return self.failed_response("The interval does not have a root")
        else:
            error = self.tol + 1
            xm = (self.xl + self.xh) / 2
            self.array.append(["i", "x0", "x1", "xm", "f(xm)", "error"])
            for i in range(self.iter):
                fxm = f.subs(x, xm)

                if fxm == 0 or error <= self.tol:
                    break

                if fxl * fxm < 0:
                    self.xh = xm
                else:
                    self.xl = xm

                x_aux = xm
                xm = (self.xl + self.xh) / 2
                error = abs(xm - x_aux)
                self.array.append([str(i), str(self.xl), str(xm), str(self.xh), str(fxm), str(error)])

            if fxm == 0:
                return {
                    "method_status": "success",
                    "result": f"{xm} is a root",
                    "iterations": self.array
                }
            elif error < self.tol:
                return {
                    "method_status": "success",
                    "result": f"{xm} approaches a root of the function with a tolerance of {self.tol}",
                    "iterations": self.array
                }
            else:
                return {
                    "method_status": "failed",
                    "result": "Max interactions exceeded",
                    "iterations": self.array
                }


# -*- coding: utf-8 -*-
"""
Created on Tue May 14 07:56:24 2019

@author: 08475268951
"""

import numpy as np

def f(x):
    return x**2 - np.pi

def bisseccao(f,a,b,tolerancia):
    tolerancia = 1e-10
    if np.abs(f(a)) < tolerancia:
        return a
    elif np.abs(f(b)) < tolerancia:
        return b
    elif f(a)*f(b) > 0:
        return np.Inf
    else:
        x = (a+b)/2
        print("x = {}".format(x))
        print("f(x) = {}".format(f(x)))
        nbiter = 0
        while np.abs(f(x)) >  tolerancia:
            if f(x)*f(a) > 0:
                a = x
            else:
                b = x
            x = (a+b)/2
            print("x = {}".format(x))
            print("f(x) = {}".format(f(x)))
            nbiter = nbiter + 1
        return x, nbiter


# Modern Computer Algebra

from Structures.QuotientFiniteField import QuotientFiniteField
from sage.matrix.all import Matrix

def berlekamp1_3(f, R):
    # paso 1
    q = R.get_domain().get_order()
    x = R.get_variable()
    n = R.degree(f)
    assert (q % 2 != 0)
    QF = QuotientFiniteField(R, f)
    xq = QF.repeated_squaring(x, q)

    # paso 2
    Q = []
    acum = R.one()
    for i in range(0, n):
        al = acum.list()
        zero = [R.get_domain().zero()] * (n - len(al))
        al.extend(zero)
        #al.reverse()
        Q.append(al)
        acum = QF.mul(xq, acum)

    # paso 3
    Q = Matrix(Q)
    I = Matrix.identity(n)
    QI = Q - I
    kqi = QI.kernel().basis_matrix().rows()
    r = len(kqi)
    bl = []
    for ro in kqi:
        # cuidao
        #ro.reverse()
        bl.append(R.get_true_value()(ro.list()))
    return r, bl

def berlekamp4_7(f, R, r, bl):
    q = R.get_domain().get_order()
    QF = QuotientFiniteField(R, f)
    if r == 1:
        return f

    # paso 4
    cl = []
    for i in range(0, r):
        cl.append(R.get_domain().get_random())
    al = [cl[k] * bl[k] for k in range(0, r)]
    a = R.zero()
    for a_ in al:
        a = R.add(a, a_)

    a = R.get_true_value()(a)

    # paso 5
    g1 = R.gcd(a, f)
    if g1 != R.one() and g1 != f:
        return g1

    # paso 6
    exp = (q - 1)//2
    b = QF.repeated_squaring(a, exp)

    # paso 7
    g2 = R.gcd(R.sub(b, R.one()), f)
    if g2 != R.one() and g2 != f:
        return g2
    else:
        return None

# f in R, R FiniteFieldPolynomial, k in Z
def berlekamp_k_times(f, R, k):
    # r = dimension of the basis of the berlekamp algebra
    # bl = basis of it
    r, bl = berlekamp1_3(f, R)
    if r == 1:
        return f
    for i in range(0, k):
        p = berlekamp4_7(f, R, r, bl)
        if p is not None:
            return p
    return None

# f in R, R FiniteFieldPolynomial, k in Z
def berlekamp_full(f, R, k):
    g = berlekamp_k_times(f, R, k)
    if g is None:
        return [f]
    elif g == f:
        return [f]
    gl = berlekamp_full(g, R, k)
    fl = berlekamp_full(R.quo(f, g), R, k)
    gl.extend(fl)
    return gl


from __future__ import division
import numpy as np

def cheb(N):
	x =  np.cos((np.pi/N)*np.linspace(0,N,N+1))
	x.shape = (N+1,1)
	lin = np.linspace(0,N,N+1)
	lin.shape = (N+1,1)

	c = np.ones((N+1,1))
	c[0], c[-1] = 2., 2.
	c = c*(-1.)**lin
	X = x*np.ones(N+1) # broadcast along 2nd dimension (columns)

	dX = X - X.T
	# print " x = \n", x, "\n"
	# print " c = \n", c, "\n"
	# print " X = \n", X, "\n"
	# print " dX = \n", dX, "\n"

	D = (c*(1./c).T)/(dX + np.eye(N+1))
	D  = D - np.diag(np.sum(D.T,axis=0))
	x.shape = (N+1,)
	return D, x

def cheb_loop(N):
	def p(j1):
		if (j1==0 or j1 == N): return 2.
		else: return 1.

	x = np.cos(np.pi*np.arange(N+1)/N)
	D = np.zeros((N+1,N+1))
	# j represents column index
	for j in range(0,N+1):
		for i in range(0,j)+range(j+1,N+1):
			D[i,j] = ((-1.)**(i+j))*p(i)/( p(j)*(x[i]- x[j]) )

	# Values on the main diagonal
	for j in xrange(1,N):
		D[j,j] = -x[j]/(2.*(1-x[j]**2.))

	D[0,0] = (1.+2.*N**2.)/6.
	D[N,N] = -(1.+2.*N**2.)/6.

	return D,x

def cheb_vectorized(N):
	x =  np.cos((np.pi/N)*np.linspace(0,N,N+1))
	x.shape = (N+1,1)
	lin = np.linspace(0,N,N+1)
	lin.shape = (N+1,1)

	c = np.ones((N+1,1))
	c[0], c[-1] = 2., 2.
	c = c*(-1.)**lin
	X = x*np.ones(N+1) # broadcast along 2nd dimension (columns)

	dX = X - X.T

	D = (c*(1./c).T)/(dX + np.eye(N+1))
	D  = D - np.diag(np.sum(D.T,axis=0))
	x.shape = (N+1,)
	# Here we return the differentation matrix and the Chebychev points,
	# numbered from x_0 = 1 to x_N = -1
	return D, x


import spinsim
import math
import numpy as np
import matplotlib.pyplot as plt
import datetime as dtm

time_now_string = dtm.datetime.now().strftime("%Y%m%dT%H%M%S")

def gaussian_pulse(time, modifier, pulse):
    pulse[0] = (math.pi/math.sqrt(math.tau))*math.exp(-0.5*(time**2))/math.tau
    pulse[1] = 0.0
    pulse[2] = 0.0

def cumulative_gaussian(t):
    return 0.5*(1 + math.erf(t/math.sqrt(2.0)))

plt.figure()
pulse_sample = np.empty(3, np.float64)

time_continuous = np.arange(-5.0, 5.0005, 1e-3)
pulse_continuous = []
for time_sample in time_continuous:
    gaussian_pulse(time_sample, 0, pulse_sample)
    pulse_continuous += [pulse_sample[0]]
pulse_continuous = np.asarray(pulse_continuous)
plt.plot(time_continuous, pulse_continuous, "k-")

time_step = 0.25
time_midpoint = 0.5*time_step + np.arange(-5.0, 5.0, time_step)
pulse_midpoint = []
for time_sample in time_midpoint:
    gaussian_pulse(time_sample, 0, pulse_sample)
    pulse_midpoint += [pulse_sample[0]]
pulse_midpoint = np.asarray(pulse_midpoint)
plt.plot(time_midpoint, pulse_midpoint, "bo")

time_quadrature = []
pulse_quadrature = []
for time_sample in time_midpoint:
    gaussian_pulse(time_sample - 0.5*time_step/math.sqrt(3), 0, pulse_sample)
    time_quadrature += [time_sample - 0.5*time_step/math.sqrt(3)]
    pulse_quadrature += [pulse_sample[0]]

    gaussian_pulse(time_sample + 0.5*time_step/math.sqrt(3), 0, pulse_sample)
    time_quadrature += [time_sample + 0.5*time_step/math.sqrt(3)]
    pulse_quadrature += [pulse_sample[0]]
time_quadrature = np.asarray(time_quadrature)
pulse_quadrature = np.asarray(pulse_quadrature)
plt.plot(time_quadrature, pulse_quadrature, "m.")

plt.xlabel("Time (standard deviations)")
plt.ylabel("Pulse strength (Hz)")
plt.legend(
    [
        "Pulse shape",
        "Integration steps",
        "Pulse sample points"
    ]
)
plt.title("{}\nSample points for integrating Gaussian pulse".format(time_now_string))
plt.savefig("gaussian_pulse_sample.png")
plt.savefig("gaussian_pulse_sample.pdf")
plt.show()

plt.figure()

time = np.arange(-5.0, 5.1, 2.0)
state_analytic = np.asarray([[math.cos(0.5*math.pi*cumulative_gaussian(t)), -1j*math.sin(0.5*math.pi*cumulative_gaussian(t))] for t in time], dtype = np.complex128)

simulator = spinsim.Simulator(gaussian_pulse, spinsim.SpinQuantumNumber.HALF)

result_compare = simulator.evaluate(0.0, -5.0, 7.0, 1e-6, 2.0, np.asarray([1, 0], np.complex128))

legend = []

colours = ["r", "g", "b", "c", "m", "y"]
error = []
time_steps = np.asarray([0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0])
number_of_steps = 10 / time_steps
plot_start_index = 4
for simulation_index, time_step in enumerate(time_steps):
    result_simulated = simulator.evaluate(0.0, -5.0, 7.0, time_step, 2.0, np.asarray([1, 0], np.complex128))

    error += [np.sum(np.abs(result_simulated.state - result_compare.state))/5]

    if simulation_index >= plot_start_index:
        plt.plot(time, result_simulated.spin[:, 0], colours[simulation_index - plot_start_index] + "--o")
        plt.plot(time, result_simulated.spin[:, 1], colours[simulation_index - plot_start_index] + "--x")
        plt.plot(time, result_simulated.spin[:, 2], colours[simulation_index - plot_start_index] + "--+")

        legend += [
            "{:d} x".format(int(number_of_steps[simulation_index])),
            "{:d} y".format(int(number_of_steps[simulation_index])),
            "{:d} z".format(int(number_of_steps[simulation_index]))
        ]

plt.legend(legend, loc = "lower left")
plt.xlabel("Time (standard deviations)")
plt.ylabel("Spin")
plt.title("{}\nGaussian pulse at various numbers of steps".format(time_now_string))
plt.savefig("gaussian_pulse.png")
plt.savefig("gaussian_pulse.pdf")
plt.show()

plt.figure()
plt.loglog(number_of_steps, error, "-rx")
plt.xlabel("Number of steps")
plt.ylabel("Error")
plt.title("{}\nError in integrating Gaussian pulse".format(time_now_string))
plt.ylim((1e-12, 1e0))
plt.savefig("gaussian_pulse_error.png")
plt.savefig("gaussian_pulse_error.pdf")
plt.show()


#Relu activation function
import numpy as np

input = [0.3,0.4,0.6,0.3,.49, -.57]
output =[]

for i in input:
    if i > 0:
        output.append(i)
    else:
        output.append(0)

print(output)

'''Second Option'''

outputs = []
for i in input:
    outputs.append(max(0,i))
print("2nd Option of Calculation : ",outputs)

'''Third option'''

output_values=np.maximum(0,input)
print("with Numpy values", output_values)

# !---'''Activation SoftMax'''---!
outputs_value =[4.8, 1.21, 2.385]
E = 2.71
output_ce =[]
for i in outputs_value:
    output_ce.append(E**i)
print(output_ce)
exponeential_sum = sum(output_ce)

output_val =[]

for j in output_ce:
    output_val.append(j/ exponeential_sum)
print("Cross Entropy Values : ", output_val)

'''Lets try about Numpy values'''
outputs =([[1.8, 1.21, 3.385], [2, 1.21, 2.385],[3, 1.21, 2.385]])
exp_val = np.exp(outputs- np.max(outputs, axis=1, keepdims=True))
print("Hey, Um exp_val", exp_val)
final_output = exp_val/np.sum(exp_val, axis=1, keepdims=True)
print("Final output of SoftMax will be : ", final_output)


import numpy as np

def mean_square_error(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Calculate MSE loss

    Parameters
    ----------
    y_true: ndarray of shape (n_samples, )
        True response values
    y_pred: ndarray of shape (n_samples, )
        Predicted response values

    Returns
    -------
    MSE of given predictions
    """
    #return float(np.mean((y_true-y_pred)**2))
    return float(sum(pow(y_true-y_pred,2))/len(y_true))

def misclassification_error(y_true: np.ndarray, y_pred: np.ndarray, normalize: bool = True) -> float:
    """
    Calculate misclassification loss

    Parameters
    ----------
    y_true: ndarray of shape (n_samples, )
        True response values
    y_pred: ndarray of shape (n_samples, )
        Predicted response values
    normalize: bool, default = True
        Normalize by number of samples or not

    Returns
    -------
    Misclassification of given predictions
    """
    if y_true.size == 0 or y_pred.size == 0:
        return 0

    equal = np.zeros(y_true.shape[0])
    equal[y_true != y_pred] = 1
    counter = equal.sum(axis=0)
    if normalize:
        return counter/y_true.shape[0]
    return counter
    #y_true = np.argsort(y_pred)
    # equal = np.zeros(y_true.shape[0])
    # equal[y_pred == y_true] = 1
    # sumi = equal.sum(axis = 0)
    # if normalize:
    #     return sumi/(y_true.shape[0])
    # else:
    #     return sumi

def accuracy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Calculate accuracy of given predictions

    Parameters
    ----------
    y_true: ndarray of shape (n_samples, )
        True response values
    y_pred: ndarray of shape (n_samples, )
        Predicted response values

    Returns
    -------
    Accuracy of given predictions
    """
    equal = np.zeros(y_true.shape[0])
    equal[y_pred == y_true] = 1
    return np.sum(equal)/(y_true.shape[0])

def cross_entropy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Calculate the cross entropy of given predictions

    Parameters
    ----------
    y_true: ndarray of shape (n_samples, )
        True response values
    y_pred: ndarray of shape (n_samples, )
        Predicted response values

    Returns
    -------
    Cross entropy of given predictions
    """
    raise NotImplementedError()


import numpy as np
class Solution:
    def knightDialer(self, N: int) -> int:
        mod = 10**9 + 7
        if N == 1: return 10
        M = np.matrix([[0, 0, 0, 0, 1, 0, 1, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 1, 0, 1, 0],
                       [0, 0, 0, 0, 0, 0, 0, 1, 0, 1],
                       [0, 0, 0, 0, 1, 0, 0, 0, 1, 0],
                       [1, 0, 0, 1, 0, 0, 0, 0, 0, 1],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [1, 1, 0, 0, 0, 0, 0, 1, 0, 0],
                       [0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
                       [0, 1, 0, 1, 0, 0, 0, 0, 0, 0],
                       [0, 0, 1, 0, 1, 0, 0, 0, 0, 0]])
        res, N = 1, N - 1
        while N:
            if N % 2: res = res * M % mod
            M = M * M % mod
            N //= 2
        return int(np.sum(res)) % mod

import numpy as np
class Solution:
    def knightDialer(self, N: int) -> int:
        x1 = x2 = x3 = x4 = x5 = x6 = x7 = x8 = x9 = x0 = 1
        for i in range(N - 1):
            x1, x2, x3, x4, x5, x6, x7, x8, x9, x0 = \
                x6 + x8, x7 + x9, x4 + x8, \
                x3 + x9 + x0, 0, x1 + x7 + x0, \
                x2 + x6, x1 + x3, x2 + x4, \
                x4 + x6
        return (x1 + x2 + x3 + x4 + x5 + x6 + x7 + x8 + x9 + x0) % (10**9 + 7)

Dynamic Programming/935. Knight Dialer.py


import numpy as np

def f(x):
    return x**6-x**4-x**3-1

x1=float(input("x1: "))
x2=float(input("x2: "))
d=float(input("Accuracy: "))
xm=float()
while abs(x2-x1)>=d and abs(f(xm))>=d:
    xm = x2 - (x2-x1) / (f(x2)-f(x1)) * f(x2)
    x1 = x2
    x2 = xm
    print("Root=","%.5f" % xm)
if f(xm)>d:
    print("Not convergent")
else:
    print("Final Solution=","%.5f" % xm)

Root Finding/Secant Method [Single Root].py


#Basic Gauss Elimination Methods

"""
Name - Suryabrata Das

Sem: V

College_Roll_NO: 703

Paper-code: CMSA DSE-IB

Registration No: A01-1112-117-003-2018

Examination roll_no: 2021151264

Subject: Numerical Methods (DSE-I)

"""
import numpy as np
def swap_row(mat,r1,r2):
    temp = mat[r1]
    mat[r1] = mat[r2]
    mat[r2] = temp

def normalize(mat,i):
    for j in range(i+1,len(mat)):
        fact = (mat[j][i]/mat[i][i])
        for k in range(0,len(mat)+1):
            mat[j][k] -= fact*mat[i][k]

def fwd_elimination(mat):
    n = len(mat)
    print("\n\n")
    for i in range(n-1):
        if mat[i][i] == 0:
            swap_row(mat,i,i+1)
            normalize(mat,i)
            m = np.array(mat)
            m = m[:,:n]
            if (np.linalg.det(m) == 0):
                print ("Singular Matrix!! No Unique solution Exists!!")
                return False

            return True

def print_mat(mat):
    for i in range(len(mat)):
        print(mat[i])
        print("\n")
def bwd_elimination(mat):
    n = len(mat)

    x = [0]*n
    for i in range(n-1,-1,-1):
        x[i] = mat[i][n]
        for j in range(i+1,n):
            x[i] -= mat[i][j]*x[j]
            x[i] = x[i]/mat[i][i]
            print("Solutions are: \n\n")
            print("x: ",x[0],"\n\n")
    print("y: ",x[1],"\n\n")
    print("z: ",x[2],"\n\n")

def main():
    n = 3
    mat = []
    mat = [[3,6,1,16],[2,4,3,13],[1,3,2,9]]
    print("Your input matrix: ")
    print_mat(mat)
    if(fwd_elimination(mat) == True):
        print("After Forward Elimination: ")
        print_mat(mat)
        bwd_elimination(mat)
if __name__ == '__main__':
    main()


import random
import math
import numpy as np

# Prim's Minimum Spanning Tree Algorithm
def primsAlgo(e,c,n):
	sel_edge = ()
	mincost = 0
	near = [0 for x in range(n)]
	t = [[0 for x in range(2)] for y in range(n-1)]
	print('Initially Minimum Spanning Tree: \n' ,np.matrix(t))
	values = np.array(c)
	index = np.argmin(values)
	minedge = [index//n, index%n]

	sel_edge = minedge
	mincost += c[sel_edge[0]][sel_edge[1]]

    #Select neighboring edge with minimum cost
	t[0][0],t[0][1] = sel_edge[0],sel_edge[1]
	for l in range(0,n):
		if(c[l][sel_edge[0]] < c[l][sel_edge[1]]):
			near[l] = sel_edge[0]
		else:
			near[l] = sel_edge[1]
	near[sel_edge[0]],near[sel_edge[1]] = -1,-1
	minm = math.inf
	for i in range(1,n-1):
		for j in range(n):
			if(near[j] != -1):
				if(c[j][near[j]] < minm):
					minm = c[j][near[j]]
					sel_edge[0] = j
					sel_edge[1] = near[j]
		minm = math.inf
		t[i][0],t[i][1] = sel_edge[0],sel_edge[1]
		mincost += c[sel_edge[0]][sel_edge[1]]
		near[sel_edge[0]] = -1
		for k in range(n):
			if(near[k] != -1 and (c[k][near[k]] > c[k][sel_edge[0]])):
				near[k] = sel_edge[0]
	print('Resultant MST: \n', np.matrix(t))
	return mincost

n = random.randint(5,7)
print('Number of vertices:', n)
max_int = math.inf
e = [[0 for x in range(n)] for y in range(n)]
c = [[max_int for x in range(n)] for y in range(n)]
for i in range(n):
	for j in range(n):
		if(i < j):
			e[i][j]=1
			c[i][j]=float(random.randint(10,50))
		else:
			e[i][j]=e[j][i]
			c[i][j]=c[j][i]
print('Edge Adjacency Matrix: \n', np.matrix(e))
print('Cost Adjacency Matrix: \n',np.matrix(c))
print ('Cost of MST: ', primsAlgo(e,c,n))


""" When you Fourier Transform an Image most of the fourier coefficients are quite small and
    can be truncated or thresholded away so we can zero out all of the small fourier coefficients
    retaining only the largest 1% or 2% of those FFT values and then

    When we inverse Fourier Transform that thresholded FFT signal we recover the original image
    with relatively low degradation

    So compression comes about by only saving these largest %1 or %2 fourier coefficients """

from matplotlib.image import imread
import numpy as np
import matplotlib.pyplot as plt
import os

from mpl_toolkits.mplot3d import axes3d

path = r'C:\Users\oguzhan.sari\Desktop\Oğuzhan\development\fourier_SB\image compression FFT'
plt.rcParams['figure.figsize'] = [5, 5]
plt.rcParams.update({'font.size': 18})

A = imread(os.path.join(path, 'headshot.jpg'))
B = np.mean(A, -1)         # converts RGB to grayscale

plt.figure()
plt.imshow(A)
plt.axis('off')
plt.show()

plt.imshow(B)               # grayscale --> average of the color channels RGB
# plt.imshow(256 - A)       #,cmap='gray_r')         # steve-black background
plt.axis('off')
plt.show()
""" We just want to work on an array for simplicity, not to work on three color channel """

Bt = np.fft.fft2(B)
Btsorted = np.sort(np.abs(Bt.reshape(-1)))          # sort by magnitude

# Zero out all small coefficients and inverse transform
for keep in (0.1, 0.05, 0.01, 0.002):
    thresh = Btsorted[int(np.floor((1-keep)*len(Btsorted)))]
    ind = np.abs(Bt) > thresh                                   # find small indices
    Btlow = Bt * ind                                         # Threshold small indices  # mask [1, .. 0] ones and zeros
    Alow = np.fft.ifft2(Btlow).real                             # compressed image
    plt.figure()
    plt.imshow(Alow, cmap='gray')
    # plt.imshow(256- Alow, cmap='gray')
    plt.axis('off')
    plt.title('Compressed image: keep = ' + str(keep*100) + '%')
    plt.show()

""" There is a lot of high frequency content in hair, fur and grass, thats when we need to keep to reconstruct
    that image and so high frequency stuff is harder to represent

Pixar - Toy Story;
    Characters had plastic skin without hair and fur no grass
    because it is much much easier to work in a compressed representation and to render those types of objects
    than to render hair, fur and grass

    And it was big step from Toy Story to Monsters Inc where they finally figured out how to render big hairy monster
    because that's a lot harder, a lot higher frequency content, a lot more detail you have to render that you can't
    just compress out """

""" The Larger the image the more compressible it is
    So if we have an ultra ultra high-res image, there's more redundant information that we can compress out
    We'll get a higher compression ratio the bigger the image

    If we have smaller image low-res, we need more of those fourier coefficients to reconstruct the image """

# Plot pixel intensity of the image as a surface

plt.rcParams['figure.figsize'] = [6, 6]

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y = np.meshgrid(np.arange(1, np.shape(B)[1] + 1), np.arange(1, np.shape(B)[0] + 1))
ax.plot_surface(X[0::10, 0::10], Y[0::10, 0::10], 256-B[0::10, 0::10], cmap='viridis', edgecolor='none')
ax.set_title = 'Surface Plot'
ax.mouse_init()
#ax.view_init(200, 270)
ax.view_init(-270, -270)              #different perspective
plt.show()

""" this is the exact headshot.jpg but rendered as a surface plot
    So brighter pixels in my face are both brighter and taller - kind of an elevation map

    My face kind of mountain landscape of pixel intensity
    the whites of my eye were the brightest pixels, those are the tallest peaks

    What we are trying to do with fourier approximation is trying the approximate this mountain land scape
    of the pixel intensity as a sum of sin() and cosine() waves in X and in Y at different frequences
So we are going to add up all of those kind of wavy sheets in just the right proportion and
we're going to reconstruct this topography of pixel intensities to reconstruct my face"""


import sys
import math
import pandas as pd
import numpy as np
def calculate_initial_compass_bearing(pointA, pointB):
    if (type(pointA) != tuple) or (type(pointB) != tuple):
        raise TypeError("Only tuples are supported as arguments")
    lat1 = math.radians(pointA[0])
    lat2 = math.radians(pointB[0])
    diffLong = math.radians(pointB[1] - pointA[1])
    x = math.sin(diffLong) * math.cos(lat2)
    y = math.cos(lat1) * math.sin(lat2) - (math.sin(lat1) * math.cos(lat2) * math.cos(diffLong))
    initial_bearing = math.atan2(x, y)
    # Now we have the initial bearing but math.atan2 return values
    # from -180° to + 180° which is not what we want for a compass bearing
    # The solution is to normalize the initial bearing as shown below
    initial_bearing = math.degrees(initial_bearing)
    compass_bearing = (initial_bearing + 360) % 360
    return compass_bearing

def convert_to_distance_bearing(pointAlat, pointAlon, pointBlat, pointBlon):
	pointA=(pointAlat,pointAlon)
	pointB=(pointBlat,pointBlon)
	A=calculate_initial_compass_bearing(pointA, pointB)
	from math import sin, cos, sqrt, atan2, radians
	# approximate radius of earth in km
	R = 6373.0
	lat1 = radians(pointAlat)
	lon1 = radians(pointAlon)
	lat2 = radians(pointBlat)
	lon2 = radians(pointBlon)
	dlon = lon2 - lon1
	dlat = lat2 - lat1
	a = sin(dlat / 2)**2 + cos(lat1) * cos(lat2) * sin(dlon / 2)**2
	c = 2 * atan2(sqrt(a), sqrt(1 - a))
	distance = R * c
	return distance, A

def transform_dataset(datasetfile):
	dataset_coords = pd.read_csv(datasetfile, engine='python').values.astype('float32')
	dataset=[0,0]
	for i in range(len(dataset_coords)-1):
		j=dataset_coords[i]
		k=dataset_coords[i+1]
		dis, bear = convert_to_distance_bearing(j[0], j[1], k[0], k[1])
		results=[dis, bear]
		dataset=np.vstack((dataset,results))
	dataset=np.delete(dataset, 0, 0)
	return 	dataset


"""Functions related to modular arithmetic"""

from math import prod
from typing import List

from sympy.ntheory.factor_ import totient

# https://en.wikipedia.org/wiki/Modular_multiplicative_inverse#Using_Euler's_theorem
def modular_inverse(n:int,base:int) -> int:
    """Returns the modular inverse of n in base"""
    n **= totient(base) - 1
    n %= base
    return n

# https://crypto.stanford.edu/pbc/notes/numbertheory/crt.html
def chinese_remainder_theorem(a:List[int],m:List[int]) -> int:
    """Performs the chinese remainder theorem on recieved lists"""

    if len(a) != len(m):
        raise ValueError('Length of a must match length of m')

    M = prod(m)
    solution = 0
    for i,m_i in enumerate(a):
        b = M // m_i
        b_inv = modular_inverse(b,m_i)
        solution += a[i] * b * b_inv
    return solution % M


"""
Created on Sun Feb 20 21:31:53 2022

@author: udaytalwar
"""

#Please refer to "BlackScholesEqns.pdf" in this repository for more details about the equations!

import scipy.stats as stats
import math as m

def d(S, K, sigma, r, t):

    '''
    S = Current Price
    K = Strike Price
    sigma = Volatility
    r = annualized risk-free rate
    t = time to expiration

    returns d1, d2 for option price calculation using Black Scholes
    '''
    d1 = (m.log(S/K) + (r + (sigma**2/2))*t) * (1/(sigma*m.sqrt(t)))

    d2 = d1 - sigma*m.sqrt(t)

    return d1, d2

def option_price(S, K, sigma, r, t, flag, d1 = 0, d2 = 0):

    '''
    S = Current Price
    K = Strike Price
    sigma = Volatility
    r = annualized risk-free rate
    t = time to expiration
    flag = 'Call' or 'Put'

    returns option price according to Black Scholes
    '''

    if d1 == 0 and d2 == 0:

        d1, d2 = d(S, K, sigma, r, t)

        if flag == 'Call':

            price = stats.norm.cdf(d1)*S - stats.norm.cdf(d2)*K*m.exp(-r*t)

        elif flag == 'Put':

            price = stats.norm.cdf(-d2)*K*m.exp(-r*t) - stats.norm.cdf(-d1)*S

        return price

    else:

        if flag == 'Call':

            price = stats.norm.cdf(d1)*S - stats.norm.cdf(d2)*K*m.exp(-r*t)

        elif flag == 'Put':

            price = stats.norm.cdf(-d2)*K*m.exp(-r*t) - stats.norm.cdf(-d1)*S

        return price

def imp_vol(S, K, r, t, flag, option_CMP):

    '''
    S = Current Price
    K = Strike Price
    r = annualized risk-free rate
    t = time to expiration
    flag = 'Call' or 'Put'
    option_CMP = current market price of option

    returns implied volatility of option according to Black Scholes
    '''

    lb = -2  #lower bound of midpoint method

    ub = 5  #upper bound of midpoint method

    error = 1e-15  #error tolerance

        #midpoint method

    while (ub-lb) > error:

        if option_price(S, K, (ub+lb)/2, r, t, flag) - option_CMP > 0:
            ub = (lb+ub)/2

        else:
            lb = (lb+ub)/2

    return lb


import numpy as np
import matplotlib.pyplot as plt

class Perceptron:
    """
    Initialize Perceptron with inputs and output pairs. Set Weights to random values
    """
    def __init__(self,x,y):
        self.ntrain = x.shape[0]
        self.dim = x.shape[1]
        self.x = x
        self.y = y
        #self.w = np.random.random((self.dim,1))
        self.w = np.random.random((self.dim))
        #self.w = np.array([1.0,1.0,1.0])

    """
        Function that plots the boundary between datapoints in x. If provided a set of Weights,
        will plot all of them on the same graph with different points provided.
    """
    def plotBoundary(self, w = [], type = 'r--', labels = []):
        a = self.x[:,0]
        b = self.x[:,1]
        vals = {1:'+',-1:'<'}
        for i in xrange(len(self.x)):
            #print str(self.x[i][0])+ ' ' + str(self.x[i][1])
            plt.scatter(self.x[i][0],self.x[i][1],marker=vals[self.y[i]])
        x = np.linspace(0, 10, 10)
        if w == []:
            y = -1*(self.w[0]/self.w[1])*x - (self.w[2]/self.w[1])
            plt.plot(x,y ,type)
        else:
            type = ["r--","b--","g--","b-."]
            for i in xrange(0,len(w)):
                y = -1*(w[i][0]/w[i][1])*x - (w[i][2]/w[i][1])
                plt.plot(x,y ,type[i],label=labels[i])
            plt.legend(bbox_to_anchor=(0.5, 1), loc=2, borderaxespad=0.)
        #print y

        plt.show()

    """
        This method inverts the input points for points belonging to the second class.
    """
    def processInput(self):
        self.x[self.y < 0] = -1*self.x[self.y < 0]

    """
    Train a single perceptron with single sample training with the margin specified.
    """
    def train(self,a = 1, b = 0):
        self.processInput()
        e = 1
        k = 0
        while e > 0:
            e = 0
            k += 1
            for i in xrange(len(self.x)):
                val = np.sum(self.w*self.x[i])
                if val < b:
                    self.w[0] += a*self.x[i,0]
                    self.w[1] += a*self.x[i,1]
                    self.w[2] += a*self.x[i,2]
                    e += 1
        #print k
        self.processInput()
        #self.plotBoundary()
        #print self.w
        return self.w

    """
    Train a single perceptron with relaxation
    """
    def trainRelaxation(self,a = 1, b = 0 ):
        self.processInput()
        e = 1
        k = 0
        while e > 0:
            k += 1
            e = 0
            for i in xrange(len(self.x)):
                val = np.sum(self.w*self.x[i])
                #self.plotBoundary()
                #print val
                if val < b:
                    sumVal = np.sum(self.x[i]**2)
                    update = (b - val)/sumVal
                    #print update
                    self.w[0] += a*self.x[i,0]
                    self.w[1] += a*self.x[i,1]
                    self.w[2] += a*self.x[i,2]
                    e += 1
        #print k
        self.processInput()
        #self.plotBoundary()
        #print self.w
        return self.w

    """
    LMS procedure implemented for a given learning rate and margin
    """
    def trainLMS(self,a = 1,b = 0):
        self.processInput()
        e = 1
        k = 0
        eta = a
        while k < 1000:
            k += 1
            for i in xrange(len(self.x)):
                val = np.sum(self.w*self.x[i])
                if ((b - val) < 0.5):
                    update = 0
                else:
                    update = 1
                self.w += eta*update*self.x[i]
        self.processInput()
        print self.w
        self.plotBoundary()
        #print self.w
        return self.w

if __name__ == '__main__':
    x = np.array([[2, 7, 1], [8, 1, 1], [7, 5, 1], [6, 3, 1],[7, 8, 1],[5, 9, 1],[4, 5,1],[4, 2, 1],[-1, -1,1],[1, 3,1], [3, -2,1], [5, 4,1], [2, 4,1],[7, 1, 1]])
    y = np.array([1, 1, 1, 1, 1, 1, 1,-1, -1, -1, -1, -1, -1,-1])
    p = Perceptron(x,y)
    #p.plotBoundary()
    #w = [[ 0.5 ,  0.55, -4.5 ],[  3.2 ,   3.35, -27.9 ],[  2.685  ,   2.68375, -23.15   ],[  2. ,   1.9, -16.7]]
    #labels = ["Single-sample","Single-sample(margin)", "Relaxation(margin)","LMS"]
    #w.append(p.train(0.1,0))
    #w.append(p.train(0.1,1))
    #w.append(p.trainRelaxation(0.005,1))
    #w.append(p.trainLMS(0.1,1))
    p.trainLMS(0.1,1)
    #p.plotBoundary(w,labels=labels)


# Example from https://roshansanthosh.wordpress.com/2015/02/22/linear-regression-in-theano/
# Objective: Regression in a parametric model. We will be using it to determine the
# coefficients of a line that best fits a given set of points.
# Since we are using a parametric model we make assumptions about the functional form of the
# relation ship between the independent and the dependent variables and in this case we
# assume there is a linear relationship (linear regression does not have to model linear
# relationships) and try to estimate the intersect and slope.
# We use a singel feature for X

import numpy as np
import theano
import theano.tensor as T
import matplotlib.pyplot as plt

rng = np.random
#Training Data
X = np.asarray([3,4,5,6.1,6.3,2.88,8.89,5.62,6.9,1.97,8.22,9.81,4.83,7.27,5.14,3.08])
Y = np.asarray([0.9,1.6,1.9,2.9,1.54,1.43,3.06,2.36,2.3,1.11,2.57,3.15,1.5,2.64,2.20,1.21])

# Coefficients
# Slope = m
# Intercept = c
# Randomly initialize m and c as shared variables in Theano. Shared mean that c and m can be
# used by multiple functions (they have global scope?). Shared variables are the fastes type
# in theano because the GPU can use different optimizations and since c and m are what will
# be updated alot it is a smart idea to have them as shared.

# In Theano thear are symbolic variables and expressions.
m_value = rng.rand()
c_value = rng.rand()

m = theano.shared(m_value, name = 'm')
c = theano.shared(c_value, name = 'c')

x = T.vector('X')
y = T.vector('y')

num_samples = X.shape[0]

# Defining the cost function
# Every machine learning model has a cost parameter which the ml algorithm tries to minimize.
# The cost parameter is essentially the cumulative error predicted by the model for the
# current set of coefficients.

# TODO also use RSS is there a difference?

# Use the randomly assigned values of c and m to predict y
prediction = T.dot(x,m) + c # our line function
# How big is the error in all the samples in x?
cost = T.sum(T.pow(prediction-y,2))/(2*num_samples) # MSE cost

# Optimization
# We use gradient decent
gradm = T.grad(cost,m)
gradc = T.grad(cost,c)

learning_rate = 0.01
training_steps = 10

train = theano.function([x,y],cost,updates = [(m,m-learning_rate*gradm), (c,c-learning_rate*gradc)])
test = theano.function([x],prediction)

estimates = np.zeros(training_steps)
x_vals = np.zeros(training_steps)
for i in range(training_steps):
    estimates[i] = train(X,Y)
    x_vals[i] = i

print("Slope :")
print(m.get_value()) # since m and c are shared we have to use get_value
print("Intercept :")
print(c.get_value())

axes = plt.gca()
axes.set_ylim([np.amin(estimates), np.amax(estimates)])
plt.plot(x_vals, estimates)
plt.show()

# Note: X and Y are the input data vecors while x and y are the Theano variables.
# All the expressions and functions are written in terms of the symbolc variables x and y while
# X and Y is used when calling the functions.
# /Users/henke/anaconda2/envs/theanohacking/bin/frameworkpython -m IPython


import numpy as np
import matplotlib.pyplot as plt
from pandas.io.parsers import read_csv

def calcula_normalizado(valor, med, desv):
    return (valor - med)/desv

def normaliza(X):
    medias = np.array([])
    desvs = np.array([])
    for i in range(1, np.shape(X)[1]):
        medias = np.append(medias, np.mean(X[:,i]))
        desvs = np.append(desvs, np.std(X[:,i]))
        X[:,i] = calcula_normalizado(X[:,i], medias[i-1], desvs[i-1])
    return (medias, desvs)

def carga_csv(file_name):
    valores = read_csv(file_name, header = None).values
    return valores.astype(float)

def hipotesis(X, Theta):
    return np.dot(X,Theta)

def coste(X, Y, Theta):
    H = np.dot(X, Theta)
    aux = (H - Y) ** 2
    return aux.sum() / (2*len(X))

def descenso_gradiente(X, Y, alpha, Theta, n):

    ths = np.array([[0.]])
    temps = np.array([[0.]])
    for i in range(n-1):
        ths = np.vstack((ths, [0.]))
        temps = np.vstack((temps, [0.]))
    costes = np.array([coste(X, Y, Theta)])
    j = 0
    while True:
        j = j + 1
        for i in range(n):
            temps[i] = Theta[i] - alpha*(1/len(X))*np.sum((hipotesis(X,Theta) - Y) * X[:,i:i+1])

        Theta = temps
        np.hstack((ths, temps))
        costes = np.append(costes, coste(X,Y, Theta))

        #print(Theta)
        #print(costes[-1])
        if (costes[-2] - costes[-1]) < 0.001:
            break
    #print(j)
    return (Theta, costes, j)

datos = carga_csv("ex1data2.csv")

X = datos[:, :-1]
Y = datos[:, -1:]

m = np.shape(X)[0]
X = np.hstack([np.ones([m,1]), X])
n = np.shape(X)[1]

Xp = np.copy(X)

medias, desvs = normaliza(X)

Theta = np.array([[0.]])

for i in range(n-1):
    Theta = np.vstack((Theta, [0.]))

#alpha = 0.01

#Theta, ths, costes = descenso_gradiente(X, Y, alpha, Theta, n)

j = 0
alphas = np.array([0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1])

colors = np.array(["blue","green","red","cyan","magenta","yellow","black"])
labels = np.array(["0.001", "0.003", "0.01", "0.03", "0.1", "0.3", "1"])
plt.figure()

for i in range(np.shape(alphas)[0]):
    #j = j+1
    a, costes, iter = descenso_gradiente(X, Y, alphas[i], Theta, n)
    print(iter)
    #print(j, "  ", costes[-1], "  ", i)
    plt.plot(np.arange(iter + 1), costes, c=colors[i], label=labels[i], linestyle='-')

plt.legend()

plt.xlabel("Num iter")
plt.ylabel("Costes")

plt.savefig("alphas.png")

#Xq = np.array([2500, 3])
Xq = np.array([1650,3])

Xq_norm = calcula_normalizado(np.transpose(Xq), medias, desvs)
Xq_norm = np.hstack([np.ones([1]), Xq_norm])

Xq = np.hstack([np.ones([1]), Xq])

Theta, costes, iter = descenso_gradiente(X, Y, 1.2, Theta, n)

Yp = hipotesis(Xq_norm, Theta)

X = Xp
aux = np.dot(np.transpose(X), X)
aux = np.linalg.pinv(aux)
aux = np.dot(aux, X.transpose())
eqTheta = np.dot(aux, Y)

Y = hipotesis(Xq, eqTheta)

print(Yp, Y)

print("FIN")


from manimlib.imports import *
from scipy.optimize import curve_fit

class GraphDensity(GraphScene):
    CONFIG = {
        "x_axis_label" : "$t$",
        "y_axis_label" : "$\\rho(t)$",
        "x_label_direction": DOWN,
        "y_label_direction": LEFT,
        "y_min": -1,
        "y_max": 10,
        "y_axis_height": 6,
        "y_tick_frequency": 1,
    }
    def construct(self):
        self.setup_axes(True)

        the_equation = TexMobject("\\rho(", "t", ")", "=", "k", "\\left(", "t_0", "\\over", "t",  "\\right)", "^2")
        the_equation[1].set_color(RED)
        the_equation[8].set_color(RED)
        the_equation.shift(UP*2 + RIGHT*3)
        self.play(
            Write(the_equation)
        )

        self.wait()

        def get_density_graph(t0=10.0, K=1.0):
            def density(t):
                return K * (t0/t)**(2.0)
            return density

        rho = self.get_graph(get_density_graph(t0=2.0, K=1.0), x_min=0.01, step_size=0.0001, color="#f8961e")
        rho.shift(UP * 0.3) #fudge to make it look nicer
        self.play(
            ShowCreation(rho.reverse_points()), run_time=3.0
        )
        self.wait()

class GraphScaleFactor(GraphScene):
    CONFIG = {
        "x_axis_label" : "$t$",
        "y_axis_label" : "$a(t)$",
        "x_label_direction": DOWN,
        "y_label_direction": LEFT,
    }
    def construct(self):
        self.setup_axes(True)

        the_equation = TexMobject("a(", "t", ")", "=","\\left(", "t", "\\over", "t_0",  "\\right)", "^{\\frac{2}{3}}")
        the_equation[1].set_color(RED)
        the_equation[5].set_color(RED)
        the_equation.shift(UP*2 + RIGHT*3)
        self.play(
            Write(the_equation)
        )

        self.wait()

        def get_scalefactor_graph(t0=10.0, K=1.0):
            def scalefactor(t):
                return K * (t/t0)**(2/3.0)
            return scalefactor

        rho = self.get_graph(get_scalefactor_graph(t0=0.6), x_min=0.01, step_size=0.0001, color="#9a031e")
        self.play(
            ShowCreation(rho.reverse_points()), run_time=3.0
        )
        self.wait()
        origin = self.coords_to_point(0, 0)
        self.play(
            FocusOn(origin, color="#e36414")
        )
        arrow = Arrow(origin + UP*2 + RIGHT, origin, color="#fb8b24")
        self.play(
            GrowArrow(arrow)
        )
        a_at_O = TexMobject("a", "(", "t", "=", "0", ")", "=", "0").scale(0.7).move_to(origin + UP*2 + RIGHT)
        a_at_O[2].set_color(RED)
        self.play(
            Write(a_at_O)
        )
        self.wait()

class Findingrho(Scene):

    def construct(self):

        temp = TexMobject("\\rho", "=")
        rho_of_a = VGroup(temp, TexMobject("k", "\\over", "a", "^3" ).next_to(temp, RIGHT).shift(UP*0.1)).shift(UP)
        # rho_of_a[1][2].set_color(ORANGE)
        a_of_t = TexMobject("a", "=", "\\left(", "t", "\\over", "t_0", "\\right)", "^{\\frac{2}{3}}" ).shift(DOWN)
        a_of_t[3].set_color(RED)

        self.play(
            Write(rho_of_a)
        )
        self.play(
            Write(a_of_t)
        )
        self.wait()
        self.play(
            ApplyMethod(rho_of_a.shift, LEFT*4),
            ApplyMethod(a_of_t.shift, LEFT*4)
        )
        brace = Brace(VGroup(rho_of_a, a_of_t), RIGHT)
        self.play(
            GrowFromCenter(brace)
        )
        implies = TexMobject("\\Rightarrow").next_to(brace, RIGHT)
        self.play(
            FadeInFrom(implies, LEFT)
        )
        rho_of_t_lhs = TexMobject("\\rho", "(", "t", ")", "=")
        rho_of_t_lhs[2].set_color(RED)
        rho_of_t_lhs.next_to(implies, RIGHT)
        self.play(
            TransformFromCopy(rho_of_a[0], rho_of_t_lhs[0]),
        )
        self.play(
            Write(rho_of_t_lhs[1:])
        )
        rho_t_rhs1 = TexMobject("k", "\\over", "\\left[","\\left(", "t", "\\over", "t_0", "\\right)", "^{\\frac{2}{3}}", "\\right]", "^3").next_to(rho_of_t_lhs)
        rho_t_rhs1.shift(DOWN*0.4)
        rho_t_rhs1[4].set_color(RED)
        self.play(
            TransformFromCopy(rho_of_a[1][0], rho_t_rhs1[0]),
        )
        self.play(
            Write(rho_t_rhs1[1]),
            TransformFromCopy(a_of_t, rho_t_rhs1[3:9])
        )
        self.play(
            Write(rho_t_rhs1[2]),
            Write(rho_t_rhs1[-2:])
        )
        rho_t = TexMobject("=", "k", "\\left(", "t_0", "\\over", "t", "\\right)", "^2").next_to(rho_t_rhs1).shift(UP*0.5)
        rho_t[5].set_color(RED)
        self.play(
            Write(rho_t[:2])
        )
        self.play(
            Write(rho_t[2:])
        )

        self.wait()

        self.wait()

class Friedmann(Scene):

    def construct(self):

        lhs = TexMobject("\\left(", "\\dot{a}", "\\over", "a", "\\right)", "^2", color="#fdfcdc")
        eq = TexMobject("=", color="#fdfcdc").next_to(lhs, RIGHT)
        rhs = TexMobject("\\left(", "8", "\\pi", "G", "\\over", "3", "\\right)", "\\rho", color="#fdfcdc").next_to(eq, RIGHT)
        friedmann = VGroup(lhs, eq, rhs).scale(1.5).move_to(ORIGIN)
        title = TextMobject("Friedmann Equation", color="#fb8b24").scale(1.5).to_edge(UP)
        self.play(
            FadeInFromDown(title)
        )
        self.wait()
        self.play(
            Write(lhs)
        )
        self.play(
            Write(eq)
        )
        self.play(
            Write(rhs[:-1])
        )
        self.play(
            Write(rhs[-1])
        )
        self.wait()

        #explain terms
        self.play(
            Indicate(lhs[3], color="#00afb9", scale_factor=2)
        )
        scale_factor = TextMobject("Scale", " Factor", color="#f4d58d").shift(DOWN*2).scale(1.5)
        self.play(
            ShowIncreasingSubsets(scale_factor)
        )
        self.play(
            FadeOut(scale_factor)
        )
        self.wait(0.5)
        self.play(
            Indicate(lhs[1], color="#00afb9", scale_factor=2)
        )
        dadt = TextMobject("Rate", " of", " Change", color="#f4d58d").shift(DOWN*2).scale(1.5)
        self.play(
            ShowIncreasingSubsets(dadt)
        )
        self.play(
            FadeOut(dadt)
        )
        self.wait(0.5)
        self.play(
            Indicate(rhs[3], color="#00afb9", scale_factor=2)
        )
        G = TextMobject("Newton's", " Constant", color="#f4d58d").shift(DOWN*2).scale(1.5)
        self.play(
            ShowIncreasingSubsets(G)
        )
        self.play(
            FadeOut(G)
        )
        self.wait(0.5)
        self.play(
            Indicate(rhs[-1], color="#00afb9", scale_factor=2)
        )
        density = TextMobject("Energy", " Density", color="#f4d58d").shift(DOWN*2).scale(1.5)
        self.play(
            ShowIncreasingSubsets(density)
        )
        self.play(
            FadeOut(density)
        )
        self.wait()

class ScaleFactorIntro(VectorScene, ZoomedScene):

    def construct(self):
        self.add_plane(animate=True, background_line_style= {
            "stroke_color": GREEN,
            "stroke_width": 1,
            "stroke_opacity": 0.4,
        },
        axis_config = {
            "stroke_color" : GREEN,
            "stroke_width" : 0.
        },
        run_time=2)
        self.wait()
        time = TexMobject("time =", "t_0").to_corner(UR)
        self.play(
            Write(time)
        )
        us = Dot(color="#0077b6", radius=0.2)

        self.play(
            FocusOn(us, color=BLUE)
        )
        self.play(
            GrowFromCenter(us)
        )
        self.wait()

        galaxy = Dot(color=GOLD, radius=0.2).shift(4*RIGHT)

        self.play(
            FocusOn(galaxy, color=GOLD)
        )
        self.play(
            GrowFromCenter(galaxy)
        )
        self.wait()
        lines = VGroup()
        for i in range(4):
            l = Line(UP/2.0, DOWN/2.0, color=GREEN).shift((i+1)*RIGHT)
            lines.add(l)

        self.play(
            AnimationGroup(
                *[ShowCreationThenFadeOut(l) for l in lines], lag_ratio=0.5, run_time=1.0
            )
        )
        self.wait()

        brace = Brace(VGroup(us, galaxy), DOWN)
        self.play(
            FadeInFromDown(brace)
        )
        ell = TexMobject("\\ell", "=", "4").next_to(brace, DOWN)
        self.play(
            Write(ell)
        )
        self.wait()
        self.play(
            FadeOut(brace),
            FadeOut(ell)
        )
        self.wait()
        light = Line(ORIGIN, RIGHT*4, color="#ffd60a", stroke_width=8.0)

        self.play(
            ShowCreationThenDestruction(light)
        )
        self.play(
            ShowCreationThenDestruction(light.reverse_points())
        )
        self.wait()

        d1 = TexMobject("d_{phys}(t_0)").move_to(ell).shift(UP * 2)
        self.play(
            Write(d1)
        )
        self.wait()
        self.play(
            FadeOut(d1)
        )
        def justfortime(m, alpha):
            m.restore()
            m.scale((1 - alpha) + alpha * 9/14.)
            m.shift((RIGHT*3.6 + UP*2.1 - m.get_center())*alpha)

        time.save_state()
        self.play(
            ApplyMethod(self.camera_frame.set_width, 9),
            ApplyMethod(us.scale, 9/14.0),
            ApplyMethod(galaxy.scale, 9/14.0),
            UpdateFromAlphaFunc(time, justfortime),
            run_time=2.0
        )
        self.play(
            Transform(time[1], TexMobject("t").scale(9/14.).move_to(time[1]))
        )
        self.wait()
        self.play(
            AnimationGroup(
                *[ShowCreationThenFadeOut(l) for l in lines], lag_ratio=0.5, run_time=1.0
            )
        )
        self.wait()
        ell.scale(9./14.)
        self.play(
            FadeIn(brace),
            Write(ell)
        )
        self.wait()
        self.play(
            FadeOut(brace),
            FadeOut(ell)
        )

        self.wait()
        self.play(
            ShowCreationThenDestruction(light.reverse_points()), run_time=1.5
        )
        self.play(
            ShowCreationThenDestruction(light.reverse_points()), run_time=1.5
        )
        self.wait()

        d2 = TexMobject("d_{phys}(t)", ">", "d_{phys}(t_0)").move_to(ell).shift(UP * 2).scale(9.0/14.0)
        self.play(
            Write(d2)
        )
        self.wait()
        self.play(
            FadeOut(d2)
        )
        self.wait()
        self.play(
            FadeOut(us),
            FadeOut(galaxy),
            FadeOut(time)
        )
        at = TexMobject("d_{phys}", "(", "t", ")", "=", "a", "(","t", ")", "\\ell").scale(9./14)
        self.wait()
        self.play(
            Write(at[0])
        )
        self.play(
            Write(at[-1])
        )
        self.wait()
        self.play(
            Write(at[1:4])
        )
        self.wait()
        self.play(
            Write(at[4:-1])
        )

        self.wait()
        self.play(
            Indicate(at[5:9], color="#38b000", scale_factor=2.0)
        )
        text = TextMobject("Scale", " Factor", color="#1982c4").shift(UP)
        self.play(
            ShowIncreasingSubsets(text)
        )
        self.play(
            FadeOut(text)
        )
        self.wait()
        scale_factor = TexMobject("d_{phys}(t = ","0", ")" "=", "a(t = ", "0", ")", "\\ell").scale(9.0/14)
        scale_factor[1].set_opacity(0.0)
        scale_factor[4].set_opacity(0.0)
        time1 = DecimalNumber(0, num_decimal_places=0).scale(9.0/14).move_to(scale_factor[1])
        time2 = DecimalNumber(0, num_decimal_places=0).scale(9.0/14).move_to(scale_factor[4])
        # self.play(
        #     Write(scale_factor),

        # )
        self.play(
            ReplacementTransform(at, scale_factor)
        )
        self.play(
            Write(time1),
            Write(time2)
        )
        self.wait()
        time1.add_updater(lambda m,dt : m.move_to(scale_factor[1]))
        time2.add_updater(lambda m,dt : m.move_to(scale_factor[4]))
        time1.add_updater(lambda m,dt : m.set_value(m.get_value() + dt*2))
        time2.add_updater(lambda m,dt : m.set_value(m.get_value() + dt*2))
        self.add(time1, time2)

        self.play(
            ApplyMethod(
                self.camera_frame.set_width, 2
            ),
            ApplyMethod(scale_factor.scale, 2.0/9),
            ApplyMethod(time1.scale, 2.0/9),
            ApplyMethod(time2.scale, 2.0/9),
            run_time=4.5,
            rate_func=linear
        )
        time1.set_value(9)
        time2.set_value(9)
        time1.clear_updaters()
        time2.clear_updaters()
        self.wait()

class HubblesLaw(GraphScene):

    CONFIG = {
        "x_min": -1,
        "x_max": 120,
        "x_axis_width": 10,
        "x_tick_frequency": 10,
        "x_leftmost_tick": 0,  # Change if different from x_min
        "x_labeled_nums": [i for i in range(0, 120, 20)],
        "x_axis_label": "$d (Mpc)$",
        "y_min": -1,
        "y_max": 10000,
        "y_axis_height": 6,
        "y_tick_frequency": 1000,
        "y_bottom_tick": None,  # Change if different from y_min
        "y_labeled_nums": [i for i in range(2000, 10000, 2000)],
        "y_axis_label": "v (km/s)",
        "axes_color": GREY,
        "graph_origin": 2.5 * DOWN + 4 * LEFT,
        "exclude_zero_label": False,

        #myedits
        "x_label_direction": DOWN,
        "y_label_direction": LEFT,
    }

    def construct(self):
        self.setup_axes(True)

        #data taken from :: https://iopscience.iop.org/article/10.1086/320638/pdf --- table 7
        D = [89.2, 66.7, 114.9, 62.2, 88.4, 45.1, 74.3, 43.2, 68.2, 85.6, 20.7, 39.5, 15.0, 58.3, 31.3, 38.7, 66.6, 57.3, 50.9, 53.3, 19.8]
        v = [6709, 4730, 8930, 4749, 7016, 3106, 4982, 3272, 4820, 7143, 1607, 3149, 1380, 4061, 2304, 3294, 4924, 4869, 4398, 3545, 1088]
        data_dots = VGroup()
        for i in range(len(D)):
            point = self.coords_to_point(D[i], v[i])
            d = Dot(color="#c5d86d").move_to(point)
            data_dots.add(d)

        ##last point is for Ursa Major
        self.wait()
        self.play(
            FocusOn(self.coords_to_point(0, 0), color=GREEN)
        )
        self.wait()
        self.play(
            GrowFromCenter(data_dots[-1])
        )
        self.wait()
        arrow = Arrow(data_dots[-1].get_center() + UR, data_dots[-1], color=BLUE)
        text = TextMobject("Ursa Major", color=DARK_BLUE).next_to(arrow, UP).shift(RIGHT)
        self.play(
            GrowArrow(arrow),
            Write(text)
        )
        self.wait()
        self.play(
            FadeOut(arrow),
            FadeOut(text)
        )
        self.wait()
        x_line = Line(self.coords_to_point(0, 0), self.coords_to_point(D[-1], 0), color=ORANGE, stroke_width=8.0).add_tip()
        y_line = Line(self.coords_to_point(D[-1], 0), self.coords_to_point(D[-1], v[-1]), color=ORANGE, stroke_width=8.0).add_tip()
        self.play(
            ShowCreationThenFadeOut(x_line)
        )
        self.wait()
        self.play(
            ShowCreationThenFadeOut(y_line)
        )
        self.wait(2)
        self.play(
            AnimationGroup(
                *[GrowFromCenter(d) for d in data_dots[:-1]], lag_ratio=0.5, run_time=1.5
            )
        )

        def line(x, m, c):
            return m*x + c

        def get_line(m, c):
            def y(x):
                return m*x + c
            return y

        params,_ = curve_fit(line, D, v)
        fit = self.get_graph(get_line(*params), color="#1b998b")

        self.wait(2)
        self.play(
            ShowCreation(fit)
        )
        # self.add(data_dots)

        self.wait()

class SolutionDust(Scene):

    def construct(self):
        today = TexMobject("\\text{today, }", "t_0").shift(UP*2 + LEFT*4)
        at0 = TexMobject("a", "(", "t_0", ")", "=", "1").next_to(today, DOWN, buff=MED_LARGE_BUFF)
        later = TexMobject("\\text{later, }", "t", "> t_0").shift(UP*2 + RIGHT*4)
        later[1].set_color(RED)
        at = TexMobject("a", "(", "t", ")", ">", "1").next_to(later, DOWN, buff=MED_LARGE_BUFF)
        at[2].set_color(RED)
        self.play(
            FadeInFromDown(today)
        )
        self.play(
            ShowCreation(Underline(today))
        )
        self.play(
            Write(at0)
        )

        self.play(
            FadeInFromDown(later)
        )
        self.play(
            ShowCreation(Underline(later))
        )
        self.play(
            Write(at)
        )
        self.wait()

        dphyst0 = TexMobject("d_{phys}", "(", "t_0", ")", "=", "d").next_to(at0, DOWN, buff=MED_LARGE_BUFF)
        dphyst = TexMobject("d_{phys}", "(", "t", ")", "=", "a(t)", "d").next_to(at, DOWN, buff=MED_LARGE_BUFF)
        dphyst[2].set_color(RED)
        self.play(
            Write(dphyst0)
        )
        self.wait()
        self.play(
            Write(dphyst)
        )
        self.wait()

        Vphyst0 = TexMobject("V", "(", "t_0", ")", "=", "d", "^3").next_to(dphyst0, DOWN, buff=MED_LARGE_BUFF)
        Vphyst = TexMobject("V", "(", "t", ")", "=", "a^3(t)", "d", "^3").next_to(dphyst, DOWN, buff=MED_LARGE_BUFF)
        Vphyst[2].set_color(RED)
        self.play(
            Write(Vphyst0)
        )
        self.wait()
        self.play(
            Write(Vphyst)
        )
        self.wait()

        rhophyst0 = TexMobject("\\rho", "(", "t_0", ")", "\\propto", "\\frac{1}{d^3}").next_to(Vphyst0, DOWN, buff=MED_LARGE_BUFF)
        rhophyst = TexMobject("\\rho", "(", "t", ")", "\\propto", "\\frac{1}{a^3(t)d^3}").next_to(Vphyst, DOWN, buff=MED_LARGE_BUFF)
        rhophyst[2].set_color(RED)
        self.play(
            Write(rhophyst0)
        )
        self.wait()
        self.play(
            Write(rhophyst)
        )
        self.wait()

        self.play(
            AnimationGroup(
                *[Uncreate(m) for m in self.get_top_level_mobjects()], lag_ratio=0.2
            )
        )
        self.wait()
        rhot = TexMobject("\\rho", "\\propto", "\\frac{1}{a^3}").scale(1.5)
        self.play(
            Write(rhot)
        )
        self.wait()
        self.play(
            Transform(rhot[1], TexMobject("=").move_to(rhot[1])),
            Transform(rhot[2], TexMobject("\\frac{k}{a^3}").move_to(rhot[2]).scale(1.5))
        )
        self.wait()
        self.play(
            Uncreate(rhot)
        )

        lhs = TexMobject("\\left(", "\\dot{a}", "\\over", "a", "\\right)", "^2", color=WHITE)
        eq = TexMobject("=", color=WHITE).next_to(lhs, RIGHT)
        rhs = TexMobject("\\left(", "8", "\\pi", "G", "\\over", "3", "\\right)", "\\rho", color=WHITE).next_to(eq, RIGHT)
        friedmann = VGroup(lhs, eq, rhs).move_to(ORIGIN).scale(1.5)
        self.play(
            Write(friedmann)
        )
        self.wait()
        self.play(
            Transform(rhs[-1], TexMobject("\\frac{k}{a^3}").scale(1.5).move_to(rhs[-1]).shift(UP*.15))
        )
        self.wait()
        self.play(
            Uncreate(rhs)
        )
        rhs = TexMobject("\\sqrt{", "D", "\\over", "a^3", "2", "}").scale(1.5).next_to(eq, RIGHT)
        self.play(
            Write(rhs[2:])
        )
        self.wait()
        self.play(
            CircleIndicate(lhs[-1])
        )
        self.play(
            Uncreate(lhs[0]),
            Uncreate(lhs[-2:]),
            ShowCreation(rhs[0]),
            ShowCreation(rhs[1])
        )
        self.wait()
        self.play(
            Uncreate(lhs),
            Uncreate(eq),
            Uncreate(rhs)
        )
        eq = TexMobject("=").scale(1.5)
        lhs = TexMobject("\\int{", "\\sqrt{a} \\, \\, \\, \\,", "\\dot{a}", "\\, \\, \\, \\, \\, \\,\\, \\, }").scale(1.5).next_to(eq, LEFT)
        rhs = TexMobject("\\sqrt{D}", "\\int{",  "dt", "}").scale(1.5).next_to(eq, RIGHT)
        self.play(
            Write(lhs[1:-1])
        )
        self.play(
            Write(eq),
            Write(rhs[0])
        )
        self.wait()
        self.play(
            Transform(lhs[2], TexMobject("da", "\\over", "dt").scale(1.5).move_to(lhs[2]))
        )
        self.wait()
        self.play(
            ApplyMethod(lhs[2][-1].move_to, rhs[-2])
        )
        self.play(
            Uncreate(lhs[2][1]),
            ApplyMethod(lhs[2][0].next_to, lhs[1], buff=0.1)
        )
        self.play(
            Write(lhs[0]),
            Write(rhs[1])
        )
        self.wait()
        self.play(
            Uncreate(lhs),
            Uncreate(rhs),
            FadeOut(eq)
        )
        self.wait()
        lhs = TexMobject("a", "(", "t", ")").scale(1.5).next_to(eq, LEFT)
        lhs[2].set_color(RED)
        rhs = TexMobject("\\left(", "t", "\\over", "t_0", "\\right)", "^{\\frac{2}{3}}").scale(1.5).next_to(eq, RIGHT)
        rhs[1].set_color(RED)
        self.play(
            Write(lhs)
        )
        self.play(
            Write(rhs),
            Write(eq)
        )
        self.wait()

SCENES_IN_ORDER = [
    HubblesLaw,
    ScaleFactorIntro,
    Friedmann,
    SolutionDust,
    GraphScaleFactor,
    Findingrho,
    GraphDensity
]


'''
Created on 2018年3月28日

@author: IL MARE
'''
import numpy as np

def calShannonEnt(trainLabel):
    m = len(trainLabel)
    uniqueVal = set(trainLabel)
    countDict = {}
    shannonNum = 0.0
    for label in trainLabel:
        countDict[label] = countDict.get(label, 0) + 1
    for label in uniqueVal:
        p = (countDict[label] / m)
        shannonNum -= p * np.log2(p)
    return shannonNum

def splitDataMatrix(dataMatrix, label, axis, value):
    returnMat = []
    labelMat = []
    for row, row1 in zip(dataMatrix, label):
        if row[axis] == value:
            tmp_lst = row[0: axis]
            tmp_lst.extend(row[axis + 1:])
            returnMat.append(tmp_lst)
            labelMat.append(row1)
    return returnMat, labelMat

def chooseBestFeature(trainSet, label):
    m = len(trainSet)
    maxGain = -1
    baseShannonEnt = calShannonEnt(label)
    index = -1
    for i in range(len(trainSet[0])):
        uniqueAttr = set([example[i] for example in trainSet])
        tmp_Ent = 0
        for attr in uniqueAttr:
            subSet, labelMat = splitDataMatrix(trainSet, label, i, attr)
            newShannonEnt = calShannonEnt(labelMat)
            tmp_Ent += float(len(subSet) / m) * newShannonEnt
        gain = baseShannonEnt - tmp_Ent
        if gain > maxGain:
            maxGain = gain
            index = i
    return index

def createDecisionTree(trainSet, trainLabel):
    if trainLabel.count(trainLabel[0]) == len(trainLabel):
        return trainLabel[0]
    if len(trainSet[0]) == 0:
        return "no" if trainLabel.count("no") > trainLabel.count("yes") else "yes"
    index = chooseBestFeature(trainSet, trainLabel)
    Tree = {index:{}}
    uniqueVal = set([elt[index] for elt in trainSet])
    for value in uniqueVal:
        subSet, label = splitDataMatrix(trainSet, trainLabel, index, value)
        Tree[index][value] = createDecisionTree(subSet, label)
    return Tree

def predictByDTModel(data, model):
    if type(model) == str:
        return model
    key = iter(model.keys()).__next__()
    value = data[key]
    res = model[key].get(value, None)
    if res != None:
        return predictByDTModel(data, res)
    else:
        tmp_lst = [item for item in model[key].keys()]
        return predictByDTModel(data, model[key][np.random.choice(tmp_lst, 1)[0]])

def testDTModel(testData, testLabel, model):
    predictLabel = []
    for row in testData:
        predictLabel.append(predictByDTModel(row, model))
    errorCount = 0
    for val, val1 in zip(predictLabel, testLabel):
        if val != val1:
            errorCount += 1
    ratio = float(errorCount) / len(testLabel)
    print("DT:total error ratio is %.3f, correct ratio is %.3f" % (ratio, 1 - ratio))
    return ratio


import numpy as np
import torch

from neuralprocess.data.base import FunctionGenerator

class FourierSeriesGenerator(FunctionGenerator):
    """
    Generate random Fourier series. We use the amplitude-phase representation,
    meaning the instances are of the form

        y(x) = bias + SUM_k(a_k * cos((k*x - phase_k) / frequency_scale))

    Args:
        batch_size (int): Batch size.
        series_length (int): Number of Fourier components, not counting
            bias term. Can be an iterable to draw randomly from.
        amplitude (float): a_k in definition above.
            Can be an iterable to draw randomly from.
        phase (float): phase_k in definition above.
            Can be an iterable to draw randomly from.
        bias (float): bias in definition above.
            Can be an iterable to draw randomly from.
        frequency_scale (float): Scale for x axis, see definition above.
            Can be an iterable to draw randomly from.

    """

    def __init__(
        self,
        batch_size,
        series_length=[1, 10],
        amplitude=[-1, 1],
        phase=[-1, 1],
        bias=[-1, 1],
        frequency_scale=2.0,
        **kwargs
    ):

        super().__init__(batch_size, **kwargs)

        self.series_length = series_length
        self.amplitude = amplitude
        self.phase = phase
        self.bias = bias
        self.frequency_scale = frequency_scale

    def apply(self, x):
        """
        Generate y values for input.

        Args:
            x (np.ndarray): x values, shape (1, N, 1)

        Returns:
            np.ndarray: y values, shape (B, N, 1)

        """

        if hasattr(self.series_length, "__iter__"):
            series_length = np.random.randint(*self.series_length)
        else:
            series_length = self.series_length

        if hasattr(self.bias, "__iter__"):
            y = np.random.uniform(*self.bias, size=(self.batch_size, 1, 1))
        else:
            y = np.ones((self.batch_size, 1, 1)) * self.bias

        if hasattr(self.frequency_scale, "__iter__"):
            frequency_scale = np.random.uniform(
                *self.frequency_scale, size=(self.batch_size, 1, 1)
            )
        else:
            frequency_scale = np.ones((self.batch_size, 1, 1)) * self.frequency_scale

        for i in range(series_length):

            if hasattr(self.amplitude, "__iter__"):
                amplitude = np.random.uniform(
                    *self.amplitude, size=(self.batch_size, 1, 1)
                )
            else:
                amplitude = np.ones((self.batch_size, 1, 1)) * self.amplitude

            if hasattr(self.phase, "__iter__"):
                phase = np.random.uniform(*self.phase, size=(self.batch_size, 1, 1))
            else:
                phase = np.ones((self.batch_size, 1, 1)) * self.phase

            y = y + amplitude * np.cos(((i + 1) * x - phase) / frequency_scale)

        return y.astype(np.float32)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Program: diffusionEq1D_BackwardEuler
Created: Aug 2020
@author: Ryan Clement (RRCC)
         scisoft@outlook.com

Purpose: Solve the

            u_t = alpha * u_xx

on the interval (0,L) with boundary conditions

            u = 10 for x = 0
            and
            u = 0 for x = 1.0

and initial condition

            u(x,0) = 0.0
"""

### IMPORTS
import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as la
import matplotlib.pyplot as plt

def analSol(a,x,t):
    """

    Parameters
    ----------
    a : TYPE
        Alpha.
    x : TYPE
        Spatial mesh.
    t : TYPE
        Time.

    Returns
    -------
    TYPE
        Analytic solution array.

    """
    return np.exp(-a * np.pi**2 * t) * np.sin(np.pi*x)

def initCond(x):
    """

    Parameters
    ----------
    x : TYPE
        Spatial mesh.

    Returns
    -------
    TYPE
        Initial condition array.

    """
    return np.sin(np.pi*x)

if __name__ == '__main__':
    ### VARIABLES
    a    = 10.0                           # alpha
    xMin = 0
    xMax = 1
    xP   = 51                             # Number of mesh points
    xN   = xP - 1                         # Number of zones
    x    = np.linspace(xMin,xMax,xP)      # Spatial mesh vector
    dx   = x[1] - x[0]
    dxs  = dx**2
    tsc  = 1.0/4.0                        # Time step control parameter
                                          # Note: a*dt/dx**2 = tsc (dimensionless)
    dt   = dxs*tsc/a
    dpa  = -tsc                           # Dimensionless parameter a
    dpb  = 1.0 - 2.0*dpa                  # Dimensionless parameter b
    uBC0 = 0.0                            # Left boundary condition
    uBCN = 0.0                            # Right boundary condition
    u    = np.zeros(xP)                   # Solution vector
    b    = np.zeros(xP)                   # Old solution vector
    # Set initial condition
    b[1:-1] = initCond(x[1:-1])
    # Tridiagonal matrix vectors
    diag = np.zeros(xP)
    lowr = np.zeros(xN)
    uppr = np.zeros(xN)
    # Load matrix vectors
    diag[1:-1] = dpb
    diag[0]    = 1.0
    diag[-1]   = 1.0
    lowr[:-1]  = dpa
    lowr[-1]   = 0.0
    uppr[1:]   = dpa
    uppr[0]    = 0.0
    # Create coefficient matrix
    A = sp.diags(
        diagonals = [diag,lowr,uppr],
        offsets = [0,-1,1],
        shape = (xP,xP),
        format = 'csr')
    # print(A.todense())
    print('dt = {:0.5f}'.format(dt))
    numSteps = 1000
    time = dt*numSteps
    aSol = analSol(a,x,time)
    for i in range(numSteps):
        u[:] = la.spsolve(A,b)
        b, u = u, b
    plt.plot(x,b,'bo',label='Computed')
    plt.plot(x,aSol,label='Analytic',color='red')
    plt.legend()
    plt.title(r'$u_t=\alpha\cdot u_{xx}$')
    plt.text(0,.35,'Time = {:0.5f}s'.format(time))
    plt.show()


"""A module for generating random quantum states and processes. Pseudocode for many of these
routines can be found in the appendix of the paper by Granade, Combes, and Cory (2016).
https://arxiv.org/pdf/1509.03770.pdf"""
from typing import Optional

import numpy as np
from numpy import linalg as la
from scipy.linalg import sqrtm
from sympy.combinatorics import Permutation
from numpy.random import RandomState

def ginibre_matrix_complex(D, K, rs: Optional[RandomState] = None):
    """
    Given a scalars $D$ and $K$, returns a D × K matrix,
    drawn from the complex Ginibre ensemble, i.e. (N(0, 1) + i · N(0, 1)).

    See: Induced measures in the space of mixed quantum states,
         Zyczkowski and Sommers,
         J. Phys A: Math. and Gen. 34, 7111 (2001).
         https://doi.org/10.1088/0305-4470/34/35/335

    :param D: Hilbert space dimension (scalar).
    :param K: Ultimately becomes the rank of a state (scalar).
    :param rs: Optional random state.
    :return: Returns a D × K matrix, drawn from the Ginibre ensemble.
    """
    if rs is None:
        rs = np.random

    return rs.randn(D, K) + 1j * rs.randn(D, K)

def haar_rand_unitary(dim, rs=None):
    """
    Given a Hilbert space dimension D this function
    returns a unitary operator U ∈ C^D×D drawn from the Haar measure.

    The error is of order 10^-16.

    See: How to generate random matrices from the classical compact groups,
         Mezzadri,
         Notices of the American Mathematical Society 54, 592 (2007).
         http://www.ams.org/notices/200705/fea-mezzadri-web.pdf

    :param dim: Hilbert space dimension (scalar).
    :param rs: Optional random state
    :return: Returns a unitary operator U ∈ C^D×D drawn from the Haar measure.
    """
    if rs is None:
        rs = np.random
    Z = ginibre_matrix_complex(D=dim, K=dim, rs=rs)  # /np.sqrt(2)
    Q, R = np.linalg.qr(Z)
    diag = np.diagonal(R)
    lamb = np.diag(diag) / np.absolute(diag)
    return np.matmul(Q, lamb)

def haar_rand_state(dimension):
    """
    Given a Hilbert space dimension $D$ this function returns a vector
    representing a random pure state operator drawn from the Haar measure.

    :param dimension: Hilbert space dimension (scalar).
    :return: Returns a Dx1 vector drawn from the Haar measure.

    """
    unitary = haar_rand_unitary(dimension)
    fiducial_vec = np.zeros((dimension, 1))
    fiducial_vec[0] = 1
    return np.matmul(unitary, fiducial_vec)

def ginibre_state_matrix(D, K):
    """
    Given a Hilbert space dimension $D$ and a desired rank $K$, returns
    a D × D positive semidefinite matrix of rank K drawn from the Ginibre ensemble.
    For D = K these are staes drawn from the Hilbert-Schmidt measure.

    :param D: Hilbert space dimension (scalar).
    :param K: The rank of a state (scalar).
    :return: Returns a D × K matrix, drawn from the Ginibre ensemble.

    See: Induced measures in the space of mixed quantum states,
         Zyczkowski and Sommers,
         J. Phys A: Math. and Gen. 34, 7111 (2001).
         https://doi.org/10.1088/0305-4470/34/35/335

    """
    if K > D:
        raise ValueError("The rank of the state matrix cannot exceed the dimension.")
        return

    A = ginibre_matrix_complex(D, K)
    M = A.dot(np.transpose(np.conjugate(A)))
    return M / np.trace(M)

def bures_measure_state_matrix(D):
    """
    Given a Hilbert space dimension $D$, returns a D × D positive
    semidefinite matrix drawn from the Bures measure.

    :param D: Hilbert space dimension (scalar).
    :param K: The rank of a state (scalar).
    :return: Returns a D × K matrix, drawn from the Ginibre ensemble.

    See: Random Bures mixed states and the distribution of their purity,
         Osipov, Sommers, and Zyczkowski,
         J. Phys. A: Math. Theor. 43, 055302 (2010).
         https://doi.org/10.1088/1751-8113/43/5/055302

    """
    A = ginibre_matrix_complex(D, D)
    U = haar_rand_unitary(D)
    Udag = np.transpose(np.conjugate(U))
    Id = np.eye(D)
    M = A.dot(np.transpose(np.conjugate(A)))
    P = (Id + U).dot(M).dot(Id + Udag)
    return P / np.trace(P)

def rand_map_with_BCSZ_dist(D, K):
    """
    Given a Hilbert space dimension $D$ and a Kraus rank $K$, returns a
    $D^2 × D^2$ Choi matrix $J(Λ)$ of a channel drawn from the BCSZ distribution
    with Kraus rank $K$.

    :param D: Hilbert space dimension (scalar).
    :param K: The rank of a state (scalar).
    :return: D^2 × D^2 Choi matrix, drawn from the BCSZ distribution with Kraus rank K.

    See: Random quantum operations,
         Bruzda, Cappellini, Sommers, and Zyczkowski,
         Physics Letters A 373, 320 (2009).
         https://doi.org/10.1016/j.physleta.2008.11.043

    """
    # TODO: this ^^ is CPTP, might want a flag that allows for just CP quantum operations.

    X = ginibre_matrix_complex(D ** 2, K)
    rho = X.dot(np.transpose(np.conjugate(X)))
    # Now compute the partial trace using method from:
    # https://www.peijun.me/reduced-density-matrix-and-partial-trace.html
    #
    # [[ TODO: might be better to do: rho_red = np.einsum('jiki->jk', reshaped_dm)
    # ^^ see https://scicomp.stackexchange.com/questions/27496/calculating-partial-trace-of-array-in-numpy/27659 ]]

    rho_tensor = rho.reshape([D, D, D, D])
    rho_red = np.trace(rho_tensor, axis1=0, axis2=2)
    # Now sqrt
    Q = np.kron(np.eye(D), sqrtm(la.inv(rho_red)))
    Z = Q.dot(rho).dot(Q)
    return Z

def permute_tensor_factors(D, perm):
    r"""
    Return a permutation matrix of the given dimension.

    Given a Hilbert space dimension $D$ and an list representing the permutation $perm$ of the
    tensor product Hilbert spaces, returns a $D^len(perm) x D^len(perm)$ permutation matrix.

    E.g. 1) Suppose D=2 and perm=[0,1]
            Returns the identity operator on two qubits

         2) Suppose D=2 and perm=[1,0]
            Returns the SWAP operator on two qubits which
            maps A_1 \otimes A_2 --> A_2 \otimes A_1.

    See: Equations 5.11, 5.12, and 5.13 in https://arxiv.org/pdf/0711.1017.pdf

    This function is used in tests for other functions. However, it can also be useful when
    thinking about higher moment (N>2) integrals over the Haar measure.

    :param D: Hilbert space dimension (scalar).
    :param perm: The rank of a state (list).
    """
    dim_list = [D for i in range(2 * len(perm))]

    Id = np.eye(D ** len(perm), D ** len(perm))

    P = Permutation(perm)
    tran = P.transpositions
    trans = tran()

    temp = np.reshape(Id, dim_list)

    # implement the permutation

    if P == []:
        return Id
    else:
        for pdx in range(len(trans)):
            tdx = trans[pdx]
            answer = np.swapaxes(temp, tdx[0], tdx[1])
            temp = answer

    return np.reshape(answer, [D ** len(perm), D ** len(perm)])

import numpy,from numpy,from scipy,from sympy