plot_heat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e699727f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import glob\n",
    "import netCDF4 as nc\n",
    "import datetime as dt\n",
    "import sys\n",
    "import gsw as sw\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib\n",
    "matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.path as mpath\n",
    "import cftime\n",
    "import coast\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c57df360",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "var = 'thetao' # thetao, so, uo, vo, siconc, siage, sivol, sithick, siu, siv, \n",
    "i_o = 'O' # SI or O for sea ice or ocean\n",
    "freq = 'mon' # mon or day\n",
    "time_s = 'highres-future' # 'highres-future' or 'hist-1950'\n",
    "\n",
    "def make_path(var, i_o, freq, time_s):\n",
    "    if 'future' in time_s:\n",
    "        ddir = 'MOHC'\n",
    "    else:\n",
    "        ddir = 'NERC'\n",
    "    root = '/badc/cmip6/data/CMIP6/HighResMIP/' + ddir + '/HadGEM3-GC31-HH/' + time_s + '/r1i1p1f1/'\n",
    "    return root + i_o + freq + '/' + var + '/gn/latest/' + var + '_' + i_o + freq + '_HadGEM3-GC31-HH_' + time_s + '_r1i1p1f1_gn_*.nc'\n",
    "\n",
    "fn_nemo_dat_t = make_path('thetao', i_o, freq, time_s)\n",
    "fn_nemo_dat_s = make_path('so', i_o, freq, time_s)\n",
    "\n",
    "domain_root = '/gws/nopw/j04/nemo_vol5/acc/eORCA12-N512/domain/'\n",
    "fn_nemo_dom1 = domain_root + 'eORCA12_coordinates.nc'\n",
    "fn_nemo_dom = domain_root + 'mesh_mask_eORCA12_v2.4.nc'\n",
    "fn_nemo_bathy = domain_root + 'eORCA12_bathymetry_v2.4.nc'\n",
    "fn_config_t_grid = './config/gc31_nemo_grid_t.json'\n",
    "\n",
    "out_file = './Processed/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d47c6799",
   "metadata": {},
   "outputs": [],
   "source": [
    "flist_t = sorted(glob.glob(fn_nemo_dat_t))\n",
    "flist_s = sorted(glob.glob(fn_nemo_dat_s))\n",
    "nemo_t = coast.Gridded(fn_data = flist_t[:3], fn_domain = fn_nemo_dom, config=fn_config_t_grid, multiple=True)\n",
    "nemo_s = coast.Gridded(fn_data = flist_s[:3], fn_domain = fn_nemo_dom, config=fn_config_t_grid, multiple=True)\n",
    "nemo_t.dataset['salinity'] = nemo_s.dataset.salinity\n",
    "\n",
    "print(nemo_t.dataset.longitude.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "debf5187-857b-47bd-90c4-22c1d3cd7cf8",
   "metadata": {},
   "source": [
    "Subset data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dec2799",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "#print(nemo_t.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c5164b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "yi1 = 2800\n",
    "yi2 = nemo_t.dataset.longitude.shape[0]\n",
    "ind_y = np.arange(yi1, yi2).astype(int)\n",
    "#print(yi2, ind_y)\n",
    "nemo_t_subset = nemo_t.isel(y_dim=ind_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c543ee-0efc-48cc-86a5-7139ab51579c",
   "metadata": {},
   "source": [
    "Time slice. Notes dates are 360 day years so use cftime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8967cfd1-5c17-4c46-bb59-b187eb9ad5ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Does not work\n",
    "#st_date = cftime.datetime(2050, 12, 1)\n",
    "#en_date = cftime.datetime(2051, 1, 1) \n",
    "#nemo_t_subset = nemo_t_subset.time_slice(st_date, en_date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f169cf3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(nemo_t_subset.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "481b58fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#nemo_t_subset.construct_density(eos='EOS10')#, pot_dens=True)\n",
    "\n",
    "def calc_rho(sp, tp, depth, lon, lat):\n",
    "    pres = sw.p_from_z(depth * -1, lat)\n",
    "    sa = sw.SA_from_SP(sp, pres, lon, lat)\n",
    "    ct = sw.CT_from_pt(sa, tp)\n",
    "    rho = sw.rho(sa, ct, 0)\n",
    "    return rho, ct\n",
    "\n",
    "def calc_heat(rho, ct):\n",
    "    cap = 3991.86795711963 # 3985 # J kg-1 K-1\n",
    "    #cap = sw.cp0 # use with conservative temp\n",
    "    return ct * rho * cap\n",
    "\n",
    "# change this to decrease resolution but decrease run time\n",
    "sub = 5\n",
    "\n",
    "ds_dom = xr.open_dataset(fn_nemo_dom).squeeze().rename({\"z\": \"z_dim\", \"x\": \"x_dim\", \"y\": \"y_dim\"})\n",
    "e1t = ds_dom.e1t[yi1+1:-1:sub, 1:-1:sub] # y, x\n",
    "e2t = ds_dom.e2t[yi1+1:-1:sub, 1:-1:sub]\n",
    "e3t = ds_dom.e3t_0[:, yi1+1:-1:sub, 1:-1:sub] # z, y, x\n",
    "e1t = np.tile(e1t, (e3t.shape[0], 1, 1))\n",
    "e2t = np.tile(e2t, (e3t.shape[0], 1, 1))\n",
    "print(e3t.shape, e1t.shape)\n",
    "volume = e1t * e2t * e3t\n",
    "\n",
    "\n",
    "lat = nemo_t_subset.dataset.latitude.values[::sub, ::sub]\n",
    "lon = nemo_t_subset.dataset.longitude.values[::sub, ::sub]\n",
    "depth = nemo_t_subset.dataset.depth_0.values[:]\n",
    "depth_g = np.tile(depth, (lon.shape[1], lon.shape[0], 1)).T\n",
    "\n",
    "mask_south = lat < 70\n",
    "heat_time = np.ma.zeros((nemo_t_subset.dataset.t_dim.shape[0]))\n",
    "for i in range(nemo_t_subset.dataset.t_dim.shape[0]):\n",
    "    temp = nemo_t_subset.dataset.temperature[i, :, ::sub, ::sub].to_masked_array() # time, lev, j, i\n",
    "    sal = nemo_t_subset.dataset.salinity[i, :, ::sub, ::sub].to_masked_array()\n",
    "    rho, ct = calc_rho(sal, temp, depth_g, lon, lat)\n",
    "\n",
    "    heat_cont = np.ma.sum(calc_heat(rho, ct) * volume, axis=0) # vertically integrated\n",
    "    heat_pole = np.ma.masked_where(mask_south, heat_cont)\n",
    "    heat_time[i] = np.ma.sum(heat_pole)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee167fa5-d9b1-46e5-831a-418c4666e2b7",
   "metadata": {},
   "source": [
    "Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d0616d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def set_circle(ax):\n",
    "    # Compute a circle in axes coordinates, which we can use as a boundary\n",
    "    # for the map. We can pan/zoom as much as we like - the boundary will be\n",
    "    # permanently circular.\n",
    "    theta = np.linspace(0, 2 * np.pi, 100)\n",
    "    center, radius = [0.5, 0.5], 0.5\n",
    "    verts = np.vstack([np.sin(theta), np.cos(theta)]).T\n",
    "    circle = mpath.Path(verts * radius + center)\n",
    "    ax.set_boundary(circle, transform=ax.transAxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0230b4a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "data_crs = ccrs.PlateCarree()\n",
    "mrc = ccrs.NorthPolarStereo(central_longitude=0.0)\n",
    "\n",
    "my_cm = plt.cm.plasma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d7f202",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "fig1 = plt.figure(figsize=(12, 6))\n",
    "ax1 = fig1.add_axes([0.03, 0.03, 0.3, 0.85], projection=mrc)\n",
    "ax2 = fig1.add_axes([0.36, 0.03, 0.6, 0.85])\n",
    "\n",
    "cax1 = fig1.add_axes([0.03, 0.94, 0.3, 0.02])\n",
    "\n",
    "cs1 = ax1.pcolormesh(lon, lat, heat_cont, transform=data_crs, cmap=my_cm)\n",
    "ax2.plot(heat_time)\n",
    "\n",
    "ax1.add_feature(cfeature.LAND, zorder=100)\n",
    "ax1.gridlines()\n",
    "ax1.set_extent([-180, 180, 60, 90], crs=data_crs)\n",
    "set_circle(ax1)\n",
    "\n",
    "ax2.set_ylabel('Heat Content (Jm$^{-2}$)')\n",
    "\n",
    "ax1.annotate('(a)', (0.05, 0.95), xycoords='axes fraction', bbox=dict(boxstyle=\"round\", fc=\"w\"), zorder=105)\n",
    "ax2.annotate('(b)', (0.05, 0.95), xycoords='axes fraction', bbox=dict(boxstyle=\"round\", fc=\"w\"), zorder=105)\n",
    "\n",
    "\n",
    "fig1.colorbar(cs1, cax=cax1, orientation='horizontal')\n",
    "cax1.set_xlabel('Heat Content (Jm$^{-3}$)')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3278c989",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "fig1.savefig('./Figures/heat_content.png')\n"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "encoding": "# coding: utf-8",
   "executable": "/usr/bin/env python",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "Python 3 + Jaspy",
   "language": "python",
   "name": "jaspy"
  },
  "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}