notebooks/Regrid_S2S_Models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n\\nThis code is part of the SIPN2 project focused on improving sub-seasonal to seasonal predictions of Arctic Sea Ice. \\nIf you use this code for a publication or presentation, please cite the reference in the README.md on the\\nmain page (https://github.com/NicWayand/ESIO). \\n\\nQuestions or comments should be addressed to nicway@uw.edu\\n\\nCopyright (c) 2018 Nic Wayand\\n\\nGNU General Public License v3.0\\n\\n\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "# S2S and C3S Model Regrid\n",
    "\n",
    "# - Loads in all daily forecasts of sea ice extent\n",
    "# - Regrids to polar stereographic,\n",
    "# - Saves to netcdf files grouped by year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\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 datetime\n",
    "#import cartopy.crs as ccrs\n",
    "#from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "import xesmf as xe\n",
    "import os\n",
    "import re\n",
    "import glob\n",
    "import seaborn as sns\n",
    "import warnings\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "import dask\n",
    "\n",
    "# ESIO Imports\n",
    "from esio import EsioData as ed\n",
    "from esio import import_data\n",
    "from esio import ice_plot\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from dask.distributed import Client\n",
    "# client = Client()\n",
    "# client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "E = ed.EsioData.load()\n",
    "# Directories\n",
    "all_models = ['ecmwfsipn','ukmetofficesipn','bom', 'ncep', 'ukmo', \n",
    "              'eccc', 'kma', 'cma', 'ecmwf', 'hcmr', 'isaccnr',\n",
    "              'jma', 'metreofr'] \n",
    "# all_models = [ 'ukmetofficesipn']\n",
    "updateall = False\n",
    "cvar = 'sic'\n",
    "stero_grid_file = E.obs['NSIDC_0051']['grid']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Regridding Options\n",
    "method='nearest_s2d' # ['bilinear', 'conservative', 'nearest_s2d', 'nearest_d2s', 'patch']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store dictionary to convert each model variable names to sipn syntax\n",
    "var_dic = {}\n",
    "\n",
    "var_dic['cma'] = {'.*ICEC':'sic'}\n",
    "# 'initial_time0_hours':'init_time',\n",
    "#                  'lat_0':'lat', 'lon_0':'lon',\n",
    "#                  'forecast_time0':'fore_time',\n",
    "# guess it looks like cma\n",
    "for model in all_models:\n",
    "    var_dic[model] = var_dic['cma']\n",
    "# Set models that are different\n",
    "#var_dic['bom'] = {'ICEC_P11_L1_GGA0_avg24h':'sic'}\n",
    "# 'initial_time0_hours':'init_time',\n",
    "#                  'lat_0':'lat', 'lon_0':'lon',\n",
    "#                  'forecast_time0':'fore_time',\n",
    "# C3S models\n",
    "var_dic['ukmetofficesipn'] = {'.*CI_GDS0_SFC':'sic'}\n",
    "var_dic['ecmwfsipn'] = {'.*CI_GDS0_SFC':'sic'}\n",
    "\n",
    "# list of models that have month init times\n",
    "monthly_init_model = ['ecmwfsipn', 'ukmetofficesipn', 'metreofr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO\n",
    "# - Get lat lon bounds "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_plot():\n",
    "    # Set up plotting info\n",
    "    cmap_sic = matplotlib.colors.ListedColormap(sns.color_palette(\"Blues\", 10))\n",
    "    cmap_sic.set_bad(color = 'red')\n",
    "\n",
    "    # Plot original projection\n",
    "    plt.figure(figsize=(20,10))\n",
    "    ax1 = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ds_p = ds.sic.isel(init_time=0).isel(fore_time=0).isel(ensemble=0)\n",
    "    ds_p.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "                                     vmin=0, vmax=1,\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",
    "    plt.title(model)\n",
    "    \n",
    "    # Plot SIC on target projection\n",
    "    (f, ax1) = ice_plot.polar_axis()\n",
    "    f.set_size_inches((10,10))\n",
    "    ds_p = ds.sic.isel(init_time=0).isel(fore_time=0).isel(ensemble=0)\n",
    "    ds_p.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "                                         transform=ccrs.PlateCarree(),\n",
    "                                         cmap=cmap_sic)\n",
    "    ax1.set_title('Orginal Grid')\n",
    "\n",
    "    # Plot SIC on target projection\n",
    "    (f, ax1) = ice_plot.polar_axis()\n",
    "    ds_p2 = var_out.isel(init_time=0).isel(fore_time=0).isel(ensemble=0)\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": {},
   "outputs": [],
   "source": [
    "# For both forecast and reforecast\n",
    "for runType in ['forecast','reforecast']: \n",
    "    print('Working on...',runType)\n",
    "    # For each model\n",
    "    for model in all_models:\n",
    "        print('Regridding ', model, '...')\n",
    "\n",
    "        data_dir = E.model[model][runType]['native']\n",
    "        data_out = E.model[model][runType]['sipn_nc']\n",
    "        model_grid_file = E.model[model]['grid']\n",
    "\n",
    "        all_files = glob.glob(os.path.join(data_dir, '*.grib'))\n",
    "        print(\"Found \",len(all_files),\" files.\")\n",
    "        if updateall:\n",
    "            print(\"Updating all files...\")\n",
    "        else:\n",
    "            print(\"Only updating new files\")\n",
    "\n",
    "        # Remove any \"empty\" files (sometimes happends with ecmwf downloads)\n",
    "        all_files_new = []\n",
    "        for cf in all_files:\n",
    "            if os.stat(cf).st_size > 0:\n",
    "                all_files_new.append(cf)\n",
    "            else:\n",
    "                print(\"Found empty file: \",cf,\". Consider deleting or redownloading.\")\n",
    "        all_files = sorted(all_files_new) # Replace and sort\n",
    "\n",
    "        weights_flag = False # Flag to set up weights have been created\n",
    "\n",
    "\n",
    "\n",
    "        # Load land/sea mask file\n",
    "        if model_grid_file.split('/')[-1]!='MISSING':\n",
    "            ds_mask = xr.open_mfdataset(model_grid_file, autoclose=True)\n",
    "        else:\n",
    "            ds_mask = None\n",
    "\n",
    "        # For each file\n",
    "        for cf in all_files:\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, os.path.basename(cf).split('.')[0]+'_Stereo.nc') # netcdf file out \n",
    "            if not updateall:\n",
    "                if (os.path.isfile(f_out)) & (cf not in all_files[-2:]):\n",
    "                    print(\"Skipping \", os.path.basename(cf), \" already imported.\")\n",
    "                    continue # Skip, file already imported\n",
    "\n",
    "            ds = xr.open_dataset(cf, engine='pynio')\n",
    "\n",
    "            # Some grib files do not have a init_time dim, because its assumed for the month\n",
    "            if model in monthly_init_model:\n",
    "                c_coords = list(ds.coords.dims.keys())\n",
    "                tar_coords = list(filter(re.compile('.*initial_time').match, c_coords))\n",
    "                if len(tar_coords)==0: # Check if we have no initial_time* coordinate\n",
    "                    print('Adding init_time as decoder failed to get it.....')\n",
    "                    ds.coords['initial_time1_hours'] = datetime.datetime(int(cf.split('.')[0].split('_')[1]), \n",
    "                                                                  int(cf.split('.')[0].split('_')[2]), 1)\n",
    "                    ds = ds.expand_dims('initial_time1_hours')\n",
    "\n",
    "            # Test we have initial_time0_hours or initial_time1_hours\n",
    "            if ('initial_time0_hours' not in ds.coords) & ('initial_time1_hours' not in ds.coords):\n",
    "                print('initial_time... not found in file: ',cf,' Skipping it, need to FIX!!!!!!!!')\n",
    "                continue\n",
    "\n",
    "            # Rename variables per sipn guidelines\n",
    "            #ds.rename(var_dic[model], inplace=True);\n",
    "            ds = import_data.rename_vars(ds=ds, var_dict=var_dic[model])\n",
    "            # Rename coords\n",
    "            ds = import_data.rename_coords(ds)\n",
    "\n",
    "            # Apply masks (if available)\n",
    "            if ds_mask:\n",
    "                # land_mask is the fraction of native grid cell that is land\n",
    "                # (1-land_mask) is fraction ocean\n",
    "                # Multiply sic by fraction ocean to get actual native grid cell sic\n",
    "                # Also mask land out where land_mask==1\n",
    "                ds[cvar] = ds[cvar] * (1 - ds_mask.land_mask.where(ds_mask.land_mask<0.5)) # Use 50% thresshold for ocean area vs land area\n",
    "\n",
    "\n",
    "        #     # Set sic below 0 to 0\n",
    "        #     if X.sic.min().values < 0:\n",
    "        #         print(\"Some negative SIC \"+str(X.sic.min().values)+\" found in input PIOMAS, setting to 0\")\n",
    "        #         X = X.where(X>=0, other=0)\n",
    "\n",
    "        #     # Apply model mask\n",
    "        #     X = X.where(X.imask)\n",
    "\n",
    "            # Check only data from one month (download bug)\n",
    "            cm = pd.to_datetime(ds.init_time.values).month\n",
    "            if model not in monthly_init_model:\n",
    "                if np.diff(cm).max() > 0:\n",
    "                    fm = int(cf.split('.')[0].split('_')[2]) # Target month in file\n",
    "                    print(\"Found dates outside month, removing...\")\n",
    "                    ds = ds.where(xr.DataArray(pd.to_datetime(ds.init_time.values).month==fm,\n",
    "                                           dims=ds.init_time.dims, coords=ds.init_time.coords), drop=True)\n",
    "\n",
    "            # Calculate regridding matrix\n",
    "            regridder = xe.Regridder(ds, obs_grid, method, periodic=True, 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",
    "                print('Removing edge cells that contain source NaN cells, should probably check here')\n",
    "                regridder = import_data.add_matrix_NaNs(regridder)\n",
    "\n",
    "            # Regrid variable\n",
    "            var_out = regridder(ds[cvar])\n",
    "\n",
    "            # Expand dims\n",
    "            var_out = import_data.expand_to_sipn_dims(var_out)\n",
    "\n",
    "            #test_plot()\n",
    "            #xr.exit()\n",
    "\n",
    "            # # Save regridded to netcdf file\n",
    "            var_out.to_netcdf(f_out)\n",
    "            var_out = None # Memory clean up\n",
    "            ds = None\n",
    "            print('Saved ', f_out)\n",
    "\n",
    "\n",
    "        # End of all files for current model\n",
    "        # Clean up data from this model\n",
    "        ds_mask = None\n",
    "\n",
    "        # Clean up\n",
    "        if weights_flag:\n",
    "            regridder.clean_weight_file()  # clean-up    "
   ]
  },
  {
   "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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cmap_sic = matplotlib.colors.ListedColormap(sns.color_palette(\"Blues\", 10))\n",
    "# cmap_sic.set_bad(color = 'red')\n",
    "\n",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# f.set_size_inches((10,10))\n",
    "# ds_mask.land_mask.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "#                                      transform=ccrs.PlateCarree(),\n",
    "#                                      cmap='Blues')"
   ]
  },
  {
   "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": [
    "# cmap_sic = matplotlib.colors.ListedColormap(sns.color_palette(\"Blues\", 10))\n",
    "# cmap_sic.set_bad(color = 'red')\n",
    "\n",
    "\n",
    "# # # Plot SIC on target projection\n",
    "# # (f, ax1) = ice_plot.polar_axis()\n",
    "# # f.set_size_inches((10,10))\n",
    "# # ds_mask.land_mask.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "# #                                      transform=ccrs.PlateCarree(),\n",
    "# #                                      cmap='Blues')\n",
    "\n",
    "# # # Plot SIC on target projection\n",
    "# # (f, ax1) = ice_plot.polar_axis()\n",
    "# # f.set_size_inches((10,10))\n",
    "# # mask_out.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "# #                                      transform=ccrs.PlateCarree(),\n",
    "# #                                      cmap='Blues')\n",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# f.set_size_inches((10,10))\n",
    "# ds_p = ds.sic.isel(init_time=0).isel(fore_time=0).isel(ensemble=0)\n",
    "# ds_p.plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "#                                      transform=ccrs.PlateCarree(),\n",
    "#                                      cmap=cmap_sic)\n",
    "# ax1.set_title('Orginal Grid')\n",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# f.set_size_inches((10,10))\n",
    "# ds_p2 = var_out.isel(init_time=0).isel(fore_time=0).isel(ensemble=0)\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": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# f.set_size_inches((10,10))\n",
    "# (1 - ds_mask.land_mask.where(ds_mask.land_mask<1)).plot.pcolormesh(ax=ax1, x='lon', y='lat', \n",
    "#                                      transform=ccrs.PlateCarree(),\n",
    "#                                      cmap=cmap_sic)\n",
    "# ax1.set_title('mask we use')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\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": []
  },
  {
   "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": []
  },
  {
   "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": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pynioNew",
   "language": "python",
   "name": "pynionew"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}