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

# # Step 4: Burger's Equation
# 
# Next, we combine what we have learned about convection and diffusion and apply it to the Burger's Equation. This equation looks like —and is— the direct combination of both of the PDE's we had been working on earlier.
# 
# $$ \frac{\partial u}{\partial t} + \frac{\partial u}{\partial x} = \nu \frac{\partial^2 u}{\partial x^2} $$
# 
# We can discretize it using the methods we have developed previously in steps 1-3. It will take forward difference for the time component, backward difference for space and our 2nd order combination method for hte second derivatives. This yields:
# 
# $$ \frac{u^{n+1}_i - u^n_i}{\Delta t} + u_i^n \frac{u^{n}_i - u^n_{i-1}}{\Delta x} =
# \nu \frac{u^{n}_{i+1} -2u^n_i + u^n_{i-1}}{\Delta x^2}
# $$
# 
# Given that we have full initial conditions as before we can solve for our only unknown $u^{n+1}_i$ and iterate through the equation that follows:
# 
# $$ u^{n+1}_i = u^n_i - u^n_i \frac{\Delta t}{\Delta x} (u^n_i - u^n_{i-1})  + \frac{\nu \Delta t}{\Delta x^2}(u^{n}_{i+1} - 2u^n_i + u^n_{i-1})
# $$
# 
# The above equation will now allow us to write a program to advance our solution in time and perform our simulation. As before, we need initial conditions, and we shall continue to use the one we obtained in the previous two steps.
# 
# ## Initial and Boundary Conditions
# 
# The Burger's equation is way more interesting than the previous ones. To have a better feel for its properties it is helpful to use different initial and boundary conditions than what we have been using for the previous steps.
# 
# \begin{eqnarray}
# u &=& -\frac{2 \nu}{\phi} \frac{\partial \phi}{\partial x} + 4 \\\
#  \phi &=& \exp \bigg(\frac{-(x-4t)^2}{4 \nu (t+1)} \bigg) + \exp \bigg(\frac{-(x-4t -2 \pi)^2}{4 \nu(t+1)} \bigg)
# \end{eqnarray}
# 
# This has an analytical solution, given by:
# \begin{eqnarray}
# u &=& -\frac{2 \nu}{\phi} \frac{\partial \phi}{\partial x} + 4 \\\
# \phi &=& \exp \bigg(\frac{-(x-4t)^2}{4 \nu (t+1)} \bigg) + \exp \bigg(\frac{-(x-4t -2 \pi)^2}{4 \nu(t+1)} \bigg)
# \end{eqnarray}
# 
# Our boundary conditions will be:
# 
# $$ u(0) = u(2 \pi)
# $$
# 
# This is a periodic boundary condition which we must be careful with.
# 
# ## Obtaining $ \frac{\partial \phi}{\partial x} $
# 
# Evaluating this initial condition by hand would be relatively painful, to avoid this we can calculate the derivative using sympy. This is basically mathematica but can be used to output the results back into Python calculations. 
# 
# We shall start by loading all of the python libraries that we will need for hte project along with a fix to make sure sympy prints our functions in latex.
# 
# 

# In[1]:


# Adding inline command to make plots appear under comments
import numpy as np
import sympy
import matplotlib.pyplot as plt
import time, sys
get_ipython().run_line_magic('matplotlib', 'inline')

sympy.init_printing(use_latex =True)


# We shall start by defining the symbolic variables in our initial conditions and then typing out the full equation.

# In[2]:


x, nu, t = sympy.symbols('x nu t')
phi = (sympy.exp(-(x - 4 * t) **2 / (4 * nu * (t+1))) +
       sympy.exp(-(x - 4 *t - 2 * np.pi)**2 / (4 * nu * (t + 1))))
phi


# In[3]:


phiprime = phi.diff(x)
phiprime


# In python code:

# In[4]:


print(phiprime)


# ## Lambdifying
# 
# Now that we have the expression for $ \frac{\partial \phi}{\partial x} $ we can finish writing the full initial condition equation and then translating it into a usable python expression. To do this we use the lambdify function which takes a sympy simbolic equation and turns it into a callable function.

# In[5]:


u = -2 * nu * (phiprime / phi) + 4
print(u)


# In[6]:


ufunc = sympy.utilities.lambdify((t,x,nu), u)
print(ufunc(1,4,3))


# Pretty neat right?!
# 
# ## Solving the Burger's Equation
# 
# Now that we can set up the initial conditions we can finish up the problem. We can generate the plot of intiial conditions using the lambifyied function.

# In[7]:


