Calculating_heat_fluxes_across_ITF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating the ITF heat flux transport using WMT diagnostics"
   ]
  },
  {
   "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:43985</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>8</li>\n",
       "  <li><b>Memory: </b>33.51 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:43985' processes=4 threads=8, memory=33.51 GB>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "base = '/g/data/e14/mv7494/access-om2/archive/'\n",
    "from dask.distributed import Client\n",
    "client = Client(n_workers=4)         # careful: do not run with 16 workers here as it will crash\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define transects\n",
    "ITH_loc={\"lat_slice\":127,\"lon_slice\":slice(24,54)}\n",
    "ITV_loc={\"lat_slice\":slice(87,119),\"lon_slice\":62}\n",
    "\n",
    "# Pre-process functions per transect\n",
    "def pre_ITH(ds):\n",
    "    dsnew = ds.isel(grid_yu_ocean=ITH_loc[\"lat_slice\"], grid_xt_ocean=ITH_loc[\"lon_slice\"]).sum(['neutral']).sum(['grid_xt_ocean']) # grid_yu_ocean=71 -> 35°S, =63 -> 42.5°S\n",
    "    ww = dsnew.temp_yflux_adv_on_nrho + dsnew.temp_yflux_submeso_on_nrho + dsnew.temp_yflux_gm_on_nrho + dsnew.temp_yflux_ndiffuse_on_nrho\n",
    "    dsnew = ww.to_dataset(name=\"wmass\")\n",
    "    return dsnew\n",
    "\n",
    "def pre_ITV(ds):\n",
    "    dsnew = ds.isel(grid_yt_ocean=ITV_loc[\"lat_slice\"], grid_xu_ocean=ITV_loc[\"lon_slice\"]).sum(['neutral']).sum(['grid_yt_ocean'])\n",
    "    ww = dsnew.temp_xflux_adv_on_nrho + dsnew.temp_xflux_submeso_on_nrho + dsnew.temp_xflux_gm_on_nrho + dsnew.temp_xflux_ndiffuse_on_nrho\n",
    "    dsnew = ww.to_dataset(name=\"wmass\")\n",
    "    return dsnew\n",
    "\n",
    "def time_calc(wmass, dt):\n",
    "    throughflow = (wmass*dt).cumsum(['time']).sel(time=slice('1972','2017'))\n",
    "    throughflow = wmass.groupby('time.year').mean('time')\n",
    "    return throughflow    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wmass calculated for ITF_horizontal\n",
      "<xarray.DataArray 'wmass' (time: 600)>\n",
      "array([-5.20645726e+14, -4.68143650e+14, -8.74013443e+14, -9.30505471e+14,\n",
      "       -7.37994739e+14, -8.95364602e+14, -9.22914752e+14, -1.33591317e+15,\n",
      "       -1.23812521e+15, -1.15727206e+15, -8.97031748e+14, -4.97358607e+14,\n",
      "       -7.11665223e+14, -7.29670844e+14, -8.54847068e+14, -4.85116096e+14,\n",
      "       -7.61539200e+14, -8.85938592e+14, -1.07184535e+15, -1.07267362e+15,\n",
      "       -1.04195108e+15, -9.69845993e+14, -7.38512346e+14, -8.06371475e+14,\n",
      "       -4.03414204e+14, -3.14523082e+14, -4.77445536e+14, -9.07189961e+14,\n",
      "       -1.22198615e+15, -8.36401214e+14, -1.13532408e+15, -1.29515952e+15,\n",
      "       -1.48756715e+15, -1.06740509e+15, -5.57186043e+14, -1.87275413e+14,\n",
      "       -4.37747825e+14, -5.91808670e+14, -3.60336838e+14, -4.71117250e+14,\n",
      "       -9.22040609e+14, -8.81466715e+14, -8.26036656e+14, -8.99858316e+14,\n",
      "       -9.11942733e+14, -7.45725356e+14, -4.98390917e+14, -3.06524037e+14,\n",
      "       -6.84940356e+13, -4.95458825e+14, -5.36160913e+14, -5.60620457e+14,\n",
      "       -4.36558116e+14, -8.17593973e+14, -1.01284118e+15, -1.23196169e+15,\n",
      "       -1.22303043e+15, -1.12939420e+15, -7.32964324e+14, -4.81527777e+14,\n",
      "       -2.79971190e+14, -1.16177168e+14, -2.94316955e+14, -5.41551834e+14,\n",
      "       -7.66449986e+14, -9.22141972e+14, -1.01056513e+15, -9.55318846e+14,\n",
      "       -1.00513808e+15, -1.04813581e+15, -7.41657610e+14, -7.25686675e+14,\n",
      "       -5.13577305e+14, -2.08482431e+14, -4.11292177e+14, -7.45152135e+14,\n",
      "       -7.66332885e+14, -9.64995486e+14, -9.89793946e+14, -1.00501877e+15,\n",
      "...\n",
      "       -9.73625714e+14, -9.57257680e+14, -9.36453467e+14, -9.23598055e+14,\n",
      "       -9.70816569e+14, -8.48744227e+14, -5.99822926e+14, -3.73734680e+14,\n",
      "       -1.61784018e+14, -5.20540880e+14, -5.84852173e+14, -6.48820920e+14,\n",
      "       -5.29299990e+14, -9.60754507e+14, -1.04903648e+15, -1.31510306e+15,\n",
      "       -1.32468422e+15, -1.16825719e+15, -7.77285953e+14, -5.05298290e+14,\n",
      "       -3.33598036e+14, -1.48068552e+14, -3.86243889e+14, -6.55676738e+14,\n",
      "       -8.24210568e+14, -1.01792738e+15, -1.02064292e+15, -1.01021324e+15,\n",
      "       -1.03170498e+15, -1.08770642e+15, -8.26199088e+14, -7.82863477e+14,\n",
      "       -7.86283795e+14, -4.77774117e+14, -7.18900142e+14, -7.60585361e+14,\n",
      "       -6.60297378e+14, -8.37034884e+14, -8.58686760e+14, -1.25953983e+15,\n",
      "       -1.17849193e+15, -1.11791604e+15, -8.63010629e+14, -4.57714561e+14,\n",
      "       -6.93037293e+14, -6.96854196e+14, -8.22980499e+14, -4.57040781e+14,\n",
      "       -7.81495012e+14, -8.90521977e+14, -1.08171863e+15, -1.07642024e+15,\n",
      "       -1.06884050e+15, -9.89960801e+14, -7.61303494e+14, -8.33527275e+14,\n",
      "       -4.24393429e+14, -3.35280883e+14, -5.42735679e+14, -9.37818929e+14,\n",
      "       -1.23119624e+15, -8.45814648e+14, -1.13227890e+15, -1.29388027e+15,\n",
      "       -1.51231335e+15, -1.09796811e+15, -5.84906250e+14, -2.16325732e+14,\n",
      "       -4.75354011e+14, -6.25618352e+14, -3.84189515e+14, -5.06870492e+14,\n",
      "       -9.66476123e+14, -9.12771368e+14, -8.58752690e+14, -9.48129982e+14,\n",
      "       -9.55018838e+14, -7.70277849e+14, -5.36731249e+14, -3.72014247e+14])\n",
      "Coordinates:\n",
      "  * time           (time) object 1972-01-16 12:00:00 ... 2021-12-16 12:00:00\n",
      "    grid_yu_ocean  float64 -3.0\n",
      "iteration done\n",
      "--------------\n",
      "wmass calculated for ITF_vertical\n",
      "<xarray.DataArray 'wmass' (time: 600)>\n",
      "array([ 4.35895229e+13,  1.51920869e+14,  6.68036568e+12, -7.75215799e+13,\n",
      "       -1.08172258e+14, -1.45033033e+14, -1.67523473e+14, -1.56245531e+14,\n",
      "       -8.78420880e+13, -7.47823668e+13, -2.10800524e+13,  5.77038785e+12,\n",
      "        5.94045914e+13,  6.52599380e+13,  7.18778911e+13, -1.50715634e+13,\n",
      "       -1.20592540e+14, -1.20434475e+14, -1.52436530e+14, -1.38010994e+14,\n",
      "       -1.45577941e+14, -1.66909471e+14, -1.11851664e+14, -4.52995952e+13,\n",
      "        5.57812945e+13,  4.57454841e+13,  2.36080710e+12, -1.07103198e+14,\n",
      "       -1.08516826e+14, -1.25092002e+14, -1.20929552e+14, -1.40412258e+14,\n",
      "       -1.29545388e+14, -3.67504495e+13, -2.74725110e+13,  3.13549832e+13,\n",
      "        6.40280746e+13,  1.38571418e+13,  2.10697881e+13, -1.07435621e+14,\n",
      "       -1.46840901e+14, -1.61950967e+14, -1.84577170e+14, -1.45613147e+14,\n",
      "       -1.51051517e+14, -1.35641281e+14, -9.82382040e+13, -3.39961670e+13,\n",
      "        3.27934274e+12,  1.66857310e+13, -6.46091739e+13, -1.43477117e+14,\n",
      "       -1.68072855e+14, -1.83300440e+14, -1.78583519e+14, -1.79330713e+14,\n",
      "       -1.26014089e+14, -1.31575534e+14, -8.00106226e+13, -7.37859695e+13,\n",
      "        8.41692338e+13,  6.48300177e+13,  6.80780901e+13, -1.13966450e+14,\n",
      "       -1.35655440e+14, -1.29499628e+14, -1.49153883e+14, -1.32445819e+14,\n",
      "       -1.37474090e+14, -1.20174911e+14, -7.25920505e+13,  5.17904490e+12,\n",
      "        1.08510785e+14,  1.43677536e+14, -2.48406574e+13, -6.20526111e+13,\n",
      "       -7.37487906e+13, -1.40629213e+14, -1.41585674e+14, -1.11467074e+14,\n",
      "...\n",
      "       -1.54384193e+14, -1.67722062e+14, -1.88319291e+14, -1.49542799e+14,\n",
      "       -1.55544857e+14, -1.40896370e+14, -1.02791747e+14, -3.95687363e+13,\n",
      "       -4.12977477e+12,  8.66235888e+12, -7.27384527e+13, -1.49322437e+14,\n",
      "       -1.69483181e+14, -1.84706594e+14, -1.79073116e+14, -1.80938334e+14,\n",
      "       -1.29151352e+14, -1.34924177e+14, -8.32607546e+13, -7.75354296e+13,\n",
      "        8.12786855e+13,  6.18124696e+13,  6.35241726e+13, -1.17084205e+14,\n",
      "       -1.39261104e+14, -1.35095280e+14, -1.53814956e+14, -1.37553253e+14,\n",
      "       -1.40649125e+14, -1.23932991e+14, -7.60422329e+13,  1.96430004e+12,\n",
      "       -2.82644594e+12,  1.27180021e+14, -1.60308136e+13, -9.77641352e+13,\n",
      "       -1.26268459e+14, -1.59953493e+14, -1.77615047e+14, -1.65783957e+14,\n",
      "       -9.75989368e+13, -8.23335050e+13, -2.60173212e+13,  3.80489241e+12,\n",
      "        5.63409701e+13,  6.07777233e+13,  6.66650212e+13, -2.01560957e+13,\n",
      "       -1.24311576e+14, -1.22958177e+14, -1.54295096e+14, -1.39357798e+14,\n",
      "       -1.45536281e+14, -1.66682687e+14, -1.12285735e+14, -4.71946817e+13,\n",
      "        5.18608604e+13,  3.72872397e+13, -2.44639041e+12, -1.12041457e+14,\n",
      "       -1.13097901e+14, -1.30096663e+14, -1.25259884e+14, -1.44211813e+14,\n",
      "       -1.33954183e+14, -4.19749337e+13, -3.15287112e+13,  2.72667067e+13,\n",
      "        5.95354466e+13,  9.18497565e+12,  1.46264852e+13, -1.12584506e+14,\n",
      "       -1.50868079e+14, -1.64986905e+14, -1.86815203e+14, -1.47898799e+14,\n",
      "       -1.52828296e+14, -1.37272358e+14, -1.00316319e+14, -3.57798387e+13])\n",
      "Coordinates:\n",
      "    grid_xu_ocean  float64 -217.0\n",
      "  * time           (time) object 1972-01-16 12:00:00 ... 2021-12-16 12:00:00\n",
      "iteration done\n",
      "--------------\n",
      "CPU times: user 15.5 s, sys: 2.04 s, total: 17.6 s\n",
      "Wall time: 3min\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "average_DT = xr.open_mfdataset(base + '1deg_jra55_iaf_branch2000_tropics_40SN/' + 'output*/ocean/ocean_wmass.nc').average_DT\n",
    "average_DT = np.mean(average_DT.values.astype('timedelta64[s]')).astype(int) # convert to seconds\n",
    "# [seconds, seconds2]\n",
    "\n",
    "round_accuracy = 1\n",
    "\n",
    "for transect in [\"ITF_horizontal\",\"ITF_vertical\"]:\n",
    "    archive = '1deg_jra55_iaf_branch2000_extratropics_40SN/' # use the correct time integration here for the output in monthly resolution\n",
    "\n",
    "    # open wmass.nc file\n",
    " \n",
    "    if transect == \"ITF_horizontal\":\n",
    "        wmass = xr.open_mfdataset(base + archive + 'output*/ocean/ocean_wmass.nc', parallel=True, preprocess = pre_ITH).load()\n",
    "        wmass = wmass[\"wmass\"]\n",
    "\n",
    "        print(f\"wmass calculated for {transect}\")\n",
    "        print(wmass)\n",
    "        # add together the different diagnostics from the advection, submeso, gm and diffusive mixing schemes\n",
    "        # these fluxes are cp*diagnostic*dxt*rho_dzt*temp binned to neutral density (Watt)\n",
    "        #temp_flux_on_nrho = (wmass.temp_yflux_adv_on_nrho + wmass.temp_yflux_submeso_on_nrho + wmass.temp_yflux_gm_on_nrho + wmass.temp_yflux_ndiffuse_on_nrho)\n",
    "        ITF_horizontal = time_calc(wmass, average_DT)\n",
    "\n",
    "    if transect == \"ITF_vertical\":\n",
    "        wmass = xr.open_mfdataset(base + archive + 'output*/ocean/ocean_wmass.nc', parallel=True, preprocess = pre_ITV).load()\n",
    "        wmass = wmass[\"wmass\"]\n",
    "        print(f\"wmass calculated for {transect}\")\n",
    "        print(wmass)\n",
    "        #temp_flux_on_nrho = (wmass.temp_xflux_adv_on_nrho + wmass.temp_xflux_submeso_on_nrho + wmass.temp_xflux_gm_on_nrho + wmass.temp_xflux_ndiffuse_on_nrho)\n",
    "        ITF_vertical = time_calc(wmass, average_DT)\n",
    "\n",
    "#    if transect == 0:\n",
    "#        ITF_horizontal = temp_flux_on_nrho\n",
    "#    else:\n",
    "#        ITF_vertical = temp_flux_on_nrho\n",
    "\n",
    "    print('iteration done')\n",
    "    print('--------------')\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-20.07]",
   "language": "python",
   "name": "conda-env-analysis3-20.07-py"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}