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

# In[1]:


from IPython.display import HTML 
s="""<h1>2D Navier-Stokes for Channel Flow</h1></br><div id="toc"></div>
<script src="https://kmahelona.github.io/ipython_notebook_goodies/ipython_notebook_toc.js" defer></script>""";
h=HTML(s); h


# ***
# # Understand the Problem
# 
# ## Question
# 
# * What is the 2D velocity **and** pressure field for the **2D Navier Stokes Equation** for flow within an **open channel**?
# 
# ## Initial Conditions
# 
# * The velocity and pressure fields are zero everywhere
# 
# ## Boundary Conditions
# 
# * `x = 0m` and `x = 2m` $\Rightarrow$ ` u, v ` and ` p ` are **periodic** 
# 
# * `y = 0m` and `y = 2m` $\Rightarrow$ `u = 0m/s` and `v = 0m/s` i.e. **no slip wall**
# 
# * `y = 0m` and `y = 2m` $\Rightarrow$ `dp/dy = 0` i.e. **impermeable wall**
# 
# **NOTE: where velocity is specified, pressure cannot also be specified, because in this case:**
# 
# * velocity is a function of pressure
# * pressure is a function of velocity
# 
# **Specifiying the pressure and the velocity would result in an overspecified problem**
# 
# <img src="https://github.com/en9apr/6_thevisualroom/blob/master/_images/navier_stokes_BCS.png" style="width:25%; height:25%">
# 
# ## Governing Equations
# 
# * The Navier Stokes Momentum Equation is described as follows:
# 
# $$ {\partial u \over \partial t} + u {\partial u \over \partial x} + v {\partial u \over \partial y} = F_x -{1 \over \rho} {{\partial p} \over {\partial x}} + \nu \left ( {\partial^2 u \over \partial x^2}+  {\partial^2 u \over \partial y^2}  \right ) $$
# 
# where: $ F_x = 1$ (a source term)
# 
# $$ {\partial v \over \partial t} + u {\partial v \over \partial x} + v {\partial v \over \partial y} = -{1 \over \rho} {{\partial p} \over {\partial y}} + \nu \left ( {\partial^2 v \over \partial x^2}+  {\partial^2 v \over \partial y^2}  \right ) $$
# 
# * The Poisson Equation to link pressure and velocity is:
# 
# $$ \nabla^2 p^{n+1} = \rho {{\nabla \cdot \mathbf{u}^n} \over {\Delta t}}-
#    \rho \nabla \cdot (\mathbf{u}^n \cdot \nabla \mathbf{u}^n)+
#    \mu \nabla^2 (\nabla \cdot \mathbf{u}^n)
# $$
# 
# Where:
# 
# $$ \nabla \cdot (\mathbf{u} \cdot \nabla \mathbf{u}) = 
# \begin{bmatrix} {{\partial} \over {\partial x} } \ {{\partial} \over {\partial y} } \end{bmatrix}
# \cdot \begin{bmatrix} (u {{\partial u} \over {\partial x} } + v {{\partial u} \over {\partial y}}) \ \mathbf{i} \\ (u {{\partial v} \over {\partial x} } + v {{\partial v} \over {\partial y} }) \ \mathbf{j} \end{bmatrix} = \left( {{\partial u} \over {\partial x}} \right)^2 + 
# 2 {{\partial u} \over {\partial y}} {{\partial v} \over {\partial x}} +
# \left( {{\partial v} \over {\partial y}} \right)^2
# $$
# 
# Hence:
# 
# * **Assuming viscosity is small**, the Poisson Equation to link pressure and velocity is as follows:
# 
# $$ {{{\partial^2 p} \over {\partial x^2}} + {{\partial^2 p} \over {\partial y^2}} =
# {\rho \over \Delta t} \left( {{\partial u} \over {\partial x}} + {{\partial v} \over {\partial y}} \right)
# - \rho \left[ \left( {{\partial u} \over {\partial x}} \right)^2 + 
# 2 {{\partial u} \over {\partial y}} {{\partial v} \over {\partial x}} +
# \left( {{\partial v} \over {\partial y}} \right)^2 \right] }$$
# 
# (this equation applies in the discrete domain)
# 
# ***

# # Formulate the Problem
# 
# ## Input Data:
# 
# The Poisson Equation has **no temporal component**, so we use a number of iterations `niter`
# 
# The Navier Stokes Momentum Equation **does have a temporal component**, so we use `nt`
# 
# * `niter` = 51 (maximum number of iterations - for Poisson Equation)
# * `nt` = 51 (number of temporal points)
# * `nx` = 21 (number of x spatial points)
# * `ny` = 21 (number of y spatial points)
# * `tmax` = 0.5
# * `xmax` = 2
# * `ymax` = 2
# * `nu` = 0.1
# * `rho` = 1
# 
# ## Initial Conditions:
# 
# * $\forall (x, y, t) \quad t = 0 \rightarrow u(x,y,t) = 0 \land v(x,y,t) = 0 \land p(x,y,t) = 0$
# 
# ## Velocity Boundary Conditions:
# 
# * $\forall (x, y, t) \quad x = 0 \lor x = 2 \rightarrow u(0,y,t) = u(2,y,t) \land v(0,y,t) = v(2,y,t)$
# 
# * $\forall (x, y, t) \quad y = 0 \lor y = 2 \rightarrow u(x,0,t) = 0 \land u(x,2,t) = 0 \land v(x,0,t) = 0 \land v(x,2,t) = 0$
# 
# ## Pressure Boundary Conditions:
# 
# * $\forall (x, y, t) \quad x = 0 \lor x = 2 \rightarrow p(0,y,t) = p(2,y,t)$
# 
# * $\forall (x, y, t) \quad y = 0 \lor y = 2 \rightarrow {{\partial p(x,0,t)} / {\partial y}} = 0 \land {{\partial p(x,2,t)} / {\partial y}} = 0$
# 
# ## Output Data:
# 
# * $\forall (x, y, t) \quad \ p(x,y,t) = ? \land u(x,y,t) = ? \land v(x,y,t) = ?$
# 
# ***
# 

