Skip to content

Instantly share code, notes, and snippets.

@ccarouge

ccarouge/Calculating_heat_fluxes_across_ITF.ipynb

Created March 22, 2021 04:08
Show Gist options
  • Save ccarouge/934e1212f638fece55439ae1e58e2086 to your computer and use it in GitHub Desktop.
Save ccarouge/934e1212f638fece55439ae1e58e2086 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment