Presentation_1.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook to produce all plots for first year presentations (EGU2023, CPT2023, M2Lines2023)"
   ]
  },
  {
   "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",
    "\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/NEVERWORLD/\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data (lazy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_restarts = 8      # Number of restarts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = xr.open_dataset(path + 'restart0/NEVERWORLD_grid_V.nc')\n",
    "U = xr.open_dataset(path + 'restart0/NEVERWORLD_grid_U.nc')\n",
    "T = xr.open_dataset(path + 'restart0/NEVERWORLD_grid_T.nc')\n",
    "for i in range(1,n_restarts):\n",
    "    _path = path + 'restart{}/'.format(i)\n",
    "    _V = xr.open_dataset(_path + 'NEVERWORLD_grid_V.nc')\n",
    "    _U = xr.open_dataset(_path + 'NEVERWORLD_grid_U.nc')\n",
    "    _T = xr.open_dataset(_path + 'NEVERWORLD_grid_T.nc')\n",
    "    V = xr.concat([V.chunk(), _V.chunk()], \"time_counter\")\n",
    "    U = xr.concat([U.chunk(), _U.chunk()], \"time_counter\")\n",
    "    T = xr.concat([T.chunk(), _T.chunk()], \"time_counter\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_restarts_hr = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V_hr = xr.open_dataset(path + 'restart10/NEVERWORLD_grid_V.nc')\n",
    "U_hr = xr.open_dataset(path + 'restart10/NEVERWORLD_grid_U.nc')\n",
    "T_hr = xr.open_dataset(path + 'restart10/NEVERWORLD_grid_T.nc')\n",
    "for i in range(n_restarts + 1,n_restarts + n_restarts_hr):\n",
    "    _path = path + 'restart{}/'.format(i)\n",
    "    _V_hr = xr.open_dataset(_path + 'NEVERWORLD_grid_V.nc')\n",
    "    _U_hr = xr.open_dataset(_path + 'NEVERWORLD_grid_U.nc')\n",
    "    _T_hr = xr.open_dataset(_path + 'NEVERWORLD_grid_T.nc')\n",
    "    V_hr = xr.concat([V_hr.chunk(), _V_hr.chunk()], \"time_counter\")\n",
    "    U_hr = xr.concat([U_hr.chunk(), _U_hr.chunk()], \"time_counter\")\n",
    "    T_hr = xr.concat([T_hr.chunk(), _T_hr.chunk()], \"time_counter\")"
   ]
  },
  {
   "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_no_ridge = Path(path + 'NO_RIDGE_init/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domcfg_no_ridge = xn.open_domain_cfg(datadir=datadir_no_ridge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir_1deg = Path(path + 'restart9')\n",
    "datadir_1_4deg_old = Path(path + 'restart10_1')\n",
    "datadir_1_4deg = Path(path + 'restart11')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domcfg_1deg = xn.open_domain_cfg(datadir=datadir_1deg)\n",
    "domcfg_1_4deg = xn.open_domain_cfg(datadir=datadir_1_4deg)\n",
    "domcfg_1_4deg_old = xn.open_domain_cfg(datadir=datadir_1_4deg_old)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bathymetry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cartopy.crs as ccrs\n",
    "\n",
    "bathy = domcfg_1_4deg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Barotropic streamfunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg = xn.open_nemo(domcfg=domcfg_1deg, files=[\n",
    "        datadir_1deg / 'NEVERWORLD_grid_T.nc',\n",
    "        datadir_1deg / 'NEVERWORLD_grid_U.nc',\n",
    "        datadir_1deg / 'NEVERWORLD_grid_V.nc'\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolating uoce on y_f"
   ]
  },
  {
   "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_1deg = xgcm.Grid(domcfg_1deg,\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": [
    "Data_1deg['u_f'] = grid_1deg.interp(Data_1deg.uoce, 'Y')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrate over depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg['U'] = grid_1deg.integrate(Data_1deg.u_f, 'Z')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cummulative integral in i-direction at last time step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Psi = (Data_1deg.U[-1,::-1,:] * domcfg_1deg.e2f[::-1,:]).cumsum('y_f') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# U_s_grid = U_s_grid.assign_coords({'x_f': U_s_grid.x_f.values - 60.})\n",
    "fig, axs = plt.subplots(1,1,figsize=(8,6))\n",
    "a = Psi.plot.contourf(x='x_f', y='y_f', 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('')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Barotropic streamfunction (1/4 degree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1_4deg = xn.open_nemo(domcfg=domcfg_1_4deg, files=[\n",
    "        datadir_1_4deg / 'NEVERWORLD_grid_T.nc',\n",
    "        datadir_1_4deg / 'NEVERWORLD_grid_U.nc',\n",
    "        datadir_1_4deg / 'NEVERWORLD_grid_V.nc'\n",
    "    ])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolating uoce on y_f"
   ]
  },
  {
   "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_1_4deg = xgcm.Grid(domcfg_1_4deg,\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": [
    "Data_1_4deg['u_f'] = grid_1_4deg.interp(Data_1_4deg.uoce, 'Y')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrate over depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1_4deg['U'] = grid_1_4deg.integrate(Data_1_4deg.u_f, 'Z')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cummulative integral in i-direction at last time step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Psi_1_4deg = (Data_1_4deg.U[0,::-1,:] * domcfg_1_4deg.e2f[::-1,:]).cumsum('y_f') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# U_s_grid = U_s_grid.assign_coords({'x_f': U_s_grid.x_f.values - 60.})\n",
    "fig, axs = plt.subplots(1,1,figsize=(8,6))\n",
    "a = Psi_1_4deg.plot.contourf(x='x_f', y='y_f', 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('')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ACC strength time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg_500y = Data_1deg.copy()\n",
    "for i in range(8,-1, -1):\n",
    "    _datadir = Path(path + 'restart{}/'.format(i))\n",
    "    _Data = xn.open_nemo(domcfg=domcfg_1deg, files=[\n",
    "        _datadir / 'NEVERWORLD_grid_T.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_U.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "    Data_1deg_500y = xr.concat([_Data.chunk(), Data_1deg_500y.chunk()], \"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC_1deg = Data_1deg_500y.isel(x_f=0).uocetr_eff.sum(['y_c', 'z_c']) / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC_1deg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10,9, -1):\n",
    "    _datadir = Path(path + 'restart{}/'.format(i))\n",
    "    _Data = xn.open_nemo(domcfg=domcfg_1_4deg, files=[\n",
    "        _datadir / 'NEVERWORLD_grid_T.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_U.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "    Data_1_4deg = xr.concat([_Data.chunk(), Data_1_4deg.chunk()], \"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC_1_4deg = Data_1_4deg.isel(x_f=0).uocetr_eff.sum(['y_c', 'z_c']) / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC = xr.concat([ACC_1deg.chunk(), ACC_1_4deg.chunk()], \"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cftime as cft\n",
    "ACC = ACC.assign_coords(dict({'t':cft.to_tuple(ACC.indexes['t'])[0]}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,1, figsize=(12,6))\n",
    "a = plt.plot(ACC.t.values[:500],ACC.data[:500],color='black', label='1° resolution')\n",
    "a = plt.plot(ACC.t.values[500:],ACC.data[500:],color='darkred', label='1/4° resolution')\n",
    "axs.set_ylabel('ACC transport [Sv]')\n",
    "plt.title('')\n",
    "plt.xlabel('time [years]')\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "\n",
    "axes = plt.axes([.45, .4, .40, .2])\n",
    "axes.plot(ACC.t.values[494:500], ACC.data[494:500], c='black', lw=1)\n",
    "axes.plot(ACC.t.values[500:], ACC.data[500:], c='darkred', lw=1)\n",
    "axes.plot(ACC.t.values[499:501], ACC.data[499:501], c='grey', linestyle=':', lw=1)\n",
    "\n",
    "axs.set_facecolor('none')\n",
    "fig.savefig('ACC.png', transparent=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Meridional overturning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gsw as gsw"
   ]
  },
  {
   "attachments": {},
   "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_1deg_top = Data_1deg_500y.isel(z_c=slice(0,-1)).copy()\n",
    "\n",
    "level_outer_data = cf_xarray.bounds_to_vertices(Data_1deg_top.deptht_bounds.isel(t=-1), 'axis_nbounds').load().data\n",
    "\n",
    "Data_1deg_top = Data_1deg_top.assign_coords({'z_f': level_outer_data})\n",
    "\n",
    "grid_1deg_top = xgcm.Grid(Data_1deg_top,\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute potential density, referenced to 2000m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg_top['sigma0'] = gsw.sigma0(Data_1deg_top.soce, gsw.CT_from_pt(Data_1deg_top.soce, Data_1deg_top.toce))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolate sigma2 on the cell faces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_v = grid_1deg_top.interp(Data_1deg_top.sigma0, 'Y')\n",
    "sigma_outer = grid_1deg_top.interp(sigma_v, 'Z',  boundary='extend')"
   ]
  },
  {
   "attachments": {},
   "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])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the meridional transport on density coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vtra_transformed = grid_1deg_top.transform(Data_1deg_top.vocetr_eff,\n",
    "        'Z',\n",
    "        sigma_target,\n",
    "        method='conservative',\n",
    "        target_data=sigma_outer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the meridional overturning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Over_1deg    = vtra_transformed.isel(t=-1).sum(dim='x_c')[:,::-1].cumsum('sigma0') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Over_1deg = Over_1deg.assign_coords(dict({'y_f': domcfg_1deg.gphif.isel(x_f=10).values}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,1,figsize=(8,6))\n",
    "a = (-Over_1deg).plot.contourf(x='y_f', y='sigma0', cmap='RdBu_r', levels=30, add_colorbar=False)\n",
    "cbar1 = fig.colorbar(a, ax=axs, label=r'$\\psi$ [Sv]')\n",
    "plt.xlabel('latitude [°N]')\n",
    "axs.invert_yaxis()\n",
    "plt.ylabel(r'$\\sigma_{0}$ [ $kg$ / $m^3$ - 1000 ]')\n",
    "plt.title('')\n",
    "fig.savefig('MOC.png', transparent=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Meridional overturning (1/4 degree)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cutting the bottom cell and create a new grid object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1_4deg_top = Data_1_4deg.isel(z_c=slice(0,-1)).copy()\n",
    "\n",
    "level_outer_data = cf_xarray.bounds_to_vertices(Data_1_4deg_top.deptht_bounds, 'axis_nbounds').load().data\n",
    "\n",
    "Data_1_4deg_top = Data_1_4deg_top.assign_coords({'z_f': level_outer_data})\n",
    "\n",
    "grid_1_4deg_top = xgcm.Grid(Data_1_4deg_top,\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute potential density, referenced to 2000m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1_4deg_top['sigma0'] = gsw.sigma0(Data_1_4deg_top.soce, gsw.CT_from_pt(Data_1_4deg_top.soce, Data_1_4deg_top.toce))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolate sigma2 on the cell faces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_v_1_4 = grid_1_4deg_top.interp(Data_1_4deg_top.sigma0, 'Y')\n",
    "sigma_outer_1_4 = grid_1_4deg_top.interp(sigma_v_1_4, 'Z',  boundary='extend')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the meridional transport on density coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vtra_transformed_1_4 = grid_1_4deg_top.transform(Data_1_4deg_top.vocetr_eff,\n",
    "        'Z',\n",
    "        sigma_target,\n",
    "        method='conservative',\n",
    "        target_data=sigma_outer_1_4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the meridional overturning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Over_1_4deg    = vtra_transformed_1_4.isel(t=-1).sum(dim='x_c')[:,::-1].cumsum('sigma0') / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,1,figsize=(8,6))\n",
    "a = (-Over_1_4deg).plot.contourf(x='y_f', y='sigma0', cmap='RdBu_r', levels=30, add_colorbar=False)\n",
    "cbar1 = fig.colorbar(a, ax=axs, label=r'$\\psi$ [Sv]')\n",
    "plt.xlabel('latitude [°N]')\n",
    "axs.invert_yaxis()\n",
    "plt.ylabel(r'$\\sigma_{0}$ [ $kg$ / $m^3$ - 1000 ]')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum overturning streamfunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SEOS(ds):\n",
    "    \"\"\"Simplified equation of state.\"\"\"\n",
    "    params = dict({'a_0': 0.165, 'b_0': 0.76554, 'lambda_1': 0.06, 'mu_1': 0.0001497 ,'rho_0': 1.026 })\n",
    "    T_a = ds['toce'] - 10.\n",
    "    S_a = ds['soce'] - 35. \n",
    "    z   = ds.e3t.cumsum('z_c')\n",
    "    rho_a = (((-params['a_0'] * (1 + 0.5 * params['lambda_1'] * T_a + params['mu_1'] * z) * T_a + params['b_0'] * S_a) / params['rho_0'] ) + 1 ) * params['rho_0'] \n",
    "    return rho_a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_a = SEOS(Data_1deg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_a.isel(t=-1, x_c = 25, y_c=slice(1,-1), z_c=slice(0,-1)).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = dict({'a_0': 0.165, 'b_0': 0.76554, 'lambda_1': 0.06, 'mu_1': 0.0001497 ,'rho_0': 1.026 })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_1deg_top['sigma0'] = gsw.sigma0(Data_1deg_top.soce, gsw.CT_from_pt(Data_1deg_top.soce, Data_1deg_top.toce))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Low scottia arc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domcfg_SA_low.e3t_0.sum('z_c')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_no_ridge = xn.open_nemo(domcfg=domcfg_SA_low, files=[\n",
    "        datadir_SA_low / 'NEVERWORLD_grid_T.nc',\n",
    "        datadir_SA_low / 'NEVERWORLD_grid_U.nc',\n",
    "        datadir_SA_low / 'NEVERWORLD_grid_V.nc'\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1,n_restarts, 1):\n",
    "    _datadir = Path(path + 'restart{}/'.format(i))\n",
    "    _Data = xn.open_nemo(domcfg=domcfg_SA_low, files=[\n",
    "        _datadir / 'NEVERWORLD_grid_T.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_U.nc',\n",
    "        _datadir / 'NEVERWORLD_grid_V.nc'\n",
    "    ])\n",
    "    Data_SA_low = xr.concat([Data_SA_low.chunk(), _Data.chunk()], \"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cftime as cft\n",
    "Data_SA_low = Data_SA_low.assign_coords(dict({'t':cft.to_tuple(Data_SA_low.indexes['t'])[0]}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC_low = Data_SA_low.isel(x_f=0).uocetr_eff.sum(['y_c', 'z_c']) / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC_low.plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nemo_env",
   "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
}