SIMPLE_LDC.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.sparse.linalg import gmres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def exchange_rows(A):\n",
    "    r, c = A.shape\n",
    "    A1 = np.zeros((r, c))\n",
    "    for i in range(0, r):\n",
    "        A1[i] = A[-i -1]\n",
    "    \n",
    "    return A1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tdma_gmres(c_diag, t_diag, b_diag, B):\n",
    "    \n",
    "    tdma = np.diag(c_diag, k=0) + np.diag(b_diag[1:], k=-1) + np.diag(t_diag[:-1], k=1)\n",
    "    sol = gmres(tdma, B)[0]\n",
    "    \n",
    "    return sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tdma_solver(c_diag, t_diag, b_diag, B):\n",
    "    \n",
    "    #print(c_diag, '\\n')\n",
    "    #print(t_diag, '\\n')\n",
    "    #print(b_diag, '\\n')\n",
    "    \n",
    "    tdma = np.diag(c_diag, k=0) + np.diag(b_diag[1:], k=-1) + np.diag(t_diag[:-1], k=1)\n",
    "    sol = np.linalg.solve(tdma, B)\n",
    "    \n",
    "    return sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def solve_X(u, aP, aE, aW, aN, aS, p, r_sweeps, c_sweeps, alpha_uv, del_y, del_x, u_lid, mu):\n",
    "    \n",
    "    sol_dim = aP.shape\n",
    "    \n",
    "    rows = np.flip(np.arange(sol_dim[0]), axis = 0)\n",
    "    #rows = np.arange(sol_dim[0])\n",
    "    columns = np.arange(sol_dim[1])\n",
    "    \n",
    "    rc_sweeps = r_sweeps + c_sweeps\n",
    "    sweep = 0\n",
    "    \n",
    "    # row-wise sweeps.....................\n",
    "    \n",
    "    while sweep < rc_sweeps:\n",
    "        \n",
    "        for i in rows:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[1])\n",
    "            t_diag = np.zeros(aP.shape[1])\n",
    "            b_diag = np.zeros(aP.shape[1])\n",
    "            Sc = np.zeros(aP.shape[1])\n",
    "            Sp = np.zeros(aP.shape[1])\n",
    "\n",
    "            if i == 0: # top row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                        if j == 0: # top left corner u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aS[i, j]*u[i + 1, j + 1] - aW[i, j]*u[i, j]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] =  Sc[j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "\n",
    "\n",
    "                        elif j >= 1 and j < sol_dim[1] - 1: # top inner u cells\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "                            b_diag[j] = aW[i, j]\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aS[i, j]*u[i + 1, j + 1]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] =  Sc[j] - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "                            \n",
    "\n",
    "                        elif j == sol_dim[1] - 1: # top right corner u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            b_diag[j] = aW[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aS[i, j]*u[i + 1, j + 1] - aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aW[i, j]*u[i, j] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                u[i, 1:-1] = u[i, 1:-1] + U[:] # update in a correction form way\n",
    "\n",
    "\n",
    "            elif i == sol_dim[0] - 1: # bottom row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                        if j == 0: # bottom left corner u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1] - aW[i, j]*u[i, j]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "\n",
    "\n",
    "                        elif j >= 1 and j < sol_dim[1] - 1: # bottom inner u cells\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "                            b_diag[j] = aW[i, j]\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "                            \n",
    "\n",
    "                        elif j == sol_dim[1] - 1: # bottom right corner u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            b_diag[j] = aW[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1] - aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aW[i, j]*u[i, j] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                u[i, 1:-1] = u[i, 1:-1] + U[:] # update in a correction form way\n",
    "\n",
    "\n",
    "            else: # all interior rows\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                        if j == 0: # left u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1] - \\\n",
    "                                                                      aS[i, j]*u[i + 1, j + 1] - aW[i, j]*u[i, j]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "\n",
    "\n",
    "                        elif j >= 1 and j < sol_dim[1] - 1: # inner u cells\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            t_diag[j] = aE[i, j]\n",
    "                            b_diag[j] = aW[i, j]\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1] - \\\n",
    "                                                                      aS[i, j]*u[i + 1, j + 1]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "\n",
    "                        elif j == sol_dim[1] - 1: # right u cell\n",
    "\n",
    "                            c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                            b_diag[j] = aW[i, j]\n",
    "\n",
    "                            # standard source\n",
    "                            Sc[j] = del_y * (p[i, j] - p[i, j + 1]) - aN[i, j]*u[i - 1, j - 1] - \\\n",
    "                                                                      aS[i, j]*u[i + 1, j + 1] - \\\n",
    "                                                                      aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                            # source with correction form\n",
    "                            Sc[j] = Sc[j] - aW[i, j]*u[i, j] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                            # extra source from B.C.\n",
    "                            Sp[j] = 0.\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                u[i, 1:-1] = u[i, 1:-1] + U[:] # update in a correction form way\n",
    "                \n",
    "        \n",
    "        # column-wise sweep....................\n",
    "        \n",
    "        for j in columns:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[0])\n",
    "            t_diag = np.zeros(aP.shape[0])\n",
    "            b_diag = np.zeros(aP.shape[0])\n",
    "            Sc = np.zeros(aP.shape[0])\n",
    "            Sp = np.zeros(aP.shape[0])\n",
    "\n",
    "            if j == 0: # left column\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top left corner u cell\n",
    "                        \n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # left inner u cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom left corner u cell\n",
    "                        \n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                u[:, j + 1] = u[:, j + 1] + np.flip(U[:],0) # update in a correction form way\n",
    "\n",
    "\n",
    "            elif j == sol_dim[1] - 1: # right column\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top right u cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # right inner u cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom right u cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                u[:, j + 1] = u[:, j + 1] + np.flip(U[:],0) # update in a correction form way\n",
    "\n",
    "            else: # all interior columns\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top u cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 8.*u_lid*mu*del_x/(3.*del_y)\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # inner u cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aE[i, j]*u[i, j + 2] - aW[i, j]*u[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*u[i + 1, j + 1] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom u cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i, j] - p[i, j + 1]) - aW[i, j]*u[i, j] - aE[i, j]*u[i, j + 2]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*u[i - 1, j - 1] - aP[i, j]*u[i, j + 1]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for u\n",
    "                U = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                u[:, j + 1] = u[:, j + 1] + np.flip(U[:],0) # update in a correction form way\n",
    "                \n",
    "        \n",
    "        sweep += 1\n",
    "    \n",
    "    \n",
    "    return u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def solve_Y(v, aP, aE, aW, aN, aS, p, r_sweeps, c_sweeps, alpha_uv, del_y, del_x, mu):\n",
    "    \n",
    "    sol_dim = aP.shape\n",
    "    #rows = np.arange(sol_dim[0])\n",
    "    rows = np.flip(np.arange(sol_dim[0]), axis = 0)\n",
    "    columns = np.arange(sol_dim[1])\n",
    "    \n",
    "    rc_sweeps = r_sweeps + c_sweeps\n",
    "    sweep = 0\n",
    "    \n",
    "    # row wise sweeps..........\n",
    "    \n",
    "    while sweep < rc_sweeps:\n",
    "        \n",
    "        for i in rows:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[1])\n",
    "            t_diag = np.zeros(aP.shape[1])\n",
    "            b_diag = np.zeros(aP.shape[1])\n",
    "            Sc = np.zeros(aP.shape[1])\n",
    "            Sp = np.zeros(aP.shape[1])\n",
    "\n",
    "            if i == 0: # top row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # top left corner v cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        #standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j]\n",
    "                            \n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "                            \n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # top inner v cells\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # top right corner u cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                v[i + 1, :] = v[i + 1, :] + V[:] # update in a correction form way\n",
    "\n",
    "\n",
    "            elif i == sol_dim[0] - 1: # bottom row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # bottom left corner v cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # bottom inner v cells\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # bottom right corner v cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                v[i + 1, :] = v[i + 1, :] + V[:] # update in a correction form way\n",
    "\n",
    "\n",
    "            else: # all interior rows\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # left v cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # inner v cells\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aE[i, j]*v[i + 1, j + 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # right v cell\n",
    "\n",
    "                        c_diag[j] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[j] = Sc[j] - aW[i, j]*v[i + 1, j - 1] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[j] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                v[i + 1, :] = v[i + 1, :] + V[:] # update in a correction form way\n",
    "                \n",
    "        \n",
    "        # column wise sweeps.................\n",
    "        \n",
    "        for j in columns:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[0])\n",
    "            t_diag = np.zeros(aP.shape[0])\n",
    "            b_diag = np.zeros(aP.shape[0])\n",
    "            Sc = np.zeros(aP.shape[0])\n",
    "            Sp = np.zeros(aP.shape[0])\n",
    "\n",
    "            if j == 0: # left column\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top left corner v cell\n",
    "                        \n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aN[i, j]*v[i, j] - aE[i, j]*v[i + 1, j + 1]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # left inner v cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aE[i, j]*v[i + 1, j + 1]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom left corner v cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aE[i, j]*v[i + 1, j + 1] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*v[i, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                v[1:-1, j] = v[1:-1, j] + np.flip(V[:],0) # update in a correction form way\n",
    "\n",
    "\n",
    "            elif j == sol_dim[1] - 1: # right column\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top right v cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aW[i, j]*v[i + 1, j - 1] - aN[i, j]*v[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # right inner v cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aW[i, j]*v[i + 1, j - 1]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom right v cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aW[i, j]*v[i + 1, j - 1] - aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*v[i, j] - aP[i, j]*v[i +1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                v[1:-1, j] = v[1:-1, j] + np.flip(V[:],0) # update in a correction form way\n",
    "\n",
    "            else: # all interior columns\n",
    "\n",
    "                for i in rows:\n",
    "\n",
    "                    if i == 0: # top v cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aE[i, j]*v[i + 1, j + 1] - \\\n",
    "                                                                  aW[i, j]*v[i + 1, j - 1] - \\\n",
    "                                                                  aN[i, j]*v[i, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "\n",
    "                    elif i >= 1 and i < sol_dim[0] - 1: # inner v cells\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        b_diag[i] = aS[i, j]\n",
    "                        t_diag[i] = aN[i, j]\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aE[i, j]*v[i + 1, j + 1] - aW[i, j]*v[i + 1, j - 1]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aS[i, j]*v[i + 2, j] - aN[i, j]*v[i, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                    elif i == sol_dim[0] - 1: # bottom v cell\n",
    "\n",
    "                        c_diag[i] = (1. + alpha_uv) * aP[i, j] # correction form\n",
    "                        t_diag[i] = aN[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[i] = del_y * (p[i + 1, j] - p[i, j]) - aW[i, j]*v[i + 1, j - 1] - \\\n",
    "                                                                  aE[i, j]*v[i + 1, j + 1] - \\\n",
    "                                                                  aS[i, j]*v[i + 2, j]\n",
    "\n",
    "                        # source with correction form\n",
    "                        Sc[i] = Sc[i] - aN[i, j]*v[i, j] - aP[i, j]*v[i + 1, j]\n",
    "\n",
    "                        # extra source from B.C.\n",
    "                        Sp[i] = 0.\n",
    "\n",
    "                # solve for v\n",
    "                V = tdma_solver(np.flip(c_diag,0), np.flip(t_diag,0), np.flip(b_diag,0), np.flip(Sc+Sp,0))\n",
    "\n",
    "                v[1:-1, j] = v[1:-1, j] + np.flip(V[:],0) # update in a correction form way\n",
    "                \n",
    "        \n",
    "        sweep += 1\n",
    "    \n",
    "    \n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def assemble_X_links(rho, mu, u, v, u_dim, p_dim, del_x, del_y):\n",
    "    \n",
    "    # X-Momentum data structure\n",
    "    aP = np.zeros(p_dim)\n",
    "    aE = np.zeros(p_dim)\n",
    "    aW = np.zeros(p_dim)\n",
    "    aN = np.zeros(p_dim)\n",
    "    aS = np.zeros(p_dim)\n",
    "    \n",
    "    # contructing coefficients aE, aW, aN, aS, aP\n",
    "    for index in np.ndindex(u_dim):\n",
    "    \n",
    "        i, j = index\n",
    "\n",
    "        if i >= 1 and i < u_dim[0] - 1 and j >= 1 and j < u_dim[1] - 1: # interior cells\n",
    "\n",
    "            Fe = rho * (u[i, j] + u[i, j + 1]) / 2.\n",
    "            Fw = rho * (u[i, j] + u[i, j - 1]) / 2.\n",
    "            Fn = rho * (v[i, j] + v[i, j - 1]) / 2.\n",
    "            Fs = rho * (v[i + 1, j - 1] + v[i + 1, j]) / 2.\n",
    "\n",
    "            aE[i, j] = - del_y * (np.fmax(-Fe, 0) + mu / del_x)\n",
    "            aW[i, j] = - del_y * (np.fmax(Fw, 0) + mu / del_x)\n",
    "            aN[i, j] = - del_x * (np.fmax(-Fn, 0) + mu / del_y)\n",
    "            aS[i, j] = - del_x * (np.fmax(Fs, 0) + mu / del_y)\n",
    "\n",
    "            aP[i, j] = del_y * (np.fmax(Fe, 0) + np.fmax(-Fw, 0)) + del_x * (np.fmax(Fn, 0) + np.fmax(-Fs, 0)) + \\\n",
    "                                                                    2.*mu*del_y/del_x + 2.*mu*del_x/del_y\n",
    "            \n",
    "\n",
    "        elif i == 0 and j >= 1 and j < u_dim[1] - 1 : # top boundary cells\n",
    "\n",
    "            Fe = rho * (u[i, j] + u[i, j + 1]) / 2.\n",
    "            Fw = rho * (u[i, j] + u[i, j - 1]) / 2.\n",
    "            Fs = rho * (v[i + 1, j - 1] + v[i + 1, j]) / 2.\n",
    "\n",
    "            aE[i, j] = - del_y * (np.fmax(-Fe, 0) + mu / del_x)\n",
    "            aW[i, j] = - del_y * (np.fmax(Fw, 0) + mu / del_x)\n",
    "            aN[i, j] = 0.\n",
    "            aS[i, j] = - del_x * (np.fmax(Fs, 0) + mu / del_y) - mu*del_x/(3.*del_y)\n",
    "\n",
    "            aP[i, j] = del_y * (np.fmax(Fe, 0) + np.fmax(-Fw, 0)) + del_x * np.fmax(-Fs, 0) + 4.*mu*del_x/del_y + \\\n",
    "                                                                                        2.*mu*del_y/del_x\n",
    "\n",
    "\n",
    "        elif i == u_dim[0] - 1 and j >= 1 and j < u_dim[1] - 1: # bottom boundary cells\n",
    "\n",
    "            Fe = rho * (u[i, j] + u[i, j + 1]) / 2.\n",
    "            Fw = rho * (u[i, j] + u[i, j - 1]) / 2.\n",
    "            Fn = rho * (v[i, j] + v[i, j - 1]) / 2.\n",
    "\n",
    "            aE[i, j] = - del_y * (np.fmax(-Fe, 0) + mu / del_x)\n",
    "            aW[i, j] = - del_y * (np.fmax(Fw, 0) + mu / del_x)\n",
    "            aN[i, j] = - del_x * (np.fmax(-Fn, 0) + mu / del_y) - mu*del_x/(3.*del_y)\n",
    "            aS[i, j] = 0.\n",
    "\n",
    "            aP[i, j] = del_y * (np.fmax(Fe, 0) + np.fmax(-Fw, 0)) + del_x * np.fmax(Fn, 0) + 2.*mu*del_y/del_x + \\\n",
    "                                                                                        4.*mu*del_x/del_y\n",
    "            \n",
    "    return aP[:, 1:], aE[:, 1:], aW[:, 1:], aN[:, 1:], aS[:, 1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def assemble_Y_links(rho, mu, u, v, v_dim, p_dim, del_x, del_y):\n",
    "    \n",
    "    # Y-Momentum coefficients\n",
    "    aP = np.zeros(p_dim)\n",
    "    aE = np.zeros(p_dim)\n",
    "    aW = np.zeros(p_dim)\n",
    "    aN = np.zeros(p_dim)\n",
    "    aS = np.zeros(p_dim)\n",
    "    \n",
    "    # contructing coefficients aE, aW, aN, aS, aP\n",
    "    for index in np.ndindex(v_dim):\n",
    "    \n",
    "        i, j = index\n",
    "        \n",
    "        if i >= 1 and i < v_dim[0] - 1 and j >= 1 and j < v_dim[1] - 1: # interior cells\n",
    "\n",
    "            Fe = rho * (u[i - 1, j + 1] + u[i, j + 1]) / 2.\n",
    "            Fw = rho * (u[i - 1, j] + u[i, j]) / 2.\n",
    "            Fn = rho * (v[i, j] + v[i - 1, j]) / 2.\n",
    "            Fs = rho * (v[i, j] + v[i + 1, j]) / 2.\n",
    "\n",
    "            aE[i, j] = - del_y * (np.fmax(-Fe, 0) + mu / del_x)\n",
    "            aW[i, j] = - del_y * (np.fmax(Fw, 0) + mu / del_x)\n",
    "            aN[i, j] = - del_x * (np.fmax(-Fn, 0) + mu / del_y)\n",
    "            aS[i, j] = - del_x * (np.fmax(Fs, 0) + mu / del_y)\n",
    "\n",
    "            aP[i, j] = del_y * (np.fmax(Fe, 0) + np.fmax(-Fw, 0)) + del_x * (np.fmax(Fn, 0) + np.fmax(-Fs, 0)) + \\\n",
    "                                                                    2.*mu*del_y/del_x + 2.*mu*del_x/del_y    \n",
    "            \n",
    "            \n",
    "        if i >= 1 and i < v_dim[0] - 1 and j == 0: # left boundary cells\n",
    "\n",
    "            Fe = rho * (u[i - 1, j + 1] + u[i, j + 1]) / 2.\n",
    "            Fn = rho * (v[i, j] + v[i - 1, j]) / 2.\n",
    "            Fs = rho * (v[i, j] + v[i + 1, j]) / 2.\n",
    "\n",
    "            aE[i, j] = - del_y * (np.fmax(-Fe, 0) + mu / del_x) - mu*del_y/(3.*del_x)\n",
    "            aW[i, j] = 0.\n",
    "            aN[i, j] = - del_x * (np.fmax(-Fn, 0) + mu / del_y)\n",
    "            aS[i, j] = - del_x * (np.fmax(Fs, 0) + mu / del_y)\n",
    "\n",
    "            aP[i, j] = del_y * np.fmax(Fe, 0) + del_x * (np.fmax(Fn, 0) + np.fmax(-Fs, 0)) + \\\n",
    "                                                                    4.*mu*del_y/del_x + 2.*mu*del_x/del_y    \n",
    "            \n",
    "            \n",
    "        if i >= 1 and i < v_dim[0] - 1 and j == v_dim[1] - 1: # right boundary cells\n",
    "            \n",
    "            Fw = rho * (u[i - 1, j] + u[i, j]) / 2.\n",
    "            Fn = rho * (v[i, j] + v[i - 1, j]) / 2.\n",
    "            Fs = rho * (v[i, j] + v[i + 1, j]) / 2.\n",
    "\n",
    "            aE[i, j] = 0.\n",
    "            aW[i, j] = - del_y * (np.fmax(Fw, 0) + mu / del_x) - mu*del_y/(3.*del_x)\n",
    "            aN[i, j] = - del_x * (np.fmax(-Fn, 0) + mu / del_y)\n",
    "            aS[i, j] = - del_x * (np.fmax(Fs, 0) + mu / del_y)\n",
    "\n",
    "            aP[i, j] = del_y * np.fmax(-Fw, 0) + del_x * (np.fmax(Fn, 0) + np.fmax(-Fs, 0)) + \\\n",
    "                                                                    4.*mu*del_y/del_x + 2.*mu*del_x/del_y    \n",
    "    \n",
    "    return aP[1:, :], aE[1:, :], aW[1:, :], aN[1:, :], aS[1:, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def solve_pressure_correction(aP, aE, aW, aN, aS, u, v, p_sweeps, rho, del_y, del_x):\n",
    "    \n",
    "    pc = np.zeros(aP.shape)\n",
    "    sol_dim = aP.shape\n",
    "    m_dot = np.zeros(aP.shape)\n",
    "    \n",
    "    #rows = np.flip(np.arange(sol_dim[0]), axis = 0)\n",
    "    rows = np.arange(sol_dim[0])\n",
    "    columns = np.arange(sol_dim[1])\n",
    "    \n",
    "    sweep = 0\n",
    "    \n",
    "    # row-wise sweeps..................\n",
    "    \n",
    "    while sweep < p_sweeps:\n",
    "        \n",
    "        for i in rows:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[1])\n",
    "            t_diag = np.zeros(aP.shape[1])\n",
    "            b_diag = np.zeros(aP.shape[1])\n",
    "            Sc = np.zeros(aP.shape[1])\n",
    "            Sp = np.zeros(aP.shape[1])\n",
    "\n",
    "            if i == 0: # top row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # top left corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        #standard source\n",
    "                        \n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j]\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # top inner p cells\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j]\n",
    "                        \n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # top right corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j]\n",
    "                        \n",
    "\n",
    "                # solve for p\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                pc[i, :] = pp[:] # update\n",
    "\n",
    "\n",
    "            elif i == sol_dim[0] - 1: # bottom row\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # bottom left corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aN[i, j]*pc[i - 1, j]\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # bottom inner p cells\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aN[i, j]*pc[i - 1, j]\n",
    "                        \n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # bottom right corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aN[i, j]*pc[i - 1, j]\n",
    "                        \n",
    "\n",
    "                # solve for p\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                pc[i, :] = pp[:] # update\n",
    "\n",
    "\n",
    "            else: # all interior rows\n",
    "\n",
    "                for j in columns:\n",
    "\n",
    "                    if j == 0: # left corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j] - aN[i, j]*pc[i - 1, j]\n",
    "\n",
    "\n",
    "                    elif j >= 1 and j < sol_dim[1] - 1: # inner p cells\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        t_diag[j] = aE[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j] - aN[i, j]*pc[i - 1, j]\n",
    "                        \n",
    "\n",
    "                    elif j == sol_dim[1] - 1: # right corner p cell\n",
    "\n",
    "                        c_diag[j] = aP[i, j]\n",
    "                        b_diag[j] = aW[i, j]\n",
    "\n",
    "                        # standard source\n",
    "                        Sc[j] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[j] = - aS[i, j]*pc[i + 1, j] - aN[i, j]*pc[i - 1, j]\n",
    "                        \n",
    "\n",
    "                # solve for v\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                pc[i, :] = pp[:] # update\n",
    "                \n",
    "        \n",
    "        # column wise sweeps..................\n",
    "    \n",
    "        for j in columns:\n",
    "            \n",
    "            c_diag = np.zeros(aP.shape[0])\n",
    "            t_diag = np.zeros(aP.shape[0])\n",
    "            b_diag = np.zeros(aP.shape[0])\n",
    "            Sc = np.zeros(aP.shape[0])\n",
    "            Sp = np.zeros(aP.shape[0])\n",
    "            \n",
    "            if j == 0: # left column\n",
    "                \n",
    "                for i in rows:\n",
    "                    \n",
    "                    if i == 0: # top left p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1]\n",
    "                        \n",
    "                    elif i >=1 and i < sol_dim[0] - 1: # inner left p cells\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1]\n",
    "                        \n",
    "                    elif i == sol_dim[0] - 1: # bottom left p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1]\n",
    "                        \n",
    "                # solve for v\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "                \n",
    "                pc[:, j] = pp[:] # update\n",
    "                \n",
    "                \n",
    "            elif j == sol_dim[1] - 1: # right column\n",
    "                \n",
    "                for i in rows:\n",
    "                    \n",
    "                    if i == 0: # top right p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                    elif i >=1 and i < sol_dim[0] - 1: # inner right p cells\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                    elif i == sol_dim[0] - 1: # bottom right p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                # solve for v\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                pc[:, j] = pp[:] # update\n",
    "                \n",
    "                \n",
    "            else: # all interior columns\n",
    "                \n",
    "                for i in rows:\n",
    "                    \n",
    "                    if i == 0: # top p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1] - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                    elif i >=1 and i < sol_dim[0] - 1: # inner p cells\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        t_diag[i] = aS[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1] - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                    elif i == sol_dim[0] - 1: # bottom p cell\n",
    "                        \n",
    "                        c_diag[i] = aP[i, j]\n",
    "                        b_diag[i] = aN[i, j]\n",
    "                        \n",
    "                        # standard source\n",
    "                        Sc[i] = - rho * ((u[i, j + 1] - u[i, j]) * del_y + \\\n",
    "                                        (v[i, j] - v[i + 1, j]) * del_x)\n",
    "                        \n",
    "                        # extra source\n",
    "                        Sp[i] = - aE[i, j]*pc[i, j + 1] - aW[i, j]*pc[i, j - 1]\n",
    "                        \n",
    "                # solve for v\n",
    "                pp = tdma_solver(c_diag, t_diag, b_diag, Sc+Sp)\n",
    "\n",
    "                pc[:, j] = pp[:] # update\n",
    "                \n",
    "        sweep += 1\n",
    "                \n",
    "                \n",
    "        \n",
    "    # calculating residue\n",
    "    m_dot[:, :] = - rho * ((u[:, 1:] - u[:, :-1]) * del_y + (v[:-1, :] - v[1:, :]) * del_x)\n",
    "    residue = np.sqrt(np.sum(np.square(m_dot[:, :])))\n",
    "    print(' ----- p reisidue : ', residue, ' ----- ', '\\n')    \n",
    "        \n",
    "    return pc, residue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def assemble_P_links(rho, aPu, aPv, u, v, del_y, del_x, ny, nx, p_dim):\n",
    "    \n",
    "    # pressure coefficients\n",
    "    aP = np.zeros(p_dim)\n",
    "    aE = np.zeros(p_dim)\n",
    "    aW = np.zeros(p_dim)\n",
    "    aN = np.zeros(p_dim)\n",
    "    aS = np.zeros(p_dim)\n",
    "    Sc = np.zeros(p_dim)\n",
    "    Sp = np.zeros(p_dim)\n",
    "    \n",
    "    # contructing coefficients aE, aW, aN, aS, aP and sources Sc and Sp\n",
    "    for index in np.ndindex(p_dim):\n",
    "        \n",
    "        i, j = index\n",
    "        \n",
    "        if i >= 1 and i < p_dim[0] - 1 and j >= 1 and j < p_dim[1] - 1: # interior cells\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "            \n",
    "        elif i >= 1 and i < p_dim[0] - 1 and j == 0: # left boundary cells\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = 0.\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "        \n",
    "        elif i >= 1 and i < p_dim[0] - 1 and j == p_dim[1] - 1: # right boundary cells\n",
    "            \n",
    "            aE[i, j] = 0.\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "        \n",
    "        elif i == 0 and j >= 1 and j < p_dim[1] - 1: # top boundary cells\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = 0.\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "        \n",
    "        elif i == p_dim[0] - 1 and j >= 1 and j < p_dim[1] - 1: # bottom boundary cells\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = 0.\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "        \n",
    "        elif i == 0 and j == 0: # top left corner cell\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = 0.\n",
    "            aN[i, j] = 0.\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "            \n",
    "        \n",
    "        elif i == 0 and j == p_dim[1] - 1: # top right corner cell\n",
    "            \n",
    "            aE[i, j] = 0.\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = 0.\n",
    "            aS[i, j] = - rho * del_x**2 / aPv[i, j]\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "        \n",
    "        elif i == p_dim[0] - 1 and j == 0: # bottom left corner cell\n",
    "            \n",
    "            aE[i, j] = - rho * del_y**2 / aPu[i, j]\n",
    "            aW[i, j] = 0.\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = 0.\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "        \n",
    "        elif i == p_dim[0] - 1 and j == p_dim[1] - 1: # bottom right corner cells\n",
    "            \n",
    "            aE[i, j] = 0.\n",
    "            aW[i, j] = - rho * del_y**2 / aPu[i, j - 1]\n",
    "            aN[i, j] = - rho * del_x**2 / aPv[i - 1, j]\n",
    "            aS[i, j] = 0.\n",
    "\n",
    "            aP[i, j] = - (aE[i, j] + aW[i, j] + aN[i, j] + aS[i, j])\n",
    "    \n",
    "    return aP[:, :], aE[:, :], aW[:, :], aN[:, :], aS[:, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def correct_u_v_p(u_new, aPu, omega_u, v_new, aPv, omega_v, p_old, pc, alpha_p, del_y, del_x, u_dim, v_dim, p_dim):\n",
    "    \n",
    "    p = np.zeros(p_dim)\n",
    "    u = np.zeros(u_dim)\n",
    "    v = np.zeros(v_dim)\n",
    "    \n",
    "    # u velocity correction\n",
    "    for index in np.ndindex(u_dim):\n",
    "    \n",
    "        i, j = index\n",
    "        \n",
    "        if j >= 1 and j < u_dim[1] - 1:\n",
    "            \n",
    "            u[i, j] = u_new[i, j] + omega_u * (pc[i, j - 1] - pc[i, j]) * del_y / aPu[i, j - 1]\n",
    "    \n",
    "    # v velocity correction\n",
    "    for index in np.ndindex(v_dim):\n",
    "    \n",
    "        i, j = index\n",
    "        \n",
    "        if i >= 1 and i < u_dim[0] - 1:\n",
    "            \n",
    "            v[i, j] = v_new[i, j] + omega_v * (pc[i - 1, j] - pc[i, j]) * del_x / aPv[i - 1, j]\n",
    "    \n",
    "    # p correction\n",
    "    \n",
    "    p[:, :] = p_old[:, :] + alpha_p * pc[:, :]\n",
    "    \n",
    "    return u, v, p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def SIMPLE_LDC():\n",
    "    \n",
    "    # Physical properties\n",
    "    rho = 1e3 # kg/m3\n",
    "    mu = 1e-3 # kg/m.s\n",
    "    u_lid = 0.05 # m/sec.\n",
    "    l = 0.01 # m\n",
    "    b = 0.01 # m\n",
    "    \n",
    "    # grid dimensions\n",
    "    nx = 40 # columns of a matrix\n",
    "    ny = 40 # rows of a matrix\n",
    "    del_x = l / nx\n",
    "    del_y = b / ny\n",
    "    p_dim = (ny, nx)\n",
    "    u_dim = (ny, nx + 1)\n",
    "    v_dim = (ny + 1, nx)\n",
    "    \n",
    "    # initial values\n",
    "    p = np.zeros(p_dim)\n",
    "    u = np.zeros(u_dim)\n",
    "    v = np.zeros(v_dim)\n",
    "    \n",
    "    # from BCs\n",
    "    u[:, 0] = 0. # u = 0 at left boundary\n",
    "    u[:, -1] = 0. # u = 0 at right boundary\n",
    "    v[0, :] = 0. # v = 0 at top boundary\n",
    "    v[-1, :] = 0. # v = 0 at bottom boundary\n",
    "    \n",
    "    # under-relaxation factors\n",
    "    alpha_uv = 0.2\n",
    "    omega_uv = 0.8\n",
    "    alpha_p = 0.2\n",
    "    \n",
    "    # iteration parameters\n",
    "    tot_iter = 2000\n",
    "    n_iter = 0 \n",
    "    e = 1e-5\n",
    "    \n",
    "    r_sweeps = 1\n",
    "    c_sweeps = 1\n",
    "    p_sweeps = 1\n",
    "    \n",
    "    while(n_iter < tot_iter):\n",
    "        \n",
    "        # assemble u links\n",
    "        aPu,aEu,aWu,aNu,aSu = assemble_X_links(rho, mu, u, v, u_dim, p_dim, del_x, del_y)\n",
    "    \n",
    "        # assemble v links\n",
    "        aPv,aEv,aWv,aNv,aSv = assemble_Y_links(rho, mu, u, v, v_dim, p_dim, del_x, del_y)\n",
    "        \n",
    "        # x-momentum solution\n",
    "        u1 = solve_X(u, aPu, aEu, aWu, aNu, aSu, p, r_sweeps, c_sweeps, alpha_uv, del_y, del_x, u_lid, mu)\n",
    "        \n",
    "        # y-momentum solution\n",
    "        v1 = solve_Y(v, aPv, aEv, aWv, aNv, aSv, p, r_sweeps, c_sweeps, alpha_uv, del_y, del_x, mu)\n",
    "        \n",
    "        # assemble p links\n",
    "        aPp,aEp,aWp,aNp,aSp = assemble_P_links(rho, aPu, aPv, u1, v1, del_y, del_x, ny, nx, p_dim)\n",
    "                            \n",
    "        # pressure correction solution\n",
    "        pc, resM = solve_pressure_correction(aPp, aEp, aWp, aNp, aSp, u1, v1, p_sweeps, rho, del_y, del_x)\n",
    "        \n",
    "        # correcting u1, v1 and p\n",
    "        u_new, v_new, p_new = correct_u_v_p(u1, aPu, omega_uv, v1, aPv, omega_uv, p, pc, \\\n",
    "                                       alpha_p, del_y, del_x, u_dim, v_dim, p_dim)\n",
    "        \n",
    "        # assigning new values\n",
    "        u[:, :] = u_new[:, :]\n",
    "        v[:, :] = v_new[:, :]\n",
    "        p[:, :] = p_new[:, :]\n",
    "        \n",
    "        if resM < e:\n",
    "\n",
    "            #print('################# Convergence ###############', '\\n')\n",
    "            print('Total no. of iterations = ', n_iter)\n",
    "            break\n",
    "        \n",
    "        n_iter += 1\n",
    "        \n",
    "    return u, v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "u, v = SIMPLE_LDC()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = np.ones((1,41)) / 100.\n",
    "U = np.append(U, u, axis = 0)\n",
    "U = np.append(U, np.zeros((1,41)), axis = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#print(U[:, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "y = np.linspace(0., 0.01, 42)\n",
    "x = np.linspace(0., 0.01, 41)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYwAAAD8CAYAAABkbJM/AAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAHTNJREFUeJzt3X+wHeV93/H3B4Eg+AcyEWBZEpGgUqjkcRihSEwnUDsO\ntkQyuW7SuCJuwdhTjRyRThg7HmHsZEqrGYLdekyDUW9jNTB1rJJg4ztFjizT2k5mKiPBAEZYCley\nMFeWjAUeySmgi9C3f+xesTo695zn3D17ft3Pa+aMzu4+z+4+uufu5z777NlVRGBmZtbMWd3eATMz\n6w8ODDMzS+LAMDOzJA4MMzNL4sAwM7MkDgwzM0uSFBiSVknaK2lU0oY6yyXp7nz5U5KWFZZtlvSC\npKdr6lwoabukZ/N/31a+OWZmvamTx1FJCyS9IumJ/LWpHW1oGhiSZgD3AKuBJcANkpbUFFsNLMpf\na4F7C8v+ElhVZ9UbgEciYhHwSD5tZjZwunQc3RcRV+avde1oR0oPYwUwGhH7I2Ic2AIM1ZQZAu6P\nzA5glqQ5ABHxXeClOusdAu7L398HfGAqDTAz6wMDcRw9O6HMXOD5wvQYsDKhzFzgUIP1XhIRE8sP\nA5fUKyRpLVnacv75umrh5Sm7bGbT3e7vv3YkIi4qs45r3n1e/Oylkynb2g28Wpg1HBHDheluHEcX\nSnoCOAp8OiL+rmlDmuiJo29EhKS69yjJ/9OHAeYunRU3fOXXTi37l299vDM7aGY972+OLTtt+tPv\nfPi5suv82UsnefDh2U3LXXHpoVcjYnnZ7ZVRcxw9BFwaES9Kugp4SNLSiDhWZhspgXEQmF+YnpfP\na7VMrZ9ImhMRh/Ju1wvNduTYifPYfviKU9PbD1/BdW/f06xa2zmozNqn9kA/FcXjwhseLr3eNuro\ncTQijgPH8/ePSdoHLAZ2Tb0JaYGxE1gkaSHZzq8Bfr+mzAhwi6QtZN2so4Vu0mRGgJuAO/N/v95s\nR8bHz+bA2Ok9zO0JDZjQrnBpxwe8Kg4z66Yqfjfqh8Hpao8LPaijx1FJFwEvRcTrki4jG0jfX7YR\nTQMjIk5IugXYBswANkfEbknr8uWbgK3A9cAo8DJw80R9SV8B3g3MljQG/GlEfClv4AOSPgo8B3yw\n2b5oXJz7o5mnzTvA5B+UBfN+etp0ygevqBu9l7J6Ocxa4eDrHZ36TJUJhtrjQq/pwnH0WuAOSa8B\nJ4F1EVFv0Lwl6qfbm7/povnxT4duPW3ez39JZ5Q7ful403XVhkm79GPIWHv0esj10h8TqX+8pQbE\nW547/Th2wb7j/O9vf+qxsuMK73zXzEgcwyi9rX7QE4PeqWa8Glyw7/ip6aOXn3vGB+Xnv3RmL2RC\nMUhSurBTCZVWezHNOID6Ry8dkHtFq78PrfQg6oWEVauvAkOvjjNzzxjjV8wD3viAHL383FNl6gXI\nhEbd1nq9kqmeF21n76XdAdQODjGrNdXPabPfsZSgAIdFp/RVYEyYuWcM4IzgmNAoQCbUnspKPQea\ncrqrnQNwVZ06K6MXQ6zIgVaNdv3cpxISMPnvssOic/oyMCYUextFk32AUoKkqN74yFQG11JCZjJV\nX/3Ri4FUVq8GWi8FWaf+j1I/v41+rxwUvaOvAwPe6G3UUxsmjT5gxTCZkBIqUD9YispewVEmcJrp\n1OWIgxhMrerVIGunVj5PzX4vHBS9p+8Do5FGYQKnB0rKh7BeqEB6sEDzcKmnHZcMVhk6KRxMg2Wq\nP8+phgQ4KHrBQAdGM7VjIc00GitJ1SxcphIoKcqETrfDphVVBdN0DqJ2/J+mfP4cFr1vWgfGhGJP\nJDU8IG2spFVV91amYhB6OGUN6lhSle1K/dw0+8w7LHqHA6NG7WmsVgJkQrMPeJlAKWolXKBzAVNP\nu76J2+/BM5k+uLVFslZ+1g6L/uLAaKLeOMhUQqSoU4FSq9WAge6GTD0Ont7U6s8l5bPosOg9Dowp\naGUwfSpa+UWpKlwmTCVkJvRa2BT5NNvUlf2/c1j0LwdGBdpxWitVt3orKfrplNlU9PLl0u3Szpv6\npX4ephoWzf6Qs/IcGB3QyQCp1cuBUmvQA6ZWr99htV3K9FJTOCg6x4HRBVWMi0xVme+fdNt0C5h+\nM5WgaLV34bDoLAdGj5jsg9+tIClK/SXu1WCZ4ICpXpnehMctep8Do8dVPcDeTr00WN8Og3BVWdXa\ncbrJQdE/HBh9rpV7afWSfj4V1sggh0wVYxEOi/7iwBhg3Rxsb4d+GrAvoxcuXa56YLpWO4LC4xed\n58CYRnp5nGQqpkugNNLpA307TNdehaRVwBfInun9FxFxZ81y5cuvJ3um94cj4vFGdSX9CrAJeDNw\nAPhQRBzLl90GfBR4Hfh3EbGtbBscGDZwQTJhUE979aPpGhITJM0A7gGuA8aAnZJGIuKZQrHVwKL8\ntRK4F1jZpO5fAJ+IiO9I+gjwx8BnJC0B1gBLgXcA35K0OCJeL9MOB4ZNqp8G3KfKoVKtqoKiD09H\nrQBGI2I/gKQtwBBQDIwh4P6ICGCHpFmS5gALGtRdDHw3r78d2AZ8Jl++JSKOAz+UNJrvw/8t0wgH\nhk1Zvw64t2pQLivulKp7E30YFgBzgecL02NkvYhmZeY2qbubLBweAn4PmF9Y14466yrFgWGV6PcB\n96lo9YmOg6KTp5s6HRY/e/18/ubYsoSSD8+WtKswYzgihqvar4KPAHdL+gwwAlR6vxkHhnVEL327\nvRsG4dRXN8ch+qBXcSQiljdYfpA3/voHmJfPSylzzmR1I2IP8D4ASYuB32xhey1zYFjXDOpg+1RN\n9YBcNmh6eUC6D4Ii1U5gkaSFZAfuNcDv15QZAW7JxyhWAkcj4pCkn05WV9LFEfGCpLOAT5NdMTWx\nrr+S9J/JBr0XAY+WbYQDw3qOg6Q1vXzAn4oBColTIuKEpFvIBqVnAJsjYrekdfnyTcBWsktqR8ku\nq725Ud181TdIWp+//yrw3/M6uyU9QDYwfgJYX/YKKXBgWB9xkAy2QQyKoojYShYKxXmbCu8DWF9b\nb7K6+fwvkH0/o16djcDGErt8BgeG9b3pOMA+SAY9KAaJA8MGjgOk9zkk+pMDwwaeA6Q3OCT6nwPD\nph0HSGc4IAaPA8Omven+HZF2cUAMPgeGWR0OkeYcENOPA8MsUfEAOR3DwwFhDgyzKRj0cRCHg9WT\nFBgVPfjjSrKvsZ9H9k3EP4iI0l9dN+uGfv1SoYPBWtE0MCp88MddwL+PiG9Iuj6ffnfbWmbWA7p5\nC3iHgbVbSg+jqgd/BPDWvP4FwI/LN8esf/iAbv0mJTCqevDHHwHbJH0OOAv4Z/U2LmktsBbgvLPe\nnLC7ZmZWhbO6uO2PAbdGxHzgVuBL9QpFxHBELI+I5TPP+oWO7qCZmb0hJTDKPPijUd2byG7HC/DX\nZKe+zMysR6UExqkHf0iaSfbwjpGaMiPAjcpcTf7gjyZ1fwz88/z9rwPPlmyLmZlVqOkYRoUP/vi3\nwBcknQ28Sj5OYWZmvSnpexgVPfjj74GrWtlZMzPrnm4OepuZWR9xYJiZWRIHhplZB0haJWmvpFFJ\nG+osl6S78+VPSVrWrK6k/ynpifx1QNIT+fwFkl4pLNtUu72p8M0HzcwqVtUtliLiXxW28Z+Ao4X1\n7YuIK9vZDvcwzMyqd+oWSxExDkzcJqno1C2WImIHMHGLpaZ18xvAfhD4SpWNcA/DzGwSx06cx/bD\nVySUfHi2pF2FGcMRMVyYruoWSxOuAX4SEcXvsy3MT1EdBT4dEX+X0JCGHBhmZuUdiYjlXdz+DZze\nuzgEXBoRL0q6CnhI0tKIOFZmIw4MM7PqlbnF0jmN6uZffv4dCt9ri4jjwPH8/WOS9gGLgWIvqGUe\nwzAzq15Vt1gC+A1gT0Scul++pIvywXIkXUY2kL6/bCPcwzAzq1iFt1iCLEBqB7uvBe6Q9BpwElgX\nES+VbYcDw8ysA6q4xVK+7MN15j0IPFhid+vyKSkzM0viwDAzsyQODDMzS+LAMDOzJA4MMzNL4sAw\nM7MkDgwzM0viwDAzsyQODDMzS+LAMDOzJA4MMzNL4sAwM7MkDgwzM0viwDAzsyQODDMzS+LAMDOz\nJA4MMzNL4sAwM7MkDgwzsw6QtErSXkmjkjbUWS5Jd+fLn5K0LKWupD+UtEfSbkl3FebflpffK+n9\n7WiDn+ltZlYxSTOAe4DrgDFgp6SRiHimUGw1sCh/rQTuBVY2qivpPcAQ8CsRcVzSxfn2lgBrgKXA\nO4BvSVocEa+XaYd7GGZm1VsBjEbE/ogYB7aQHeiLhoD7I7MDmCVpTpO6HwPujIjjABHxQmFdWyLi\neET8EBjN11OKexhmZpMYHz+bA2MXpRSdLWlXYXo4IoYL03OB5wvTY2S9CJqUmduk7mLgGkkbgVeB\nT0TEzrzOjjrrKsWBYWZW3pGIWN6F7Z4NXAhcDfwq8ICky6raWNIpqU4P1piZDZiDwPzC9Lx8XkqZ\nRnXHgK/mp7EeBU4CsxO317KmgVEYcFkNLAFuyAdUioqDNWvJBmsa1q0ZrFkKfK5sY8zMetROYJGk\nhZJmkg1Ij9SUGQFuzP8Avxo4GhGHmtR9CHgPgKTFwEzgSL58jaRzJS0kOzY/WrYRKaekTg245Ds1\nMeBSHN0/NVgD7JA0MVizoEHdyQZrzMwGSkSckHQLsA2YAWyOiN2S1uXLNwFbgevJBqhfBm5uVDdf\n9WZgs6SngXHgpvw4vFvSA2TH2hPA+rJXSEFaYHR6sOY0ktaS9Vo476w3J+yumVnviYitZKFQnLep\n8D6A9al18/njwL+epM5GYGOJXT5DNwe96w7W5P9pp+RXGgwDXHDOxXHGWszMrCNSAqPMYM05Deqe\nGqwBHpU0MVjz0+S9NzOzjkm5SqrTgzVmZtaDmvYwujBYY2ZmPShpDKPTgzVmZtZ7fC8pMzNL4sAw\nM7MkDgwzM0viwDAzsyQODDMzS+LAMDOzJA4MMzNL4sAwM7MkDgwzM0viwDAzsyQODDMzS+LAMDOz\nJA4MMzNL4sAwM+sASask7ZU0KmlDneWSdHe+/ClJy1qo+3FJIWl2Pr1A0iuSnshfm2rrTEU3H9Fq\nZjYtSJoB3ANcR/a00Z2SRiLimUKx1cCi/LUSuBdY2ayupPnA+4Af1Wx2X0Rc2c52uIdhZla9FcBo\nROzPnwW0BRiqKTME3B+ZHcAsSXMS6n4e+CRQ+QPoHBhmZuXNlrSr8Fpbs3wu8Hxheiyfl1Jm0rqS\nhoCDEfFknX1amJ+O+o6ka1pv0pl8SsrMbBIaF+f+aGZK0SMRsbzq/SmSdD7wKbLTUbUOAZdGxIuS\nrgIekrQ0Io6V2aZ7GGZm1TsIzC9Mz8vnpZSZbP7lwELgSUkH8vmPS3p7RByPiBcBIuIxYB+wuGwj\nHBhmZtXbCSyStFDSTGANMFJTZgS4Mb9a6mrgaEQcmqxuRHw/Ii6OiAURsYDsVNWyiDgs6aJ8sBxJ\nl5ENpO8v2wifkjIzq1hEnJB0C7ANmAFsjojdktblyzcBW4HrgVHgZeDmRnWbbPJa4A5JrwEngXUR\n8VLZdjgwzMw6ICK2koVCcd6mwvsA1qfWrVNmQeH9g8CDJXa3Lp+SMjOzJA4MMzNL4sAwM7MkDgwz\nM0viwDAzsyQODDMzS+LAMDOzJA4MMzNL4sAwM7MkDgwzM0viwDAzsyQODDMzS+LAMDOzJEmBIWmV\npL2SRiVtqLNcku7Olz8laVkLdT8uKSTNLtcUMzOrUtPAyB/CcQ+wGlgC3CBpSU2x1WQP6FgErAXu\nTakraT7Z4wV/VLolZmZWqZQexgpgNCL2R8Q4sAUYqikzBNwfmR3ALElzEup+HvgkEGUbYmZm1UoJ\njLnA84XpsXxeSplJ60oaAg5GxJONNi5praRdknaNn3wlYXfNzKwKXXninqTzgU+RnY5qKCKGgWGA\nC8652D0RM7MuSelhHATmF6bn5fNSykw2/3JgIfCkpAP5/Mclvb2VnTcz6xdVXDwk6T/kZZ+Q9E1J\n7ygsuy0vv1fS+9vRhpTA2AkskrRQ0kxgDTBSU2YEuDFv8NXA0Yg4NFndiPh+RFwcEQvy59COAcsi\n4nA7GmVm1ksqvHjosxHxroi4EvhfwJ/kdZaQHW+XAquAL+brKaVpYETECeAWYBvwA+CBiNgtaZ2k\ndXmxrcB+YBT4b8AfNKpbdqfNzPpMJRcPRcSxQv038cYFREPAlog4HhE/JDs2ryjbiKQxjIjYShYK\nxXmbCu8DWJ9at06ZBSn7YWbWo2ZL2lWYHs7HXyfUuwBoZc06Wrl46FRdSRuBG4GjwHsK69pRZ12l\ndGXQ28ysH8wYh7c8l3StzZGIWF71/tQTEbcDt0u6jeyMzp9WtS3fGsTMrHpVXDxU68vA77awvZY5\nMMzMqtf2i4cAJC0q1B8C9hTWtUbSuZIWkg2kP1q2ET4lZWZWsYg4IWniAqAZwOaJi4fy5ZvIxnqv\nJxugfhm4uVHdfNV3Svpl4CTwHDCxvt2SHgCeAU4A6yPi9bLtcGCYmXVAFRcPRcTv1ik+sWwjsHGq\n+1uPT0mZmVkSB4aZmSVxYJiZWRIHhpmZJXFgmJlZEgeGmZklcWCYmVkSB4aZmSVxYJiZWRIHhpmZ\nJXFgmJlZEgeGmZklcWCYmVkSB4aZmSVxYJiZWRIHhpmZJXFgmJlZEgeGmZklcWCYmXWApFWS9koa\nlbShznJJujtf/pSkZc3qSvo9SbslnZS0vDB/gaRXJD2RvzbVbm8q/ExvM7OKSZoB3ANcB4wBOyWN\nRMQzhWKrgUX5ayVwL7CySd2ngd8B/mudze6LiCvb2Q73MMzMqrcCGI2I/RExDmwBhmrKDAH3R2YH\nMEvSnEZ1I+IHEbG3U41wYJiZlTdb0q7Ca23N8rnA84XpsXxeSpmUuvUszE9HfUfSNUmtaMKnpMzM\nJjHj1eCCfcdTih6JiOXNi3XMIeDSiHhR0lXAQ5KWRsSxMit1D8PMrHoHgfmF6Xn5vJQyKXVPExHH\nI+LF/P1jwD5g8ZT2vMCBYWZWvZ3AIkkLJc0E1gAjNWVGgBvzq6WuBo5GxKHEuqeRdFE+WI6ky8gG\n0veXbYRPSZmZVSwiTki6BdgGzAA2R8RuSevy5ZuArcD1wCjwMnBzo7oAkv4F8F+Ai4CHJT0REe8H\nrgXukPQacBJYFxEvlW2HA8PMrAMiYitZKBTnbSq8D2B9at18/teAr9WZ/yDwYMldPoNPSZmZWRIH\nhpmZJUkKjIq+0v5ZSXvy8l+TNKs9TTIzsyo0DYzC19JXA0uAGyQtqSlW/Er7WrKvtDerux14Z0S8\nC/gH4LbSrTEzs8qk9DCq+kr7NyPiRF5/B9m1xWZm1qNSAqMTX2n/CPCNehuXtHbi6/bjJ19J2F0z\nM6tC1we9Jd0OnAC+XG95RAxHxPKIWD7zrF/o7M6ZmdkpKd/DKPOV9nMa1ZX0YeC3gPfm1yCbmVmP\nSulhVPKVdkmrgE8Cvx0RL7epPWZmVpGmPYyqvtIO/DlwLrBdEsCOiFjXzsaZmVn7JN0apKKvtP+T\nlvbUzMy6quuD3mZm1h8cGGZmlsSBYWZmSRwYZmaWxIFhZmZJHBhmZpbEgWFmZkkcGGZmHVDRc4Uu\nlLRd0rP5v28rLLstL79X0vvb0QYHhplZxSp8rtAG4JGIWAQ8kk+TL18DLAVWAV/M11OKA8PMrHqV\nPFco//e+/P19wAcK87dExPGI+CHZbZtWlG2EA8PMrHpVPVfokvxGrwCHgUta2F7Lku4lZWY2HenV\ncWbuGUspOlvSrsL0cEQMV7RbdUVESKr0MREODDPrC+NXTP4U58SDepWORMTyBsureq7QTyTNiYhD\n+emrF1rYXst8SsrMetL4FfNOezUr2+Mqea5Q/u9N+fubgK8X5q+RdK6khWQD6Y+WbYR7GGbWE8oe\n9Cfq90Bv4wwVPlfoTuABSR8FngM+mNfZLekB4BmyR2Cvj4jXy7bDgWFmXdXu3sH4FfN6NTSqeK7Q\ni8B7J6mzEdhYYpfP4MAws66o8jTS+BXzsmuGrK0cGGbWMX0w1mANODDMrHIOisHgwDCzyjgoBosD\nw8zazkExmPw9DDNrK4fF4HIPw8zawkEx+NzDMLPSHBbTg3sYZjZlDorpxT0MM5sSh8X04x6GmbXE\nQTF9uYdhZskcFtObA8PMkjgszIFhZk05LAwcGGbWhMPCJjgwzGxSDgsrcmCYmVkSB4aZTaoXn1xn\n3ePAMLOG+i00Zu4Z67t97hdJgSFplaS9kkYlbaizXJLuzpc/JWlZs7qSLpS0XdKz+b9va0+TzKzd\n+uUA3C/7WSv1eNjq8VTSL0r6P5L+UdKf16zr2/m6nshfFzfbz6aBIWkGcA+wGlgC3CBpSU2x1cCi\n/LUWuDeh7gbgkYhYBDyST5tZj+r1g3Gv718TTY+HUzyevgp8BvjEJNv9UERcmb9eaLaTKT2MFcBo\nROyPiHFgCzBUU2YIuD8yO4BZkuY0qTsE3Je/vw/4QMK+mFkX9epBuVf3qwUpx8OWj6cR8f8i4u/J\ngqO0lHtJzQWeL0yPASsTysxtUveSiDiUvz8MXFJv45LWkvVaAI7/7eEvPp2wz/1mNnCk2zvRZoPY\nJhjMdrXWpsPV7Uib/XLZFRw78dNtf3v4i7MTip4naVdhejgihlvYVMrxsPTxtI77JL0GPAj8x4iI\nRoV74uaDERGS6u5o/p8+DCBpV0Qs7+jOdcAgtmsQ2wSD2a5BbBNk7Sq7johY1Y59AZD0LeDtdRbd\nXrPNSY+HKVqo/6GIOCjpLWSB8W+A+xtVSAmMg8D8wvS8fF5KmXMa1P2JpDkRcSg/fdX0/JmZWb+K\niN+YbJmklONho2Nxy8fTiDiY//tzSX9FdsqrYWCkjGHsBBZJWihpJrAGGKkpMwLcmF8tdTVwNO8e\nNao7AtyUv78J+HrCvpiZDaKU42HbjqeSzpY0O39/DvBbQPPT/RHR9AVcD/wDsA+4PZ+3DliXvxfZ\n6P0+4PvA8kZ18/m/SDaa/yzwLeDChP1Ym7K//fYaxHYNYpsGtV2D2KZ+a9dkx0PgHcDWQrmWj6fA\nAeAl4B/Jxj2WAG8CHgOeAnYDXwBmNNtP5Ss0MzNryN/0NjOzJA4MMzNL0rXAGNTbjVTUrs9K2pOX\n/5qkWZ1qT6P9KixvuU2F5R+XFBMDcJ1UVbsk/WH+89ot6a5OtCVlv/LlU/n8XSlph7LbR+yStKJT\n7Wm2b4Xljdq1WdILkp6uqdP140Xf6dIAzwyyQZvLgJnAk8CSmjLXA98gG1C/Gvhes7rAXcCG/P0G\n4M8GpF3vA87O3/9ZJ9tVVZvy5fOBbcBzwOwB+Vm9h2zQ8dx8+uIBaNM3gdWF+t/ul59VvuxaYBnw\ndE2drh4v+vHVrR7GoN5upJJ2RcQ3I+JEXn8H2fXXnVLVzwrg88AngW5ceVFVuz4G3BkRxwEi4f48\nbVRVmwJ4a/7+AuDHVTekRpl2ERHfJbtKqFa3jxd9p1uBMdmtRFLKNKo71a/Ht0tV7Sr6CNlfUp1S\nSZskDQEHI+LJdu9woqp+VouBayR9T9J3JP1qW/e6sara9EfAZyU9D3wOuK2N+5yiTLsa6fbxou8M\n7KB3RATd+cu1MpJuB04AX+72vpQh6XzgU8CfdHtfKnA2cCHZaZE/Bh6QpO7uUmkfA26NiPnArcCX\nurw/bTeIx4sqdCswytxupOnX4wHUnduNVNUuJH2Y7NuYH8o/3J1SRZsuBxYCT0o6kM9/XFK9++xU\npaqf1Rjw1fzUyKPASbKb+3VCVW26Cfhq/v6vyU4RdVKZdjXS7eNF/+nGwAnZX2H7yQ4aE4NYS2vK\n/CanD2I92qwu8FlOH8S6a0DatQp4BrhoUH5WNfUP0PlB76p+VuuAO/L3i8lOk6jP2/QD4N35+/cC\nj/XLz6qwfAFnDnp39XjRj6/ubbhHbjfSJ+0azQ88T+SvTf3eppr1H6DDgVHhz2om8D/I7svzOPDr\nA9CmXyO7jcSTwPeAq/rsZ/UV4BDwGlkP8KP5/K4fL/rt5VuDmJlZkoEd9DYzs/ZyYJiZWRIHhpmZ\nJXFgmJlZEgeGmZklcWCYmVkSB4aZmSX5/4uvfxRiz0hEAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x10af8deb8>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cl = plt.contourf(x, y, exchange_rows(U))\n",
    "plt.colorbar(cl)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}