# coding: utf-8

# In[1]:

import numpy as np
get_ipython().magic('matplotlib inline')
from matplotlib import pyplot as plt

# Adapted from http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-4
#
# ## Plane Poiseuille flow - BVP solve by shooting method
#
# One approach to solving BVPs is to use the shooting method. The reason we cannot use an initial value solver for a BVP is that there is not enough information at the initial value to start. In the shooting method, we take the function value at the initial point, and guess what the function derivatives are so that we can do an integration. If our guess was good, then the solution will go through the known second boundary point. If not, we guess again, until we get the answer we need. In this example we repeat the pressure driven flow example, but illustrate the shooting method.
#
# In the pressure driven flow of a fluid with viscosity $\mu$ between two stationary plates separated by distance $d$ and driven by a pressure drop $\Delta P/\Delta x$, the governing equations on the velocity $u$ of the fluid are (assuming flow in the x-direction with the velocity varying only in the y-direction):
#
# $$\frac{\Delta P}{\Delta x} = \mu \frac{d^2u}{dy^2}$$
#
# with boundary conditions $u(y=0) = 0$ and $u(y=d) = 0$, i.e. the no-slip condition at the edges of the plate.
#
# we convert this second order BVP to a system of ODEs by letting $u_1 = u$, $u_2 = u_1'$ and then $u_2' = u_1''$. This leads to:
#
# $\frac{d u_1}{dy} = u_2$
#
# $\frac{d u_2}{dy} = \frac{1}{\mu}\frac{\Delta P}{\Delta x}$
#
# with boundary conditions $u_1(y=0) = 0$ and $u_1(y=d) = 0$.
#
# for this problem we let the plate separation be d=0.1, the viscosity $\mu = 1$, and $\frac{\Delta P}{\Delta x} = -100$.

# ## Shooting method
# To integrate the differential equations we need starting values $u_1(0)$ and $u_2(0)$, but we only have $u_1(0)=0$. We need to guess a value for $u_2(0)$ and see if the solution goes through the $u_2(d)=0$ boundary value.
#
#

# In[2]:

from scipy.integrate import odeint

d = 0.1 # plate thickness
mu = 1
Pdrop = -100

def odefun(U, y):
    u1, u2 = U
    du1dy = u2
    du2dy = 1.0/mu * Pdrop
    return [du1dy, du2dy]

u1_0 = 0 # known
u2_0 = 1 # guessed

dspan = np.linspace(0, d)

U = odeint(odefun, [u1_0, u2_0], dspan)

def plot_results(dspan, U):
    plt.plot(dspan, U[:,0], label="Shooting method")
    plt.plot([0,d],[0,0], 'ro') # the boundary values as red circles
    plt.xlim(0,d)
    plt.xlabel('d')
    plt.ylabel('$u_1$')

plot_results(dspan, U)

# Too low! Guess higher

# In[3]:

u1_0 = 0 # known
u2_0 = 4 # guessed higher than last time

U = odeint(odefun, [u1_0, u2_0], dspan)

plot_results(dspan, U)

# Rather than keep guessing, let's use a nonlinear solver to converge upon the best value for `u2_0`

# In[4]:

from scipy.optimize import fsolve

u1_0 = 0 # known

def objective(u2_0):
    U = odeint(odefun, [u1_0, u2_0], dspan)
    return U[-1,0]

u2_0, = fsolve(objective, 1.0)
print("Optimized u2 is {}".format(u2_0))

# now solve with optimal u2_0
U = odeint(odefun, [u1_0, u2_0], dspan)

plot_results(dspan, U)

# Also plot an analytical solution
u = -(Pdrop) * d**2 / 2 / mu * (dspan / d - (dspan / d)**2)
plt.plot(dspan, u, 'r--', label='Analytical solution')
plt.legend(loc='best')

plt.show()

# ## Finite difference (second order, one dimension)
# Adapted from http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-4-2 which was
# adapted from http://www.physics.arizona.edu/~restrepo/475B/Notes/sourcehtml/node24.html which is a dead link.
#
#
# For this example, we solve the plane poiseuille flow problem using a finite difference approach. An advantage of the approach we use here is we do not have to rewrite the second order ODE as a set of coupled first order ODEs, nor do we have to provide guesses for the solution. We do, however, have to discretize the derivatives and formulate a linear algebra problem.
#
#
# > A reminder:
# In the pressure driven flow of a fluid with viscosity $\mu$ between two stationary plates separated by distance $d$ and driven by a pressure drop $\Delta P/\Delta x$, the governing equations on the velocity $u$ of the fluid are (assuming flow in the x-direction with the velocity varying only in the y-direction):
# $$\frac{\Delta P}{\Delta x} = \mu \frac{d^2u}{dy^2}$$
# with boundary conditions $u(y=0) = 0$ and $u(y=d) = 0$, i.e. the no-slip condition at the edges of the plate.
# As before we let the plate separation be $d=0.1$, the viscosity $\mu = 1$, and $\frac{\Delta P}{\Delta x} = -100$.
#
#
# The idea behind the finite difference method is to approximate the derivatives by finite differences on a grid.
#
# Using central differences
# $$\frac{d^2u}{dx^2} = \frac{1}{\mu} \frac{\Delta P}{\Delta x}$$
# becomes
# $$\frac{u_{i-1} - 2 u_{i} + u_{i+1}}{h^2} = \frac{1}{\mu} \frac{\Delta P}{\Delta x} = -100$$
#
#
# By discretizing the ODE, we arrive at a set of linear algebra equations of the form $A y = b$, where $A$ and $b$ are defined as follows.
#
# $$
# A = \left [ \begin{array}{ccccc} %
#  2         & -1  & 0                    & 0 & 0 \\
# -1  & 2          & -1 & 0 & 0 \\
# 0                    & \ddots               & \ddots               & \ddots & 0 \\
# 0                    & 0                    & -1& 2 & -1  \\
# 0                    & 0                    & 0  & -1  & 2 \end{array} \right ] $$
#
#
# $$ y = \left [ \begin{array}{c} y_i \\ \vdots \\ y_N \end{array} \right ] $$
#
# $$ b = \left [ \begin{array}{c} 100 h^2   \\
#  100 h^2 \\
# \vdots \\
#  100 h^2 \\
#  100 h^2 \end{array} \right]
# $$
#
#
#
#
# If we want to solve a **linear boundary value problem** of the **general form**:
# $$y'' = p(x)y' + q(x)y + r(x)$$
# with boundary conditions $y(x_1) = \alpha$ and $y(x_2) = \beta$.
# then we arrive at a set of linear algebra equations of the form $A y = b$, where $A$ and $b$ are defined as follows.
#
# $$
# A = \left [ \begin{array}{ccccc} %
#  2 + h^2 q_1         & -1 + \frac{h}{2} p_1 & 0                    & 0 & 0 \\
# -1 - \frac{h}{2} p_2 & 2 + h^2 q_2          & -1 + \frac{h}{2} p_2 & 0 & 0 \\
# 0                    & \ddots               & \ddots               & \ddots & 0 \\
# 0                    & 0                    & -1 - \frac{h}{2} p_{N-1} & 2 + h^2 q_{N-1} & -1 + \frac{h}{2} p_{N-1} \\
# 0                    & 0                    & 0  & -1 - \frac{h}{2} p_N & 2 + h^2 q_N \end{array} \right ] $$
#
#
# $$ y = \left [ \begin{array}{c} y_i \\ \vdots \\ y_N \end{array} \right ] $$
#
# $$ b = \left [ \begin{array}{c} -h^2 r_1 + ( 1 + \frac{h}{2} p_1) \alpha \\
# -h^2 r_2 \\
# \vdots \\
# -h^2 r_{N-1} \\
# -h^2 r_N + (1 - \frac{h}{2} p_N) \beta \end{array} \right]
# $$
#
# Mapping our $u'' = \frac{1}{\mu} \frac{\Delta P}{\Delta x}$ onto the general form $$y'' = p(x)y' + q(x)y + r(x)$$ with boundary conditions $y(x_1) = \alpha$ and $y(x_2) = \beta$, we see $p(x)=0$, $q(x)=0$, $r(x)=-100$, $\alpha=0$, $\beta=0$.

# In[5]:

"""
This is set up for general p, q, r, alpha, beta
"""
# we use the notation for y'' = p(x)y' + q(x)y + r(x)
def p(x):
    return 0

def q(x):
    return 0

def r(x):
    return -100

#we use the notation y(x1) = alpha and y(x2) = beta
x1 = 0; alpha = 0.0
x2 = 0.1; beta = 0.0

npoints = 1000

# compute interval width
h = (x2-x1)/npoints;

# preallocate and shape the b vector and A-matrix
b = np.zeros((npoints - 1, 1));
A = np.zeros((npoints - 1, npoints - 1));
X = np.zeros((npoints - 1, 1));

#now we populate the A-matrix and b vector elements
for i in range(npoints - 1):
    X[i,0] = x1 + (i + 1) * h

    # get the value of the BVP Odes at this x
    pi = p(X[i])
    qi = q(X[i])
    ri = r(X[i])

    if i == 0:
        # first boundary condition
        b[i] = -h**2 * ri + (1 + h / 2 * pi)*alpha;
    elif i == npoints - 1:
        # second boundary condition
        b[i] = -h**2 * ri + (1 - h / 2 * pi)*beta;
    else:
        b[i] = -h**2 * ri # intermediate points

    for j in range(npoints - 1):
        if j == i: # the diagonal
            A[i,j] = 2 + h**2 * qi
        elif j == i - 1: # left of the diagonal
            A[i,j] = -1 - h / 2 * pi
        elif j == i + 1: # right of the diagonal
            A[i,j] = -1 + h / 2 * pi
        else:
            A[i,j] = 0 # off the tri-diagonal

# solve the equations A*y = b for Y
Y = np.linalg.solve(A,b)

x = np.hstack([x1, X[:,0], x2])
y = np.hstack([alpha, Y[:,0], beta])

plt.plot(x, y, label='finite difference')

mu = 1
d = 0.1
x = np.linspace(0,0.1);
Pdrop = -100 # this is DeltaP/Deltax
u = -(Pdrop) * d**2 / 2.0 / mu * (x / d - (x / d)**2)
plt.plot(x,u,'r--', label='analytical solution')

plt.xlabel('distance between plates')
plt.ylabel('fluid velocity')
plt.legend(loc='best')

# ## Boundary conditions
# This example had a *Dirichlet boundary condition* (fixed value). What if you had a *Neumann boundary condition* (fixed derivative)??
# An example is solving for temperature $T(x)$ you may specify thermal insulation ($dT/dx=0$) at one boundary instead of a fixed temperature.

# In[6]:

# Good question! (...for homework!!!)

# # PDEs - method of lines
#
# Adapted from http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-5-1
#
# The PDE that describes the transient behavior of a plug flow reactor with constant volumetric flow rate is:
# $$ \frac{\partial C_A}{\partial t} = -v_0 \frac{\partial C_A}{\partial V} + r_A $$
#
# To solve this numerically in python, we will utilize the method of lines. The idea is to discretize the reactor in volume, and approximate the spatial derivatives by finite differences. Then we will have a set of coupled ordinary differential equations that can be solved in the usual way. Let us simplify the notation with $C = C_A$, and let $r_A = -k C^2$. Graphically this looks like this:
#
# ![pde-method-of-lines](http://kitchingroup.cheme.cmu.edu/pycse/images/pde-method-of-lines.png)
#
# This leads to the following set of equations:
#
# \begin{eqnarray}
# \frac{dC_0}{dt} &=& 0 \text{ (entrance concentration never changes)} \\
# \frac{dC_1}{dt} &=& -v_0 \frac{C_1 - C_0}{V_1 - V_0} - k C_1^2 \\
# \frac{dC_2}{dt} &=& -v_0 \frac{C_2 - C_1}{V_2 - V_1} - k C_2^2 \\
# \vdots \\
# \frac{dC_4}{dt} &=& -v_0 \frac{C_4 - C_3}{V_4 - V_3} - k C_4^2
# \end{eqnarray}
#
# Last, we need initial conditions for all the nodes in the discretization. Let us assume the reactor was full of empty solvent, so that $C_i = 0$ at $t=0$. In the next block of code, we get the transient solutions, and the steady state solution.

# In[7]:

from scipy.integrate import odeint

Ca0 = 2     # Entering concentration
vo = 2      # volumetric flow rate
volume = 20 # total volume of reactor, spacetime = 10
k = 1       # reaction rate constant

N = 100     # number of points to discretize the reactor volume on

init = np.zeros(N)    # Concentration in reactor at t = 0
init[0] = Ca0         # concentration at entrance

V = np.linspace(0, volume, N) # discretized volume elements
tspan = np.linspace(0, 25)    # time span to integrate over

def method_of_lines(C, t):
    'coupled ODES at each node point'
    D = np.zeros_like(C) # C0 is constant at entrance, so D[0] = 0
    D[1:] = -vo * np.diff(C) / np.diff(V) - k * C[1:]**2
    return D
sol = odeint(method_of_lines, init, tspan)

plt.plot(tspan, sol[:, -1])
plt.xlabel('time')
plt.ylabel('$C_A$ at exit')
plt.show()

# In[ ]:

# After approximately one space time, the steady state solution is reached at the exit. For completeness, we also examine the steady state solution.

# In[8]:

# steady state solution
def pfr(C, V):
    return 1.0 / vo * (-k * C**2)
ssol = odeint(pfr, Ca0, V)
plt.plot(V, ssol, label='Steady state')
plt.plot(V, sol[-1], label='t = {}'.format(tspan[-1]))
plt.xlabel('Volume')
plt.ylabel('$C_A$')
plt.legend(loc='best')
plt.show()

# There is some minor disagreement between the final transient solution and the steady state solution. That is due to the approximation in discretizing the reactor volume. In this example we used 100 nodes. You get better agreement with a larger number of nodes, say 200 or more. Of course, it takes slightly longer to compute then, since the number of coupled odes is equal to the number of nodes.
#
#

# In[9]:

"""
Note that you will need to install ffmpeg or mencoder. Simplest way is probably:
    conda install -c rmg ffmpeg
"""
from matplotlib import animation
import matplotlib
matplotlib.rc('animation', html='html5')
# make empty figure
fig = plt.figure()
ax = plt.axes(xlim=(0, 20), ylim=(0, 2))
line, = ax.plot(V, init, lw=2)

def animate(i):
    line.set_xdata(V)
    line.set_ydata(sol[i])
    ax.set_title('t = {0:.1f}'.format(tspan[i]))
    ax.figure.canvas.draw()
    return line,

anim = animation.FuncAnimation(fig, animate, frames=50,  blit=True)
anim

# ## Transient heat conduction
# Adapted from http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-5-2 which is
# adapated from http://msemac.redwoods.edu/~darnold/math55/DEproj/sp02/AbeRichards/slideshowdefinal.pdf
# which is a dead link
#
# We examine the transient behavior of a rod at constant T put between two heat reservoirs at different temperatures, T1 = 100, and T2 = 200.
# The rod will start at 150. Over time, we should expect a solution that approaches the steady state solution: a linear temperature profile from one side of the rod to the other.
#
# $$\frac{\partial u}{\partial t} = k \frac{\partial^2 u}{\partial x^2}$$
#
# at $t=0$, in this example we have $u_0(x) = 150$ as an initial condition, with boundary conditions $u(0,t)=100$ and $u(L,t)=200$.
#
# In Matlab there is the pdepe command. There is not yet a PDE solver in scipy. Instead, we will utilze the method of lines to solve this problem. We discretize the rod into segments, and approximate the second derivative in the spatial dimension as
# $$\frac{\partial^2 u}{\partial x^2} = \frac{(u(x + h) - 2 u(x) + u(x-h))}{ h^2}$$
# at each node. This leads to a set of coupled ordinary differential equations that is easy to solve.
#
# Let us say the rod has a length of 1, $k=0.02$, and solve for the time-dependent temperature profiles.