# # Design Algorithm to Solve Problem
# 
# ## Space-time discretisation:
# 
# * i $\rightarrow$ index of grid in x
# * j $\rightarrow$ index of grid in y
# * n $\rightarrow$ index of time
# * m $\rightarrow$ index of iterations
# 
# ## Numerical schemes for each Momentum Equation
# 
# * For the **one** first derivative of velocity in time: 1st order FD in time
# * For the **two** first derivatives of velocity in space: 1st order BD in space
# * For the **one** first derivative of pressure in space: 2nd order CD in space
# * For the **two** second derivatives of velocity in space: 2nd order CD in space
# 
# ## Numerical schemes for the Poisson Equation
# 
# * For the **two** second derivatives of pressure in space: 2nd order CD in space
# * The the **four** first derivatives of velocity in space: 2nd order CD in space
# 
# ## Discrete equation for u-Momentum Equation
# 
# $$ {{u_{i,j}^{n+1} - u_{i,j}^n} \over {\Delta t}} + 
#    u_{i,j}^n {{u_{i,j}^n - u_{i-1,j}^n} \over \Delta x} + 
#    v_{i,j}^n {{u_{i,j}^n - u_{i,j-1}^n} \over \Delta y} = \\
#    1 -{1 \over \rho} {{p_{i+1,j}^n - p_{i-1,j}^n} \over {2 \Delta x}} +
#    \nu {{u_{i-1,j}^n - 2u_{i,j}^n + u_{i+1,j}^n} \over \Delta x^2} + 
#    \nu {{u_{i,j-1}^n - 2u_{i,j}^n + u_{i,j+1}^n} \over \Delta y^2}
# $$
# 
# ## Transpose
# 
# Assume $ \Delta x = \Delta y = h$
# 
# $$ {u_{i,j}^{n+1} = u_{i,j}^n - {{\Delta t} \over h} \left[ u_{i,j}^n(u_{i,j}^n - u_{i-1,j}^n) + v_{i,j}^n(u_{i,j}^n - u_{i,j-1}^n)\right] \\
# + {\Delta t} - {{\Delta t} \over {2 \rho h}} (p_{i+1,j}^n - p_{i-1,j}^n) + {{\Delta t \nu} \over {h^2}}
# (u_{i-1,j}^n + u_{i+1,j}^n + u_{i,j-1}^n + u_{i,j+1}^n - 4 u_{i,j}^n )}
# $$
# 
# ## Discrete equation for v-Momentum Equation
# 
# $$ {{v_{i,j}^{n+1} - v_{i,j}^n} \over {\Delta t}} + 
#    u_{i,j}^n {{v_{i,j}^n - v_{i-1,j}^n} \over \Delta x} + 
#    v_{i,j}^n {{v_{i,j}^n - v_{i,j-1}^n} \over \Delta y} = \\
#    -{1 \over \rho} {{p_{i,j+1}^n - p_{i,j-1}^n} \over {2 \Delta y}} +
#    \nu {{v_{i-1,j}^n - 2v_{i,j}^n + v_{i+1,j}^n} \over \Delta x^2} + 
#    \nu {{v_{i,j-1}^n - 2v_{i,j}^n + v_{i,j+1}^n} \over \Delta y^2}
# $$
# 
# ## Transpose
# 
# Assume $ \Delta x = \Delta y = h$
# 
# $$ v_{i,j}^{n+1} = v_{i,j}^n - {{\Delta t} \over h} \left[ u_{i,j}^n(v_{i,j}^n - v_{i-1,j}^n) + v_{i,j}^n(v_{i,j}^n - v_{i,j-1}^n)\right] - \\
# {{\Delta t} \over {2 \rho h}} (p_{i,j+1}^n - p_{i,j-1}^n) + {{\Delta t \nu} \over {h^2}}
# (v_{i-1,j}^n + v_{i+1,j}^n + v_{i,j-1}^n + v_{i,j+1}^n - 4 v_{i,j}^n )
# $$
# 
# ## Discrete equation for Poisson Equation must be Divergence Free
# 
# Since no terms have a differential temporal component - bring source term to other side and equate to forward differencing (as we did for Laplace Equation). Then take steady state after a certain number of iterations.
# 
# $$ {{\partial p} \over {\partial t}} = {\partial^2 p \over \partial x^2} + {\partial^2 p \over \partial y^2} - 
# \left[ {\rho \over \Delta t} \left( {{\partial u} \over {\partial x}} + {{\partial v} \over {\partial y}} \right)
# - \rho \left[ \left( {{\partial u} \over {\partial x}} \right)^2 + 
# 2 {{\partial u} \over {\partial y}} {{\partial v} \over {\partial x}} +
# \left( {{\partial v} \over {\partial y}} \right)^2 \right] \right] $$
# 
# $$ {{\partial p} \over {\partial t}} = {\partial^2 p \over \partial x^2} + {\partial^2 p \over \partial y^2} - 
# b $$
# 
# $$ {{p_{i,j}^{m+1}-p_{i,j}^m} \over {\Delta \tau}} = {{p_{i+1,j}^m -2p_{i,j}^m + p_{i-1,j}^m} \over \Delta x^2} + {{p_{i,j+1}^m -2p_{i,j}^m + p_{i,j-1}^m} \over \Delta y^2} - b_{i,j}^m $$
# 
# Assume that $ \Delta x = \Delta y = h $
# 
# $$ {{p_{i,j}^{m+1}-p_{i,j}^m} \over {\Delta \tau}} = {{p_{i+1,j}^m -2p_{i,j}^m + p_{i-1,j}^m} \over h^2} + {{p_{i,j+1}^m -2p_{i,j}^m + p_{i,j-1}^m} \over h^2} - b_{i,j}^m $$
# 
# 
# $$ p_{i,j}^{m+1} = p_{i,j}^m + {{\Delta \tau} \over {h^2}} \left( 
# p_{i+1,j}^m + p_{i-1,j}^m + p_{i,j+1}^m + p_{i,j-1}^m - 4p_{i,j}^m \right) - b_{i,j}^m \Delta \tau $$
# 
# For the fastest convergence $ r = {{\Delta \tau} \over {h^2}} = {1 \over 4} $ and $ \Delta \tau = {{h^2} \over 4} $
# 
# Hence:
# 
# $$ p_{i,j}^{m+1} = {1 \over 4} \left( p_{i+1,j}^m + p_{i-1,j}^m + p_{i,j+1}^m + p_{i,j-1}^m - b_{i,j}^m h^2 \right) $$
# 
# So now we need to define $ b_{i,j}^m $
# 
# $$
# b = {\rho \over \Delta t} \left( {{\partial u} \over {\partial x}} + {{\partial v} \over {\partial y}} \right)
# - \rho \left[ \left( {{\partial u} \over {\partial x}} \right)^2 + 
# 2 {{\partial u} \over {\partial y}} {{\partial v} \over {\partial x}} +
# \left( {{\partial v} \over {\partial y}} \right)^2 \right] 
# $$
# 
# 
# 
# $$ b_{i,j}^{m} = 
# {\rho \over \Delta t} \left( {{u_{i+1,j}^n - u_{i-1,j}^n} \over {2 \Delta x}} + {{v_{i,j+1}^n - v_{i,j-1}^n} \over {2 \Delta y}} \right) +
# \rho \left[ \left( {{u_{i+1,j}^n - u_{i-1,j}^n} \over {2 \Delta x}} \right)^2 + 2 \left( {{u_{i,j+1}^n - u_{i,j-1}^n} \over {2 \Delta y}} {{v_{i+1,j}^n - v_{i-1,j}^n} \over {2 \Delta x}}\right) +
# \left( {{v_{i,j+1}^n - v_{i,j-1}^n} \over {2 \Delta y}} \right)^2 \right] $$
# 
# Also assume that $ \Delta x = \Delta y = h $
# 
# 
# $$ b_{i,j}^{m} = 
# {\rho \over {2 h \Delta t}} ( {{u_{i+1,j}^n - u_{i-1,j}^n}} + {{v_{i,j+1}^n - v_{i,j-1}^n}} ) +
# {\rho \over {4h^2}} \left[ (u_{i+1,j}^n - u_{i-1,j}^n)^2 + 4h ( u_{i,j+1}^n - u_{i,j-1}^n) (v_{i+1,j}^n - v_{i-1,j}^n) + (v_{i,j+1}^n - v_{i,j-1}^n)^2 \right] $$
# 
# ## Question: What governs the stability of these schemes?
# 
# * The term that grows quickest with reducing spatial step size?
# * The CD term for diffusion contains the coefficient $ {\nu \Delta t} / h^2 $
# * Based on the previous von Neumann stability analysis $ {\nu \Delta t} / h^2 \le 0.5$

