AnnualCycle.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook to test the fixed annual cycle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import xgcm\n",
    "from matplotlib import pyplot as plt\n",
    "import nc_time_axis\n",
    "import numpy as np\n",
    "\n",
    "#plt.rcParams['axes.facecolor']='black'\n",
    "#plt.style.use('dark_background')\n",
    "plt.style.use('default')\n",
    "\n",
    "plt.rcParams.update({'font.size': 20})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path   = \"/data/dkamm/nemo_output/DINO/\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load domain (not lazy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import xnemogcm as xn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir_semi    = Path(path + '../NEVERWORLD/GAUS_RIDGE_init/restart7')\n",
    "datadir_annual  = Path(path + 'AnnualCycle/restart7')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domcfg_semi = xn.open_domain_cfg(datadir=datadir_semi)\n",
    "domcfg_annual   = xn.open_domain_cfg(datadir=datadir_annual  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domcfg_semi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual = xn.open_nemo(domcfg=domcfg_annual, files=[\n",
    "        datadir_annual / '../restart7/NEVERWORLD_grid_T.nc',\n",
    "        datadir_annual / '../restart7/NEVERWORLD_grid_U.nc',\n",
    "        datadir_annual / '../restart7/NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "\n",
    "Data_semi = xn.open_nemo(domcfg=domcfg_semi, files=[\n",
    "        datadir_semi / '../restart7/NEVERWORLD_grid_T.nc',\n",
    "        datadir_semi / '../restart7/NEVERWORLD_grid_U.nc',\n",
    "        datadir_semi / '../restart7/NEVERWORLD_grid_V.nc'\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual_mo = xn.open_nemo(domcfg=domcfg_annual, files=[\n",
    "        datadir_annual / '../restart7/NEVERWORLD_grid_T_mo.nc'\n",
    "    ])\n",
    "\n",
    "Data_semi_mo = xn.open_nemo(domcfg=domcfg_semi, files=[\n",
    "        datadir_semi / '../restart7/NEVERWORLD_grid_T_mo.nc'\n",
    "    ])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monthly averaged MLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual_mo.mldr10_1.isel(t=slice(-120,-1)).groupby(\"t.month\").mean().plot(col='month', col_wrap=4)\n",
    "plt.suptitle('Monthly average over the last ten years: Annual cycle', y=1.04)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_semi_mo.mldr10_1.isel(t=slice(-120,-1)).groupby(\"t.month\").mean().plot(col='month', col_wrap=4)\n",
    "plt.suptitle('Monthly average over the last ten years: Semi-annual cycle', y=1.04)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barotropic Streamfunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = {\n",
    "    ('X',): ['e1t', 'e1u', 'e1v', 'e1f'], # X distances\n",
    "    ('Y',): ['e2t', 'e2u', 'e2v', 'e2f'], # Y distances\n",
    "    ('Z',): ['e3t_0', 'e3u_0', 'e3v_0', 'e3f_0'] # Z distances\n",
    "}\n",
    "\n",
    "grid_annual = xgcm.Grid(domcfg_annual,\n",
    "        coords={\n",
    "            \"X\": {\"right\": \"x_f\", \"center\": \"x_c\"},\n",
    "            \"Y\": {\"right\": \"y_f\", \"center\":\"y_c\"},\n",
    "            \"Z\": {\"right\": \"z_f\", \"center\": \"z_c\"}\n",
    "        },\n",
    "        periodic=False,\n",
    "        metrics=metrics\n",
    ")\n",
    "\n",
    "grid_semi = xgcm.Grid(domcfg_semi,\n",
    "        coords={\n",
    "            \"X\": {\"right\": \"x_f\", \"center\": \"x_c\"},\n",
    "            \"Y\": {\"right\": \"y_f\", \"center\":\"y_c\"},\n",
    "            \"Z\": {\"right\": \"z_f\", \"center\": \"z_c\"}\n",
    "        },\n",
    "        periodic=False,\n",
    "        metrics=metrics\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Interpolating on u_f\n",
    "Data_annual['u_f'] = grid_annual.interp(Data_annual.uoce, 'Y')\n",
    "Data_semi['u_f'] = grid_semi.interp(Data_semi.uoce, 'Y')\n",
    "\n",
    "#Integrating over depth\n",
    "Data_annual['U'] = grid_annual.integrate(Data_annual.u_f, 'Z')\n",
    "Data_semi['U'] = grid_semi.integrate(Data_semi.u_f, 'Z')\n",
    "\n",
    "#Cumulative integral over x\n",
    "Psi_annual = (Data_annual.U[-1,::-1,:] * domcfg_annual.e2f[::-1,:]).cumsum('y_f') / 1e6\n",
    "Psi_semi = (Data_semi.U[-1,::-1,:] * domcfg_semi.e2f[::-1,:]).cumsum('y_f') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(16,10))\n",
    "a = Psi_annual.plot.contourf(x='x_f', y='y_f', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "cbar1 = fig.colorbar(a, ax=axs[:2], location='bottom', label=r'$\\psi$ [Sv]')\n",
    "b = Psi_semi.plot.contourf(x='x_f', y='y_f', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "axs[0].set_ylabel('latitude [°N]')\n",
    "axs[1].set_ylabel('latitude [°N]')\n",
    "axs[0].set_xlabel('longitude [°E]')\n",
    "axs[1].set_xlabel('longitude [°E]')\n",
    "axs[0].set_title('Semi-annual Cycle')\n",
    "axs[1].set_title('Annual Cycle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Psi_semi = Psi_semi.assign_coords(dict(\n",
    "    x_sh=([\"y_f\", \"x_f\"], Psi_semi.glamf.values - 60.)\n",
    "    ))\n",
    "fig, axs = plt.subplots(1,1,figsize=(8,6))\n",
    "a = Psi_semi.plot.contourf(x='x_sh', y='gphif', cmap='RdBu_r', levels=30, add_colorbar=False)\n",
    "cbar1 = fig.colorbar(a, ax=axs, label=r'$\\psi$ [Sv]')\n",
    "axs.set_ylabel('latitude [°N]')\n",
    "axs.set_xlabel('longitude [°E]')\n",
    "plt.title('')\n",
    "axs.set_facecolor('none')\n",
    "#fig.savefig('BSF.png', transparent=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stratification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gsw as gsw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual['sigma0'] = gsw.sigma0(Data_annual.soce, gsw.CT_from_pt(Data_annual.soce, Data_annual.toce))\n",
    "Data_semi['sigma0'] = gsw.sigma0(Data_semi.soce, gsw.CT_from_pt(Data_semi.soce, Data_semi.toce))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual = Data_annual.assign_coords(dict(\n",
    "    depth_v=([\"z_c\", \"y_f\", \"x_c\"], - Data_annual.isel(t=-1).e3v.cumsum('z_c').values),\n",
    "    depth_t=([\"z_c\", \"y_c\", \"x_c\"], - Data_annual.isel(t=-1).e3t.cumsum('z_c').values),\n",
    "    lon_f=([\"y_c\", \"x_f\"], domcfg_annual.glamf.values-60.),\n",
    "    lon_t=([\"y_c\", \"x_c\"], domcfg_annual.glamt.values-60.),\n",
    "    lat_f=([\"y_f\", \"x_c\"], domcfg_annual.gphif.values),\n",
    "    lat_t=([\"y_c\", \"x_c\"], domcfg_annual.gphit.values),\n",
    "    )\n",
    ")\n",
    "\n",
    "Data_semi = Data_semi.assign_coords(dict(\n",
    "    depth_v=([\"z_c\", \"y_f\", \"x_c\"], - Data_semi.isel(t=-1).e3v.cumsum('z_c').values),\n",
    "    depth_t=([\"z_c\", \"y_c\", \"x_c\"], - Data_semi.isel(t=-1).e3t.cumsum('z_c').values),\n",
    "    lon_f=([\"y_c\", \"x_f\"], domcfg_semi.glamf.values-60.),\n",
    "    lon_t=([\"y_c\", \"x_c\"], domcfg_semi.glamt.values-60.),\n",
    "    lat_f=([\"y_f\", \"x_c\"], domcfg_semi.gphif.values),\n",
    "    lat_t=([\"y_c\", \"x_c\"], domcfg_semi.gphit.values),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual.sigma0.isel(t=-1, x_c=15, z_c=slice(0,25), y_c=slice(1,-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(12,8), sharey=True)\n",
    "a = (Data_annual.isel(t=-1, x_c=15, z_c=slice(0,25), y_c=slice(1,-1))).sigma0.plot.contourf(x='lat_t', y='depth_t', cmap='Blues_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "cbar1 = fig.colorbar(a, ax=axs[:2], location='bottom', label=r'$\\sigma$ [g / m^3 - 1000]')\n",
    "b = (Data_semi.isel(t=-1, x_c=15, z_c=slice(0,25), y_c=slice(1,-1))).sigma0.plot.contourf(x='lat_t', y='depth_t', cmap='Blues_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "axs[0].set_ylabel('depth [m]')\n",
    "axs[1].set_ylabel('depth [m]')\n",
    "axs[0].set_facecolor('grey')\n",
    "axs[1].set_facecolor('grey')\n",
    "axs[0].set_xlabel('latitude [°N]')\n",
    "axs[1].set_xlabel('latitude [°N]')\n",
    "axs[0].set_title('Semi-annual Cycle', fontsize=15)\n",
    "axs[1].set_title('Annual Cycle', fontsize=15)\n",
    "plt.suptitle('Density of the top 1000m', fontsize=17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(12,8), sharey=True)\n",
    "a = (Data_annual.sigma0.isel(t=-1, x_c=15, z_c=slice(24,-1), y_c=slice(1,-1))).plot.contourf(x='lat_t', y='depth_t', cmap='Blues_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "cbar1 = fig.colorbar(a, ax=axs[:2], location='bottom', label=r'$\\sigma$ [g / m^3 - 1000]')\n",
    "b = (Data_semi.sigma0.isel(t=-1, x_c=15, z_c=slice(24,-1), y_c=slice(1,-1))).plot.contourf(x='lat_t', y='depth_t', cmap='Blues_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "axs[0].set_ylabel('depth [m]')\n",
    "axs[1].set_ylabel('depth [m]')\n",
    "axs[0].set_facecolor('grey')\n",
    "axs[1].set_facecolor('grey')\n",
    "axs[0].set_xlabel('latitude [°N]')\n",
    "axs[1].set_xlabel('latitude [°N]')\n",
    "axs[0].set_title('Semi-annual Cycle', fontsize=15)\n",
    "axs[1].set_title('Annual Cycle', fontsize=15)\n",
    "plt.suptitle('Density below 1000m', fontsize=17)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cross-equatorial meridional velocities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_contours = list(np.linspace(24.25,28.0,16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(12,8), sharey=True)\n",
    "a = (Data_annual.vocetr_eff.isel(t=-1, y_f=100, z_c=slice(0,25)) / 10e6).plot.contourf(x='x_c', y='depth_v', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "Data_annual.sigma0.isel(t=-1, y_c=100, x_c=slice(1,-1), z_c=slice(0,25)).plot.contour(ax=axs[1], y='depth_t', colors='grey', levels=sigma_contours)\n",
    "cbar1 = fig.colorbar(a, ax=axs[:2], location='bottom', label=r'$\\psi$ [Sv]')\n",
    "b = (Data_semi.vocetr_eff.isel(t=-1, y_f=100, z_c=slice(0,25)) / 10e6).plot.contourf(x='x_c', y='depth_v', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "Data_semi.sigma0.isel(t=-1, y_c=100, x_c=slice(1,-1), z_c=slice(0,25)).plot.contour(ax=axs[0], y='depth_t', colors='grey', levels=sigma_contours)\n",
    "axs[0].set_ylabel('depth [m]')\n",
    "axs[1].set_ylabel('depth [m]')\n",
    "axs[0].set_facecolor('grey')\n",
    "axs[1].set_facecolor('grey')\n",
    "axs[0].set_xlabel('longitude [°E]')\n",
    "axs[1].set_xlabel('longitude [°E]')\n",
    "axs[0].set_title('Semi-annual Cycle', fontsize=15)\n",
    "axs[1].set_title('Annual Cycle', fontsize=15)\n",
    "plt.suptitle('Cross-equatorial velocities in the top 1000m', fontsize=17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_contours = list(np.linspace(28.12,28.19,16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(12,8), sharey=True)\n",
    "a = (Data_annual.vocetr_eff.isel(t=-1, y_f=100, z_c=slice(24,-1)) / 10e6).plot.contourf(x='x_c', y='depth_v', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "Data_annual.sigma0.isel(t=-1, y_c=100, x_c=slice(1,-1), z_c=slice(24,-1)).plot.contour(ax=axs[1], y='depth_t', colors='grey', levels=sigma_contours)\n",
    "cbar1 = fig.colorbar(a, ax=axs[:2], location='bottom', label=r'$\\psi$ [Sv]')\n",
    "b = (Data_semi.vocetr_eff.isel(t=-1, y_f=100, z_c=slice(24,-1)) / 10e6).plot.contourf(x='x_c', y='depth_v', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "Data_semi.sigma0.isel(t=-1, y_c=100, x_c=slice(1,-1), z_c=slice(24,-1)).plot.contour(ax=axs[0], y='depth_t', colors='grey', levels=sigma_contours)\n",
    "axs[0].set_ylabel('depth [m]')\n",
    "axs[1].set_ylabel('depth [m]')\n",
    "axs[0].set_facecolor('grey')\n",
    "axs[1].set_facecolor('grey')\n",
    "axs[0].set_xlabel('longitude [°E]')\n",
    "axs[1].set_xlabel('longitude [°E]')\n",
    "axs[0].set_title('Semi-annual Cycle', fontsize=15)\n",
    "axs[1].set_title('Annual Cycle', fontsize=15)\n",
    "plt.suptitle('Cross-equatorial velocities in the top 1000m', fontsize=17)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monthly averaged MLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_semi_mo.qsr.isel(y_c=150, x_c=30, t=slice(-13,-1)).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual_mo.qsr.isel(t=slice(-13,-1), x_c=30, y_c=175).plot()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Meridional overturning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cutting the bottom cell and create a new grid object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cf_xarray\n",
    "\n",
    "Data_annual_cp = Data_annual.isel(z_c=slice(0,-1)).copy()\n",
    "Data_semi_cp = Data_semi.isel(z_c=slice(0,-1)).copy()\n",
    "\n",
    "level_outer_annual = cf_xarray.bounds_to_vertices(Data_annual_cp.deptht_bounds, 'axis_nbounds').load().data\n",
    "level_outer_semi = cf_xarray.bounds_to_vertices(Data_semi_cp.deptht_bounds, 'axis_nbounds').load().data\n",
    "\n",
    "Data_annual_cp = Data_annual_cp.assign_coords({'z_f': level_outer_annual})\n",
    "Data_semi_cp = Data_semi_cp.assign_coords({'z_f': level_outer_semi})\n",
    "\n",
    "grid_annual_cp = xgcm.Grid(Data_annual_cp,\n",
    "    coords={\n",
    "        \"X\": {\"center\": \"x_c\"},\n",
    "        \"Y\": {\"right\": \"y_f\", \"center\":\"y_c\"},\n",
    "        \"Z\": {\"center\": \"z_c\", \"outer\": \"z_f\"}\n",
    "    },\n",
    "    periodic=False)\n",
    "\n",
    "grid_semi_cp = xgcm.Grid(Data_semi_cp,\n",
    "    coords={\n",
    "        \"X\": {\"center\": \"x_c\"},\n",
    "        \"Y\": {\"right\": \"y_f\", \"center\":\"y_c\"},\n",
    "        \"Z\": {\"center\": \"z_c\", \"outer\": \"z_f\"}\n",
    "    },\n",
    "    periodic=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute potential density, referenced to 2000m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_annual_cp['sigma0'] = gsw.sigma0(Data_annual_cp.soce, gsw.CT_from_pt(Data_annual_cp.soce, Data_annual_cp.toce))\n",
    "Data_semi_cp['sigma0'] = gsw.sigma0(Data_semi_cp.soce, gsw.CT_from_pt(Data_semi_cp.soce, Data_semi_cp.toce))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolate sigma2 on the cell faces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_v_annual = grid_annual_cp.interp(Data_annual_cp.sigma0, 'Y')\n",
    "sigma_outer_annual = grid_annual_cp.interp(sigma_v_annual, 'Z',  boundary='extend')\n",
    "\n",
    "sigma_v_semi = grid_semi_cp.interp(Data_semi_cp.sigma0, 'Y')\n",
    "sigma_outer_semi = grid_semi_cp.interp(sigma_v_semi, 'Z',  boundary='extend')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Target values for the new density coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "sigma_target = np.array([24, 24.2, 24.4, 24.65, 24.9, 25.15, 25.4, 25.65,\n",
    "                        25.9, 26.15, 26.4, 26.57, 26.75, 26.90, 27.05, 27.17,\n",
    "                        27.30, 27.37, 27.45, 27.52, 27.58,\n",
    "                        27.63, 27.68, 27.72, 27.75, 27.78, 27.80, 27.815,\n",
    "                        27.83, 27.845, 27.86, 27.875, 27.89,\n",
    "                        27.905, 27.92, 27.935, 27.95, 27.965,\n",
    "                        27.98, 27.99, 28.01, 28.025,\n",
    "                        28.04, 28.055, 28.07, 28.085, 28.1, 28.35,\n",
    "                        28.6, 28.8, 29.1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the meridional transport on density coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vtra_annual = grid_annual_cp.transform(Data_annual_cp.vocetr_eff,\n",
    "        'Z',\n",
    "        sigma_target,\n",
    "        method='conservative',\n",
    "        target_data=sigma_outer_annual)\n",
    "\n",
    "vtra_semi = grid_semi_cp.transform(Data_semi_cp.vocetr_eff,\n",
    "        'Z',\n",
    "        sigma_target,\n",
    "        method='conservative',\n",
    "        target_data=sigma_outer_semi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the meridional overturning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Over_annual    = vtra_annual.isel(t=-1).sum(dim='x_c')[:,::-1].cumsum('sigma0') / 1e6\n",
    "Over_semi    = vtra_semi.isel(t=-1).sum(dim='x_c')[:,::-1].cumsum('sigma0') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Over_annual = Over_annual.assign_coords(dict({'y_f': domcfg_annual.gphif.isel(x_f=10).values}))\n",
    "Over_semi = Over_semi.assign_coords(dict({'y_f': domcfg_semi.gphif.isel(x_f=10).values}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(12,8), sharey=True)\n",
    "a = (-Over_annual).plot.contourf(x='y_f', y='sigma0', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[1])\n",
    "cbar1 = fig.colorbar(a, ax=axs, label=r'$\\psi$ [Sv]')\n",
    "b = (-Over_semi).plot.contourf(x='y_f', y='sigma0', cmap='RdBu_r', levels=30, add_colorbar=False, ax=axs[0], vmin=cbar1.vmin, vmax=cbar1.vmax)\n",
    "axs[0].invert_yaxis()\n",
    "axs[0].set_ylabel(r'$\\sigma_{0}$ [ $kg$ / $m^3$ - 1000 ]')\n",
    "axs[1].set_ylabel(r'$\\sigma_{0}$ [ $kg$ / $m^3$ - 1000 ]')\n",
    "axs[0].set_xlabel('latitude [°N]')\n",
    "axs[1].set_xlabel('latitude [°N]')\n",
    "axs[0].set_title('Semi-annual Cycle', fontsize=15)\n",
    "axs[1].set_title('Annual Cycle', fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!jupyter nbconvert --to html ./AnnualCycle.ipynb --HTMLExporter.theme=dark "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ACC strength"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_restarts = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Combine restart files\n",
    "for i in range(n_restarts-2, 0, -1):\n",
    "    _datadir = Path(path + 'AnnualCycle/restart{}/'.format(i))\n",
    "    _Data = xn.open_nemo(domcfg=domcfg_annual, files=[\n",
    "        _datadir / 'NEVERWORLD_grid_T.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_U.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "    Data_annual = xr.concat([_Data.chunk(), Data_annual.chunk()], \"t\")\n",
    "\n",
    "for i in range(n_restarts-2, 0, -1):\n",
    "    _datadir = Path(path + '../NEVERWORLD/GAUS_RIDGE_init/restart{}/'.format(i))\n",
    "    _Data = xn.open_nemo(domcfg=domcfg_semi, files=[\n",
    "        _datadir / 'NEVERWORLD_grid_T.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_U.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "    Data_semi = xr.concat([_Data.chunk(), Data_semi.chunk()], \"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Compute ACC transport\n",
    "ACC_annual = Data_annual.isel(x_f=0).uocetr_eff.sum(['y_c', 'z_c']) / 1e6\n",
    "ACC_semi = Data_semi.isel(x_f=0).uocetr_eff.sum(['y_c', 'z_c']) / 1e6\n",
    "\n",
    "#OR!!!!\n",
    "#ACC_annual = (Data_annual.isel(x_f=0).uoce * Data_annual.isel(x_f=0).e2u * Data_annual.isel(x_f=0).e3u).sum(['y_c', 'z_c']) / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cftime as cft\n",
    "ACC_annual = ACC_annual.assign_coords(dict({'t':cft.to_tuple(ACC_annual.indexes['t'])[0]}))\n",
    "ACC_semi = ACC_semi.assign_coords(dict({'t':cft.to_tuple(ACC_semi.indexes['t'])[0]}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,1, figsize=(12,6))\n",
    "ACC_annual.plot(color='midnightblue', label='Annaul Cycle')\n",
    "ACC_semi.plot(color='darkred', label='Semi-annual cycle')\n",
    "axs.set_ylabel('ACC transport [Sv]')\n",
    "plt.title('')\n",
    "plt.xlabel('time [years]')\n",
    "plt.legend()\n",
    "plt.grid()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xgcm",
   "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.9.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}