notebooks/Regrid_NESM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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 pandas as pd\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",
    "import datetime\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "# ESIO Imports\n",
    "\n",
    "from esio import EsioData as ed\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": [
    "E = ed.EsioData.load()\n",
    "# Directories\n",
    "all_models=['usnavygofs','usnavyncep','usnavysipn']\n",
    "runType='forecast'\n",
    "updateall = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stero_grid_file = E.obs['NSIDC_0051']['grid']\n",
    "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='conservative_normed' # ['bilinear', 'conservative', 'nearest_s2d', 'nearest_d2s', 'patch']\n",
    "method = 'nearest_s2d'"
   ]
  },
  {
   "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": [
    "# Set models that are different\n",
    "var_dic = {'aice':'sic'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    # Files are stored as per time step (about 45 per init_time)\n",
    "    # First parse files to see what unique init_times we have\n",
    "    # ARCu0.08_121_2018042112_t0300.nc\n",
    "    prefix = 'ARCu0'\n",
    "    all_files = glob.glob(os.path.join(data_dir, '*'+prefix+'*.nc'))\n",
    "    if model=='usnavygofs':\n",
    "        init_N = 4\n",
    "    else:\n",
    "        init_N = 2\n",
    "    init_times = list(set([s.split('_')[init_N] for s in all_files]))\n",
    "    \n",
    "    print(\"Found \",len(init_times),\" initialization times.\")\n",
    "    if updateall:\n",
    "        print(\"Updating all files...\")\n",
    "    else:\n",
    "        print(\"Only updating new files\")\n",
    "\n",
    "\n",
    "    weights_flag = False # Flag to set up weights have been created\n",
    "\n",
    "    # Load land/sea mask file\n",
    "    if os.path.basename(model_grid_file)!='MISSING':\n",
    "        ds_mask = xr.open_mfdataset(model_grid_file)\n",
    "    else:\n",
    "        ds_mask = None\n",
    "\n",
    "    for cf in sorted(init_times):\n",
    "        # Check if already imported and skip (unless updateall flag is True)\n",
    "        f_out = os.path.join(data_out, prefix+'_'+cf+'_Stereo.nc') # netcdf file out \n",
    "        if not updateall:\n",
    "            # TODO: Test if the file is openable (not corrupted)\n",
    "            if os.path.isfile(f_out):\n",
    "                print(\"Skipping \", cf, \" already imported.\")\n",
    "                continue # Skip, file already imported\n",
    "\n",
    "        c_files = sorted(glob.glob(os.path.join(data_dir, '*'+prefix+'*_'+cf+'*.nc')))\n",
    "                    \n",
    "        \n",
    "        # Some files have a \"tau\" variable that is hours since analysis\n",
    "        try:\n",
    "            ds = xr.open_mfdataset(c_files, concat_dim='time', decode_times=False, autoclose=True)\n",
    "\n",
    "            # Format times\n",
    "            ds.coords['init_time'] = np.datetime64(ds.tau.attrs['time_origin'])\n",
    "            ds.coords['tau'] = ds.tau\n",
    "            ds.swap_dims({'time':'tau'}, inplace=True)\n",
    "            ds.rename({'tau':'fore_time'}, inplace=True)\n",
    "            ds.fore_time.attrs['units'] = 'Forecast offset from initial time'\n",
    "            ds = ds.drop(['time'])\n",
    "            ds.coords['fore_time'] = ds.fore_time.astype('timedelta64[h]') \n",
    "            \n",
    "        # Some do not\n",
    "        except AttributeError:\n",
    "            \n",
    "            ds = xr.open_mfdataset(c_files, concat_dim='time', decode_times=True, autoclose=True)\n",
    "            dt_mod = ds.time.values[1] - ds.time.values[0]\n",
    "            ds.coords['init_time'] = ds.time.values[0] - dt_mod\n",
    "            ds.coords['fore_time'] = ds.time - ds.init_time\n",
    "            ds.swap_dims({'time':'fore_time'}, inplace=True);\n",
    "            ds = ds.drop('time')\n",
    "            \n",
    "        # Rename variables per esipn guidelines\n",
    "        ds.rename(var_dic, inplace=True);\n",
    "        \n",
    "        # Apply masks (if available)\n",
    "        if ds_mask:\n",
    "            print('found 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 = ds * (1 - ds_mask.land_mask.where(ds_mask.land_mask<1))\n",
    "            \n",
    "        # Add mask variable so conservative regridding works as expected\n",
    "        # DOESN\"T WORK WITH OTHER METHODS!!\n",
    "        #ds['mask'] = ds.sic.isel(fore_time=0).notnull() # Hardcoded variable choice\n",
    "                        \n",
    "        # Calculate regridding matrix\n",
    "        regridder = xe.Regridder(ds, obs_grid, method, periodic=False, reuse_weights=weights_flag)\n",
    "\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 = import_data.add_matrix_NaNs(regridder)\n",
    "\n",
    "        # Regrid variables\n",
    "\n",
    "        var_list = []\n",
    "        for cvar in ds.data_vars:\n",
    "            # 0 to NaN hack\n",
    "            #offset = 10\n",
    "            #da_coarse = regridder(ds[cvar]+10)\n",
    "            #da_coarse = da_coarse.where(da_coarse>(offset)) - offset\n",
    "            #var_list.append(da_coarse)\n",
    "            \n",
    "            # When doing nearest neighbor\n",
    "            da_coarse = regridder(ds[cvar])\n",
    "            var_list.append(da_coarse)\n",
    "            \n",
    "        ds_out = xr.merge(var_list)\n",
    "\n",
    "        # Expand dims\n",
    "        ds_out = import_data.expand_to_sipn_dims(ds_out)\n",
    "                \n",
    "        # # Save regridded to netcdf file\n",
    "        ds_out.to_netcdf(f_out)\n",
    "        \n",
    "        ds_out = None # Memory clean up\n",
    "        print('Saved ', f_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": "markdown",
   "metadata": {},
   "source": [
    "# Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# sic_all = xr.open_mfdataset(f_out)\n",
    "\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(fore_time=79)\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",
    "\n",
    "# # Plot SIC on target projection\n",
    "# (f, ax1) = ice_plot.polar_axis()\n",
    "# ds_p2 = sic_all.sic.isel(init_time=0).isel(fore_time=79).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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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
}