# ## Pseudo-code
# 
#     # Inputs:    
#     niter = 51 
#     nx =  21
#     xmax = 2
#     tmax = 0.5
#     r = 0.25
#     nu = 0.1
#     rho = 1
#     
#     # Calculated values - nt is a function of r and dx (should be int) - WHAT VALUE OF r IS STABLE? 0.5?
#     nt = int{ ( (nu*tmax) / (r*(dx)**2) )+1 }
#     dt = tmax/(nt-1)
#     ymax = xmax
# 
#     # Increments in x and y are the same - h, dy & dx should be rational numbers:
#     h = dy = dx = xmax/(nx-1)
# 
#     # Compute ny from dy and ymax - ny should be an integer
#     ny = int{ (ymax/dy) + 1 }
# 
#     # Initial Conditions:
#         p[:,:] = 0
#         pm[:,:] = 0
#         u[:,:] = 0
#         v[:,:] = 0
#         un[:,:] = 0
#         vn[:,:] = 0
#         b[:,:] = 0
#         
#     # Velocity Boundary Conditions:
#         u[:,0] = u[:,ny-1] = v[:,0] = v[:,ny-1] = 0
#        
#     for n between 0 and nt-1
#     
#         # Numerical Computation - u and v at i=0 and i=nx-1 are unknown
#                                 - u and v at j=0 and j=ny-1 are known
#                                 
#         for i between 0 and nx-1
#             for j between 1 and ny-2
#             
#                 # Periodic Boundary Conditions
#                 if(i==0) 
#                     u[i-1,:] = u[nx-1,:]
#                     v[i-1,:] = v[nx-1,:]
#                 else if(i==nx-1)
#                     u[i+1,:] = u[0, :]
#                     v[i+1,:] = v[0, :]
#                 
#                 # Value of Poisson RHS
#                 b[i,j] = { rho / (2*h*dt) } * { u[i+1,j] - u[i-1,j] + v[i,j+1] - v[i,j-1] } +
#                          { rho / (4*h^2) } * { (u[i+1,j] - u[i-1,j])^2 +
#                                            4*h*(u[i,j+1] - u[i,j-1])*(v[i+1,j] - v[i-1,j]) + 
#                                                (v[i,j+1] - v[i,j-1])^2 } 
#      
#         for m between 0 and niter-1
#         
#             pm = copyof(p)
#             
#             for i between 0 and nx-1
#                 for j between 1 and ny-2
# 
#                     # Periodic Boundary Conditions
#                     if(i==0) 
#                         pm[i-1,:] = pm[nx-1,:]
#                     else if(i==nx-1)
#                         pm[i+1,:] = pm[0, :]
#                         
#                     p[i,j] = 0.25*( pm[i+1,j] + pm[i-1,j] + pm[i,j+1] + pm[i,j-1] - b[i,j]*h**2 )
#                     
#             un = copyof(u)
#             vn = copyof(v)
#             
#             for i between 0 and nx-1
#                 for j between 1 and ny-2
# 
#                     # Periodic Boundary Conditions
#                     if(i==0) 
#                         un[i-1,:] = un[nx-1,:]
#                         vn[i-1,:] = vn[nx-1,:]
#                     else if(i==nx-1)
#                         un[i+1,:] = un[0, :]
#                         vn[i+1,:] = vn[0, :]
# 
#                     u[i,j] = un[i,j] - {dt / h} * { un[i,j] * ( un[i,j] - un[i-1,j] ) + 
#                                                     vn[i,j] * ( un[i,j] - un[i,j-1] ) } + 
#                                   {dt} - {dt / (2 * rho * h)} * { p[i+1,j] - p[i-1,j] } + 
#                        {(dt * nu) / (h^2)} * { un[i-1,j] + un[i+1,j] + un[i,j-1] + un[i,j+1] - 4 * un[i,j] } 
# 
#                     v[i,j] = vn[i,j] - {dt / h} * { un[i,j] * ( vn[i,j] - vn[i-1,j] ) + 
#                                                     vn[i,j] * ( vn[i,j] - vn[i,j-1] ) } - 
#                                        {dt / (2 * rho * h)} * { p[i,j+1] - p[i,j-1] } + 
#                        {(dt * nu) / (h^2)} * { vn[i-1,j] + vn[i+1,j] + vn[i,j-1] + vn[i,j+1] - 4 * vn[i,j] }       
#        
# ***

# # Implement Algorithm in Python
# 
# ## Question: How does the value of r affect the pressure in the **Poisson** Equation?

# Try to implement a Parabolic velocity in u and v so that b is non-zero

# In[91]:


def poisson_initialisation(niter, r, nx_or_ny, tmax, xmax_or_ymax, nu):
    """
    Initialise conditions for Poisson Equation
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int(((nu * tmax) / (r*(dx)**2))+1)
    dt = tmax/(nt-1)

    # Initialise data structures:
    import numpy as np
    
    # Initial conditions
    p = np.zeros((nx,ny))
    u = np.zeros((nx,ny))
    v = np.zeros((nx,ny))
    
    # x and y range
    x = np.zeros(nx)
    y = np.zeros(ny)

    # X Loop
    x = np.linspace(0.0,2.0,nx)
    
    # Y Loop
    y = np.linspace(0.0,2.0,ny)
     
    # Velocity Initial Conditions:
        
    #u[:,0:ny] = [(1.0/nu) * (y[j] - ((y[j])**2.0 / 2.0)) for j in range(ny) ]  
    for i in range(nx):
        for j in range(ny):
            
            u[i,j] = x[i] * y[j] * (2.0-x[i])*(2.0-y[j])
            v[i,j] = x[i] * y[j] * (2.0-x[i])*(2.0-y[j])
            
    return p, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, r


# In[92]:


(p30, x30, y30, u30, v30, nx, ny, nt, 
    dx, dy, dt, niter, r) = poisson_initialisation(50, 0.25, 51, 0.5, 2.0, 0.1)  


# In[93]:


def plot_3D(u,x,y,title,label):
    """
    Plots the 2D velocity field
    """

    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    fig=plt.figure(figsize=(11,7),dpi=100)
    ax=fig.gca(projection='3d')
    ax.set_xlabel('x (m)')
    ax.set_ylabel('y (m)')
    ax.set_zlabel(label)
    ax.set_xlim(0,2)
    ax.set_ylim(0,1)
    ax.view_init(30,225)
    Y,X=np.meshgrid(y,x) #note meshgrid uses y,x not x,y!!!
    surf=ax.plot_surface(X,Y,u[:,:], rstride=1, cstride=1)
    plt.title(title)
    plt.show()


# In[94]:


plot_3D(u30,x30,y30,'Figure 1: Initial Conditions','u (m/s)')
plot_3D(v30,x30,y30,'Figure 2: Initial Conditions','v (m/s)')
plot_3D(p30,x30,y30,'Figure 3: Initial Conditions','p (Pa)')


# In[95]:


def poisson_equation(rho, nu, r):
                  
    (p, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, r) = poisson_initialisation(50, r, 51, 0.5, 2.0, 0.1) 
    
    # Increments
    h = dx
    
    import numpy as np
    
    # Intermediate copies:
    un = np.zeros((nx, ny))
    vn = np.zeros((nx, ny))
    pm = np.zeros((nx, ny))
    bm = np.zeros((nx, ny)) # bm needs to be exactly zero at the boundaries  
        
    # First points for b. We know the velocity at j=0 and j=ny-1. We don't know it at i=0 and i=nx-1
    # However, we will re-write the equation for b at i=0 and i=nx-1
    
     # Loop - use decimal points for all floating point numbers
    for n in range(nt):
    
        i=1
        j=1    
        bm[i:nx-1, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1] 
                                                        + v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1])**2.0 + 
                             4.0*h*(u[i:nx-1, j+1:ny] - u[i:nx-1, j-1:ny-2])*
                                   (v[i+1:nx, j:ny-1] - v[i-1:nx-2, j:ny-1]) + 
                                   (v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2])**2.0 ) )
        i=0
        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1, j:ny-1] - u[nx-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1, j:ny-1] - u[nx-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[i+1, j:ny-1] - v[nx-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )

        i=nx-1
        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[0, j:ny-1] - u[i-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[0, j:ny-1] - u[i-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[0, j:ny-1] - v[i-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )

        # First points for p. We don't know the pressure at i=0 and i=nx-1. We DO know the pressure at j=ny-1
        for m in range(niter):
            pm = np.copy(p)
            i=1
            j=1
            p[i:nx-1, j:ny-1] = 0.25*( pm[i+1:nx, j:ny-1] + pm[i-1:nx-2, j:ny-1] 
                                     + pm[i:nx-1, j+1:ny] + pm[i:nx-1, j-1:ny-2]
                                     - bm[i:nx-1, j:ny-1]*h**2.0 )        
            i=0
            # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
            p[i, j:ny-1] = 0.25*( pm[i+1, j:ny-1] + pm[nx-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )  
            i=nx-1
            # Periodic at i = nx-1 Replace all the i=i+1 with i=0
            p[i, j:ny-1] = 0.25*( pm[0, j:ny-1] + pm[i-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )

            # Set zero gradient boundary conditions
            p[:, 0] = p[:, 1]
            p[:, ny-1] = p[:, ny-2]
        
    return u, v, p, x, y, bm 


# In[96]:


u40, v40, p40, x40, y40, bm40 = poisson_equation(1, 0.1, 0.25)


# In[97]:


plot_3D(u40,x40,y40,'Figure 1: Final Conditions','u (m/s)')
plot_3D(v40,x40,y40,'Figure 2: Final Conditions','v (m/s)')
plot_3D(p40,x40,y40,'Figure 3: Final Conditions','p (Pa)')
plot_3D(bm40,x40,y40,'Figure 4: Final Conditions','b')


# ## Question: What is the effect of r on the magnitude of p?

# In[98]:


u1, v1, p1, x1, y1, bm1 = poisson_equation(1, 0.1, 0.25)
u2, v2, p2, x2, y2, bm2 = poisson_equation(1, 0.1, 0.50)
u3, v3, p3, x3, y3, bm3 = poisson_equation(1, 0.1, 1.00)
u4, v4, p4, x4, y4, bm4 = poisson_equation(1, 0.1, 1.50)
u5, v5, p5, x5, y5, bm5 = poisson_equation(1, 0.1, 2.00)


# In[99]:


plot_3D(p1,x1,y1,'Figure 1: Final Conditions','p (Pa)')
plot_3D(p2,x2,y2,'Figure 2: Final Conditions','p (Pa)')
plot_3D(p3,x3,y3,'Figure 3: Final Conditions','p (Pa)')
plot_3D(p4,x4,y4,'Figure 4: Final Conditions','p (Pa)')
plot_3D(p5,x5,y5,'Figure 5: Final Conditions','p (Pa)')


# ## Conclusion about the value of r in the equation for b

# r does affect the magnitude of p - therefore it is not a question of stability, but of dimensional consistency - the same r must be used so that the units of pressure are maintained.
# 
# In other words, pressure and velocity must step forward in time using the same timestep

# ## So, now this is established - what is the effect of r on the stability of the Burgers Equation for this problem?

# See whether starting the values at zero and then allowing convection would make a difference?

# In[108]:


def burgers_initialisation(niter, r, nx_or_ny, tmax, xmax_or_ymax, nu):
    """
    Initialise conditions for Poisson Equation
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int(((nu * tmax) / (r*(dx)**2))+1)
    dt = tmax/(nt-1)

    # Initialise data structures:
    import numpy as np
    
    # Initial conditions
    p = np.zeros((nx,ny))
    u = np.zeros((nx,ny))
    v = np.zeros((nx,ny))
    
    # x and y range
    x = np.zeros(nx)
    y = np.zeros(ny)

    # X Loop
    x = np.linspace(0.0,2.0,nx)
    
    # Y Loop
    y = np.linspace(0.0,2.0,ny)
            
    return p, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, r