# In[10]:

N = 100  # number of points to discretize
L = 1.0
X = np.linspace(0, L, N) # position along the rod
h = L / (N - 1)

k = 0.02

def odefunc(u, t):
    dudt = np.zeros(X.shape)

    dudt[0] = 0 # constant at boundary condition
    dudt[-1] = 0

    # now for the internal nodes
    for i in range(1, N-1):
        dudt[i] = k * (u[i + 1] - 2*u[i] + u[i - 1]) / h**2

    return dudt

init = 150.0 * np.ones(X.shape) # initial temperature
init[0] = 100.0  # one boundary condition
init[-1] = 200.0 # the other boundary condition

tspan = np.linspace(0.0, 5.0, 100)
sol = odeint(odefunc, init, tspan)

for i in range(0, len(tspan), 5):
    plt.plot(X, sol[i], label='t={0:1.2f}'.format(tspan[i]))

# put legend outside the figure
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
plt.xlabel('X position')
plt.ylabel('Temperature')

# adjust figure edges so the legend is in the figure
plt.subplots_adjust(top=0.89, right=0.77)

# In[11]:

# Try it again using np.diff and see if it's better

# In[12]:

"""
You'll neeed to
    conda install seaborn
"""
import seaborn as sns
plt.figure(figsize=(5, 4))
num_lines = int(len(tspan)/5)+1
sns.set_palette(sns.color_palette("coolwarm",num_lines))
for i in range(0, len(tspan), 5):
    plt.plot(X, sol[i], label='t={0:1.2f}'.format(tspan[i]))

# put legend outside the figure
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
plt.xlabel('X position')
plt.ylabel('Temperature')
plt.show()

# In[13]:

from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

SX, ST = np.meshgrid(X, tspan)
ax.plot_surface(SX, ST, sol, cmap='jet')
ax.set_xlabel('X')
ax.set_ylabel('time')
ax.set_zlabel('T')
ax.view_init(elev=15, azim=-124) # adjust view so it is easy to see

# ## Transient diffusion, PDEs, Finite Difference discretization
# Adapted from http://kitchingroup.cheme.cmu.edu/pycse/pycse.html#sec-10-5-3
# See also: http://coast.nd.edu/jjwteach/www/www/441/PdfNotes/lecture16.pdf
#
# We want to solve for the concentration profile of component that diffuses into a 1D rod, with an impermeable barrier at the end. The PDE governing this situation is:
#
# $$\frac{\partial C}{\partial t} = D \frac{\partial^2 C}{\partial x^2}$$
#
# In this example we have $C_0(x) = 0$ at $t=0$ as an initial condition, with boundary conditions $C(0,t)=0.1$ and $\frac{\partial C}{\partial x}(L,t)=0$.
#
# We are going to discretize this equation in both time and space to arrive at the solution. We will let $i$ be the index for the spatial discretization, and $j$ be the index for the temporal discretization. The discretization looks like this.
#
# ![pde-diffusion-discretization-scheme](http://kitchingroup.cheme.cmu.edu/pycse/images/pde-diffusion-discretization-scheme.png)
#
# Note that we cannot use the method of lines as we did before because we have the derivative-based boundary condition at one of the boundaries.
#
# We approximate the time derivative as:
#
# $$\frac{\partial C}{\partial t} \bigg| _{i,j} \approx \frac{C_{i,j+1} - C_{i,j}}{\Delta t} $$
#
# $$\frac{\partial^2 C}{\partial  x^2} \bigg| _{i,j} \approx \frac{C_{i+1,j} - 2 C_{i,j} + C_{i-1,j}}{h^2} $$
#
# We define $\alpha = \frac{D \Delta t}{h^2}$, and from these two approximations and the PDE, we solve for the unknown solution at a later time step as:
#
# $$C_{i, j+1} = \alpha C_{i+1,j} + (1 - 2 \alpha) C_{i,j}  + \alpha C_{i-1,j} $$
#
# We know $C_{i,j=0}$ from the initial conditions, so we simply need to iterate to evaluate $C_{i,j}$, which is the solution at each time step.
#
#

# In[14]:

N = 20  # number of points to discretize
L = 1.0
X = np.linspace(0, L, N) # position along the rod
h = L / (N - 1)          # discretization spacing

C0t = 0.1  # concentration at x = 0
D = 0.02

tfinal = 50.0
Ntsteps = 1000
dt = tfinal / (Ntsteps - 1)
t = np.linspace(0, tfinal, Ntsteps)

alpha = D * dt / h**2
print(alpha)

C_xt = [] # container for all the time steps

# initial condition at t = 0
C = np.zeros(X.shape)
C[0] = C0t

C_xt += [C]

for j in range(1, Ntsteps):
    N = np.zeros(C.shape)
    N[0] =  C0t
    N[1:-1] = alpha*C[2:] + (1 - 2 * alpha) * C[1:-1] + alpha * C[0:-2]
    N[-1] = N[-2]  # derivative boundary condition flux = 0
    C[:] = N
    C_xt += [N]

    # plot selective solutions
    sns.set_palette(sns.color_palette("coolwarm",10))
    if j in [1,2,5,10,20,50,100,200,500]:
        plt.plot(X, N, label='t={0:1.2f}'.format(t[j]))

plt.xlabel('Position in rod')
plt.ylabel('Concentration')
plt.title('Concentration at different times')
plt.legend(loc='best')
plt.show()

C_xt = np.array(C_xt)
plt.figure()
sns.set_palette(sns.color_palette("coolwarm",4))
plt.plot(t, C_xt[:,5], label='x={0:1.2f}'.format(X[5]))
plt.plot(t, C_xt[:,10], label='x={0:1.2f}'.format(X[10]))
plt.plot(t, C_xt[:,15], label='x={0:1.2f}'.format(X[15]))
plt.plot(t, C_xt[:,19], label='x={0:1.2f}'.format(X[19]))
plt.legend(loc='best')
plt.xlabel('Time')
plt.ylabel('Concentration')
plt.title('Concentration at different locations')
plt.show()

# Now copy the code and increase the diffusivity from $D=0.020$ to $D=0.028$ or $D=0.029$.
# Or increase the number of points in the $x$ direction to discretize from 20 to something bigger. What happens? Can you get it stable again?

# In[15]:

N = 20  # number of points to discretize
L = 1.0
X = np.linspace(0, L, N) # position along the rod
h = L / (N - 1)          # discretization spacing

C0t = 0.1  # concentration at x = 0
D = 0.028

tfinal = 50.0
Ntsteps = 1000
dt = tfinal / (Ntsteps - 1)
t = np.linspace(0, tfinal, Ntsteps)

alpha = D * dt / h**2
print(alpha)

C_xt = [] # container for all the time steps

# initial condition at t = 0
C = np.zeros(X.shape)
C[0] = C0t

C_xt += [C]

for j in range(1, Ntsteps):
    N = np.zeros(C.shape)
    N[0] =  C0t
    N[1:-1] = alpha*C[2:] + (1 - 2 * alpha) * C[1:-1] + alpha * C[0:-2]
    N[-1] = N[-2]  # derivative boundary condition flux = 0
    C[:] = N
    C_xt += [N]

    # plot selective solutions
    sns.set_palette(sns.color_palette("coolwarm",10))
    if j in [1,2,5,10,20,50,100,200,500]:
        plt.plot(X, N, label='t={0:1.2f}'.format(t[j]))

plt.xlabel('Position in rod')
plt.ylabel('Concentration')
plt.title('Concentration at different times')
plt.legend(loc='best')
plt.show()

# # Exercise
# Try the rabbits and foxes again (start with simple ODEs not KMC) but have a grid of farms covering a large area.
# Add terms for foxes and rabbits to migrate from one farm to a neighbouring farm.

# In[16]:

from scipy.integrate import odeint

k1 = 0.015
k2 = 0.00004
k3 = 0.0004
k4 = 0.04
end_time = 600.

def rates(variables, time):
    """
    Return the right hand side of the ODE
    """
    rabbits, foxes = variables
    rate_rabbits = (k1 * rabbits - k2 * rabbits * foxes)
    rate_foxes = (k3 * rabbits * foxes - k4 * foxes)
    return (rate_rabbits, rate_foxes)

times = np.arange(0, end_time)
initial_conditions = (400., 200.)
result = odeint(rates, initial_conditions, times)
rabbits = result[:,0]
foxes = result[:,1]
plt.plot(times, rabbits, label='rabbits')
plt.plot(times, foxes, label='foxes')
plt.legend(loc="best") # put the legend at the best location to avoid overlapping things
plt.show()

# In[ ]:

# In[ ]:


"""
Create training/validation data for discrete spectrum example.

All data comes from solutions to a simple ODE.
Training data:
    Initial conditions:
        3,500 ICs randomly selected in the square:
        [-.5, .5] x [-.5, .5]
    Solve from t = 0 to 1 in steps of 0.02
    Parameters mu = -0.05, lambda = -1
Validation data:
    Same structure as training data but with 1,000 ICs
"""

import numpy as np
from scipy.integrate import solve_ivp
from numpy.random import default_rng

def func(t, x):
    mu = -0.05
    lam = -1
    return [mu*x[0], lam*(x[1] - x[0]**2)]

def create_data(numICs, seed):
    rng = default_rng(seed)

    # select a bunch of initial conditions in the correct range
    ic_ranges = np.array([[-.5, .5], [-.5, .5]])
    ic_spans = ic_ranges[:,1] - ic_ranges[:,0]

    initial_conds = rng.random((numICs,2))
    initial_conds = initial_conds * ic_spans + ic_ranges[:,0]

    lenT = 51
    data = np.zeros((numICs, lenT, 2))

    t_span = [0, 1]
    t = np.linspace(0, 1, lenT) # tSpan = 0:0.02:1
    for j in np.arange(numICs):
        sol = solve_ivp(func, t_span, initial_conds[j,:], t_eval=t)
        data[j,:,:] = (sol.y).transpose() # shape [2, 51]
    return data

data_prefix = 'DiscreteSpectrumExample'
n_IC_train = 3500  # Number of initial conditions
n_IC_val = 1000

train_data = create_data(n_IC_train, seed=1)
np.save("DiscreteSpectrumExample_train1_x.npy", train_data, allow_pickle=False)

val_data = create_data(n_IC_val, seed=2)
np.save("DiscreteSpectrumExample_val_x.npy", val_data, allow_pickle=False)


#!python3
# -*- coding: utf-8 -*-

import numpy as np

def gen_matrix_X(xs):
    """
    Generate the matrix X  from the given input vector x.

    Example:
        xs = [1, 2, 3]

        X = [[1, 1, 1], [1, 2, 4], [1, 3, 9]]

    Parameters
    ----------
    xs : np.array(dtype=float)..shape==(n,)
        the input vector x

    Returns
    -------
    np.array(dtype=float)..shape==(n,n)
        The generated matrix X from the vector x.
    """

    # if the length of xs is 1, the return the matrix [[1]] of size 1*1
    if (len(xs) == 1):
        return np.array([[1]])

    ones = np.full(len(xs), 1)
    X = [ones, xs]
    temps = xs
    for i in range(len(xs) - 2):
        temps = np.multiply(temps, xs)
        X.append(temps)
    X = np.array(X).T

    return X

def gen_constants(X, ys):
    """Generate the constants vector of the taylor series from the matrix M and product vector y, where y = X*a.

    Parameters
    ----------
    X : np.array(dtype=float)..shape==(n,n)
        the matrix X, which is generated from the input vector x
    ys : np.array(dtype=float)..shape==(n,)
        the resulting vector of the matrix multiplication of :math:`X*a`

    Returns
    -------
    np.array(dtype=float)..shape==(n,)
        the parameter of taylor series for each variable :math:`x^n`, where :math:`a = X^-1 * y`
    """

    X_inv = np.linalg.inv(X)

    return np.dot(X_inv, ys.T)

def fcn_taylor(x, consts):
    """Taylor series function: :math:`y = f(x) = c_0*x^0 + c_1*x^1 + c_2*x^2 + c_3*x^3 + ...`

    Where consts = [c0, c1, c2, ...] and x is the input value for the talor series function

    Parameters
    ----------
    x : float
        the input variable x of the function f
    consts : np.array(dtype=float)..shape==(n,n)
        the constants value for taylor series

    Returns
    -------
    float
        the result of the taylor series function
    """

    y = consts[0]
    temp = x
    for const in consts[1:]:
        y += const*temp
        temp *= x

    return y

def fcn_taylors(xs, consts):
    """
    Vector input type for taylor series functions, where
        xs = [x0, x1, x2, ...]
        consts = [c0, c1, c2, ...]

        ys = [y0, y1, y2, ...].T
           = [xs^0, xs^1, xs^2, ...].T * consts
           = [[ 1, x0, x0^2, ...],
              [ 1, x1, x1^2, ...],
              [ 1, x2, x2^2, ...],
              ...                ] * consts

    Parameters
    ----------
    xs : np.array(dtype=float)..shape==(n,)
        the input vector x
    consts : np.array(dtype=float)..shape==(n,n)
        the constants value for taylor series, where :math:`f(x) = c_0*x^0 + c_1*x^1 + c_2*x^2 + c_3*x^3 + ...`

    Returns
    -------
    np.array(dtype=float)..shape==(n,)
        the matrix multiplication of X*xs
    """

    ys = [fcn_taylor(x, consts) for x in xs]
    return np.array(ys)

#######################################################
# main
#######################################################
def _main_():

    print("Get random x_src & y_src:")
    x_src = np.arange(start=0.01, stop=0.04, step=0.01)
    y_src = np.random.uniform(75.5, 125.5, size=len(x_src))
    print("x_src: shape-", x_src.shape, "data-", x_src)
    print("y_src: shape-", y_src.shape, "data-", y_src, "\n")

    print("Generate the matrix X from x_src:")
    X = gen_matrix_X(x_src)
    print("shape: ", X.shape)
    print("data:")
    print(X, "\n")

    print("Calculate the constants of the Taylor Series from x_src & y_src:")
    consts = gen_constants(X, y_src)
    print("consts: shape-", consts.shape, "data-", consts, "\n")

    print("Check the correctness of the calculated Taylor Series constants")
    y_test = fcn_taylors(x_src, consts)
    print("y_test = fcn_taylors(x_src, consts)")
    print("y_test: shape-", y_test.shape, "data-", y_test)
    print("y_src: shape-", y_src.shape, "data-", y_src, "\n")

    print("y_test - y_src =", y_test - y_src)

if __name__ == "__main__":
    _main_()


import numpy as np
import pandas as pd
class GradientDescentLinearRegression:
    def __init__(self, learning_rate=0.01, iterations=1000):
        self.learning_rate, self.iterations = learning_rate, iterations

    def fit(self, X, y):
        b = 0
        m = 5
        n = X.shape[0]
        for _ in range(self.iterations):
            b_gradient = -2 * np.sum(y - m*X + b) / n
            m_gradient = -2 * np.sum(X*(y - (m*X + b))) / n
            b = b + (self.learning_rate * b_gradient)
            m = m - (self.learning_rate * m_gradient)
        self.m, self.b = m, b

    def predict(self, X):
        return self.m*X + self.b

