Density_stratifacation_slope.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c53deeb7-9609-489d-958d-482e1afcca9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "import glob\n",
    "from cmip6_preprocessing.preprocessing import combined_preprocessing\n",
    "import cftime\n",
    "from datetime import datetime\n",
    "xr.set_options(display_style='html')\n",
    "plt.rcParams['figure.figsize'] = 8,3\n",
    "\n",
    "import gsw as gs\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b7800527-e8ed-43fe-942c-96cf281445bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def intersection(lst1, lst2):\n",
    "    lst3 = [value for value in lst1 if value in lst2]\n",
    "    return lst3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5465f50-1109-42c5-a207-debc2af9ad63",
   "metadata": {},
   "source": [
    "### We calculate density, buoyancy frequency, and isopycnal slope from T and S.\n",
    "#### First read in temp and salinity data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d9405ea2-633f-4295-9711-920311936edd",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/gws/nopw/j04/ukfafmip/users/enewsom/regrid_from_jonathan_thetoa/cmip6/CONCATTED/'\n",
    "data_dir_so = '/gws/nopw/j04/ukfafmip/users/enewsom/regrid_from_jonathan_so/cmip6/'\n",
    "\n",
    "\n",
    "files_c6 =  glob.glob(data_dir+ 'thetao_Oyr_*_piControl*.nc')\n",
    "files_so_c6 =  glob.glob(data_dir_so+ 'so_Oyr_*_piControl*.nc')\n",
    "\n",
    "\n",
    "models=[]\n",
    "for file in files_c6:\n",
    "    model = file.split(\"_\")[5]\n",
    "    models.append(model)\n",
    "MODELS=np.unique(models)\n",
    "\n",
    "\n",
    "models=[]\n",
    "for file in files_so_c6:\n",
    "    model = file.split(\"_\")[5]\n",
    "    models.append(model)\n",
    "MODELS2=np.unique(models)\n",
    " \n",
    "MODELS_do_c6 = intersection(MODELS, MODELS2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dffffc50-3541-48ea-bb0a-1831049d63a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/gws/nopw/j04/ukfafmip/users/enewsom/regrid_from_jonathan_thetoa/cmip5/CONCATTED/'\n",
    "data_dir_so = '/gws/nopw/j04/ukfafmip/users/enewsom/regrid_from_jonathan_so/cmip5/'\n",
    "\n",
    "files_c5 =  glob.glob(data_dir+ 'thetao_*yr_*_piControl*.nc')\n",
    "files_so_c5 =  glob.glob(data_dir_so+ 'so_*yr_*_piControl*.nc')\n",
    "\n",
    "models=[]\n",
    "for file in files_c5:\n",
    "    model = file.split(\"_\")[5]\n",
    "    models.append(model)\n",
    "MODELS=np.unique(models)\n",
    "\n",
    "\n",
    "models=[]\n",
    "for file in files_so_c5:\n",
    "    model = file.split(\"_\")[5]\n",
    "    models.append(model)\n",
    "MODELS2=np.unique(models)\n",
    "\n",
    "MODELS_do_c5 = intersection(MODELS, MODELS2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17abaf38-6256-452d-b985-3bf722f8d0ba",
   "metadata": {},
   "source": [
    "### Below chooses a random model (CESM2) to grab a z-grid to interpolate to and "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8c4fb27c-0e22-403d-93b6-03b14cc87ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = '/gws/nopw/j04/ukfafmip/users/enewsom/regrid_from_jonathan_thetoa/cmip6/CONCATTED/thetao_Oyr_CESM2_piControl_r1i1p1f1_gn_regridded_050001-059912.nc'\n",
    "\n",
    "dc  = xr.open_dataset(file).assign_attrs(source_id='CESM2')\n",
    "dc=combined_preprocessing(dc)\n",
    "Tc= dc['thetao'].isel(time=0)\n",
    "ones_interp_grid = (Tc/Tc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0350ba1e-2ba7-4962-9a9d-cfe6fef6c300",
   "metadata": {},
   "source": [
    "### Next calculate metrics salinity, density, N$^2$ (called \"stratification\"), or $\\frac{\\partial\\sigma_2}{\\partial y}$ and $\\frac{\\partial\\sigma_2}{\\partial z}$ for isopycal slope <br>\n",
    "(Below could be written as a loop over each metric, but it takes a long time for a notebook and I liked using it interactively)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "41a726de-3393-48e8-88fc-d49d05a42887",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ACCESS-CM2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/coding/times.py:123: SerializationWarning: Ambiguous reference date string: 950-01-01. The first value is assumed to be the year hence will be padded with zeros to remove the ambiguity (the padded reference date string is: 0950-01-01). To remove this message, remove the ambiguity by padding your reference date strings with zeros.\n",
      "  warnings.warn(warning_msg, SerializationWarning)\n",
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/coding/times.py:526: SerializationWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using cftime.datetime objects instead, reason: dates out of range\n",
      "  dtype = _decode_cf_datetime_dtype(data, units, calendar, self.use_cftime)\n",
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/core/indexing.py:422: SerializationWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using cftime.datetime objects instead, reason: dates out of range\n",
      "  return np.asarray(array[self.key], dtype=None)\n",
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/coding/times.py:123: SerializationWarning: Ambiguous reference date string: 950-01-01. The first value is assumed to be the year hence will be padded with zeros to remove the ambiguity (the padded reference date string is: 0950-01-01). To remove this message, remove the ambiguity by padding your reference date strings with zeros.\n",
      "  warnings.warn(warning_msg, SerializationWarning)\n",
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/coding/times.py:526: SerializationWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using cftime.datetime objects instead, reason: dates out of range\n",
      "  dtype = _decode_cf_datetime_dtype(data, units, calendar, self.use_cftime)\n",
      "/home/users/enewsom/.conda/envs/nbenv2/lib/python3.10/site-packages/xarray/core/indexing.py:422: SerializationWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using cftime.datetime objects instead, reason: dates out of range\n",
      "  return np.asarray(array[self.key], dtype=None)\n"
     ]
    }
   ],
   "source": [
    "p=2000\n",
    "averaging_period_start=0\n",
    "averaging_period_end=80\n",
    "g=9.8\n",
    "rho_0 = 1028\n",
    "lat_to_meters = 1/(111*1000)\n",
    "\n",
    "\n",
    "density = 1\n",
    "stratification = 1\n",
    "slopes = 1\n",
    "\n",
    "cmip6=1\n",
    "cmip5=0\n",
    "\n",
    "if cmip6:\n",
    "    MODELS_do = MODELS_do_c6\n",
    "    files = files_c6\n",
    "    files_so = files_so_c6\n",
    "    ens = 'cmip6'\n",
    "else:\n",
    "    MODELS_do = MODELS_do_c5\n",
    "    files = files_c5\n",
    "    files_so = files_so_c5\n",
    "    ens='cmip5'\n",
    "\n",
    "\n",
    "S=[]\n",
    "ST=[]\n",
    "PX=[]\n",
    "\n",
    "\n",
    "for count,model in enumerate(MODELS_do[0:1]):\n",
    "    print(model)\n",
    "    \n",
    "    files_c = [match for match in files if model in match]\n",
    "    files_s = [match for match in files_so if model in match]\n",
    "           \n",
    "    file = files_c[0]\n",
    "    file_s = files_s[0]\n",
    "    \n",
    "    dc  = xr.open_dataset(file).assign_attrs(source_id=model)\n",
    "    dc=combined_preprocessing(dc)\n",
    "    variable = list(dc.keys())[-1] #grab variable name for temp (not all the same)\n",
    "    Tc= dc[variable].rename('thetao')#rename to the same variable name in data array\n",
    "    Tc = Tc.isel(time = slice(averaging_period_start,averaging_period_end)).mean('time')\n",
    "    Tc=Tc.where(np.abs(Tc)<1e3) #here I mask where some models have weird fill values\n",
    "    Tc = Tc.rename({Tc.dims[0]:'lev'}) #here is because not all models have the same vertical coord name\n",
    "    Tcx=Tc.assign_coords(model=model)\n",
    "        \n",
    "    ds  = xr.open_dataset(file_s).assign_attrs(source_id=model)\n",
    "    ds=combined_preprocessing(ds)\n",
    "    Sc= ds['so']\n",
    "    Sc = Sc.isel(time = slice(averaging_period_start,averaging_period_end)).mean('time')\n",
    "    Sc=Sc.where(np.abs(Tc)<1e3)\n",
    "    Sc = Sc.rename({Sc.dims[0]:'lev'}) \n",
    "    Scx=Sc.assign_coords(model=model)\n",
    "         \n",
    "    SA = gs.SA_from_SP(Scx,p,Scx.x,Scx.y)\n",
    "    CT = gs.CT_from_t(SA,Tcx,p)\n",
    "    Pc = gs.density.sigma2(SA,CT)\n",
    "\n",
    "    if stratification: \n",
    "        strat = (g/rho_0)*Pc.differentiate(\"lev\").mean('x')\n",
    "        strat = strat.assign_coords(model=model)\n",
    "        strat = strat.interp_like(ones_interp_grid)\n",
    "        ST.append(strat)\n",
    "\n",
    "    if slopes:\n",
    "        #slope = drhody/drhodz, but it's smoother to take zonal mean of each first so we save seperately\n",
    "        drhodz = Pc.differentiate(\"lev\").mean('x')\n",
    "        drhody= Pc.differentiate(\"y\").mean('x')*lat_to_meters\n",
    "        drhodz = drhodz.assign_coords(model=model)\n",
    "        drhodz = drhodz.interp_like(ones_interp_grid)\n",
    "        drhody = drhody.assign_coords(model=model)\n",
    "        drhody = drhody.interp_like(ones_interp_grid)\n",
    "        slp = drhody/drhodz\n",
    "        S.append(slp)\n",
    "\n",
    "    if density:\n",
    "        Pcx = Pc.interp_like(ones_interp_grid).mean('x')\n",
    "        Pcx = Pcx.assign_coords(model=model)\n",
    "        PX.append(Pcx)\n",
    "  \n",
    "clear_output(wait=True)\n",
    "        \n",
    "if slopes: \n",
    "    slope_control_model=xr.concat(S,dim='model',coords='minimal',compat='override')\n",
    "    #slope_control_model.to_netcdf('slope_'+ens+'.nc')\n",
    "    #slope_control_model_sig2.close()\n",
    "    \n",
    "if stratification:\n",
    "    strat_control_model_sig2=xr.concat(ST,dim='model',coords='minimal',compat='override')\n",
    "    #strat_control_model_sig2.to_netcdf('N2_'+ens+'.nc')\n",
    "    #strat_control_model_sig2.close()\n",
    "\n",
    "if density:\n",
    "    Dense_control_model_sig2=xr.concat(PX,dim='model',coords='minimal',compat='override')\n",
    "    #Dense_control_model_sig2.to_netcdf('density_'+ens+'.nc')\n",
    "    #Dense_control_model_sig2.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nbenv2",
   "language": "python",
   "name": "nbenv2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}