notebooks/Regrid_GFDL_Forecast.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GFDL FLOR Forecast\n",
    "\n",
    "- Loads in all monthly/daily forecasts of sea ice extent\n",
    "- Regrids to polar stereographic,\n",
    "- Saves to netcdf files grouped by year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "'''\n",
    "\n",
    "This code is part of the SIPN2 project focused on improving sub-seasonal to seasonal predictions of Arctic Sea Ice. \n",
    "If you use this code for a publication or presentation, please cite the reference in the README.md on the\n",
    "main page (https://github.com/NicWayand/ESIO). \n",
    "\n",
    "Questions or comments should be addressed to nicway@uw.edu\n",
    "\n",
    "Copyright (c) 2018 Nic Wayand\n",
    "\n",
    "GNU General Public License v3.0\n",
    "\n",
    "\n",
    "'''\n",
    "\n",
    "# Standard Imports\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload\n",
    "import matplotlib\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xesmf as xe\n",
    "import os\n",
    "import glob\n",
    "import seaborn as sns\n",
    "import warnings\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "# ESIO Imports\n",
    "from esio import import_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# General plotting settings\n",
    "sns.set_style('whitegrid')\n",
    "sns.set_context(\"talk\", font_scale=1.5, rc={\"lines.linewidth\": 2.5})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Directories\n",
    "model='gfdlsipn'\n",
    "runType='forecast'\n",
    "base_dir = r'/home/disk/sipn/nicway/data/'\n",
    "ftp_dir = r'/home/disk/sipn/upload/'\n",
    "data_dir = os.path.join(ftp_dir, model, runType)\n",
    "data_out = os.path.join(base_dir, 'model', model, runType, 'sipn_nc')\n",
    "model_grid_file = os.path.join(base_dir, 'grids', 'GFDLFLOR_gridinfo.nc' )\n",
    "stero_grid_file = os.path.join(base_dir, 'grids', 'stereo_gridinfo.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "updateall = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfdl_grid = import_data.load_grid_info(model_grid_file, model='GFDL')\n",
    "gfdl_grid.dims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# # Plot the input Grid\n",
    "# %matplotlib notebook\n",
    "# plt.figure(figsize=(10,8));\n",
    "# ax1 = plt.axes(projection=ccrs.PlateCarree())\n",
    "# ax1.scatter(gfdl_grid['lon_b'], gfdl_grid['lat_b'], s=6, c='r', transform=ccrs.PlateCarree())  # plot grid locations\n",
    "# ax1.scatter(gfdl_grid['lon'], gfdl_grid['lat'], s=3, c='b', transform=ccrs.PlateCarree())  # plot grid locations\n",
    "# plt.xlabel('lon')\n",
    "# plt.ylabel('lat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# # Plot the input Grid\n",
    "# %matplotlib notebook\n",
    "# plt.figure(figsize=(10,8));\n",
    "# ax1 = plt.axes(projection=ccrs.NorthPolarStereo(central_longitude=0))\n",
    "# ax1.scatter(gfdl_grid['lon_b'], gfdl_grid['lat_b'], s=6, c='r', transform=ccrs.PlateCarree())  # plot grid locations\n",
    "# ax1.scatter(gfdl_grid['lon'], gfdl_grid['lat'], s=3, c='b', transform=ccrs.PlateCarree())  # plot grid locations\n",
    "# plt.xlabel('lon')\n",
    "# plt.ylabel('lat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "obs_grid = import_data.load_grid_info(stero_grid_file, model='NSIDC')\n",
    "# Ensure latitude is within bounds (-90 to 90)\n",
    "# Have to do this because grid file has 90.000001\n",
    "obs_grid['lat_b'] = obs_grid.lat_b.where(obs_grid.lat_b < 90, other = 90)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Regridding Options\n",
    "method='nearest_s2d' # ['bilinear', 'conservative', 'nearest_s2d', 'nearest_d2s', 'patch']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_flag = False # Flag to set up weights have been created\n",
    "have_grid_file  = False\n",
    "\n",
    "# ice_CN_032018_05.nc\n",
    "# GFDL model variable names\n",
    "varnames = ['CN', 'HI']\n",
    "\n",
    "for year in np.arange(2018, 2025, 1):\n",
    "    \n",
    "    for cm in np.arange(1,13):\n",
    "        \n",
    "        # Check if already imported and skip (unless updateall flag is True)\n",
    "        # Always import the most recent two months of files (because they get updated)\n",
    "        f_out = os.path.join(data_out, 'GFDL_FLOR_'+str(year)+'_'+str(cm)+'_Stereo.nc')\n",
    "        if not updateall:\n",
    "            if os.path.isfile(f_out):\n",
    "                print(\"Skipping \", os.path.basename(f_out), \" already imported.\")\n",
    "                continue # Skip, file already imported\n",
    "    \n",
    "\n",
    "        # Check any files for this year exist:\n",
    "        if not glob.glob(os.path.join(data_dir, 'ice_daily_*'+format(cm, '02')+str(year)+'*.nc')):\n",
    "            print(\"No files found for year \", year, cm, \".\")\n",
    "            continue    \n",
    "\n",
    "        print(\"Procesing year \", year, cm)\n",
    "        have_grid_file = True # Flag to clean up regrid file\n",
    "        ds_var_list = []\n",
    "        for varname in varnames:\n",
    "            # Loop over each ensemble member\n",
    "            ds_list = []\n",
    "            for e in np.arange(1,13,1):\n",
    "                cfiles = os.path.join(data_dir, 'ice_daily_'+varname+'*'+format(cm, '02')+str(year)+'_'+format(e, '02')+'.nc')\n",
    "                ds_list.append(import_data.open_1_member(cfiles, e))\n",
    "            ds_var_list.append(xr.concat(ds_list, dim='ensemble')) # conact ensembles\n",
    "\n",
    "        # Merge variables\n",
    "        ds_all = xr.merge(ds_var_list)\n",
    "\n",
    "        # Fix GFDL grid\n",
    "        # Files comes with incorrect lat/lon at poles. Use correct separate file to update coords\n",
    "        ds_all.rename({'yt':'nj', 'xt':'ni'}, inplace=True); # Rename wrong-lat/lon to indices\n",
    "        ds_all.coords['nj'] = gfdl_grid.nj\n",
    "        ds_all.coords['ni'] = gfdl_grid.ni\n",
    "        ds_all.coords['lat'] = gfdl_grid.lat\n",
    "        ds_all.coords['lon'] = gfdl_grid.lon\n",
    "        ds_all.coords['lat_b'] = gfdl_grid.lat_b\n",
    "        ds_all.coords['lon_b'] = gfdl_grid.lon_b\n",
    "\n",
    "        # Split up grid into \"top\" and \"bottom\"\n",
    "        (ds_top, ds_bottom) = import_data.split_GFDL(ds_all, varnames=varnames)\n",
    "\n",
    "        # Calculate regridding matrix\n",
    "        regridder_top = xe.Regridder(ds_top, obs_grid, method, periodic=False, reuse_weights=weights_flag)\n",
    "        regridder_bottom = xe.Regridder(ds_bottom, obs_grid, method, periodic=False, reuse_weights=weights_flag)\n",
    "        weights_flag = True # Set true for following loops\n",
    "\n",
    "        # Add NaNs to empty rows of matrix (forces any target cell with ANY source cells containing NaN to be NaN)\n",
    "        if method=='conservative':\n",
    "            regridder_top = import_data.add_matrix_NaNs(regridder_top)\n",
    "            regridder_bottom = import_data.add_matrix_NaNs(regridder_bottom)\n",
    "\n",
    "        # Average over all sea ice categories \n",
    "        native_top = ds_top.sum(dim='ct', skipna=False)\n",
    "        native_bottom = ds_bottom.sum(dim='ct', skipna=False)\n",
    "\n",
    "        # Loop through each variable (xemsf not dataset enabled yet)\n",
    "        da_out_all_list = []\n",
    "        for cvar in native_top.data_vars:\n",
    "            da_out_all_list.append( import_data.regrid_gfdl_split_domain(ds_all, native_top[cvar], native_bottom[cvar],\n",
    "                                                                  regridder_top, regridder_bottom))\n",
    "        ds_out_all = xr.merge(da_out_all_list)\n",
    "\n",
    "        # Rename to common sipn variable names\n",
    "        ds_out_all = ds_out_all.rename({'CN':'sic','HI':'hi'})\n",
    "\n",
    "        # Expand dims\n",
    "        ds_out_all = import_data.expand_to_sipn_dims(ds_out_all)\n",
    "\n",
    "        # Save regridded to netcdf file\n",
    "        \n",
    "        ds_out_all.to_netcdf(f_out)\n",
    "\n",
    "        # Save regridded to multiple netcdf files by month\n",
    "#         months, datasets = zip(*ds_out_all.groupby('init_time.month'))\n",
    "#         paths = [os.path.join(data_out, 'GFDL_FLOR_'+str(year)+'_'+str(m)+'_Stereo.nc') for m in months]\n",
    "#         xr.save_mfdataset(datasets, paths)\n",
    "\n",
    "        ds_out_all = None # Memory clean up\n",
    "        print('Saved file', f_out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up\n",
    "if have_grid_file:\n",
    "    regridder_top.clean_weight_file()  # clean-up\n",
    "    regridder_bottom.clean_weight_file()  # clean-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# ds_new = xr.open_dataset(paths[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# # Set up plotting info\n",
    "# cmap_sic = matplotlib.colors.ListedColormap(sns.color_palette(\"Blues\", 10))\n",
    "# cmap_sic.set_bad(color = 'green')\n",
    "\n",
    "# # Pick forecast lead time and ensemble member to plot\n",
    "# FT=0\n",
    "# E=0\n",
    "# it=0\n",
    "# var='HI'\n",
    "\n",
    "# # Plot original projection\n",
    "# plt.figure(figsize=(20,10))\n",
    "# ax1 = plt.axes(projection=ccrs.PlateCarree())\n",
    "# ds_p = ds_all[var].isel(fore_time=FT).isel(ensemble=E) #.sum(dim='ct')\n",
    "# ds_p.plot.pcolormesh(ax=ax1, x='lon', y='lat',\n",
    "#                                  cmap=matplotlib.colors.ListedColormap(sns.color_palette(\"Blues\", 10)),\n",
    "#                     transform=ccrs.PlateCarree());\n",
    "# ax1.set_extent([-180, 180, -90, 90], crs=ccrs.PlateCarree())\n",
    "# gl = ax1.gridlines(crs=ccrs.PlateCarree(), linestyle='-')\n",
    "# gl.xlabels_bottom = True\n",
    "# gl.ylabels_left = True\n",
    "# gl.xformatter = LONGITUDE_FORMATTER\n",
    "# gl.yformatter = LATITUDE_FORMATTER\n",
    "# ax1.coastlines(linewidth=0.75, color='black', resolution='50m');\n",
    "\n",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# ds_p2 = ds_new[var].isel(fore_time=FT).isel(ensemble=E).isel(init_time=it)\n",
    "# ds_p2.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "#                                      transform=ccrs.PlateCarree(),\n",
    "#                                      cmap=cmap_sic)\n",
    "# ax1.set_title('Target Grid')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.6.4 esio",
   "language": "python",
   "name": "esio"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}