#New initial conditions
grid_length = 2
grid_points = 101
nt = 150
dx = grid_length * np.pi / (grid_points - 1) 
nu = .07 
dt = dx * nu #Dynamically scaling dt based on grid size to ensure convergence

#Initiallizing the array containing the shape of our initial conditions
x = np.linspace(0,2 * np.pi, grid_points)
un = np.empty(grid_points)
t = 0

u = np.asarray([ufunc(t,x0,nu) for x0 in x])
u


# In[8]:


plt.figure(figsize=(11, 7), dpi= 100)
plt.plot(x, u, marker='o', lw=2)
plt.xlim([0, 2 * np.pi])
plt.ylim([0, 10]);
plt.xlabel('x')
plt.ylabel('u')
plt.title('Burgers Equation at t=0');


# This new function is known as a `sawtooth` function. 

# ## Periodic boundary conditions
# 
# The biggest difference between this step and the previous ones is the use of periodic boundary conditions. If you have experimented with steps 1-2 you would have seen that eventually the wave moves out of the picture to the right and does not show up in the plot. 
# 
# With periodic BC, what happens now is that when the wave hits the end of the frame it wraps around and starts from the beginning again.
# 
# Now we will apply the discretization as outlined above and check out the final results.

# In[9]:


for n in range(nt): #Runs however many timesteps you set earlier
    un = u.copy()   #copy the u array to not overwrite values
    for i in range(1,grid_points-1):
        u[i] = un[i] - un[i] * dt/dx * (un[i]-un[i-1])  + nu * (dt/dx**2) * (un[i+1]- 2*un[i] + un[i-1]) 
        
    u[0] = un[0] - un[0] * dt / dx * (un[0] - un[-2]) + nu*(dt / dx**2) *(un[1] - 2* un[0] + un[-2])
    u[-1] = u[0]

u_anal = np.asarray([ufunc(nt* dt , xi, nu) for xi in x])


# In[10]:


plt.figure(figsize=(11, 7), dpi=100)
plt.plot(x,u, marker ='o', lw=2, label='Computational')
plt.plot(x, u_anal, label='Analytical')
plt.xlim([0, 2* np.pi])
plt.ylim([0,10])
plt.xlabel('x')
plt.ylabel('u')
plt.title('Burgers Equation at t=10');
plt.legend();


# ### Results
# 
# Looks pretty Cool! But I would like to know how this evolves over time, not just at its final timestep. As always we will attempt to animate the wave and  take a look at the results.
# 
# ## Animating the wave moving

# In[11]:


#Imports for animation and display within a jupyter notebook
from matplotlib import animation, rc 
from IPython.display import HTML

#Generating the figure that will contain the animation
fig, ax = plt.subplots()
fig.set_dpi(100)
fig.set_size_inches(9, 5)
ax.set_xlim(( 0, 2*np.pi))
ax.set_ylim((0, 10))
comp, = ax.plot([], [], marker='o', lw=2,label='Computational')
anal, = ax.plot([], [], lw=2,label='Analytical')
ax.legend();
plt.xlabel('x')
plt.ylabel('u')
plt.title('Burgers Equation time evolution from t=0 to t=10');

#Resetting the U wave back to initial conditions
u = np.asarray([ufunc(0, x0, nu) for x0 in x])


# In[12]:


#Initialization function for funcanimation
def init():
    comp.set_data([], [])
    anal.set_data([], [])
    return (comp,anal,)


# In[13]:


#Main animation function, each frame represents a time step in our calculation
def animate(j):
    un = u.copy()   #copy the u array to not overwrite values
    for i in range(1,grid_points-1):
        u[i] = un[i] - un[i] * dt/dx * (un[i]-un[i-1])  + nu * (dt/dx**2) * (un[i+1]- 2*un[i] + un[i-1]) 
        
    u[0] = un[0] - un[0] * dt / dx * (un[0] - un[-2]) + nu*(dt / dx**2) *(un[1] - 2* un[0] + un[-2])
    u[-1] = u[0]
    u_anal = np.asarray([ufunc(j * dt, xi, nu) for xi in x])
    comp.set_data(x, u)
    anal.set_data(x, u_anal)
    return (comp,anal,)


# In[14]:


anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=nt, interval=20)
anim.save('../gifs/1dBurgers.gif',writer='imagemagick',fps=60)
#HTML(anim.to_jshtml())


# ## Conclusion
# 
# This concludes our examination of 1D sims and boy oh boy was this cool! This last model in particular shines in the animation showing the behavior and properties of the burghers equation quite well. 
# 
# Next, we will start our move to 2D  but before this a quick detour on array operations on NumPy.