{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "71c6dadf-029b-4897-88e1-778d0577be4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/marcel/opt/anaconda3/envs/duplessis2021_JGR/lib/python3.9/site-packages/xarray/backends/cfgrib_.py:27: UserWarning: Failed to load cfgrib - most likely there is a problem accessing the ecCodes library. Try `import cfgrib` to get the full error message\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "from tqdm import tqdm_notebook\n",
    "import cmocean.cm as cmo\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import glidertools as gt\n",
    "import my_functions as my"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aae787cc-e583-4a1e-8e72-ebbec091f3e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "dat_saz = xr.open_dataset('../data/sg542_saz_L2.nc')\n",
    "dat_pfz = xr.open_dataset('../data/slocum_pfz_L2.nc')\n",
    "dat_miz = xr.open_dataset('../data/sg643_miz_L2.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "274b08d1-2dc4-47f9-8976-e4dd8aa6093f",
   "metadata": {},
   "outputs": [],
   "source": [
    "interp_res = '6H'\n",
    "offset     = '3H'\n",
    "dt=21600\n",
    "\n",
    "dat_saz_interp = dat_saz.resample(time=interp_res, loffset=offset).mean()\n",
    "dat_pfz_interp = dat_pfz.resample(time=interp_res, loffset=offset).mean()\n",
    "dat_miz_interp = dat_miz.resample(time=interp_res, loffset=offset).mean()\n",
    "\n",
    "dat_saz_interp=dat_saz_interp.interpolate_na(dim='time')\n",
    "dat_pfz_interp=dat_pfz_interp.interpolate_na(dim='time')\n",
    "dat_miz_interp=dat_miz_interp.interpolate_na(dim='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "199e0865-484c-47a7-9e00-22e0779c2923",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, dat in enumerate([dat_saz_interp, dat_pfz_interp, dat_miz_interp]):\n",
    "    \n",
    "    dat['mld_03'] = (('time'), my.calc_mld(dat.density, dat.depth, den_lim=0.03))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "610c3bef-2734-428a-a47c-b1ebbec98e49",
   "metadata": {},
   "outputs": [],
   "source": [
    "for dat in [dat_saz_interp, dat_pfz_interp, dat_miz_interp]:\n",
    "\n",
    "    xt=[]\n",
    "    xs=[]\n",
    "    \n",
    "    for i, m in enumerate(np.round(dat['mld_03'].astype(int))):\n",
    "        \n",
    "        xt += dat.temp.sel(depth=slice(5, m-5)).mean(dim='depth').values[i],\n",
    "        xs += dat.salt.sel(depth=slice(5, m-5)).mean(dim='depth').values[i],\n",
    "    \n",
    "    dat['ml_s'] = (('time'), xs)\n",
    "    dat['ml_t'] = (('time'), xt)\n",
    "    \n",
    "    dat['ml_s_smooth'] = (('time'), gt.cleaning.rolling_window(dat['ml_s'], func=np.nanmean, window=6).data)\n",
    "    dat['ml_t_smooth'] = (('time'), gt.cleaning.rolling_window(dat['ml_t'], func=np.nanmean, window=6).data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "25fa7c78-4bb3-45ce-a5e6-fc2980e97288",
   "metadata": {},
   "outputs": [],
   "source": [
    "era5 = xr.open_mfdataset('../data/ERA5_flux_data/*.nc')\n",
    "era5 = era5.sel(latitude=slice(-40, -65), longitude=slice(-5, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "805fd7a0-2a08-482c-a899-6866f7ff7def",
   "metadata": {},
   "outputs": [],
   "source": [
    "var_name=['slhf', 'sshf', 'str', 'ssr', 'tp', 'e']\n",
    "era5=my.convert_era5_to_Wm2(ds=era5, var_name=var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eb439060-c703-4c6c-9d0f-ea3a4ac4ebed",
   "metadata": {},
   "outputs": [],
   "source": [
    "era5_interp = era5.resample(time=interp_res, loffset=offset).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0d483a8f-74f2-491e-a7f1-ce8c390b45e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "era5['emp']  = (('time', 'latitude', 'longitude'), (era5['e']-era5['tp']).data)\n",
    "era5['qnet'] = (('time', 'latitude', 'longitude'), (era5['slhf']+era5['sshf']+era5['str' ]+era5['ssr' ]).data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ed5079dc-8f25-4915-98c5-529c84a03550",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████| 8/8 [17:09<00:00, 128.63s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "var_name = [\n",
    "            'slhf',\n",
    "            'sshf',\n",
    "            'str' ,\n",
    "            'ssr' ,\n",
    "            'tp' ,\n",
    "            'e' ,\n",
    "#             'emp',\n",
    "            'u10',\n",
    "            'v10',\n",
    "#             'qnet'\n",
    "           ]  \n",
    "\n",
    "for var in tqdm(var_name):\n",
    "  \n",
    "    dat_saz_interp[var] = (('time'), my.interp_glider_era5(dat_saz_interp, era5_interp, var=var))\n",
    "    dat_pfz_interp[var] = (('time'), my.interp_glider_era5(dat_pfz_interp, era5_interp, var=var))\n",
    "    dat_miz_interp[var] = (('time'), my.interp_glider_era5(dat_miz_interp, era5_interp, var=var))\n",
    "    \n",
    "for i, dat in enumerate([dat_saz_interp, dat_pfz_interp, dat_miz_interp]):\n",
    "    \n",
    "    dat['emp']  = (('time'), (dat['e']-dat['tp']).data)\n",
    "    dat['qnet'] = (('time'), (dat['slhf']+dat['sshf']+dat['str' ]+dat['ssr' ]).data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dc982e4b-1301-465f-a2ba-1a6c0c3e336f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getWindDir(uwind, vwind):\n",
    "    \n",
    "    data_dir = []\n",
    "    \n",
    "    for i in range(len(uwind)):\n",
    "        data_dir += math.atan2(vwind[i], uwind[i])/math.pi*180,\n",
    "\n",
    "    data_dir = np.array(data_dir)\n",
    "    \n",
    "    return data_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7dcfa7bc-c630-4da8-8de3-e22e470aaeb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def setWindDirZero360(wind_dir):\n",
    "\n",
    "    ind1 = np.nonzero((wind_dir > 0)    & (wind_dir < 90))\n",
    "    ind2 = np.nonzero((wind_dir > 90)   & (wind_dir < 190))\n",
    "    ind3 = np.nonzero((wind_dir > -180) & (wind_dir < -90))\n",
    "    ind4 = np.nonzero((wind_dir > -90)  & (wind_dir < 0))\n",
    "    \n",
    "    wind_dir[ind1] = 90 - wind_dir[ind1]\n",
    "    wind_dir[ind2] = 450 - wind_dir[ind2]\n",
    "    wind_dir[ind3] = abs(wind_dir[ind3] - 90)\n",
    "    wind_dir[ind4] = abs(wind_dir[ind4]) + 90\n",
    "    \n",
    "    return wind_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0dd656b4-063f-47af-b249-fac40ce6c957",
   "metadata": {},
   "outputs": [],
   "source": [
    "import windstress\n",
    "import math\n",
    "\n",
    "airT = [14, 0, -5]\n",
    "\n",
    "for i, dat in enumerate([dat_saz_interp, dat_pfz_interp, dat_miz_interp]):\n",
    "    \n",
    "    dat['taux'] = (('time'), windstress.stress(dat.u10, z=10., drag='largepond', rho_air=1.22, Ta=airT[i]))\n",
    "    dat['tauy'] = (('time'), windstress.stress(dat.v10, z=10., drag='largepond', rho_air=1.22, Ta=airT[i]))\n",
    "    \n",
    "    dat['wind_dir'] = (('time'), getWindDir(dat.u10, dat.v10).data)\n",
    "    dat['wind_dir'] = (('time'), setWindDirZero360(dat['wind_dir'].values).data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0cbbd52f-d19b-4d3e-a346-830577bfeaa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dat_saz_interp.to_netcdf('../data/dat_saz_6H.nc')\n",
    "dat_pfz_interp.to_netcdf('../data/dat_pfz_6H.nc')\n",
    "dat_miz_interp.to_netcdf('../data/dat_miz_6H.nc')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "duplessis2021_JGR",
   "language": "python",
   "name": "duplessis2021_jgr"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}