np.random.seed(42)
X = np.array(sorted(list(range(5))*20)) + np.random.normal(size=100, scale=0.5)
y = np.array(sorted(list(range(5))*20)) + np.random.normal(size=100, scale=0.25)
"""
data = pd.read_csv("height-weight.csv")
X = data.iloc[:,:-1]
y = data.iloc[:,-1]
"""
clf = GradientDescentLinearRegression()
clf.fit(X, y)
"""
import matplotlib.pyplot as plt
#plt.style.use('fivethirtyeight')

plt.scatter(X, y, color='black')
plt.plot(X, clf.predict(X))
plt.gca().set_title("Gradient Descent Linear Regressor")
"""
clf.b
clf.m
print("Result:")
pred = clf.predict(X)
print("y:", y)
print("y_pred:", pred)
from sklearn.metrics import r2_score
scores = r2_score(y, pred)
print("Scores:",scores*100)


import numpy as np
from numba import njit

import np_clip_fix

# Numerical methods
@njit
def sigmoid(x, k = 1):
    """Sigmoid logistic function"""
    return 1 / (1 + np.exp(-x * k))

@njit
def softmax(x):
    exps = np.exp(x)
    return exps / np.sum(exps)

@njit
def cross_entropy(predictions: np.ndarray, targets: np.ndarray):
    """ Computes cross entropy between two distributions.
    Input: x: iterabale of N non-negative values
           y: iterabale of N non-negative values
    Returns: scalar
    """

    if np.any(predictions < 0) or np.any(targets < 0):
        raise ValueError('Negative values exist.')

    if not np.any(predictions):
        predictions = np.full(predictions.shape, 1 / len(predictions))

    # Force to proper probability mass function.
    #predictions = np.array(predictions, dtype=np.float)
    #targets = np.array(targets, dtype=np.float)
    predictions /= np.sum(predictions)
    targets /= np.sum(targets)

    # Ignore zero 'y' elements.
    mask = targets > 0
    x = predictions[mask]
    y = targets[mask]
    ce = -np.sum(x * np.log(y))
    return ce

@njit
def mean_squared_error(p1, p2):
    """Calculates the mean squared error between two vectors'"""
    return np.square(np.subtract(p1, p2)).mean()
    #return 0.5 * np.sum(((p1 - p2) ** 2))

@njit
def add_noise(p: np.ndarray, noise: float):
    """Add normally distributed noise to a vector."""
    if noise > 0:
        noise_vector = np.random.normal(1, noise, p.size)
        p = p * noise_vector
        p = np.clip(p, 0., 1.)
    return p


from scipy.special import softmax
import numpy as np

"""
RNN and LSTM
https://colab.research.google.com/github/mrm8488/shared_colab_notebooks/blob/master/basic_self_attention_.ipynb

Attentions
https://www.youtube.com/watch?v=S27pHKBEp30

Position Embeddings
https://www.youtube.com/watch?v=dichIcUZfOw

@ symbol
https://www.python.org/dev/peps/pep-0465/#semantics

"""
x = np.array([
  [1, 0, 1, 0],   # Input 1
  [0, 2, 0, 2],   # Input 2
  [1, 1, 1, 1],   # Input 3
  [1, 2, 1, 2],   # Input 4
  [2, 2, 2, 2],   # Input 5
 ])
seql, emb = x.shape

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

w_key = np.array([
    [0, 0, 1, 0, 0, 1],
    [1, 1, 0, 1, 1, 0],
    [0, 1, 0, 0, 1, 0],
    [1, 1, 0, 1, 1, 0]
])
w_value = np.array([
    [0, 2, 0, 0, 2, 0],
    [0, 3, 0, 0, 3, 0],
    [1, 0, 3, 1, 0, 3],
    [1, 1, 0, 1, 1, 0]
])
key = []
query = []
value = []

# Generate Query, Key, and Value
for i in range(len(x)):
    # The out dim: 1X4 @ 4X3 = 1X3 = array(3)
    query_i = x[i] @ w_query
    key_i = x[i] @ w_key
    value_i = x[i] @ w_value
    query.append(query_i)
    key.append(key_i)
    value.append(value_i)

# print(query)
# print(key)
# print(value)
# exit()
heads = 2
head_dim = 3

# Convert list into numpy array
query = np.stack(query).reshape((seql, heads, head_dim))
key = np.stack(key).reshape((seql, heads, head_dim))
value = np.stack(value).reshape((seql, heads, head_dim))
query = np.transpose(query, (1, 0, 2))
key = np.transpose(key, (1, 0, 2))
value = np.transpose(value, (1, 0, 2))

# Transpose key again to get relevance score per head
key = np.transpose(key, (0, 2, 1))

# Generate the relevance score
relevance = query @ key
# Apply softmax to get probability scores of relevance
relevance_scores = softmax(relevance, axis=-1)
print(relevance_scores.round(decimals=2))
exit()
out = relevance_scores @ value
print(out)


