eccov4r3_budgets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xgcm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Load dataset\n",
    "ds = xr.open_zarr('/rigel/ocp/users/jt2796/eccov4r3_output/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Creating the 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))}}}\n",
    "\n",
    "grid = xgcm.Grid(ds, face_connections=face_connections)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Define terms\n",
    "# Before doing the budget calculations we need to define some terms that will be used in the budget calculations.\n",
    "\n",
    "### Number of seconds between each snapshot\n",
    "# There are no snapshots for the first and last time point. Thus, we are skipping budget calculations for first and last month of the given time period.\n",
    "dt = ds.time_snp[1:].load()\n",
    "\n",
    "# delta t in seconds. Note: devide by 10**9 to convert nanoseconds to seconds\n",
    "dt.values = [float(t)/10**9 for t in np.diff(ds.time_snp)]\n",
    "\n",
    "# Rename time (and iter) axis\n",
    "dt = dt.rename({'time_snp':'time','iter_snp':'iter'})\n",
    "\n",
    "# time axis of dt should be the same as of the monthly averages\n",
    "dt.time.values = ds.time[1:-1].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Relevant constants\n",
    "# Density kg/m^3\n",
    "rhoconst = 1029\n",
    "\n",
    "# Heat capacity (J/kg/K)\n",
    "c_p = 3994\n",
    "\n",
    "# Constants for surface heat penetration (from Table 2 of Paulson and Simpson, 1977)\n",
    "R = 0.62\n",
    "zeta1 = 0.6\n",
    "zeta2 = 20.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Ocean depth\n",
    "Depth = ds.Depth\n",
    "\n",
    "### Grid dimensions\n",
    "dxG = ds.dxG\n",
    "dyG = ds.dyG\n",
    "drF = ds.drF\n",
    "rA = ds.rA\n",
    "hFacC = ds.hFacC.load()\n",
    "\n",
    "# Volume (m^3)\n",
    "vol = (rA*drF*hFacC).transpose('face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Land masks\n",
    "\n",
    "# Make copy of hFacC\n",
    "mskC = hFacC.copy(deep=True).load()\n",
    "\n",
    "# Change all fractions (ocean) to 1. land = 0\n",
    "mskC.values[mskC.values>0] = 1\n",
    "\n",
    "# Make 2D land mask for surface (This is just for plotting/mapping purposes)\n",
    "land_mask = mskC[0]\n",
    "land_mask.values[land_mask.values==0] = np.nan\n",
    "\n",
    "# Make 3D land mask\n",
    "land_mask_3d = mskC.copy()\n",
    "land_mask_3d.values[land_mask_3d.values==0] = np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Evaluating the volume budget\n",
    "### Total tendency\n",
    "# - ETAN: Surface Height Anomaly (m)\n",
    "\n",
    "# Load snapshots for surface height anomaly from dataset\n",
    "ETANsnp = ds.ETAN_snp\n",
    "\n",
    "# Total tendency (1/s)\n",
    "tendV = xr.DataArray(50*[1],coords={'k': np.array(range(0,50))},dims=['k'])*grid.diff(ETANsnp, 'T', boundary='fill', \n",
    "                                                                                      fill_value=0.0)/Depth/dt\n",
    "tendV['Z'] = ds.Z\n",
    "tendV['hFacC'] = ds.hFacC\n",
    "tendV = xr.concat([np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[0]).expand_dims('time'),tendV,\n",
    "                   np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[-1]).expand_dims('time')],\n",
    "                  dim='time').transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Forcing\n",
    "# - oceFWflx: net surface Fresh-Water flux into the ocean (kg/m^2/s)\n",
    "\n",
    "# Load monthly averaged freshwater flux\n",
    "oceFWflx = ds.oceFWflx.assign_coords(k=0).expand_dims('k')\n",
    "\n",
    "# Sea surface forcing on volume (1/s)\n",
    "forcV = xr.concat([(oceFWflx/rhoconst)/(hFacC*drF), \n",
    "                   xr.zeros_like(((oceFWflx[0]/rhoconst)/(hFacC*drF)).transpose('time','face','k','j','i'))[:,:,1:]], \n",
    "                  dim='k').transpose('time','face','k','j','i')*land_mask_3d.transpose('face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Horizontal convergence\n",
    "# - UVELMASS: U Mass-Weighted Comp of Velocity (m/s)\n",
    "# - VVELMASS: V Mass-Weighted Comp of Velocity (m/s)\n",
    "\n",
    "# Load monthly averaged velocities\n",
    "UVELMASS = ds.UVELMASS\n",
    "VVELMASS = ds.VVELMASS\n",
    "\n",
    "# Horizontal volume transports (m^3/s)\n",
    "u_transport = UVELMASS * dyG * drF\n",
    "v_transport = VVELMASS * dxG * drF\n",
    "\n",
    "uv_diff = grid.diff_2d_vector({'X' : u_transport, 'Y' : v_transport}, boundary = 'fill')\n",
    "\n",
    "u_diffx = uv_diff['X']\n",
    "v_diffy = uv_diff['Y']\n",
    "\n",
    "# Convergence of the horizontal flow (1/s)\n",
    "hConvV = (-(u_diffx + v_diffy)/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Vertical convergence\n",
    "# - WVELMASS: Vertical Mass-Weighted Comp of Velocity (m/s)\n",
    "\n",
    "# Load monthly averaged vertical velocity\n",
    "WVELMASS = ds.WVELMASS.transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Vertical volume transport (m^3/s)\n",
    "w_transport = (WVELMASS * rA).transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Apparently, it is required to add the vertical volume flux at the air-sea interface (`oceFWflx`) to the surface layer to balance the budget.\n",
    "# Add the vertical volume flux at the air-sea interface\n",
    "w_transport = xr.concat([(w_transport.sel(k_l=0) + (forcV*vol).sel(k=0)).expand_dims('k_l').drop(['k','Z','PHrefC','drF','hFacC']),\n",
    "                         w_transport[:,:,1:]], dim='k_l').transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Convergence of the vertical flow (1/s)\n",
    "vConvV = grid.diff(w_transport, 'Z', boundary='fill')/vol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Evaluating the heat budget\n",
    "\n",
    "# ### Total tendency\n",
    "# - THETA: Potential Temperature (degC)\n",
    "\n",
    "# Load snapshots of theta\n",
    "THETAsnp = ds.THETA_snp\n",
    "\n",
    "# Calculate the s∗theta term\n",
    "HCsnp = (THETAsnp*(1+ETANsnp/Depth)).transpose('time_snp','face','k','j','i')\n",
    "\n",
    "# Total tendency (degC/s)\n",
    "tendH = grid.diff(HCsnp, 'T', boundary='fill', fill_value=0.0)/dt\n",
    "\n",
    "# Add coordinates\n",
    "tendH['Depth'] = ds.Depth\n",
    "tendH['XC'] = ds.XC\n",
    "tendH['YC'] = ds.YC\n",
    "tendH['Z'] = ds.Z\n",
    "tendH['rA'] = ds.rA\n",
    "tendH['PHrefC'] = ds.PHrefC\n",
    "tendH['drF'] = ds.drF\n",
    "tendH['hFacC'] = ds.hFacC\n",
    "tendH = xr.concat([np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[0]).expand_dims('time'),tendH,\n",
    "                   np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[-1]).expand_dims('time')],\n",
    "                  dim='time').transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Forcing\n",
    "# - TFLUX: total heat flux (match heat-content variations) (W/m^2)\n",
    "# - oceQsw: net Short-Wave radiation (+=down) (W/m^2)\n",
    "\n",
    "# Load monthly averages of heat flux and shortwave radiation\n",
    "TFLUX = ds.TFLUX\n",
    "oceQsw = ds.oceQsw\n",
    "\n",
    "#### Defining terms needed for evaluating surface heat forcing\n",
    "Z = ds.Z.load()\n",
    "RF = np.concatenate([ds.Zp1.values[:-1],[np.nan]])\n",
    "\n",
    "# **Note**: `Z` and `Zp1` are used in deriving surface heat penetration. MATLAB code uses `RF` from `mygrid` structure.\n",
    "\n",
    "q1 = R*np.exp(1.0/zeta1*RF[:-1]) + (1.0-R)*np.exp(1.0/zeta2*RF[:-1])\n",
    "q2 = R*np.exp(1.0/zeta1*RF[1:]) + (1.0-R)*np.exp(1.0/zeta2*RF[1:])\n",
    "\n",
    "# Correction for the 200m cutoff\n",
    "zCut = np.where(Z < -200)[0][0]\n",
    "q1[zCut:] = 0\n",
    "q2[zCut-1:] = 0\n",
    "\n",
    "# Save q1 and q2 as xarray data arrays\n",
    "q1 = xr.DataArray(q1,coords=[Z.k],dims=['k'])\n",
    "q2 = xr.DataArray(q2,coords=[Z.k],dims=['k'])\n",
    "\n",
    "\n",
    "#### Compute vertically penetrating flux\n",
    "\n",
    "# Surface heat flux (below the surface)\n",
    "forcH_subsurf = ((q1*(mskC==1)-q2*(mskC.shift(k=-1)==1))*oceQsw).transpose('time','face','k','j','i')\n",
    "\n",
    "# Surface heat flux (at the sea surface)\n",
    "forcH_surf = ((TFLUX - (1-(q1[0]-q2[0]))*oceQsw)*mskC[0]).transpose('time','face','j','i').assign_coords(k=0).expand_dims('k')\n",
    "\n",
    "# Full-depth forcing\n",
    "forcH = xr.concat([forcH_surf,forcH_subsurf[:,:,1:]], dim='k').transpose('time','face','k','j','i')\n",
    "\n",
    "# #### Add geothermal heat flux\n",
    "# Geothermal flux needs to be a three dimensional field since the sources are distributed along the ocean floor at various depths. This requires a three dimensional mask.\n",
    "\n",
    "# Create 3d bathymetry mask\n",
    "mskC_shifted = mskC.shift(k=-1)\n",
    "\n",
    "mskC_shifted.values[-1,:,:,:] = 0\n",
    "mskb = mskC - mskC_shifted\n",
    "\n",
    "# Create 3d field of geothermal heat flux\n",
    "geoflx3d = ds.GEOFLX * mskb.transpose('face','k','j','i')\n",
    "GEOFLX = geoflx3d.transpose('face','k','j','i')\n",
    "\n",
    "# Add geothermal heat flux to forcing field and convert from W/m^2 to degC/s\n",
    "forcH = ((forcH + GEOFLX)/(rhoconst*c_p))/(hFacC*drF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Advection\n",
    "# #### Horizontal convergence\n",
    "# - ADVx_TH: U Comp. Advective Flux of Pot.Temperature (degC m^3/s)\n",
    "# - ADVy_TH: V Comp. Advective Flux of Pot.Temperature (degC m^3/s)\n",
    "\n",
    "ADVx_TH = ds.ADVx_TH\n",
    "ADVy_TH = ds.ADVy_TH\n",
    "\n",
    "ADVxy_diff = grid.diff_2d_vector({'X' : ADVx_TH, 'Y' : ADVy_TH}, boundary = 'fill')\n",
    "\n",
    "ADVx_diffx = ADVxy_diff['X']\n",
    "ADVy_diffy = ADVxy_diff['Y']\n",
    "\n",
    "# Convergence of horizontal advection (degC/s)\n",
    "adv_hConvH = (-(ADVx_diffx + ADVy_diffy)/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #### Vertical convergence\n",
    "# - ADVr_TH: Vertical Advective Flux of Pot.Temperature (degC m^3/s)\n",
    "\n",
    "# Load monthly averages of vertical advective flux\n",
    "ADVr_TH = ds.ADVr_TH.transpose('time','face','k_l','j','i')\n",
    "\n",
    "# For `ADVr_TH`, `DFrE_TH` and `DFrI_TH`, we need to make sure that sequence of dimensions are consistent. \n",
    "# When loading the fields use `.transpose('time','face','k_l','j','i')`.\n",
    "#Otherwise, the divergences will be not correct (at least for face = 12).\n",
    "\n",
    "# Convergence of vertical advection (degC/s)\n",
    "adv_vConvH = (grid.diff(ADVr_TH, 'Z', boundary='fill')/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Diffusion\n",
    "# #### Horizontal convergence\n",
    "# - DFxE_TH: U Comp. Diffusive Flux of Pot.Temperature (degC m^3/s)\n",
    "# - DFyE_TH: V Comp. Diffusive Flux of Pot.Temperature (degC m^3/s)\n",
    "\n",
    "# Load monthly averages of diffusive fluxes\n",
    "DFxE_TH = ds.DFxE_TH\n",
    "DFyE_TH = ds.DFyE_TH\n",
    "\n",
    "DFxyE_diff = grid.diff_2d_vector({'X' : DFxE_TH, 'Y' : DFyE_TH}, boundary = 'fill')\n",
    "\n",
    "DFxE_diffx = DFxyE_diff['X']\n",
    "DFyE_diffy = DFxyE_diff['Y']\n",
    "\n",
    "# Convergence of horizontal diffusion (degC/s)\n",
    "dif_hConvH = (-(DFxE_diffx + DFyE_diffy)/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #### Vertical convergence\n",
    "# - DFrE_TH: Vertical Diffusive Flux of Pot.Temperature (Explicit part) (degC m^3/s)\n",
    "# - DFrI_TH: Vertical Diffusive Flux of Pot.Temperature (Implicit part) (degC m^3/s)\n",
    "\n",
    "# Load monthly averages of vertical diffusive fluxes\n",
    "DFrE_TH = ds.DFrE_TH.transpose('time','face','k_l','j','i')\n",
    "DFrI_TH = ds.DFrI_TH.transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Convergence of vertical diffusion (degC/s)\n",
    "dif_vConvH = ((grid.diff(DFrE_TH, 'Z', boundary='fill') + grid.diff(DFrI_TH, 'Z', boundary='fill'))/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## Evaluating the salt budget\n",
    "\n",
    "# ### Total tendency\n",
    "# - SALT: Salinity (psu)\n",
    "\n",
    "# Load salinity snapshot\n",
    "SALTsnp = ds.SALT_snp\n",
    "\n",
    "# Calculate s*S term\n",
    "sSALT = (SALTsnp*(1+ETANsnp/Depth)).transpose('time_snp','face','k','j','i')\n",
    "\n",
    "# Total tendency (psu/s)\n",
    "tendS = grid.diff(sSALT, 'T', boundary='fill', fill_value=0.0)/dt\n",
    "\n",
    "# Add coordinates\n",
    "tendS['Depth'] = ds.Depth\n",
    "tendS['XC'] = ds.XC\n",
    "tendS['YC'] = ds.YC\n",
    "tendS['Z'] = ds.Z\n",
    "tendS['rA'] = ds.rA\n",
    "tendS['PHrefC'] = ds.PHrefC\n",
    "tendS['drF'] = ds.drF\n",
    "tendS['hFacC'] = ds.hFacC\n",
    "tendS = xr.concat([np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[0]).expand_dims('time'),tendS,\n",
    "                   np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[-1]).expand_dims('time')],\n",
    "                  dim='time').transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Forcing\n",
    "# - SFLUX: total salt flux (match salt-content variations) (g/m^2/s)\n",
    "# - oceSPtnd: salt tendency due to salt plume flux (g/m^2/s)\n",
    "\n",
    "# Load monthly averaged SFLUX and oceSPtnd\n",
    "SFLUX = ds.SFLUX.assign_coords(k=0).expand_dims('k')\n",
    "oceSPtnd = ds.oceSPtnd.transpose('time','face','k','j','i')\n",
    "\n",
    "# `SFLUX` and `oceSPtnd` is given in g/m^2/s.\n",
    "# Dividing by density and corresponding vertical length scale (`drF`) results in g/kg/s, which is the same as psu/s.\n",
    "\n",
    "# Surface salt flux (psu/s)\n",
    "forcS_surf = ((SFLUX/rhoconst)/(hFacC*drF)).transpose('time','face','k','j','i')\n",
    "\n",
    "# Salt tendency (psu/s)\n",
    "forcS_tend = ((oceSPtnd/rhoconst)/(hFacC*drF)).transpose('time','face','k','j','i')\n",
    "\n",
    "forcS = xr.concat([forcS_surf+forcS_tend,forcS_tend[:,:,1:]], dim='k').transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Advection\n",
    "# #### Horizontal convergence\n",
    "# - ADVx_SLT: U Comp. Advective Flux of Salinity (psu m^3/s)\n",
    "# - ADVy_SLT: V Comp. Advective Flux of Salinity (psu m^3/s)\n",
    "\n",
    "# Load monthly averaged advective fluxes\n",
    "ADVx_SLT = ds.ADVx_SLT\n",
    "ADVy_SLT = ds.ADVy_SLT\n",
    "\n",
    "ADVxy_SLT_diff = grid.diff_2d_vector({'X' : ADVx_SLT, 'Y' : ADVy_SLT}, boundary = 'fill')\n",
    "\n",
    "ADVx_SLT_diffx = ADVxy_SLT_diff['X']\n",
    "ADVy_SLT_diffy = ADVxy_SLT_diff['Y']\n",
    "\n",
    "# Convergence of horizontal advection (psu/s)\n",
    "adv_hConvS = (-(ADVx_SLT_diffx + ADVy_SLT_diffy)/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #### Vertical convergence\n",
    "# - ADVr_SLT: Vertical Advective Flux of Salinity (psu m^3/s)\n",
    "\n",
    "# Load monthly averaged vertical advective flux\n",
    "ADVr_SLT = ds.ADVr_SLT.transpose('time','face','k_l','j','i')\n",
    "\n",
    "# The salt budget only balances when the sea surface forcing is not added to the vertical salt flux \n",
    "#(at the air-sea interface). This is different from the volume and salinity budget.\n",
    "\n",
    "# Convergence of vertical advection (psu/s)\n",
    "adv_vConvS = (grid.diff(ADVr_SLT, 'Z', boundary='fill')/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### Diffusion\n",
    "# #### Horizontal convergence\n",
    "# - DFxE_SLT: U Comp. Diffusive Flux of Salinity (psu m^3/s)\n",
    "# - DFyE_SLT: V Comp. Diffusive Flux of Salinity (psu m^3/s)\n",
    "\n",
    "# Load monthly averaged horizontal diffusive fluxes\n",
    "DFxE_SLT = ds.DFxE_SLT\n",
    "DFyE_SLT = ds.DFyE_SLT\n",
    "\n",
    "DFxyE_SLT_diff = grid.diff_2d_vector({'X' : DFxE_SLT, 'Y' : DFyE_SLT}, boundary = 'fill')\n",
    "\n",
    "DFxE_SLT_diffx = DFxyE_SLT_diff['X']\n",
    "DFyE_SLT_diffy = DFxyE_SLT_diff['Y']\n",
    "\n",
    "# Convergence of horizontal diffusion (psu/s)\n",
    "dif_hConvS = (-(DFxE_SLT_diffx + DFyE_SLT_diffy)/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #### Vertical convergence\n",
    "# - DFrE_SLT: Vertical Diffusive Flux of Salinity (Explicit part) (psu m^3/s)\n",
    "# - DFrI_SLT: Vertical Diffusive Flux of Salinity (Implicit part) (psu m^3/s)\n",
    "\n",
    "# Load monthly averaged vertical diffusive fluxes\n",
    "DFrE_SLT = ds.DFrE_SLT.transpose('time','face','k_l','j','i')\n",
    "DFrI_SLT = ds.DFrI_SLT.transpose('time','face','k_l','j','i')\n",
    "\n",
    "# Convergence of vertical diffusion (psu/s)\n",
    "dif_vConvS = ((grid.diff(DFrE_SLT, 'Z', boundary='fill') + grid.diff(DFrI_SLT, 'Z', boundary='fill'))/vol).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Evaluating the salinity budget\n",
    "#### Scale factor\n",
    "# - Depth: Ocean_depth (m)\n",
    "# - ETAN: Surface Height Anomaly (m)\n",
    "\n",
    "# Load monthly averaged surface height anomaly\n",
    "ETAN = ds.ETAN\n",
    "\n",
    "# Scale factor\n",
    "rstarfac = ((Depth + ETAN)/Depth).transpose('time','face','j','i')\n",
    "\n",
    "### Total tendency\n",
    "# Total tendency (psu/s)\n",
    "tendSln = grid.diff(SALTsnp, 'T', boundary='fill', fill_value=0.0)/dt\n",
    "\n",
    "# Add coordinates\n",
    "tendSln['Depth'] = ds.Depth\n",
    "tendSln['XC'] = ds.XC\n",
    "tendSln['YC'] = ds.YC\n",
    "tendSln['Z'] = ds.Z\n",
    "tendSln['rA'] = ds.rA\n",
    "tendSln['PHrefC'] = ds.PHrefC\n",
    "tendSln['drF'] = ds.drF\n",
    "tendSln['hFacC'] = ds.hFacC\n",
    "\n",
    "tendSln = xr.concat([np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[0]).expand_dims('time'),tendSln,\n",
    "                     np.nan*xr.zeros_like(vol).assign_coords(time=ds.time[-1]).expand_dims('time')],\n",
    "                    dim='time').transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Forcing\n",
    "# The forcing term is comprised of both salt flux (`forcS`) and volume (i.e., freshwater) flux (`forcV`).\n",
    "# - SALT: Salinity (psu)\n",
    "# Load monthly averaged salinity fields\n",
    "SALT = ds.SALT.transpose('time','face','k','j','i')\n",
    "\n",
    "### Sea surface forcing for salinity (psu/s)\n",
    "forcSln = ((-SALT*forcV + forcS)/rstarfac).transpose('time','face','k','j','i')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Advection\n",
    "### Horizontal convergence\n",
    "adv_hConvSln = ((-SALT*hConvV + adv_hConvS)/rstarfac).transpose('time','face','k','j','i')\n",
    "\n",
    "### Vertical convergence\n",
    "adv_vConvSln = ((-SALT*vConvV + adv_vConvS)/rstarfac).transpose('time','face','k','j','i')\n",
    "\n",
    "#### Diffusion\n",
    "### Horizontal convergence\n",
    "dif_hConvSln = dif_hConvS/rstarfac\n",
    "\n",
    "### Vertical convergence\n",
    "dif_vConvSln = dif_vConvS/rstarfac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Save to dataset\n",
    "\n",
    "varnames = [\n",
    "'tendV',\n",
    " 'forcV',\n",
    " 'hConvV',\n",
    " 'vConvV',\n",
    " 'tendH',\n",
    " 'forcH',\n",
    " 'adv_hConvH',\n",
    " 'adv_vConvH',\n",
    " 'dif_hConvH',\n",
    " 'dif_vConvH',\n",
    " 'tendS',\n",
    " 'forcS',\n",
    " 'adv_hConvS',\n",
    " 'adv_vConvS',\n",
    " 'dif_hConvS',\n",
    " 'dif_vConvS',\n",
    " 'tendSln',\n",
    " 'forcSln',\n",
    " 'adv_hConvSln',\n",
    " 'adv_vConvSln',\n",
    " 'dif_hConvSln',\n",
    " 'dif_vConvSln'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out = xr.Dataset(data_vars={})\n",
    "for varname in varnames:\n",
    "    ds_out[varname] = globals()[varname].chunk(chunks={'time':1,'face':13,'k':50,'j':90,'i':90})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out = ds_out.assign_coords(dt=dt)\n",
    "ds_out.dt.attrs = {'units': 's','standard_name': 'dt','coordinate': 'time','long_name': 'time span between snapshots'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out.time.encoding = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:       (face: 13, i: 90, j: 90, k: 50, time: 288)\n",
       "Coordinates:\n",
       "    Depth         (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    XC            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    YC            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "  * face          (face) int64 0 1 2 3 4 5 6 7 8 9 10 11 12\n",
       "  * i             (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "  * j             (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "    rA            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    PHrefC        (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "    Z             (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "    drF           (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "  * k             (k) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "    hFacC         (k, face, j, i) float32 dask.array<shape=(50, 13, 90, 90), chunksize=(50, 13, 90, 90)>\n",
       "    iter          (time) float64 nan 1.428e+03 2.172e+03 2.892e+03 3.636e+03 ...\n",
       "  * time          (time) datetime64[ns] 1992-01-15 1992-02-13 1992-03-15 ...\n",
       "    dt            (time) float64 nan 2.506e+06 2.678e+06 2.592e+06 2.678e+06 ...\n",
       "Data variables:\n",
       "    tendV         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcV         (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    hConvV        (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    vConvV        (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendH         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcH         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_hConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendS         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcS         (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_hConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendSln       (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcSln       (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_hConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save to zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.diagnostics import ProgressBar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/rigel/home/jt2796/miniconda/envs/default/lib/python3.6/site-packages/xarray/backends/api.py:858: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs\n",
      "  dataset.dump_to_store(store, sync=True, encoding=encoding, compute=compute)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[                                        ] | 0% Completed |  1min  8.5s"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/rigel/home/jt2796/dask/dask/local.py:255: RuntimeWarning: divide by zero encountered in true_divide\n",
      "  return func(*args2)\n",
      "/rigel/home/jt2796/dask/dask/local.py:255: RuntimeWarning: invalid value encountered in true_divide\n",
      "  return func(*args2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[                                        ] | 0% Completed |  1min  8.9s"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/rigel/home/jt2796/dask/dask/local.py:255: RuntimeWarning: invalid value encountered in multiply\n",
      "  return func(*args2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[########################################] | 100% Completed | 17min 49.0s\n"
     ]
    }
   ],
   "source": [
    "with ProgressBar():\n",
    "    ds_out.to_zarr('/rigel/ocp/users/jt2796/eccov4r3_datasets/eccov4r3_budgets')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_budg = xr.open_zarr('/rigel/ocp/users/jt2796/eccov4r3_datasets/eccov4r3_budgets')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:       (face: 13, i: 90, j: 90, k: 50, time: 288)\n",
       "Coordinates:\n",
       "    Depth         (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    PHrefC        (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "    XC            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    YC            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "    Z             (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "    drF           (k) float32 dask.array<shape=(50,), chunksize=(50,)>\n",
       "    dt            (time) int64 dask.array<shape=(288,), chunksize=(287,)>\n",
       "  * face          (face) int64 0 1 2 3 4 5 6 7 8 9 10 11 12\n",
       "    hFacC         (k, face, j, i) float32 dask.array<shape=(50, 13, 90, 90), chunksize=(50, 13, 90, 90)>\n",
       "  * i             (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "    iter          (time) int64 dask.array<shape=(288,), chunksize=(1,)>\n",
       "  * j             (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "  * k             (k) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...\n",
       "    rA            (face, j, i) float32 dask.array<shape=(13, 90, 90), chunksize=(13, 90, 90)>\n",
       "  * time          (time) datetime64[ns] 1992-01-15 1992-02-13 1992-03-15 ...\n",
       "Data variables:\n",
       "    adv_hConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_hConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_hConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    adv_vConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_hConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvH    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvS    (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    dif_vConvSln  (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcH         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcS         (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcSln       (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    forcV         (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    hConvV        (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendH         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendS         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendSln       (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    tendV         (time, face, k, j, i) float64 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>\n",
       "    vConvV        (time, face, k, j, i) float32 dask.array<shape=(288, 13, 50, 90, 90), chunksize=(1, 13, 50, 90, 90)>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_budg"
   ]
  },
  {
   "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
}