dataset_access/conus404_regrid.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# CONUS404 Regridding (Curvilinear => Rectilinear)\n",
    "Create a rectilinear grid (1D lon/lat coordinates) for a specific region. Extract spatial and temporal subset of regridded data to a netcdf file. (Extraction to netcdf may also be done for curvilinear grid.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import xarray as xr\n",
    "import xesmf as xe\n",
    "import numpy as np\n",
    "import fsspec\n",
    "import hvplot.xarray\n",
    "import geoviews as gv\n",
    "from matplotlib import path \n",
    "import intake\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open dataset from Intake Catalog\n",
    "* Select `on-prem` dataset from /caldera if running on prem (Denali/Tallgrass)\n",
    "* Select `cloud`/`osn` object store data if running elsewhere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# open the hytest data intake catalog\n",
    "hytest_cat = intake.open_catalog(\"https://raw.githubusercontent.com/hytest-org/hytest/main/dataset_catalog/hytest_intake_catalog.yml\")\n",
    "list(hytest_cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open the conus404 sub-catalog\n",
    "cat = hytest_cat['conus404-catalog']\n",
    "list(cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Select the dataset you want to read into your notebook and preview its metadata\n",
    "dataset = 'conus404-hourly-osn' \n",
    "cat[dataset]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Set Up AWS Credentials (Optional)\n",
    "\n",
    "This notebook reads data from the OSN pod by default, which is object store data on a high speed internet connection that is free to access from any environment. If you change this notebook to use one of the CONUS404 datasets stored on S3 (options ending in `-cloud`), you will be pulling data from a `requester-pays` S3 bucket. This means you have to set up your AWS credentials, else we won't be able to load the data. Please note that reading the `-cloud` data from S3 may incur charges if you are reading data outside of the us-west-2 region or running the notebook outside of the cloud altogether. If you would like to access one of the `-cloud` options, uncomment and run the following code snippet to set up your AWS credentials. You can find more info about this AWS helper function [here](https://hytest-org.github.io/hytest/environment_set_up/Help_AWS_Credentials.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment the lines below to read in your AWS credentials if you want to access data from a requester-pays bucket (-cloud)\n",
    "# os.environ['AWS_PROFILE'] = 'default'\n",
    "# %run ../environment_set_up/Help_AWS_Credentials.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallelize with Dask \n",
    "Some of the steps we will take are aware of parallel clustered compute environments\n",
    "using `dask`. We're going to start a cluster now so that future steps can take advantage\n",
    "of this ability. \n",
    "\n",
    "This is an optional step, but speed ups data loading significantly, especially \n",
    "when accessing data from the cloud.\n",
    "\n",
    "We have documentation on how to start a Dask Cluster in different computing environments [here](../environment_set_up/clusters.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ../environment_set_up/Start_Dask_Cluster_Nebari.ipynb\n",
    "## If this notebook is not being run on Nebari/ESIP, replace the above \n",
    "## path name with a helper appropriate to your compute environment.  Examples:\n",
    "# %run ../environment_set_up/Start_Dask_Cluster_Denali.ipynb\n",
    "# %run ../environment_set_up/Start_Dask_Cluster_Tallgrass.ipynb\n",
    "# %run ../environment_set_up/Start_Dask_Cluster_Desktop.ipynb\n",
    "# %run ../environment_set_up/Start_Dask_Cluster_PangeoCHS.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = cat[dataset].to_dask()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nc_outfile = 'CONUS404_DRB_rectilinear.nc'\n",
    "bbox = [-75.9, -74.45, 38.7, 42.55]\n",
    "dx = dy = 3./111.    # 3km grid\n",
    "vars_out = ['T2', 'SNOW']\n",
    "start = '2017-04-01 00:00'\n",
    "stop  = '2017-05-01 00:00'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use xESMF to regrid\n",
    "xESMF is a xarray-enabled interface to the ESMF regridder from NCAR.\n",
    "ESMF has options for regridding between curvilinear, rectilinear, and unstructured grids, with conservative regridding options, and much more"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bbox2ij(lon,lat,bbox=[-160., -155., 18., 23.]):\n",
    "    \"\"\"Return indices for i,j that will completely cover the specified bounding box.     \n",
    "    i0,i1,j0,j1 = bbox2ij(lon,lat,bbox)\n",
    "    lon,lat = 2D arrays that are the target of the subset\n",
    "    bbox = list containing the bounding box: [lon_min, lon_max, lat_min, lat_max]\n",
    "\n",
    "    Example\n",
    "    -------  \n",
    "    >>> i0,i1,j0,j1 = bbox2ij(lon_rho,[-71, -63., 39., 46])\n",
    "    >>> h_subset = nc.variables['h'][j0:j1,i0:i1]       \n",
    "    \"\"\"\n",
    "    bbox=np.array(bbox)\n",
    "    mypath=np.array([bbox[[0,1,1,0]],bbox[[2,2,3,3]]]).T\n",
    "    p = path.Path(mypath)\n",
    "    points = np.vstack((lon.ravel(),lat.ravel())).T   \n",
    "    n,m = np.shape(lon)\n",
    "    inside = p.contains_points(points).reshape((n,m))\n",
    "    ii,jj = np.meshgrid(range(m),range(n))\n",
    "    return min(ii[inside]),max(ii[inside]),min(jj[inside]),max(jj[inside])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we regrid to rectilinear, let's subset a region that covers our area of interest.  Becuase lon,lat are 2D arrays, we can't just use xarray to slice these coordinate variables.  So we have a routine that finds the i,j locations of a specified bounding box, and then slice on those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "i0,i1,j0,j1 = bbox2ij(ds['lon'].values, ds['lat'].values, bbox=bbox)\n",
    "print(i0,i1,j0,j1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds_subset = ds.isel(x=slice(i0-1,i1+1), y=slice(j0-1,j1+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_subset = ds_subset.sel(time=slice(start,stop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_subset.nbytes/1e9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da = ds_subset.T2.sel(time='2017-04-25 00:00', method='nearest')\n",
    "viz = da.hvplot.quadmesh(x='lon', y='lat', geo=True, rasterize=True, cmap='turbo')\n",
    "base = gv.tile_sources.OSM\n",
    "base * viz.opts(alpha=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_subset.nbytes/1e9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "ds_subset = ds_subset.chunk({'x':-1, 'y':-1, 'time':24})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds_out = xr.Dataset({'lon': (['lon'], np.arange(bbox[0], bbox[1], dx)),\n",
    "                     'lat': (['lat'], np.arange(bbox[2], bbox[3], dy))})\n",
    "\n",
    "regridder = xe.Regridder(ds_subset, ds_out, 'bilinear')\n",
    "regridder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds_out = regridder(ds_subset[vars_out])\n",
    "print(ds_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out['SNOW']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(ds_out.variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(ds_out.data_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out['T2'].encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out.time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding={}\n",
    "for var in ds_out.variables:\n",
    "    encoding[var] = dict(zlib=True, complevel=2, \n",
    "                         fletcher32=False, shuffle=True,\n",
    "                         _FillValue=None\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# you will need to update the filepaths and uncomment the following line to save out your data.\n",
    "ds_out.load().to_netcdf(nc_outfile, encoding=encoding, mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_nc = xr.open_dataset(nc_outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(ds_nc['T2']-273.15).hvplot(x='lon',y='lat', geo=True,\n",
    "                rasterize=True, cmap='turbo', \n",
    "                tiles='OSM', clim=(2,15))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds_outcl = ds_subset[vars_out]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(ds_outcl.data_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding={}\n",
    "for var in ds_outcl.variables:\n",
    "    encoding[var] = dict(zlib=True, complevel=2, \n",
    "                         fletcher32=False, shuffle=True,\n",
    "                         _FillValue=None\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# you will need to update the filepaths and uncomment the following line to save out your data.\n",
    "# ds_outcl.load().to_netcdf('CONUS404_DRB_curvilinear.nc', encoding=encoding, mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close(); cluster.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "85ce4eb777a47250062e69db63a13ea88790b42e9bccd9f9c9a4eb5ac7d53107"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}