Save_Vorticity_Terms.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook for computing vorticity budget terms from depth-integrated momentum budget terms. We then average these terms over few years and save into a separate netcdf file for $1/8^{\\circ}$, $1/4^{\\circ}$ simulations.\n",
    "\n",
    "In a non steady state, the following balance holds\n",
    "\n",
    "\\begin{equation}\n",
    "\\beta \\overline{\\int_{-H}^{\\eta}vdz} = \\dfrac{1}{\\rho_o}\\overline{J(p_b, H)} - \\overline{\\dfrac{fQ_m}{\\rho_o}} + f\\partial_t\\eta - \\nabla \\wedge \\int_{-H}^{\\eta} \\mathbf{u}dz + \\dfrac{1}{\\rho_o}\\nabla \\wedge\\overline{\\left(\\boldsymbol{\\tau_s - \\tau_b}\\right)} + \\overline{\\nabla \\wedge \\int_{-H}^{\\eta}\\mathbf{a}dz} + \\overline{\\nabla \\wedge \\int_{-H}^{\\eta}\\mathbf{b}dz} \\tag{1}\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "\\beta \\overline{V} = \\dfrac{1}{\\rho_o}\\overline{J(p_b, H)} - f\\overline{\\dfrac{Q_m}{\\rho_o}} + f\\partial_t\\eta - \\nabla \\wedge \\mathcal{U}_t + \\dfrac{1}{\\rho_o}\\nabla \\wedge\\overline{\\left(\\boldsymbol{\\tau_s - \\tau_b}\\right)} + \\overline{\\nabla \\wedge \\mathcal{A}} + \\overline{\\nabla \\wedge \\mathcal{B}} \\tag{2}\n",
    "\\end{equation}\n",
    "\n",
    "It is seen that, in addition to boundary stress terms (fifth term on the RHS), bottom pressure torque (first term on the RHS) and nonlinear terms (sixth term on the RHS) do contribute significantly to the overall vorticity baalnce. Here, we isolate the relative contributions of these terms in 1/8 deg MOM6 simulation.\n",
    "\n",
    "In order to compute bottom pressure torque, we use the following relation (the terms on the RHS are available as diagnostics)\n",
    "\n",
    "\\begin{equation}\n",
    "\\dfrac{1}{\\rho_o}\\overline{J(p_b, H)} = \\nabla \\wedge\\left[- f \\hat{{\\bf z}} \\wedge \\int_{-H}^{\\eta} \\mathbf{u}dz - \\frac{1}{\\rho_o}\\int_{-H}^{\\eta} \\nabla pdz \\right] + \\beta \\overline{V} + f\\overline{\\dfrac{Q_m}{\\rho_o}} - f\\partial_t\\eta \\tag{3}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://127.0.0.1:35392</li>\n",
       "  <li><b>Dashboard: </b><a href='http://127.0.0.1:8787/status' target='_blank'>http://127.0.0.1:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>4</li>\n",
       "  <li><b>Cores: </b>16</li>\n",
       "  <li><b>Memory: </b>406.01 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:35392' processes=4 threads=16, memory=406.01 GB>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "from xgcm import Grid\n",
    "import filter\n",
    "from dask.diagnostics import ProgressBar\n",
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import glob, os\n",
    "\n",
    "from dask.distributed import Client\n",
    "from dask.distributed import LocalCluster\n",
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset>\n",
      "Dimensions:             (tim: 30, xh: 1440, xq: 1440, yh: 1080, yq: 1080)\n",
      "Coordinates:\n",
      "  * xh                  (xh) float64 -299.7 -299.5 -299.2 ... 59.53 59.78 60.03\n",
      "  * yh                  (yh) float64 -80.39 -80.31 -80.23 ... 89.73 89.84 89.95\n",
      "  * xq                  (xq) float64 -299.6 -299.3 -299.1 ... 59.66 59.91 60.16\n",
      "  * yq                  (yq) float64 -80.35 -80.27 -80.19 ... 89.78 89.89 90.0\n",
      "Dimensions without coordinates: tim\n",
      "Data variables:\n",
      "    hf_CAu_2d           (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_CAv_2d           (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_PFu_2d           (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_PFv_2d           (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_diffu_2d         (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_diffv_2d         (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_du_dt_visc_2d    (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_dudt_2d          (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_dv_dt_visc_2d    (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_dvdt_2d          (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_gKEu_2d          (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_gKEv_2d          (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_rvxu_2d          (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_rvxv_2d          (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_u_BT_accel_2d    (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hf_v_BT_accel_2d    (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    pbo                 (tim, yh, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    taux                (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    taux_bot            (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    tauy                (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    tauy_bot            (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    ubt_dt              (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    vbt_dt              (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    wfo                 (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    col_height          (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_CAu_2d         (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_CAv_2d         (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_diffu_2d       (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_diffv_2d       (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_gKEu_2d        (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_gKEv_2d        (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_PFu_2d         (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_PFv_2d         (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_rvxu_2d        (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_rvxv_2d        (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_u_BT_accel_2d  (tim, yh, xq) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    intz_v_BT_accel_2d  (tim, yq, xh) float64 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    Coriolis            (tim, yq, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    areacello           (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    areacello_bu        (tim, yq, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    areacello_cu        (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    areacello_cv        (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    deptho              (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dxCu                (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dxCv                (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dxt                 (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dyCu                (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dyCv                (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    dyt                 (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolat              (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolat_c            (tim, yq, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolat_u            (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolat_v            (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolon              (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolon_c            (tim, yq, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolon_u            (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    geolon_v            (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    hfgeou              (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    sftof               (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    wet                 (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    wet_c               (tim, yq, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    wet_u               (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    wet_v               (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    vmo                 (tim, yq, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    umo                 (tim, yh, xq) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    zos                 (tim, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n"
     ]
    }
   ],
   "source": [
    "# Read Data\n",
    "\n",
    "# for 1/8 deg data\n",
    "#path1 = \"/archive/Hemant.Khatri/MOM_Budget/OM4p125/\"\n",
    "#filelist = glob.glob(path1 + \"OM4p125*.nc\")\n",
    "#save_file = \"OM4p125_Vorticity_Budget.nc\" \n",
    "\n",
    "#path_grid = \"/archive/Raphael.Dussin/FMS2019.01.03_devgfdl_20210308/CM4_piControl_c192_OM4p125_v3/\" + \\\n",
    "#\"gfdl.ncrc4-intel18-prod-openmp/pp/ocean_monthly/\"\n",
    "#ds_grid = xr.open_dataset(path_grid + \"ocean_monthly.static.nc\")\n",
    "#ds_grid = ds_grid.isel(xq = slice(0,2880), yq=slice(0,2240))\n",
    "\n",
    "# for 1/4 deg data\n",
    "path1 = \"/archive/Hemant.Khatri/MOM_Budget/OM4p25/\"\n",
    "filelist = glob.glob(path1 + \"OM4p25*.nc\")\n",
    "save_file = \"OM4p25_Vorticity_Budget.nc\"\n",
    "\n",
    "path_grid = \"/archive/Raphael.Dussin/FMS2019.01.03_devgfdl_20210308/CM4_piControl_c192_OM4p25/\" + \\\n",
    "\"gfdl.ncrc4-intel18-prod-openmp/pp/ocean_monthly/\"\n",
    "ds_grid = xr.open_dataset(path_grid + \"ocean_monthly.static.nc\")\n",
    "ds_grid = ds_grid.isel(xq = slice(1,1441), yq=slice(1,1081))\n",
    "\n",
    "#-------------------------------------------------------------\n",
    "filelist.sort()\n",
    "\n",
    "ds = []\n",
    "for i in range(0,len(filelist)):\n",
    "    \n",
    "    d = xr.open_dataset(filelist[i])\n",
    "    ds.append(d)\n",
    "    \n",
    "ds = xr.concat(ds, dim='tim')\n",
    "ds = ds.chunk({'tim': 1})\n",
    "\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset>\n",
      "Dimensions:       (xh: 1440, xq: 1440, yh: 1080, yq: 1080)\n",
      "Coordinates:\n",
      "  * xh            (xh) float64 -299.7 -299.5 -299.2 -299.0 ... 59.53 59.78 60.03\n",
      "  * xq            (xq) float64 -299.6 -299.3 -299.1 -298.9 ... 59.66 59.91 60.16\n",
      "  * yh            (yh) float64 -80.39 -80.31 -80.23 -80.15 ... 89.73 89.84 89.95\n",
      "  * yq            (yq) float64 -80.35 -80.27 -80.19 -80.11 ... 89.78 89.89 90.0\n",
      "Data variables:\n",
      "    Coriolis      (yq, xq) float32 ...\n",
      "    areacello     (yh, xh) float32 ...\n",
      "    areacello_bu  (yq, xq) float32 ...\n",
      "    areacello_cu  (yh, xq) float32 ...\n",
      "    areacello_cv  (yq, xh) float32 ...\n",
      "    deptho        (yh, xh) float32 ...\n",
      "    dxCu          (yh, xq) float32 ...\n",
      "    dxCv          (yq, xh) float32 ...\n",
      "    dxt           (yh, xh) float32 ...\n",
      "    dyCu          (yh, xq) float32 ...\n",
      "    dyCv          (yq, xh) float32 ...\n",
      "    dyt           (yh, xh) float32 ...\n",
      "    geolat        (yh, xh) float32 ...\n",
      "    geolat_c      (yq, xq) float32 ...\n",
      "    geolat_u      (yh, xq) float32 ...\n",
      "    geolat_v      (yq, xh) float32 ...\n",
      "    geolon        (yh, xh) float32 ...\n",
      "    geolon_c      (yq, xq) float32 ...\n",
      "    geolon_u      (yh, xq) float32 ...\n",
      "    geolon_v      (yq, xh) float32 ...\n",
      "    hfgeou        (yh, xh) float32 ...\n",
      "    sftof         (yh, xh) float32 ...\n",
      "    wet           (yh, xh) float32 ...\n",
      "    wet_c         (yq, xq) float32 ...\n",
      "    wet_u         (yh, xq) float32 ...\n",
      "    wet_v         (yq, xh) float32 ...\n",
      "Attributes:\n",
      "    external_variables:  areacello areacello\n"
     ]
    }
   ],
   "source": [
    "print(ds_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create grid and interpolate depth, beta\n",
    "\n",
    "OMEGA = 7.2921e-5\n",
    "RAD_EARTH = 6.378e6\n",
    "\n",
    "grid = Grid(ds, coords={'X': {'center': 'xh', 'right': 'xq'},\n",
    "                        'Y': {'center': 'yh', 'right': 'yq'} }, periodic=[ ])\n",
    "\n",
    "depth_u = grid.interp(ds['deptho'].isel(tim=0) * ds['areacello'].isel(tim=0), 'X',  boundary='fill')  / ds['areacello_cu'].isel(tim=0)\n",
    "depth_v = grid.interp(ds['deptho'].isel(tim=0) * ds['areacello'].isel(tim=0), 'Y',  boundary='fill') / ds['areacello_cv'].isel(tim=0)\n",
    "depth_q = grid.interp(depth_u * ds['areacello_cu'].isel(tim=0), 'Y',  boundary='fill') / ds['areacello_bu'].isel(tim=0)\n",
    "\n",
    "colh_u = grid.interp(ds['col_height'] * ds['areacello'].isel(tim=0), 'X',  boundary='fill') / ds['areacello_cu'].isel(tim=0)\n",
    "colh_v = grid.interp(ds['col_height'] * ds['areacello'].isel(tim=0), 'Y',  boundary='fill') / ds['areacello_cv'].isel(tim=0)\n",
    "\n",
    "beta_v = 2*OMEGA*np.cos(ds.geolat_v.isel(tim=0) * np.pi /180.)/RAD_EARTH\n",
    "beta_q = 2*OMEGA*np.cos(ds.geolat_c.isel(tim=0) * np.pi /180.)/RAD_EARTH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute terms in vorticity budget\n",
    "\n",
    "rho_0 = 1035.\n",
    "\n",
    "BPT = xr.Dataset()\n",
    "\n",
    "vmo_bv = (ds['vmo'] / (rho_0 * ds['dxCv'].isel(tim=0)))\n",
    "vmo_bv = beta_q * grid.interp(vmo_bv, 'X',  boundary='fill')\n",
    "BPT['vmo_bv'] = vmo_bv\n",
    "\n",
    "umo = (ds['umo'] / (rho_0 * ds['dyCu'].isel(tim=0)))\n",
    "umo = grid.interp(umo, 'Y',  boundary='fill')\n",
    "BPT['umo'] = umo\n",
    "\n",
    "BPT_1 = (( - grid.diff((ds['intz_PFu_2d'] + ds['intz_u_BT_accel_2d']) * ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "           + grid.diff((ds['intz_PFv_2d'] + ds['intz_v_BT_accel_2d']) * ds.dyCv.isel(tim=0), 'X', boundary='fill')) \n",
    "         / ds.areacello_bu.isel(tim=0))\n",
    "BPT['BPT'] = BPT_1\n",
    "\n",
    "BPT['depth'] = (depth_q.load())\n",
    "\n",
    "Mass_Surf = (grid.interp(grid.interp(ds['wfo'] * ds['areacello'].isel(tim=0), 'X', boundary='fill'), 'Y',  boundary='fill') \n",
    "             * ds['Coriolis'].isel(tim=0) / (rho_0)) / ds['areacello_bu'].isel(tim=0)\n",
    "BPT['Qm'] = Mass_Surf\n",
    "\n",
    "dhdt = (grid.interp(grid.interp(ds['zos'] * ds['areacello'].isel(tim=0), 'X', boundary='fill'), 'Y',  boundary='fill') \n",
    "             * ds['Coriolis'].isel(tim=0)) / ds['areacello_bu'].isel(tim=0)\n",
    "BPT['fdhdt'] = dhdt\n",
    "\n",
    "div_u = (grid.diff(ds['umo'] / (rho_0), 'X', boundary='fill') + \n",
    "         grid.diff(ds['vmo'] / (rho_0), 'Y', boundary='fill') ) / ds['areacello'].isel(tim=0)\n",
    "div_u = - (grid.interp(grid.interp(div_u * ds['areacello'].isel(tim=0), 'X', boundary='fill'), 'Y', boundary='fill')\n",
    "        * ds['Coriolis'].isel(tim=0)) / ds['areacello_bu'].isel(tim=0)\n",
    "BPT['div_u'] = div_u\n",
    "\n",
    "Curl_dudt = ( - grid.diff(ds['hf_dudt_2d'] * colh_u * ds['dxCu'].isel(tim=0), 'Y', boundary='fill')\n",
    "            + grid.diff(ds['hf_dvdt_2d'] * colh_v * ds['dyCv'].isel(tim=0), 'X', boundary='fill') ) / ds.areacello_bu.isel(tim=0)\n",
    "BPT['Curl_dudt'] = Curl_dudt\n",
    "\n",
    "Curl_taus = ( - grid.diff((ds['taux'])* ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "                   + grid.diff((ds['tauy'])* ds.dyCv.isel(tim=0), 'X', boundary='fill') )/ ds.areacello_bu.isel(tim=0) \n",
    "Curl_taus = Curl_taus / (rho_0 )\n",
    "BPT['Curl_taus'] = Curl_taus\n",
    "\n",
    "Curl_taub = ( - grid.diff((-ds['taux_bot'])* ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "                   + grid.diff(-ds['tauy_bot'] * ds.dyCv.isel(tim=0), 'X', boundary='fill') )/ ds.areacello_bu.isel(tim=0) \n",
    "Curl_taub = Curl_taub / (rho_0 )\n",
    "BPT['Curl_taub'] = Curl_taub\n",
    "\n",
    "Curl_Hrv2 = ( - grid.diff((ds['intz_rvxv_2d'] + ds['intz_gKEu_2d']) * ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "           + grid.diff((ds['intz_rvxu_2d'] + ds['intz_gKEv_2d']) * ds.dyCv.isel(tim=0), 'X', boundary='fill') )/ ds.areacello_bu.isel(tim=0) \n",
    "BPT['Curl_NL'] = Curl_Hrv2\n",
    "\n",
    "Curl_Hdiff2 = ( - grid.diff(ds['intz_diffu_2d'] * ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "             + grid.diff(ds['intz_diffv_2d'] * ds.dyCv.isel(tim=0), 'X', boundary='fill') )/ ds.areacello_bu.isel(tim=0) \n",
    "BPT['Curl_Hdiff'] = Curl_Hdiff2\n",
    "\n",
    "Curl_Cor2 = ( - grid.diff((ds['intz_CAu_2d'] - ds['intz_gKEu_2d'] - ds['intz_rvxv_2d'])* ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "           + grid.diff((ds['intz_CAv_2d'] - ds['intz_gKEv_2d'] - ds['intz_rvxu_2d'])* ds.dyCv.isel(tim=0), 'X', boundary='fill'))/ ds.areacello_bu.isel(tim=0)\n",
    "BPT['Curl_Cor'] = Curl_Cor2\n",
    "\n",
    "tmpx = (ds['hf_dudt_2d'] * colh_u - ds['intz_CAu_2d']-ds['intz_PFu_2d']-ds['intz_diffu_2d']-\n",
    "       ds['intz_u_BT_accel_2d'] - ds['taux']/rho_0 + ds['taux_bot']/rho_0)\n",
    "tmpy = (ds['hf_dvdt_2d'] * colh_v - ds['intz_CAv_2d']-ds['intz_PFv_2d']-ds['intz_diffv_2d']-\n",
    "       ds['intz_v_BT_accel_2d'] - ds['tauy'] /rho_0 + ds['tauy_bot']/rho_0)\n",
    "Curl_remap = ( - grid.diff(tmpx * ds.dxCu.isel(tim=0), 'Y', boundary='fill')\n",
    "             + grid.diff(tmpy * ds.dyCv.isel(tim=0), 'X', boundary='fill') )/ ds.areacello_bu.isel(tim=0) \n",
    "BPT['Curl_remap'] = Curl_remap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = np.linspace(2.5, len(filelist) * 5 - 2.5, len(filelist))\n",
    "\n",
    "ds_save = xr.Dataset()\n",
    "\n",
    "ds_save['ssh'] = ds['col_height'] - ds['deptho']\n",
    "ds_save['ssh'].attrs['units'] = \"m\"\n",
    "ds_save['ssh'].attrs['standard_name'] = \"sea surface height above geoid\"\n",
    "\n",
    "ds_save['beta_V'] = BPT['vmo_bv']\n",
    "ds_save['beta_V'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['beta_V'].attrs['standard_name'] = \"Meridional Coriolis gradient x depth-integrated meridional velocity\"\n",
    "\n",
    "ds_save['BPT'] = BPT['BPT'] + BPT['Curl_Cor'] + BPT['vmo_bv'] + BPT['Qm'] - BPT['fdhdt']\n",
    "ds_save['BPT'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['BPT'].attrs['standard_name'] = \"Bottom Pressure Torque\"\n",
    "\n",
    "ds_save['Curl_Adv'] = (BPT['Curl_NL'] + BPT['Curl_remap'])\n",
    "ds_save['Curl_Adv'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Curl_Adv'].attrs['standard_name'] = \"Curl of depth-integrated nonlinear advetion term\"\n",
    "\n",
    "ds_save['Curl_taus'] =  BPT['Curl_taus']\n",
    "ds_save['Curl_taus'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Curl_taus'].attrs['standard_name'] = \"Curl of Surface Wind Stress / rho_0\"\n",
    "\n",
    "ds_save['Curl_taub'] =  BPT['Curl_taub']\n",
    "ds_save['Curl_taub'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Curl_taub'].attrs['standard_name'] = \"- Curl of bottom boundary stress / rho_0\"\n",
    "\n",
    "ds_save['Curl_diff'] = BPT['Curl_Hdiff']\n",
    "ds_save['Curl_diff'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Curl_diff'].attrs['standard_name'] = \"Curl of depth-integrated horizontal diffusion\"\n",
    "\n",
    "ds_save['Mass_flux'] = (- BPT['Qm'])\n",
    "ds_save['Mass_flux'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Mass_flux'].attrs['standard_name'] = \" - Coriolis x Surface mass flux / rho_0\"\n",
    "\n",
    "ds_save['eta_dt'] = BPT['fdhdt']\n",
    "ds_save['eta_dt'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['eta_dt'].attrs['standard_name'] = \" Coriolis x d(eta)/dt\"\n",
    "\n",
    "ds_save['Curl_dudt'] = (-BPT['Curl_dudt'])\n",
    "ds_save['Curl_dudt'].attrs['units'] = \"m/s^2\"\n",
    "ds_save['Curl_dudt'].attrs['standard_name'] = \" - Curl of depth-integrated du/dt\"\n",
    "\n",
    "ds_save = xr.merge([ds_save, ds_grid])\n",
    "ds_save = ds_save.rename({'tim': 'time'})\n",
    "\n",
    "ds_save.coords['time'] = times\n",
    "ds_save.time.attrs['units'] = \"years since 0001\"\n",
    "\n",
    "ds_save = ds_save.transpose('time','yq','yh','xq','xh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset>\n",
      "Dimensions:       (time: 30, xh: 1440, xq: 1440, yh: 1080, yq: 1080)\n",
      "Coordinates:\n",
      "  * xh            (xh) float64 -299.7 -299.5 -299.2 -299.0 ... 59.53 59.78 60.03\n",
      "  * yh            (yh) float64 -80.39 -80.31 -80.23 -80.15 ... 89.73 89.84 89.95\n",
      "  * xq            (xq) float64 -299.6 -299.3 -299.1 -298.9 ... 59.66 59.91 60.16\n",
      "  * yq            (yq) float64 -80.35 -80.27 -80.19 -80.11 ... 89.78 89.89 90.0\n",
      "  * time          (time) float64 2.5 7.5 12.5 17.5 ... 132.5 137.5 142.5 147.5\n",
      "Data variables:\n",
      "    ssh           (time, yh, xh) float32 dask.array<chunksize=(1, 1080, 1440), meta=np.ndarray>\n",
      "    beta_V        (time, yq, xq) float32 dask.array<chunksize=(1, 1080, 1439), meta=np.ndarray>\n",
      "    BPT           (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Curl_Adv      (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Curl_taus     (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Curl_taub     (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Curl_diff     (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Mass_flux     (time, yq, xq) float32 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    eta_dt        (time, yq, xq) float32 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Curl_dudt     (time, yq, xq) float64 dask.array<chunksize=(1, 1079, 1439), meta=np.ndarray>\n",
      "    Coriolis      (yq, xq) float32 ...\n",
      "    areacello     (yh, xh) float32 ...\n",
      "    areacello_bu  (yq, xq) float32 ...\n",
      "    areacello_cu  (yh, xq) float32 ...\n",
      "    areacello_cv  (yq, xh) float32 ...\n",
      "    deptho        (yh, xh) float32 ...\n",
      "    dxCu          (yh, xq) float32 ...\n",
      "    dxCv          (yq, xh) float32 ...\n",
      "    dxt           (yh, xh) float32 ...\n",
      "    dyCu          (yh, xq) float32 ...\n",
      "    dyCv          (yq, xh) float32 ...\n",
      "    dyt           (yh, xh) float32 ...\n",
      "    geolat        (yh, xh) float32 ...\n",
      "    geolat_c      (yq, xq) float32 ...\n",
      "    geolat_u      (yh, xq) float32 ...\n",
      "    geolat_v      (yq, xh) float32 ...\n",
      "    geolon        (yh, xh) float32 ...\n",
      "    geolon_c      (yq, xq) float32 ...\n",
      "    geolon_u      (yh, xq) float32 ...\n",
      "    geolon_v      (yq, xh) float32 ...\n",
      "    hfgeou        (yh, xh) float32 ...\n",
      "    sftof         (yh, xh) float32 ...\n",
      "    wet           (yh, xh) float32 ...\n",
      "    wet_c         (yq, xq) float32 ...\n",
      "    wet_u         (yh, xq) float32 ...\n",
      "    wet_v         (yq, xh) float32 ...\n"
     ]
    }
   ],
   "source": [
    "print(ds_save)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 15.4 s, sys: 15.6 s, total: 30.9 s\n",
      "Wall time: 32.1 s\n"
     ]
    }
   ],
   "source": [
    "path1 = \"/work/Hemant.Khatri/\"\n",
    "%time ds_save.load().to_netcdf(path1 + save_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.close()\n",
    "ds_save.close()\n",
    "client.close()\n",
    "cluster.close()"
   ]
  }
 ],
 "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}