#   # -*- coding: UTF-8 -*-
#   trial on the : \TBD
#   Created by Gakki on 2018/3/4
#   Project name :  class 1
#   Please contact CHIH, HSIN-CHING/D0631008 when expect to refer this source code.
#   NOTE : no liability on any loss nor damage by using this source code. it is your own risk.
import numpy as np
from scipy.linalg import toeplitz

def matrix_diag():
    v = np.matrix(np.eye(16)) * -2
    w = np.matrix(np.diag([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 1))
    u = np.matrix(np.diag([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], -1))
    diag_result = v + w + u
    diag_result[0, 15] = 1
    diag_result[15, 0] = 1
    print("Diag = \n" + str(diag_result))

def matrix_toeplitz():
    c = np.matrix([-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    r = np.matrix([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1])
    b_p = toeplitz(c, r) + np.transpose(toeplitz(c, r))
    print("Toeplitz = \n" + str(b_p))

import math

# definition :
# ax^2+bx+c=0
# x1/x2= (-b(+/-)sqrt(b^2-4ac))/2a
# https://stackoverflow.com/questions/10725522/arbitrary-precision-of-square-rootsfrom decimal import *
from bigfloat import *
import cmath

# error code 01 : invalid input as 0x^2+0x+c=0 format.
# error code 02 : invalid input as 0x^2+bx+c=0 format.
def quadratic_eq_root(a, b, c):
    # input data verification
    if (a != 0):
        eigen = (b ** 2 - 4 * a * c)
        if (eigen > 0):
            # real part only
            x1 = (-b + sqrt(eigen, precision(100))) / (2 * a)
            x2 = (-b - sqrt(eigen, precision(100))) / (2 * a)
        if (eigen == 0):
            # real part only
            x1, x2 = -b / (2 * a)
        if (eigen < 0):
            # real part and imaginary part
            x1 = (-b + cmath.sqrt(eigen)) / (2 * a)
            x2 = (-b - cmath.sqrt(eigen)) / (2 * a)
    else:
        if (b == 0):
            print("Error code 01  ")
            x1, x2 = 0, 0
        else:
            print("Error code 02  ")
            x1, x2 = 0, 0

    return (x1, x2)

# https://stackoverflow.com/questions/15390807/integer-square-root-in-python
def isqrt(num):
    x = num
    res = x
    y = (x + 1) >> 1
    while y < x:
        x = y
        y = (x + floor(num, x)) >> 1
        res = num - square(x)
    print(num, x, res)
    return x, res

# https://www.geeksforgeeks.org/calculate-square-of-a-number-without-using-and-pow/
def square(n):
    result = n
    if (n > 0):
        for i in range(1, n):
            result = result + n
    return result

# http://www.bogotobogo.com/python/python_interview_questions_2.php
def floor(a, b):
    count = 0
    sign = 1
    if a < 0: sign = -1
    while True:
        if b == 1: return a
        # positive
        if a >= 0:
            a = a - b
            if a < 0: break
        # negative
        else:
            a = -a - b
            a = -a
            if a > 0:
                count += 1
                break
        count += 1
    return count * sign

import time
import numpy
import matplotlib.pyplot as plt

def main():
    print
    "let's make it happen!"
    # Task1
    if True:
        avg1 = []
        avg2 = []
        for i in range(0, 1000):
            time_1 = time.time()
            matrix_diag()
            time_2 = time.time()
            matrix_toeplitz()
            time_3 = time.time()
            Diag_method_time = time_2 - time_1
            Toeplitz_method_time = time_3 - time_2
            avg1 = np.append(avg1, Diag_method_time)
            avg1_avg = np.mean(avg1)
            avg2 = np.append(avg2, Toeplitz_method_time)
            avg2_avg = np.mean(avg2)
            print(i)
        print("Diag_method: " + str(avg1_avg) + " second. Toeplitz method: " + str(avg2_avg) + " second")
    # Task2
    if True:
        x1, x2 = (quadratic_eq_root(1., -1.e+8, 1.))
        print("Task2 -own =")
        print(x1, x2)
        x1, x2 = (numpy.roots([1., -1.e+8, 1.]))
        print("Task2 -roots =")
        print(x1, x2)
    # Task3
    if True:
        root_own = []
        root_res = []
        root_libary = []
        fig, ax1 = plt.subplots()
        x = np.arange(0, 200, 1)
        for i in range(0, 200):
            root_own_data, root_own_res = isqrt(i)
            root_own = np.append(root_own, root_own_data)
            root_res = np.append(root_res, root_own_res)
            root_libary = np.append(root_libary, cmath.sqrt(i))
        ax1.plot(x, root_own, color='r', label='own code in root')
        ax1.plot(x, root_libary, color='g', label='cmath library')
        plt.legend(bbox_to_anchor=(0.35, 1))
        plt.ylim((0, 15))
        plt.ylabel('root value')
        plt.xlabel('input number')
        ax2 = ax1.twinx()
        ax2.bar(x, root_res, color='b', label='own code in residue')
        plt.xlabel('x')
        plt.ylim((0, 25))
        plt.ylabel('residure value')
        plt.grid(True)
        plt.legend(bbox_to_anchor=(0.8, 1))
        fig.tight_layout()
        plt.show()

if __name__ == "__main__":
    main()