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

# In[1]:


from IPython.display import HTML 
s="""<h1>2D Navier-Stokes for Flow in a Cavity (Loop Method)</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 over a cavity?
# 
# ## Initial Conditions
# 
# * The velocity and pressure fields are zero everywhere
# 
# ## Boundary Conditions
# 
# * The u-boundary condition at y=2m is 1m/s (the lid) 
# 
# * The other velocity boundary conditions are zero (no slip). 
# 
# * The p-boundary condition at y=2m is 0Pa (atmospheric pressure)
# 
# * The gradient of pressure at y=0m, x=0m and x=2m is zero (no-slip)
# 
# ## 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} = -{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 ) $$
# 
# $$ {\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:
# 
# $$ {{{\partial^2 p} \over {\partial x^2}} + {{\partial^2 p} \over {\partial y^2}} = - \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 equations applies in the continous domain, so it may not ensure continuity 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 (n, x, y) \quad n = 0 \rightarrow u = 0 \land v = 0 \land p = 0$
# 
# ## Velocity Boundary Conditions:
# 
# * $\forall (n, x, y) \quad y = 2 \rightarrow u = 1$
# 
# * $\forall (n, x, y) \quad x = 0 \lor x = 2 \lor y = 0  \rightarrow u = 0$
# 
# * $\forall (n, x, y) \quad x = 0 \lor x = 2 \lor y = 0 \lor y = 2  \rightarrow v = 0$
# 
# ## Pressure Boundary Conditions:
# 
# * $\forall (n, x, y) \quad y = 2 \rightarrow p = 0$
# 
# * $\forall (n, x, y) \quad y = 0 \rightarrow {{\partial p} \over {\partial y}} = 0$
# 
# * $\forall (n, x, y) \quad x = 0 \lor x = 2 \rightarrow {{\partial p} \over {\partial x}} = 0$
# 
# ## Output Data:
# 
# * $\forall (n,x,y) \quad \ p = ? \land u = ? \land v = ?$
# 
# ***
# 
# # 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 \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} \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
# 
# 
# Since no terms have a 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} + \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) $$
# 
# 
# 
# 
# 
# $$ {{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} + \\
# \rho \left[ \left( {{u_{i+1,j}^m - u_{i-1,j}^m} \over {2 \Delta x}} \right)^2 + 2 \left( {{u_{i,j+1}^m - u_{i,j-1}^m} \over {2 \Delta y}} {{v_{i+1,j}^m - v_{i-1,j}^m} \over {2 \Delta x}}\right) +
# \left( {{v_{i,j+1}^m - v_{i,j-1}^m} \over {2 \Delta y}} \right)^2
# \right] $$
# 
# ## Transpose
# 
# Also assume that $ \Delta x = \Delta y = h $
# 
# $$ p_{i,j}^{m+1} = p_{i,j}^m + 
# r \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 + \\
# {\rho \over 4} \left( (u_{i+1,j}^m - u_{i-1,j}^m)^2 + 
# 4 h (u_{i,j+1}^m - u_{i,j-1}^m)(v_{i+1,j}^m - v_{i-1,j}^m) + 
# (v_{i,j+1}^m - v_{i,j-1}^m)^2 \right) \right]$$
# 
# where: $ r = {{\Delta \tau} \over {h^2}} $
# 
# We don't know what r is for the fastest convergence, but $ r=0.25$ has worked before
# 

# ## Pseudo-code
# 
#     # Inputs:
#     
#     niter = 51 
#     nx =  21
#     xmax = 2
#     ymax = 2
#     nt = 51
#     tmax = 0.5
#     r = 0.25
#     
#     # Calculated increments:
#     
#     dt = tmax/(nt-1)
#    
#     # 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] = 0
#         pm[:,:,0] = 0
#         u[:,:,0] = 0
#         v[:,:,0] = 0
# 
#     # Velocity Boundary Conditions:
#     
#         u[:,ny-1,:] = 1
#         u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = 0
#         v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 0
#         
#     # Pressure Boundary Conditions:
#     
#         p[:,ny-1,:] = 0
#         pm[:,ny-1,:] = 0
#         
#     # Values for correction at boundaries:
#     
#     for i between 0 and nx-1
#         ipos[i] = i + 1
#         ineg[i] = i - 1
#     
#     for j between 0 and ny-1
#         jneg[j] = j - 1
#        
#     # Set Reflection at i = 0, nx-1 and j = 0: 
#     
#         ineg[0] = 1          # i.e. index -1 = 1
#         ipos[nx-1] = nx-2    # i.e. index nx = nx-2
#         jneg[0] = 1          # i.e. index -1 = 1
#        
#     # Numerical Computation
#     
#       for n between 0 and nt-1
#       
#         # Poisson Equation - Function 1 - test with Laplace Conditions first
#         
#         pm = poisson_equation(u,v,pm,r)
#         
#         for m between 0 and niter-2
#         
#           # Poisson Equation - Function 1a: 
#           # b = poisson_uv(u, v)
#         
#           for i between 1 and nx-2
#             for j between 1 and ny-2
#         
#               b[i,j,m] = 0.25*rho*{ (u[i+1,j,n] - u[i-1,j,n])^2 + 
#                                  4*h*(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 }   
#             end
#           end 
#           
#           # Poisson Equation - Function 1b: 
#           # pm = poisson_p(pm, r, b)
#           
#           for i between 0 and nx-1
#             for j between 0 and ny-2
#          
#               pm[i,j,m+1] = pm[i,j,m] + 
#                         r*{ pm[ipos[i],j,m] + pm[ineg[i],j,m] + pm[i,j+1,m] + pm[i,jneg[j],m] - 4 * pm[i,j,m] + b[i,j,m] }
#                         
#             end
#           end
#           
#         end -- end of iteration loop
#         
#         # Equate Pressure - Function 2 - test to see if final result is stored: 
#         # p, pu, pv = pressure(pm, p, h, rho, dt)
#         
#         for i between 0 and nx-1
#           for j between 0 and ny-2
#           
#             p[i,j,n] = pm[i,j,niter-1] 
#      
#             pu[i,j,n] = {dt / (2 * rho * h)} * ( p[ipos[i],j,n] - p[ineg[i],j,n] )
#         
#             pv[i,j,n] = {dt / (2 * rho * h)} * ( p[i,j+1,n] - p[i,jneg[j],n] )
#             
#           end
#         end
#         
#         # Momentum Equation - Function 3 - test as Burgers Equation first (zero pressure term): 
#         # p, u, v = momentum(p, u, v, nu, h, dt, pu, pv)
#                         
#         for i between 1 and nx-2
#           for j between 1 and ny-2
#             
#               p[i,j,n] = p[i,j,niter-1] 
#               
#               u[i,j,n+1] = u[i,j,n] - (dt / h) * [ 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] ) - 
#                                                    pu[i,j,n] + 
#                    {(dt * nu) / (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] ) ] 
#                    
#               v[i,j,n+1] = v[i,j,n] - (dt / h) * [ 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] ) - 
#                                                    pv[i,j,n] + 
#                    {(dt * nu) / (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] ) ]      
#        
#             end
#           end
#           
#         end  -- end of time loop
#       
# ***
# 
# # Implement algorithm in Python
# 
# I initially tried this with the laplace equation from the previous case
# 

# ## Simplified Poisson Equation Computation (Actually Laplace's Equation)

# In[4]:


def initialisation2(niter, r, nx_or_ny, tmax, xmax_or_ymax):
    """
    Returns the velocity field and distance for 2D linear convection
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int((tmax / (r*(dx)**2))+1)
    dt = tmax/(nt-1)
    
    # Initialise data structures:
    import numpy as np
    p = np.zeros(((nx,ny,nt)))
    pm = np.zeros(((nx,ny,niter)))
    u = np.zeros(((nx,ny,nt)))
    v = np.zeros(((nx,ny,nt)))
    x = np.zeros(nx)
    y = np.zeros(ny)
    ipos = np.zeros(nx)
    ineg = np.zeros(nx)
    jpos = np.zeros(ny)
    jneg = np.zeros(ny)

    # X Loop
    for i in range(0,nx):
        x[i] = i*dx

    # Y Loop
    for j in range(0,ny):
        y[j] = j*dy

    # Initial conditions
    #p[:,:,0] = pm[:,:,0] = 0

    # Pressure Boundary Conditions:
    
    # p boundary, top:
    #p[:,ny-1,:] = pm[:,ny-1,:] = 1

    # Boundary conditions
    u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = u[:,ny-1,:] = 1
    v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 1

    # Initial conditions
    u[:,:,0] = v[:,:,0] = 1
    u[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2
    v[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2
        
    # Boundary conditions
    p[0,:,:] = p[nx-1,:,:] = p[:,0,:] = p[:,ny-1,:] = 1
    pm[0,:,:] = pm[nx-1,:,:] = pm[:,0,:] = pm[:,ny-1,:] = 1
    
    # Initial conditions
    p[:,:,0] = 1
    p[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2   
        
    # von Neumann Boundary Conditions:

    # Values for correction at boundaries:
    for i in range(0,nx-1):
        ipos[i] = i + 1
        ineg[i] = i - 1

    for j in range(0,ny-1):
        jpos[j] = j + 1
        jneg[j] = j - 1
        
    # Set Reflection: 
    
    ineg[0] = 1           # i.e. index -1 = 1
    ipos[nx-1] = nx-2     # i.e. index nx = nx-2
    jneg[0] = 1           # i.e. index -1 = 1
    jpos[ny-1] = ny-2     # i.e. index ny = ny-2
    
        
    return p, pm, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, ipos, ineg, jpos, jneg, r


# In[1]:


def initialisation(niter, r, nx_or_ny, tmax, xmax_or_ymax):
    """
    Returns the velocity field and distance for 2D linear convection
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int((tmax / (r*(dx)**2))+1)
    dt = tmax/(nt-1)
    
    
    # Initialise data structures:
    import numpy as np
    p = np.zeros(((nx,ny,nt)))
    pm = np.zeros(((nx,ny,niter)))
    u = np.zeros(((nx,ny,nt)))
    v = np.zeros(((nx,ny,nt)))
    x = np.zeros(nx)
    y = np.zeros(ny)
    ipos = np.zeros(nx)
    ineg = np.zeros(nx)
    jpos = np.zeros(ny)
    jneg = np.zeros(ny)

    # X Loop
    for i in range(0,nx):
        x[i] = i*dx

    # Y Loop
    for j in range(0,ny):
        y[j] = j*dy

    # Initial conditions
    p[:,:,0] = u[:,:,0] = v[:,:,0] = 0

    # Pressure Boundary Conditions:
    
    # p boundary, top:
    p[:,ny-1,:] = pm[:,ny-1,:] = 1
        
    # Velocity Boundary Conditions:
        
    u[:,ny-1,:] = 1
    u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = 0
    v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 0
        
    # von Neumann Boundary Conditions:

    # Values for correction at boundaries:
    for i in range(0,nx-1):
        ipos[i] = i + 1
        ineg[i] = i - 1

    for j in range(0,ny-1):
        jpos[j] = j + 1
        jneg[j] = j - 1
        
    # Set Reflection: 
    
    ineg[0] = 1           # i.e. index -1 = 1
    ipos[nx-1] = nx-2     # i.e. index nx = nx-2
    jneg[0] = 1           # i.e. index -1 = 1
    jpos[ny-1] = ny-2     # i.e. index ny = ny-2
    
        
    return p, pm, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, ipos, ineg, jpos, jneg, r


# In[2]:


(p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation(51, 0.5, 51, 0.5, 2.0)


# In[3]:


nt


# In[78]:


(p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation2(11, 0.25, 11, 0.5, 2.0)


# In[79]:


nt


# In[30]:


def poisson_equation_numerical(rho, n):

    """
    Returns the velocity field and distance for 2D linear convection
    """
    
    (p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation(11, 0.25, 51, 0.5, 2.0)
    
    b = np.zeros((nx,ny))
    h = dx

        # We must assume that the velocity field for the Poisson Equation uses the same boundary
        # conditions as the pressure field in order for the equation to make sense - i.e. for us
        # to have the same spatial range across all parts of the equation
        
    for i in range(0,nx):
        for j in range(0,ny-1):
            b[i,j] = 0.25*rho*( (u[ipos[i],j,n] - u[ineg[i],j,n])**2 + 
                     4.0*h*(u[i,j+1,n] - u[i,jneg[j],n])*(v[ipos[i],j,n] - v[ineg[i],j,n]) + 
                     (v[i,j+1,n] - v[i,jneg[j],n])**2 )
            
    # At the start of the iterations - the pressure is from the previous timestep        
    pm[:,:,0] = p[:,:,n] 
    
    for m in range(0,niter-1):
        for i in range(0,nx):
            for j in range(0,ny-1):
                pm[i,j,m+1] = (pm[i,j,m] + 
                              r*( pm[ipos[i],j,m] + pm[ineg[i],j,m] 
                              + pm[i,j+1,m] + pm[i,jneg[j],m] - 4 * pm[i,j,m] ))#+ b[i,j] ))
                #pm[i,j] = (pm[i,j] + 
                #              r*( pm[ipos[i],j] + pm[ineg[i],j] 
                #              + pm[i,j+1] + pm[i,jneg[j]] - 4 * pm[i,j] + b[i,j] ))
    # At the end of the iterations - the pressure is from the end of the iterations            
    p[:,:,n+1] = pm[:,:,niter-1]   
    
    return p, x, y


# In[31]:


p2, x2, y2 = poisson_equation_numerical(1, 0)


# In[24]:


def plot_3D(p,x,y,time,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,p[:,:,time], rstride=1, cstride=1)
    plt.title(title)
    plt.show()


# In[32]:


plot_3D(p2,x2,y2,0,'Figure 1: Initial Conditions: n=0, nx=51','p (Pa)')
plot_3D(p2,x2,y2,1,'Figure 2: Final Conditions: n=1, nx=51','p (Pa)')


# ## Simplified Momentum Equation (Actually Burgers Equation)

# In[40]:


def burgers_equation(rho, nu):
    
    (p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation2(11, 0.5, 51, 0.5, 2.0)

    # Increments
    h = dx
    
    # Initialise data structures
    pu = np.zeros((nx,ny))
    pv = np.zeros((nx,ny))   
        
    # Loop
    for n in range(0,nt-1):
        
        #p, x, y = poisson_equation_numerical(rho, n)
        
        for i in range(0,nx):
            for j in range(0,ny-1):
                pu[i,j] = (dt / (2 * rho * h)) * ( p[ipos[i],j,n] - p[ineg[i],j,n] )
                pv[i,j] = (dt / (2 * rho * h)) * ( p[i,j+1,n] - p[i,jneg[j],n] )
    
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                u[i,j,n+1] = ( u[i,j,n] - (dt / h) * ( 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] ) - 
                                                     pu[i,j] + 
                ((dt * nu) / (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] ) ) ) 
                v[i,j,n+1] = ( v[i,j,n] - (dt / h) * ( 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] ) - 
                                                     pv[i,j] + 
                ((dt * nu) / (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] ) ) )         
                
    return u, v, p, x, y 


# In[41]:


u10, v10, p10, x10, y10 = burgers_equation(1, 0.1)


# In[9]:


def plot_3D_2(p,x,y,time,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.view_init(30,225)
    Y,X=np.meshgrid(y,x) #note meshgrid uses y,x not x,y!!!
    surf=ax.plot_surface(X,Y,p[:,:,time], rstride=1, cstride=1)
    plt.title(title)
    plt.show()


# In[61]:


plot_3D_2(u10,x10,y10,0,'Figure 1: Initial Conditions: n=0, nx=51','u (m/s)')
plot_3D_2(u10,x10,y10,100,'Figure 3: Initial Conditions: n=0, nx=51','v (m/s)')
plot_3D_2(u10,x10,y10,200,'Figure 4: Final Conditions: n=310, nx=51','v (m/s)')
plot_3D_2(u10,x10,y10,300,'Figure 2: Final Conditions: n=310, nx=51','u (m/s)')
plot_3D_2(u10,x10,y10,400,'Figure 4: Final Conditions: n=310, nx=51','v (m/s)')
plot_3D_2(u10,x10,y10,500,'Figure 2: Final Conditions: n=310, nx=51','u (m/s)')
plot_3D_2(v10,x10,y10,0,'Figure 3: Initial Conditions: n=0, nx=51','v (m/s)')
plot_3D_2(v10,x10,y10,500,'Figure 4: Final Conditions: n=310, nx=51','v (m/s)')


# ## Navier Stokes Equations - Square Channel Boundary Conditions

# In[25]:


def initialisation3(niter, r, nx_or_ny, tmax, xmax_or_ymax):
    """
    Returns the velocity field and distance for 2D linear convection
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int((tmax / (r*(dx)**2))+1)
    dt = tmax/(nt-1)
    
    # Initialise data structures:
    import numpy as np
    p = np.zeros(((nx,ny,nt)))
    pm = np.zeros(((nx,ny,niter)))
    u = np.zeros(((nx,ny,nt)))
    v = np.zeros(((nx,ny,nt)))
    x = np.zeros(nx)
    y = np.zeros(ny)
    ipos = np.zeros(nx)
    ineg = np.zeros(nx)
    jpos = np.zeros(ny)
    jneg = np.zeros(ny)

    # X Loop
    for i in range(0,nx):
        x[i] = i*dx

    # Y Loop
    for j in range(0,ny):
        y[j] = j*dy

    # Initial conditions
    #p[:,:,0] = pm[:,:,0] = 0

    # Pressure Boundary Conditions:
    
    # p boundary, top:
    #p[:,ny-1,:] = pm[:,ny-1,:] = 1

    # Boundary conditions
    u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = u[:,ny-1,:] = 1
    v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 1

    # Initial conditions
    u[:,:,0] = v[:,:,0] = 1
    u[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2
    v[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2
        
    # Initial conditions
    p[:,:,0] = 0

    # Pressure Boundary Conditions:
    
    # p boundary, top:
    p[:,ny-1,:] = pm[:,ny-1,:] = 1
        
    # Boundary conditions
   # p[0,:,:] = p[nx-1,:,:] = p[:,0,:] = p[:,ny-1,:] = 1
   # pm[0,:,:] = pm[nx-1,:,:] = pm[:,0,:] = pm[:,ny-1,:] = 1
    
    # Initial conditions
   # p[:,:,0] = 1
   # p[(nx-1)/4:(nx-1)/2,(ny-1)/4:(ny-1)/2,0] = 2   
        
    # von Neumann Boundary Conditions:

    # Values for correction at boundaries:
    for i in range(0,nx-1):
        ipos[i] = i + 1
        ineg[i] = i - 1

    for j in range(0,ny-1):
        jpos[j] = j + 1
        jneg[j] = j - 1
        
    # Set Reflection: 
    
    ineg[0] = 1           # i.e. index -1 = 1
    ipos[nx-1] = nx-2     # i.e. index nx = nx-2
    jneg[0] = 1           # i.e. index -1 = 1
    jpos[ny-1] = ny-2     # i.e. index ny = ny-2
    
        
    return p, pm, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, ipos, ineg, jpos, jneg, r


# In[30]:


def navier_stokes_equation(rho, nu):
    
   # (p, pm, x, y, u, v, nx, ny, nt, 
  #  dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation2(11, 0.25, 11, 0.5, 2.0)
    (p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation3(11, 0.5, 11, 0.5, 2.0)
    # Increments
    h = dx
    r_factor = 0.25
    
    # Initialise data structures
    b = np.zeros((nx,ny))
    pu = np.zeros((nx,ny))
    pv = np.zeros((nx,ny))   
        
    # Loop
    for n in range(0,nt-1):
   
        # Poisson's Equation (Iteration)
        for i in range(0,nx):
            for j in range(0,ny-1):
                b[i,j] = 0.25*rho*( (u[ipos[i],j,n] - u[ineg[i],j,n])**2 + 
                    4.0*h*(u[i,j+1,n] - u[i,jneg[j],n])*(v[ipos[i],j,n] - v[ineg[i],j,n]) + 
                    (v[i,j+1,n] - v[i,jneg[j],n])**2 )  
        
        # At the start of the iterations - the pressure is from the previous timestep        
        pm[:,:,0] = p[:,:,n] 

        for m in range(0,niter-1):
            for i in range(0,nx):
                for j in range(0,ny-1):
                    pm[i,j,m+1] = (pm[i,j,m] + 
                                  r_factor*( pm[ipos[i],j,m] + pm[ineg[i],j,m] 
                                  + pm[i,j+1,m] + pm[i,jneg[j],m] - 4 * pm[i,j,m] ))#+ b[i,j] ))

        # At the end of the iterations - the pressure is from the end of the iterations            
        p[:,:,n+1] = pm[:,:,niter-1] 
        
        
        
        
        
        
        # Momentum Equation - Pressure Terms
      #  for i in range(0,nx):
      #     for j in range(0,ny-1):
      #          pu[i,j] = (dt / (2 * rho * h)) * ( p[ipos[i],j,n] - p[ineg[i],j,n] )
      #          pv[i,j] = (dt / (2 * rho * h)) * ( p[i,j+1,n] - p[i,jneg[j],n] )
    
        # Momentum Equation - Velocity Terms
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                u[i,j,n+1] = ( u[i,j,n] - (dt / h) * ( 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] ) - 
                                                    #  pu[i,j] + 
                ((dt * nu) / (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] ) ) ) 
                v[i,j,n+1] = ( v[i,j,n] - (dt / h) * ( 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] ) - 
                                                   #    pv[i,j] + 
                ((dt * nu) / (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] ) ) )         
                
    return u, v, p, x, y 


# In[31]:


(p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation3(11, 0.5, 11, 0.5, 2.0)


# In[32]:


nt


# In[33]:


u12, v12, p12, x12, y12 = navier_stokes_equation(1, 0.1)


# In[34]:


plot_3D_2(u12,x12,y12,0,'Initial u','u (m/s)')
plot_3D_2(u12,x12,y12,24,'Final u','u (m/s)')
plot_3D_2(v12,x12,y12,0,'Initial v','v (m/s)')
plot_3D_2(v12,x12,y12,24,'Final v','v (m/s)')
plot_3D_2(p12,x12,y12,0,'Initial Pressure','p (Pa)')
plot_3D_2(p12,x12,y12,24,'Final Pressure','p (Pa)')


# ## Final Navier Stokes Equations - Open Cavity Boundary Conditions

# In[4]:


def initialisation4(niter, r, nx_or_ny, tmax, xmax_or_ymax):
    """
    Returns the velocity field and distance for 2D linear convection
    """
    # Increments:
    nx = ny = nx_or_ny
    xmax = ymax = xmax_or_ymax
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    nt = int((tmax / (r*(dx)**2))+1)
    dt = tmax/(nt-1)
    
    
    # Initialise data structures:
    import numpy as np
    p = np.zeros(((nx,ny,nt)))
    pm = np.zeros(((nx,ny,niter)))
    u = np.zeros(((nx,ny,nt)))
    v = np.zeros(((nx,ny,nt)))
    x = np.zeros(nx)
    y = np.zeros(ny)
    ipos = np.zeros(nx)
    ineg = np.zeros(nx)
    jpos = np.zeros(ny)
    jneg = np.zeros(ny)

    # X Loop
    for i in range(0,nx):
        x[i] = i*dx

    # Y Loop
    for j in range(0,ny):
        y[j] = j*dy

    # Initial conditions
    p[:,:,0] = u[:,:,0] = v[:,:,0] = 0

    # Pressure Boundary Conditions:
    
    # p boundary, top:
    p[:,ny-1,:] = pm[:,ny-1,:] = 0
        
    # Velocity Boundary Conditions:
        
    u[:,ny-1,:] = 1
    u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = 0
    v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 0
        
    # von Neumann Boundary Conditions:

    # Values for correction at boundaries:
    for i in range(0,nx-1):
        ipos[i] = i + 1
        ineg[i] = i - 1

    for j in range(0,ny-1):
        jpos[j] = j + 1
        jneg[j] = j - 1
        
    # Set Reflection: 
    
    ineg[0] = 1           # i.e. index -1 = 1
    ipos[nx-1] = nx-2     # i.e. index nx = nx-2
    jneg[0] = 1           # i.e. index -1 = 1
    jpos[ny-1] = ny-2     # i.e. index ny = ny-2
    
        
    return p, pm, x, y, u, v, nx, ny, nt, dx, dy, dt, niter, ipos, ineg, jpos, jneg, r


# In[34]:


def navier_stokes_equation(rho, nu):
    
    (p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation4(11, 0.5, 31, 0.5, 2.0)
    
    # Increments
    h = dx
    r_factor = 0.25
    
    # Initialise data structures
    b = np.zeros((nx,ny))
    pu = np.zeros((nx,ny))
    pv = np.zeros((nx,ny))   
        
    # Loop
    for n in range(0,nt-1):
   
        # Poisson's Equation (Iteration) - uses an r factor of 0.25
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                b[i,j] = ( (rho / (2.0 * h * dt)) * ( u[i+1,j,n] - u[i-1,j,n] 
                                                    + v[i,j+1,n] - v[i,j-1,n] ) +
            (rho / (4.0*h**2)) * ( (u[i+1,j,n] - u[i-1,j,n])**2.0 + 
                             4.0*h*(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.0 ) )           
                
        # At the start of the iterations - the pressure is from the previous timestep        
        pm[:,:,0] = p[:,:,n] 

        for m in range(0,niter-1):
            for i in range(0,nx):
                for j in range(0,ny-1):
                    pm[i,j,m+1] = (r_factor*( pm[ipos[i],j,m] + pm[ineg[i],j,m] 
                                  + pm[i,j+1,m] + pm[i,jneg[j],m] - (h**2.0 * b[i,j]) ))

        # At the end of the iterations - the pressure is from the end of the iterations            
        p[:,:,n+1] = pm[:,:,niter-1] 
        
        # Momentum Equation - Pressure Terms - uses an r factor of 0.5
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                pu[i,j] = (dt / (2.0 * rho * h)) * ( p[i+1,j,n+1] - p[i-1,j,n+1] )
                pv[i,j] = (dt / (2.0 * rho * h)) * ( p[i,j+1,n+1] - p[i,j-1,n+1] )
    
        # Momentum Equation - Velocity Terms - uses an r factor of 0.5
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                u[i,j,n+1] = ( u[i,j,n] - (dt / h) * ( 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] ) ) - 
                                                       pu[i,j] + 
                ((dt * nu) / (h**2)) * ( u[i-1,j,n] + u[i+1,j,n] + u[i,j-1,n] + u[i,j+1,n] - 
                                        4.0 * u[i,j,n] ) )  
                v[i,j,n+1] = ( v[i,j,n] - (dt / h) * ( 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] ) ) - 
                                                       pv[i,j] + 
                ((dt * nu) / (h**2)) * ( v[i-1,j,n] + v[i+1,j,n] + v[i,j-1,n] + v[i,j+1,n] - 
                                        4.0 * v[i,j,n] ) )          
                
    return u, v, p, x, y 


# In[39]:


(p, pm, x, y, u, v, nx, ny, nt, 
    dx, dy, dt, niter, ipos, ineg, jpos, jneg, r) = initialisation4(11, 0.5, 41, 0.5, 2.0)


# In[40]:


nt


# In[41]:


u13, v13, p13, x13, y13 = navier_stokes_equation(1, 0.1)


# In[42]:


fig = plt.figure(figsize=(11,7), dpi=100)
Y,X=np.meshgrid(y13,x13) #note meshgrid uses y,x not x,y!!!
plt.contourf(X,Y,p13[:,:,-1],alpha=0.5)    ###plotting the pressure field as a contour
plt.colorbar()
plt.contour(X,Y,p13[:,:,-1])               ###plotting the pressure field outlines
plt.quiver(X[::2,::2],Y[::2,::2],u13[::2,::2,-1],v13[::2,::2,-1]) ##plotting velocity
plt.xlabel('X')
plt.ylabel('Y')


# In[ ]: