This gist contains three functions for the implementation of the predictor-corrector finite difference scheme for the solution of the incompressible Navier-Stokes equations.

FlowPy_FD_functions.py

#The first function is used to get starred velocities from u and v at timestep t
def GetStarredVelocities(space,fluid):
    #Save object attributes as local variable with explicity typing for improved readability
    rows=int(space.rowpts)
    cols=int(space.colpts)
    u=space.u.astype(float,copy=False)
    v=space.v.astype(float,copy=False)
    dx=float(space.dx)
    dy=float(space.dy)
    dt=float(space.dt)
    S_x=float(space.S_x)
    S_y=float(space.S_y)
    rho=float(fluid.rho)
    mu=float(fluid.mu)
    
    #Copy u and v to new variables u_star and v_star
    u_star=u.copy()
    v_star=v.copy()
    
    #Calculate derivatives of u and v using the finite difference scheme. 
    #Numpy vectorization saves us from using slower for loops to go over each element in the u and v matrices
    u1_y=(u[2:,1:cols+1]-u[0:rows,1:cols+1])/(2*dy)
    u1_x=(u[1:rows+1,2:]-u[1:rows+1,0:cols])/(2*dx)
    u2_y=(u[2:,1:cols+1]-2*u[1:rows+1,1:cols+1]+u[0:rows,1:cols+1])/(dy**2)
    u2_x=(u[1:rows+1,2:]-2*u[1:rows+1,1:cols+1]+u[1:rows+1,0:cols])/(dx**2)
    v_face=(v[1:rows+1,1:cols+1]+v[1:rows+1,0:cols]+v[2:,1:cols+1]+v[2:,0:cols])/4
    u_star[1:rows+1,1:cols+1]=u[1:rows+1,1:cols+1]-dt*(u[1:rows+1,1:cols+1]*u1_x+v_face*u1_y)+(dt*(mu/rho)*(u2_x+u2_y))+(dt*S_x)   

    v1_y=(v[2:,1:cols+1]-v[0:rows,1:cols+1])/(2*dy)
    v1_x=(v[1:rows+1,2:]-v[1:rows+1,0:cols])/(2*dx)
    v2_y=(v[2:,1:cols+1]-2*v[1:rows+1,1:cols+1]+v[0:rows,1:cols+1])/(dy**2)
    v2_x=(v[1:rows+1,2:]-2*v[1:rows+1,1:cols+1]+v[1:rows+1,0:cols])/(dx**2)
    u_face=(u[1:rows+1,1:cols+1]+u[1:rows+1,2:]+u[0:rows,1:cols+1]+u[0:rows,2:])/4
    v_star[1:rows+1,1:cols+1]=v[1:rows+1,1:cols+1]-dt*(u_face*v1_x+v[1:rows+1,1:cols+1]*v1_y)+(dt*(mu/rho)*(v2_x+v2_y))+(dt*S_y)
    
    #Save the calculated starred velocities to the space object 
    space.u_star=u_star.copy()
    space.v_star=v_star.copy()    
    
    
#The second function is used to iteratively solve the pressure Possion equation from the starred velocities 
#to calculate pressure at t+delta_t
def SolvePressurePoisson(space,fluid,left,right,top,bottom):
    #Save object attributes as local variable with explicity typing for improved readability
    rows=int(space.rowpts)
    cols=int(space.colpts)
    u_star=space.u_star.astype(float,copy=False)
    v_star=space.v_star.astype(float,copy=False)
    p=space.p.astype(float,copy=False)
    dx=float(space.dx)
    dy=float(space.dy)
    dt=float(space.dt)
    rho=float(fluid.rho)
    factor=1/(2/dx**2+2/dy**2)
    
    #Define initial error and tolerance for convergence (error > tol necessary initially)
    error=1
    tol=1e-3

    #Evaluate derivative of starred velocities
    ustar1_x=(u_star[1:rows+1,2:]-u_star[1:rows+1,0:cols])/(2*dx)
    vstar1_y=(v_star[2:,1:cols+1]-v_star[0:rows,1:cols+1])/(2*dy)

    #Continue iterative solution until error becomes smaller than tolerance
    i=0
    while(error>tol):
        i+=1
        
        #Save current pressure as p_old
        p_old=p.astype(float,copy=True)
        
        #Evaluate second derivative of pressure from p_old
        p2_xy=(p_old[2:,1:cols+1]+p_old[0:rows,1:cols+1])/dy**2+(p_old[1:rows+1,2:]+p_old[1:rows+1,0:cols])/dx**2
        
        #Calculate new pressure 
        p[1:rows+1,1:cols+1]=(p2_xy)*factor-(rho*factor/dt)*(ustar1_x+vstar1_y)
        
        #Find maximum error between old and new pressure matrices
        error=np.amax(abs(p-p_old))
        
        #Apply pressure boundary conditions
        SetPBoundary(space,left,right,top,bottom)
        
        #Escape condition in case solution does not converge after 500 iterations
        if(i>500):
            tol*=10
            
    
#The third function is used to calculate the velocities at timestep t+delta_t using the pressure at t+delta_t and starred velocities
def SolveMomentumEquation(space,fluid):
    #Save object attributes as local variable with explicity typing for improved readability
    rows=int(space.rowpts)
    cols=int(space.colpts)
    u_star=space.u_star.astype(float,copy=False)
    v_star=space.v_star.astype(float,copy=False)
    p=space.p.astype(float,copy=False)
    dx=float(space.dx)
    dy=float(space.dy)
    dt=float(space.dt)
    rho=float(fluid.rho)
    u=space.u.astype(float,copy=False)
    v=space.v.astype(float,copy=False)

    #Evaluate first derivative of pressure in x direction
    p1_x=(p[1:rows+1,2:]-p[1:rows+1,0:cols])/(2*dx)
    #Calculate u at next timestep
    u[1:rows+1,1:cols+1]=u_star[1:rows+1,1:cols+1]-(dt/rho)*p1_x

    #Evaluate first derivative of pressure in y direction
    p1_y=(p[2:,1:cols+1]-p[0:rows,1:cols+1])/(2*dy)
    #Calculate v at next timestep
    v[1:rows+1,1:cols+1]=v_star[1:rows+1,1:cols+1]-(dt/rho)*p1_y