eccov4r3_budgetH_anom.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xgcm\n",
    "import zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.10.7'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.18.1+38.gca8d62e9'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define face_connections for grid object\n",
    "face_connections = {'face':\n",
    "                    {0: {'X':  ((12, 'Y', False), (3, 'X', False)),\n",
    "                         'Y':  (None,             (1, 'Y', False))},\n",
    "                     1: {'X':  ((11, 'Y', False), (4, 'X', False)),\n",
    "                         'Y':  ((0, 'Y', False),  (2, 'Y', False))},\n",
    "                     2: {'X':  ((10, 'Y', False), (5, 'X', False)),\n",
    "                         'Y':  ((1, 'Y', False),  (6, 'X', False))},\n",
    "                     3: {'X':  ((0, 'X', False),  (9, 'Y', False)),\n",
    "                         'Y':  (None,             (4, 'Y', False))},\n",
    "                     4: {'X':  ((1, 'X', False),  (8, 'Y', False)),\n",
    "                         'Y':  ((3, 'Y', False),  (5, 'Y', False))},\n",
    "                     5: {'X':  ((2, 'X', False),  (7, 'Y', False)),\n",
    "                         'Y':  ((4, 'Y', False),  (6, 'Y', False))},\n",
    "                     6: {'X':  ((2, 'Y', False),  (7, 'X', False)),\n",
    "                         'Y':  ((5, 'Y', False),  (10, 'X', False))},\n",
    "                     7: {'X':  ((6, 'X', False),  (8, 'X', False)),\n",
    "                         'Y':  ((5, 'X', False),  (10, 'Y', False))},\n",
    "                     8: {'X':  ((7, 'X', False),  (9, 'X', False)),\n",
    "                         'Y':  ((4, 'X', False),  (11, 'Y', False))},\n",
    "                     9: {'X':  ((8, 'X', False),  None),\n",
    "                         'Y':  ((3, 'X', False),  (12, 'Y', False))},\n",
    "                     10: {'X': ((6, 'Y', False),  (11, 'X', False)),\n",
    "                          'Y': ((7, 'Y', False),  (2, 'X', False))},\n",
    "                     11: {'X': ((10, 'X', False), (12, 'X', False)),\n",
    "                          'Y': ((8, 'Y', False),  (1, 'X', False))},\n",
    "                     12: {'X': ((11, 'X', False), None),\n",
    "                          'Y': ((9, 'Y', False),  (0, 'X', False))}}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main disagnostic output\n",
    "ds_main = xr.open_zarr('/rigel/ocp/users/jt2796/eccov4r3_output')\n",
    "coords_main = ds_main.coords.to_dataset().reset_coords()\n",
    "ds_main = ds_main.reset_coords(drop=True)\n",
    "\n",
    "# Budget terms\n",
    "ds_budg = xr.open_zarr('/rigel/ocp/users/jt2796/eccov4r3_budgets')\n",
    "#ds_budg.time.values = ds_main.time.values\n",
    "coords_budg = ds_budg.coords.to_dataset().reset_coords()\n",
    "ds_budg = ds_budg.reset_coords(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell z size\n",
    "drF = coords_main.drF\n",
    "rA = coords_main.rA\n",
    "hFacC = coords_main.hFacC\n",
    "\n",
    "# Volume (m^3)\n",
    "vol = (rA*drF*hFacC).transpose('face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Seperate averages and snapshots\n",
    "ds_ave = ds_main[['ETAN','THETA','SALT','UVELMASS','VVELMASS','WVELMASS',\n",
    "                  'ADVx_TH','ADVx_SLT','ADVy_TH','ADVy_SLT','ADVr_TH','ADVr_SLT']]\n",
    "\n",
    "ds_snp = ds_main[['ETAN_snp','THETA_snp','SALT_snp']].rename({'time_snp':'time'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove oceFWflx from WVELMASS\n",
    "WVELMASS = ds_main.WVELMASS.transpose('time','face','k_l','j','i')\n",
    "oceFWflx = ds_main.oceFWflx.assign_coords(k_l=0).expand_dims('k_l').transpose('time','face','k_l','j','i')\n",
    "\n",
    "rhoconst = 1029\n",
    "oceFWflx = (oceFWflx/rhoconst)\n",
    "WVELMASS = xr.concat([WVELMASS.sel(k_l=0) + oceFWflx, WVELMASS[:,:,1:]], \n",
    "                     dim='k_l').transpose('time','face','k_l','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Monthly means\n",
    "ds_ave_clim = ds_ave.groupby('time.month').mean('time')\n",
    "WVELMASS_clim = WVELMASS.groupby('time.month').mean('time')\n",
    "ds_snp_clim = ds_snp.groupby('time.month').mean('time')\n",
    "ds_budg_clim = ds_budg.groupby('time.month').mean('time')\n",
    "\n",
    "#### Monthly anomalies\n",
    "ds_ave_anom = ds_ave.groupby('time.month') - ds_ave_clim\n",
    "WVELMASS_anom = WVELMASS.groupby('time.month') - WVELMASS_clim\n",
    "ds_snp_anom = ds_snp.groupby('time.month') - ds_snp_clim\n",
    "ds_budg_anom = ds_budg.groupby('time.month') - ds_budg_clim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Monthly mean terms\n",
    "grid = xgcm.Grid(ds_ave_clim, face_connections=face_connections)\n",
    "\n",
    "# Transport\n",
    "u_clim = (ds_ave_clim.UVELMASS * coords_main.dyG * coords_main.drF).transpose('month','face','k','j','i_g')\n",
    "v_clim = (ds_ave_clim.VVELMASS * coords_main.dxG * coords_main.drF).transpose('month','face','k','j_g','i')\n",
    "w_clim = (WVELMASS_clim * coords_main.rA).transpose('month','face','k_l','j','i')\n",
    "\n",
    "# Potential Temperature (degC)\n",
    "THETA_clim = ds_ave_clim.THETA.transpose('month','face','k','j','i')\n",
    "THETA_clim_at_u = grid.interp(THETA_clim, 'X', boundary='extend')\n",
    "THETA_clim_at_v = grid.interp(THETA_clim, 'Y', boundary='extend')\n",
    "THETA_clim_at_w = grid.interp(THETA_clim, 'Z', boundary='extend')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Monthly anomaly terms\n",
    "grid = xgcm.Grid(ds_ave_anom, face_connections=face_connections)\n",
    "\n",
    "# Transport\n",
    "u_anom = (ds_ave_anom.UVELMASS * coords_main.dyG * coords_main.drF).transpose('time','face','k','j','i_g')\n",
    "v_anom = (ds_ave_anom.VVELMASS * coords_main.dxG * coords_main.drF).transpose('time','face','k','j_g','i')\n",
    "w_anom = (WVELMASS_anom * coords_main.rA).transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Temperature\n",
    "THETA_anom = ds_ave_anom.THETA.transpose('time','face','k','j','i')\n",
    "THETA_anom_at_u = grid.interp(THETA_anom, 'X', boundary='extend')\n",
    "THETA_anom_at_v = grid.interp(THETA_anom, 'Y', boundary='extend')\n",
    "THETA_anom_at_w = grid.interp(THETA_anom, 'Z', boundary='extend')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monthly anomaly budget\n",
    "$$\\frac{\\partial\\theta^{\\prime}}{\\partial t} + \\overline{\\bf{u}}^m\\cdot\\nabla\\theta^{\\prime} + \\bf{u}^{\\prime} \\cdot \\overline{\\nabla\\theta}^m -\\nabla \\cdot ({\\bf{u}}^{\\prime}\\,\\theta^{\\prime}-\\overline{\\bf{u}^{\\prime}\\,\\theta^{\\prime}}^m) = -\\nabla \\cdot {\\bf{F_{diff}}}^{\\prime} - F_{\\textrm{forc}}^{\\prime}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Anomaleous heat tendency\n",
    "tendH_anom = ds_budg_anom.tendH\n",
    "\n",
    "#### Anomaleous forcing\n",
    "forcH_anom = ds_budg_anom.forcH\n",
    "\n",
    "#### Anomaleous diffusive heat convergence\n",
    "dif_hConvH = ds_budg_anom.dif_hConvH\n",
    "dif_vConvH = ds_budg_anom.dif_vConvH\n",
    "dif_ConvH_anom = dif_hConvH + dif_vConvH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Mean advection of anomaleous temperature\n",
    "uclimTanom = u_clim * THETA_anom_at_u.groupby('time.month')\n",
    "vclimTanom = v_clim * THETA_anom_at_v.groupby('time.month')\n",
    "wclimTanom = w_clim * THETA_anom_at_w.groupby('time.month')\n",
    "\n",
    "# Convergence\n",
    "ADVxy_diff = grid.diff_2d_vector({'X' : uclimTanom, 'Y' : vclimTanom}, boundary = 'fill')\n",
    "ADVx_diffx = ADVxy_diff['X']\n",
    "ADVy_diffy = ADVxy_diff['Y']\n",
    "adv_hConvH = (-(ADVx_diffx + ADVy_diffy)/vol)\n",
    "adv_vConvH = (grid.diff(wclimTanom, 'Z', boundary='fill')/vol)\n",
    "\n",
    "adv_ConvH_Uclim_Tanom = adv_hConvH + adv_vConvH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n",
      "/rigel/home/jt2796/dask/dask/array/core.py:3460: PerformanceWarning: Increasing number of chunks by factor of 24\n",
      "  **atop_kwargs)\n"
     ]
    }
   ],
   "source": [
    "#### Anomaleous advection of mean temperature\n",
    "uanomTclim = u_anom.groupby('time.month') * THETA_clim_at_u\n",
    "vanomTclim = v_anom.groupby('time.month') * THETA_clim_at_v\n",
    "wanomTclim = w_anom.groupby('time.month') * THETA_clim_at_w\n",
    "\n",
    "# Convergence\n",
    "ADVxy_diff = grid.diff_2d_vector({'X' : uanomTclim, 'Y' : vanomTclim}, boundary = 'fill')\n",
    "ADVx_diffx = ADVxy_diff['X']\n",
    "ADVy_diffy = ADVxy_diff['Y']\n",
    "adv_hConvH = (-(ADVx_diffx + ADVy_diffy)/vol)\n",
    "adv_vConvH = (grid.diff(wanomTclim, 'Z', boundary='fill')/vol)\n",
    "\n",
    "adv_ConvH_Uanom_Tclim = adv_hConvH + adv_vConvH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Anomaleous advection of anomaleous temperature\n",
    "uanomTanom = u_anom * THETA_anom_at_u\n",
    "vanomTanom = v_anom * THETA_anom_at_v\n",
    "wanomTanom = w_anom * THETA_anom_at_w\n",
    "\n",
    "uanomTanom_clim = uanomTanom.groupby('time.month').mean('time')\n",
    "vanomTanom_clim = vanomTanom.groupby('time.month').mean('time')\n",
    "wanomTanom_clim = wanomTanom.groupby('time.month').mean('time')\n",
    "\n",
    "uanomTanom_clim_anom = uanomTanom.groupby('time.month') - uanomTanom_clim\n",
    "vanomTanom_clim_anom = vanomTanom.groupby('time.month') - vanomTanom_clim\n",
    "wanomTanom_clim_anom = wanomTanom.groupby('time.month') - wanomTanom_clim\n",
    "\n",
    "# Convergence\n",
    "ADVxy_diff = grid.diff_2d_vector({'X' : uanomTanom_clim_anom, 'Y' : vanomTanom_clim_anom}, boundary = 'fill')\n",
    "ADVx_diffx = ADVxy_diff['X']\n",
    "ADVy_diffy = ADVxy_diff['Y']\n",
    "adv_hConvH = (-(ADVx_diffx + ADVy_diffy)/vol)\n",
    "adv_vConvH = (grid.diff(wanomTanom_clim_anom, 'Z', boundary='fill')/vol)\n",
    "\n",
    "adv_ConvH_Uanom_Tanom = adv_hConvH + adv_vConvH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Residual in the advective flux\n",
    "\n",
    "# Transport\n",
    "u_transport = (ds_ave.UVELMASS * coords_main.dyG * coords_main.drF)\n",
    "v_transport = (ds_ave.VVELMASS * coords_main.dxG * coords_main.drF)\n",
    "w_transport = (WVELMASS * rA)\n",
    "\n",
    "# Temperature\n",
    "THETA = ds_main.THETA\n",
    "THETA_at_u = grid.interp(THETA, 'X', boundary='extend')\n",
    "THETA_at_v = grid.interp(THETA, 'Y', boundary='extend')\n",
    "THETA_at_w = grid.interp(THETA, 'Z', boundary='extend')\n",
    "\n",
    "# Advection\n",
    "uT = (u_transport * THETA_at_u)\n",
    "vT = (v_transport * THETA_at_v)\n",
    "wT = (w_transport * THETA_at_w)\n",
    "\n",
    "# Convergence\n",
    "ADVxy_diff = grid.diff_2d_vector({'X' : uT, 'Y' : vT}, boundary = 'fill')\n",
    "ADVx_diffx = ADVxy_diff['X']\n",
    "ADVy_diffy = ADVxy_diff['Y']\n",
    "adv_hConvH = (-(ADVx_diffx + ADVy_diffy)/vol).transpose('time','face','k','j','i')\n",
    "adv_vConvH = (grid.diff(wT, 'Z', boundary='fill')/vol).transpose('time','face','k','j','i')\n",
    "\n",
    "# Reconstructed\n",
    "adv_ConvH_reco = adv_hConvH + adv_vConvH\n",
    "\n",
    "# True (diagnostic)\n",
    "adv_ConvH_true = ds_budg.adv_hConvH + ds_budg.adv_vConvH\n",
    "\n",
    "# Residual\n",
    "adv_ConvH_res = adv_ConvH_true - adv_ConvH_reco\n",
    "adv_ConvH_res_clim = adv_ConvH_res.groupby('time.month').mean('time')\n",
    "adv_ConvH_res_anom = adv_ConvH_res.groupby('time.month') - adv_ConvH_res_clim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save to dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.Dataset(data_vars={})\n",
    "ds['tendH_anom'] = tendH_anom\n",
    "ds['forcH_anom'] = forcH_anom\n",
    "ds['dif_ConvH_anom'] = dif_ConvH_anom\n",
    "ds['adv_ConvH_Uclim_Tanom'] = adv_ConvH_Uclim_Tanom\n",
    "ds['adv_ConvH_Uanom_Tclim'] = adv_ConvH_Uanom_Tclim\n",
    "ds['adv_ConvH_Uanom_Tanom'] = adv_ConvH_Uanom_Tanom\n",
    "ds['adv_ConvH_res_anom'] = adv_ConvH_res_anom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\\frac{\\partial\\theta^{\\prime}}{\\partial t} + \\overline{\\bf{u}}^m\\cdot\\nabla\\theta^{\\prime} + \\bf{u}^{\\prime} \\cdot \\overline{\\nabla\\theta}^m -\\nabla \\cdot ({\\bf{u}}^{\\prime}\\,\\theta^{\\prime}-\\overline{\\bf{u}^{\\prime}\\,\\theta^{\\prime}}^m) = -\\nabla \\cdot {\\bf{F_{diff}}}^{\\prime} - F_{\\textrm{forc}}^{\\prime}$$\n",
    "- `tendH_anom`: Anomaleous tendency ($\\frac{\\partial\\theta^{\\prime}}{\\partial t}$)\n",
    "- `forcH_anom`: Anomaleous forcing ($F_{\\textrm{forc}}^{\\prime}$)\n",
    "- `dif_ConvH_anom`: Anomaleous diffusion ($\\nabla \\cdot {\\bf{F_{diff}}}^{\\prime}$)\n",
    "- `adv_ConvH_Uclim_Tanom`: Mean advection of T anomalies ($\\overline{\\bf{u}}^m\\cdot\\nabla\\theta^{\\prime}$)\n",
    "- `adv_ConvH_Uanom_Tclim`: Anomaleous advection of mean T ($\\bf{u}^{\\prime} \\cdot \\overline{\\nabla\\theta}^m$)\n",
    "- `adv_ConvH_Uanom_Tanom`: $\\nabla \\cdot ({\\bf{u}}^{\\prime}\\,\\theta^{\\prime}-\\overline{\\bf{u}^{\\prime}\\,\\theta^{\\prime}}^m)$\n",
    "- `adv_ConvH_res_anom`: Residual in advection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}