# In[111]:


def burgers_equation(rho, nu, r):
                  
    (p, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, r) = burgers_initialisation(50, r, 51, 0.5, 2.0, nu)
    
    # Increments
    h = dx
    
    import numpy as np
    
    # Intermediate copies:
    un = np.zeros((nx, ny))
    vn = np.zeros((nx, ny))
        
    # Loop - use decimal points for all floating point numbers
    for n in range(nt):    
        
        # First points for b. We know the velocity at i=0, j=0, i=nx-1 and j=ny-1. b is zero at the boundaries.
        i=1
        j=1    
        
        # First points for u and v. We know the velocity at i=0, j=0, i=nx-1 and j=ny-1.
        # We are simply using the value of pressure here
        un = np.copy(u)
        vn = np.copy(v)
              
        u[i:nx-1, j:ny-1] = ( un[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i:nx-1, j-1:ny-2] ) ) + 
                        dt - (dt / (2.0 * rho * h)) * ( p[i+1:nx, j:ny-1] - p[i-1:nx-2, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1:nx-2, j:ny-1] + un[i+1:nx, j:ny-1] + 
                                                     un[i:nx-1, j-1:ny-2] + un[i:nx-1, j+1:ny] -
                                               4.0 * un[i:nx-1, j:ny-1] ) ) 

        v[i:nx-1, j:ny-1] = ( vn[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i:nx-1, j-1:ny-2] ) ) - 
                            (dt / (2.0 * rho * h)) * ( p[i:nx-1, j+1:ny] - p[i:nx-1, j-1:ny-2] ) + 
                            ((dt * nu) / (h**2.0)) * ( vn[i-1:nx-2, j:ny-1] + vn[i+1:nx, j:ny-1] + 
                                                     vn[i:nx-1, j-1:ny-2] + vn[i:nx-1, j+1:ny] -
                                               4.0 * vn[i:nx-1, j:ny-1] ) )          
    
        i=0
        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        
        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[nx-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                         dt - (dt / (2.0 * rho * h)) * ( p[i+1, j:ny-1] - p[nx-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[nx-1, j:ny-1] + un[i+1, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )  
        
        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[nx-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[i+1, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) ) 
        i=nx-1
        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[i-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                       dt - (dt / (2.0 * rho * h)) * ( p[0, j:ny-1] - p[i-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1, j:ny-1] + un[0, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )
        
        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[0, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) ) 

    return u, v, p, x, y


# In[115]:


u10, v10, p10, x10, y10 = burgers_equation(1, 0.1, 0.25)
u20, v20, p20, x20, y20 = burgers_equation(1, 0.1, 0.50)
u30, v30, p30, x30, y30 = burgers_equation(1, 0.1, 1.00)
u40, v40, p40, x40, y40 = burgers_equation(1, 0.1, 1.50)
u50, v50, p50, x50, y50 = burgers_equation(1, 0.1, 2.00)


# In[117]:


plot_3D(u10,x10,y10,'Figure 1: Final Conditions r = 0.25','u (m/s)')
plot_3D(u20,x20,y20,'Figure 2: Final Conditions r = 0.5','u (m/s)')
plot_3D(u30,x30,y30,'Figure 3: Final Conditions r = 1.0' ,'u (m/s)')
plot_3D(u40,x40,y40,'Figure 4: Final Conditions r = 1.5','u (m/s)')
plot_3D(u50,x50,y50,'Figure 5: Final Conditions r = 2.0','u (m/s)')


# As expected, if the value of r is greater than 0.5, then we have an unstable Burgers Equation

# ## Now try it with the full Navier Stokes Equations

# In[22]:


def navier_stokes_initialisation(niter, r, nx_or_ny, tmax, xmax_or_ymax, nu):
    """
    Initialise conditions for Poisson Equation
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int(((nu * tmax) / (r*(dx)**2))+1)
    dt = tmax/(nt-1)

    # Initialise data structures:
    import numpy as np
    
    # Initial conditions
    p = np.zeros((nx,ny))
    u = np.zeros((nx,ny))
    v = np.zeros((nx,ny))
    
    # x and y range
    x = np.zeros(nx)
    y = np.zeros(ny)

    # X Loop
    x = np.linspace(0.0,2.0,nx)
    
    # Y Loop
    y = np.linspace(0.0,2.0,ny)
            
    return p, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, r


# In[72]:


p, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, r = navier_stokes_initialisation(50, 0.5, 51, 0.5, 2.0, 0.1)


# In[74]:


dt


# In[124]:


def navier_stokes_equation(rho, nu, r):
                  
    (p, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, r) = navier_stokes_initialisation(50, r, 51, 0.5, 2.0, nu)
    
    # Increments
    h = dx
    
    import numpy as np
    
    # Intermediate copies:
    un = np.zeros((nx, ny))
    vn = np.zeros((nx, ny))
    pm = np.zeros((nx, ny))
    bm = np.zeros((nx, ny)) # bm needs to be exactly zero at the boundaries
        
    # Loop - use decimal points for all floating point numbers
    for n in range(nt):    
        
        # First points for bm. We don't know the value at i=0 and i=nx-1. b is zero at j=0 and j=ny-1.
        i=1
        j=1    
        bm[i:nx-1, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1] 
                                                        + v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1])**2.0 + 
                             4.0*h*(u[i:nx-1, j+1:ny] - u[i:nx-1, j-1:ny-2])*
                                   (v[i+1:nx, j:ny-1] - v[i-1:nx-2, j:ny-1]) + 
                                   (v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2])**2.0 ) )
        i=0
        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1, j:ny-1] - u[nx-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1, j:ny-1] - u[nx-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[i+1, j:ny-1] - v[nx-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )

        i=nx-1
        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[0, j:ny-1] - u[i-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[0, j:ny-1] - u[i-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[0, j:ny-1] - v[i-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )
        
        for m in range(niter):
            
            # First points for p. We don't know the pressure at i=0 and i=nx-1. dp/dy = 0 at y=0 and y=2
            pm = np.copy(p)
            i=1
            j=1
            p[i:nx-1, j:ny-1] = 0.25*( pm[i+1:nx, j:ny-1] + pm[i-1:nx-2, j:ny-1] 
                                     + pm[i:nx-1, j+1:ny] + pm[i:nx-1, j-1:ny-2]
                                     - bm[i:nx-1, j:ny-1]*h**2.0 )        
            i=0
            # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
            p[i, j:ny-1] = 0.25*( pm[i+1, j:ny-1] + pm[nx-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )  
            i=nx-1
            # Periodic at i = nx-1 Replace all the i=i+1 with i=0
            p[i, j:ny-1] = 0.25*( pm[0, j:ny-1] + pm[i-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )

            # Set zero gradient boundary conditions
            p[:, 0] = p[:, 1]
            p[:, ny-1] = p[:, ny-2]

        # First points for u and v. We don't know u and v at i=0 and i=nx-1. u and v are zero at j=0 and j=ny-1.
        i=1
        j=1    
        
        un = np.copy(u)
        vn = np.copy(v)
              
        u[i:nx-1, j:ny-1] = ( un[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i:nx-1, j-1:ny-2] ) ) + 
                        dt - (dt / (2.0 * rho * h)) * ( p[i+1:nx, j:ny-1] - p[i-1:nx-2, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1:nx-2, j:ny-1] + un[i+1:nx, j:ny-1] + 
                                                     un[i:nx-1, j-1:ny-2] + un[i:nx-1, j+1:ny] -
                                               4.0 * un[i:nx-1, j:ny-1] ) ) 

        v[i:nx-1, j:ny-1] = ( vn[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i:nx-1, j-1:ny-2] ) ) - 
                            (dt / (2.0 * rho * h)) * ( p[i:nx-1, j+1:ny] - p[i:nx-1, j-1:ny-2] ) + 
                            ((dt * nu) / (h**2.0)) * ( vn[i-1:nx-2, j:ny-1] + vn[i+1:nx, j:ny-1] + 
                                                     vn[i:nx-1, j-1:ny-2] + vn[i:nx-1, j+1:ny] -
                                               4.0 * vn[i:nx-1, j:ny-1] ) )          
        
        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        i=0
        
        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[nx-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                         dt - (dt / (2.0 * rho * h)) * ( p[i+1, j:ny-1] - p[nx-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[nx-1, j:ny-1] + un[i+1, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )  
        
        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[nx-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[i+1, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) )
        
        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        i=nx-1
        
        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[i-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                       dt - (dt / (2.0 * rho * h)) * ( p[0, j:ny-1] - p[i-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1, j:ny-1] + un[0, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )
        
        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[0, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) ) 

    return u, v, p, x, y


# In[125]:


u50, v50, p50, x50, y50 = navier_stokes_equation(1.0, 0.1, 0.5)


# In[126]:


plot_3D(u50,x50,y50,'Figure 1: Final Conditions','u (m/s)')
plot_3D(v50,x50,y50,'Figure 2: Final Conditions','v (m/s)')
plot_3D(p50,x50,y50,'Figure 3: Final Conditions','p (Pa)')


# ## How does this compare with the analytical solution?

# We know the answer to this problem, it's just flow between parallel plates
# 
# $$ {\partial u \over \partial t} + u {\partial u \over \partial x} + v {\partial u \over \partial y} = 1 -{1 \over \rho} {{\partial p} \over {\partial x}} + \nu \left ( {\partial^2 u \over \partial x^2}+  {\partial^2 u \over \partial y^2}  \right ) $$
# 
# In steady state:
# 
# * $ {\partial u / \partial t} = 0 $
# * $ {\partial u / \partial x} = 0 $
# * $ v = 0 $
# * $ {\partial p / \partial x} = 0 $
# * $ {\partial^2 u / \partial x^2} = 0 $
# 
# Hence:
# 
# $$ 1 + \nu {\partial^2 u \over \partial y^2} = 0 \qquad \Rightarrow \qquad {\partial^2 u \over \partial y^2} = -{1 \over \nu} \qquad \Rightarrow \qquad \int \int {d^2 u \over d y^2} dy dy = -{1 \over \nu} \int \int dy dy$$
# 
# $$ \int {d u \over d y} dy = -{1 \over \nu} \int (y + C_1 ) dy \qquad \Rightarrow \qquad u = -{1 \over \nu} \left({y^2 \over 2} + C_1y +C_2 \right) $$
# 
# * $u = 0 \text{ at } y = 0 \qquad \Rightarrow  \qquad C_2 = 0$
# * $u = 0 \text{ at } y = 2 \qquad \Rightarrow  \qquad C_2 = - 1$
# 
# $$ u = -{1 \over \nu} \left({y^2 \over 2} - y \right) = {1 \over \nu} \left(y - {y^2 \over 2} \right)$$

# In[7]:


def parallel_plates_analytical(YMAX, NY, NU):
    
    # Initialise data structures
    u_analytical = np.zeros(NY)
    y = np.zeros(NY)
    
    # Constants
    DY = YMAX/(NY-1)

    # Y Loop
    for j in range(0,NY):
        y[j] = j*DY
    
    # Analytical solution  
    for j in range(0, NY):
        u_analytical[j]= (1.0 / NU) * (y[j] - (y[j]**2.0 / 2.0))
        
    return u_analytical, y


# In[8]:


u_analytic, y_analytic = parallel_plates_analytical(2.0, 51, 0.1)


# In[59]:


def plot_diffusion(u1,u2,y1,y2,NX):
   """
   Plots the 1D velocity field
   """

   import matplotlib.pyplot as plt
   import matplotlib.cm as cm
   plt.figure()
   ax=plt.subplot(111)
   ax.plot(y1,u1[NX-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical')
   ax.plot(y2,u2[:],linestyle='-',c='r',label='Analytical')
   box=ax.get_position()
   ax.set_position([box.x0, box.y0, box.width*1.5,box.height*1.5])
   ax.legend( bbox_to_anchor=(1.02,1), loc=2)
   plt.xlabel('y (m)')
   plt.ylabel('u (m/s)')
   plt.show()


# In[150]:


plot_diffusion(u50,u_analytic,y_analytic,51)


# ## Why Doesn't the Numerical Solution Agree with the Analytical?

# * Where the gradient is shallow, the values are not changing very much. There I suspect convergence is loose
# * I'm not sure at this stage why the magnitude is about one tenth - could be a order of magnitude due to misplacement of the kinematic viscosity (which is 0.1)
# 
# * **Instead of advancing in time, look for the steady state solution when the streamwise velocity does not change by more than** $10^{-3}$

# In[95]:


def navier_stokes_equation_accurate(niter, nx, rho, nu, r, tmax):
                  
    (p, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, r) = navier_stokes_initialisation(niter, r, nx, tmax, 2.0, nu)
    
    # Increments
    h = dx
    
    import numpy as np
    
    # Intermediate copies:
    un = np.zeros((nx, ny))
    vn = np.zeros((nx, ny))
    pm = np.zeros((nx, ny))
    bm = np.zeros((nx, ny)) # bm needs to be exactly zero at the boundaries
        
    # Loop - use decimal points for all floating point numbers
    #for n in range(nt): 
    
    udiff = 1.0
    
    while udiff > 0.001:   
        
        # First points for bm. We don't know the value at i=0 and i=nx-1. b is zero at j=0 and j=ny-1.
        i=1
        j=1    
        bm[i:nx-1, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1] 
                                                        + v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1:nx, j:ny-1] - u[i-1:nx-2, j:ny-1])**2.0 + 
                             4.0*h*(u[i:nx-1, j+1:ny] - u[i:nx-1, j-1:ny-2])*
                                   (v[i+1:nx, j:ny-1] - v[i-1:nx-2, j:ny-1]) + 
                                   (v[i:nx-1, j+1:ny] - v[i:nx-1, j-1:ny-2])**2.0 ) )
        i=0
        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[i+1, j:ny-1] - u[nx-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[i+1, j:ny-1] - u[nx-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[i+1, j:ny-1] - v[nx-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )

        i=nx-1
        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        bm[i, j:ny-1] = ( (rho / (2.0 * h * dt)) * ( u[0, j:ny-1] - u[i-1, j:ny-1] 
                                                        + v[i, j+1:ny] - v[i, j-1:ny-2] ) +
            (rho / (4.0*h**2)) * ( (u[0, j:ny-1] - u[i-1, j:ny-1])**2.0 + 
                             4.0*h*(u[i, j+1:ny] - u[i, j-1:ny-2])*
                                   (v[0, j:ny-1] - v[i-1, j:ny-1]) + 
                                   (v[i, j+1:ny] - v[i, j-1:ny-2])**2.0 ) )
        
        for m in range(niter):
            
            # First points for p. We don't know the pressure at i=0 and i=nx-1. dp/dy = 0 at y=0 and y=2
            pm = np.copy(p)
            i=1
            j=1
            p[i:nx-1, j:ny-1] = 0.25*( pm[i+1:nx, j:ny-1] + pm[i-1:nx-2, j:ny-1] 
                                     + pm[i:nx-1, j+1:ny] + pm[i:nx-1, j-1:ny-2]
                                     - bm[i:nx-1, j:ny-1]*h**2.0 )        
            i=0
            # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
            p[i, j:ny-1] = 0.25*( pm[i+1, j:ny-1] + pm[nx-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )  
            i=nx-1
            # Periodic at i = nx-1 Replace all the i=i+1 with i=0
            p[i, j:ny-1] = 0.25*( pm[0, j:ny-1] + pm[i-1, j:ny-1] 
                                     + pm[i, j+1:ny] + pm[i, j-1:ny-2]
                                     - bm[i, j:ny-1]*h**2.0 )

            # Set zero gradient boundary conditions
            p[:, 0] = p[:, 1]
            p[:, ny-1] = p[:, ny-2]

        # First points for u and v. We don't know u and v at i=0 and i=nx-1. u and v are zero at j=0 and j=ny-1.
        i=1
        j=1    

        un = np.copy(u)
        vn = np.copy(v)

        u[i:nx-1, j:ny-1] = ( un[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( un[i:nx-1, j:ny-1] - un[i:nx-1, j-1:ny-2] ) ) + 
                        dt - (dt / (2.0 * rho * h)) * ( p[i+1:nx, j:ny-1] - p[i-1:nx-2, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1:nx-2, j:ny-1] + un[i+1:nx, j:ny-1] + 
                                                     un[i:nx-1, j-1:ny-2] + un[i:nx-1, j+1:ny] -
                                               4.0 * un[i:nx-1, j:ny-1] ) ) 

        v[i:nx-1, j:ny-1] = ( vn[i:nx-1, j:ny-1] - 
                 (dt / h) * ( un[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i-1:nx-2, j:ny-1] ) + 
                              vn[i:nx-1, j:ny-1] * ( vn[i:nx-1, j:ny-1] - vn[i:nx-1, j-1:ny-2] ) ) - 
                            (dt / (2.0 * rho * h)) * ( p[i:nx-1, j+1:ny] - p[i:nx-1, j-1:ny-2] ) + 
                            ((dt * nu) / (h**2.0)) * ( vn[i-1:nx-2, j:ny-1] + vn[i+1:nx, j:ny-1] + 
                                                     vn[i:nx-1, j-1:ny-2] + vn[i:nx-1, j+1:ny] -
                                               4.0 * vn[i:nx-1, j:ny-1] ) )          

        # Periodic at i = 0 Replace all the i=i-1 with i=nx-1
        i=0

        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[nx-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                         dt - (dt / (2.0 * rho * h)) * ( p[i+1, j:ny-1] - p[nx-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[nx-1, j:ny-1] + un[i+1, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )  

        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[nx-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[i+1, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) )

        # Periodic at i = nx-1 Replace all the i=i+1 with i=0
        i=nx-1

        u[i, j:ny-1] = ( un[i, j:ny-1] - 
                 (dt / h) * ( un[i, j:ny-1] * ( un[i, j:ny-1] - un[i-1, j:ny-1] ) + 
                              vn[i, j:ny-1] * ( un[i, j:ny-1] - un[i, j-1:ny-2] ) ) + 
                       dt - (dt / (2.0 * rho * h)) * ( p[0, j:ny-1] - p[i-1, j:ny-1] ) + 
                            ((dt * nu) / (h**2.0)) * ( un[i-1, j:ny-1] + un[0, j:ny-1] + 
                                                     un[i, j-1:ny-2] + un[i, j+1:ny] -
                                               4.0 * un[i, j:ny-1] ) )

        v[i, j:ny-1] = ( vn[i, j:ny-1] - 
         (dt / h) * ( un[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i-1, j:ny-1] ) + 
                      vn[i, j:ny-1] * ( vn[i, j:ny-1] - vn[i, j-1:ny-2] ) ) - 
                    (dt / (2.0 * rho * h)) * ( p[i, j+1:ny] - p[i, j-1:ny-2] ) + 
                    ((dt * nu) / (h**2.0)) * ( vn[i-1, j:ny-1] + vn[0, j:ny-1] + 
                                             vn[i, j-1:ny-2] + vn[i, j+1:ny] -
                                       4.0 * vn[i, j:ny-1] ) ) 
        
        udiff = (np.sum(u)-np.sum(un))/np.sum(u)

    return u, v, p, x, y


# In[78]:


u60, v60, p60, x60, y60 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.5, 2.0)


# In[76]:


np.max(p60)


# In[70]:


plot_diffusion(u60,u_analytic,y60,y_analytic,51)


# ## Observations about accuracy
# 
# * Iterating until the velocity doesn't change in steady state has improved the prediction
# 
# ## Observations about stability
# 
# * We are just on the limit of stability here with r=0.5.
# 
# ## Observations about maximum velocity
# 
# * However, there may now be a dependency of the maximum velocity on the value of dt because:
# 
#   - `dt x F` is the source term (like a short pulse)
#   
# ## Observations about maximum velocity
# 
# * We need to know what value of dt is going to produce the the correct pressure gradient and so the correct maximum velocity
# 
# * We need to look at the discretised equation in steady state

# $$ u_{i,j}^{n+1} = u_{i,j}^n - {{\Delta t} \over h} \left[ u_{i,j}^n(u_{i,j}^n - u_{i-1,j}^n) + v_{i,j}^n(u_{i,j}^n - u_{i,j-1}^n)\right] \\
# + {\Delta t} - {{\Delta t} \over {2 \rho h}} (p_{i+1,j}^n - p_{i-1,j}^n) + {{\Delta t \nu} \over {h^2}}
# (u_{i-1,j}^n + u_{i+1,j}^n + u_{i,j-1}^n + u_{i,j+1}^n - 4 u_{i,j}^n )
# $$

# ## What is the stable value of r?

# In[115]:


u02, v02, p02, x02, y02 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.2, 2.0)
u03, v03, p03, x03, y03 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.3, 2.0)
u04, v04, p04, x04, y04 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.4, 2.0)
u05, v05, p05, x05, y05 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.5, 2.0)
u06, v06, p06, x06, y06 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.6, 2.0)


# In[119]:


def plot_diffusion(u1,u2,u3,u4,u5,y1,y2,y3,y4,y5,NX):
   """
   Plots the 1D velocity field
   """

   import matplotlib.pyplot as plt
   import matplotlib.cm as cm
   plt.figure()
   ax=plt.subplot(111)
   ax.plot(y1,u1[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical r=0.2')
   ax.plot(y2,u2[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical r=0.3')
   ax.plot(y3,u3[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical r=0.4')
   ax.plot(y4,u4[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical r=0.5')
   ax.plot(y5,u5[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical r=0.6') 
   box=ax.get_position()
   ax.set_position([box.x0, box.y0, box.width*1.5,box.height*1.5])
   ax.legend( bbox_to_anchor=(1.02,1), loc=2)
   plt.xlabel('y (m)')
   plt.ylabel('u (m/s)')
   plt.show()


# In[120]:


plot_diffusion(u02,u03,u04,u05,u06,y02,y03,y04,y05,y06,51)


# Clearly r=0.5 is still the limit of stability, and notice that r affects the maximum velocity (via it's influence on dt)

# For the same r, reducing nx will increase dx, reduce nt and increase dt

# In[118]:


u002, v002, p002, x002, y002 = navier_stokes_equation_accurate(50, 51, 1.0, 0.1, 0.5, 2.0)
u003, v003, p003, x003, y003 = navier_stokes_equation_accurate(50, 41, 1.0, 0.1, 0.5, 2.0)
u004, v004, p004, x004, y004 = navier_stokes_equation_accurate(50, 31, 1.0, 0.1, 0.5, 2.0)
u005, v005, p005, x005, y005 = navier_stokes_equation_accurate(50, 21, 1.0, 0.1, 0.5, 2.0)
u006, v006, p006, x006, y006 = navier_stokes_equation_accurate(50, 11, 1.0, 0.1, 0.5, 2.0)


# In[122]:


def plot_diffusion_2(u1,u2,u3,u4,u5,y1,y2,y3,y4,y5,NX):
   """
   Plots the 1D velocity field
   """

   import matplotlib.pyplot as plt
   import matplotlib.cm as cm
   plt.figure()
   ax=plt.subplot(111)
   ax.plot(y1,u1[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical nx=51')
   ax.plot(y2,u2[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical nx=41')
   ax.plot(y3,u3[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical nx=31')
   ax.plot(y4,u4[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical nx=21')
   ax.plot(y5,u5[-1,:],'-', markerfacecolor='none', alpha=0.5, label='Numerical nx=11') 
   box=ax.get_position()
   ax.set_position([box.x0, box.y0, box.width*1.5,box.height*1.5])
   ax.legend( bbox_to_anchor=(1.02,1), loc=2)
   plt.xlabel('y (m)')
   plt.ylabel('u (m/s)')
   plt.show()


# In[123]:


plot_diffusion_2(u002,u003,u004,u005,u006,y002,y003,y004,y005,y006,51)


# Clearly there is also an effect or dx on stability and the maximum occurs when nx=31

# **Bascially we are using an explicit method here, which is not ideal, we keep reaching the limits of stability**
# 
# * It is possible that the SOR method could be used to solve the Poisson Equation

# This makes it difficult to compare the values between the numerical and analytical solutions - the source term is being included like a pulse in time, so this might require much more stable methods

# ## How is the source term included?
# 
# * The source term is split into steady and unsteady components:
# 
# $$ p = P + p' $$

# * The applied pressure gradient is:
#     
# $$ -{{\partial P} \over {\partial x}} = 1 $$

# * So the solution in steady state is where the fluctuating pressure is:
#     
# $$ {{\partial p'} \over {\partial x}} = 0 $$

# ### Why was this done?
# 
# * We used periodic boundary conditions. 
# * **Without the source term** this means the pressure on the left and right would have been the same and no flow would have occured
# * **With the source term** the pressure in all the equations is the fluctuating pressure, which in steady state settles out to zero

# ### When is the fluctuating pressure non-zero?
# 
# * In turbulent flow the fluctuating pressure is non-zero

# In[ ]: