NWM-transport-plot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Explore the National Water Model Reanalysis\n",
    "Use [Xarray](http://xarray.pydata.org/en/stable/), [Dask](https://dask.org) and [hvPlot](https://hvplot.holoviz.org) from the [HoloViz](https://holoviz.org) tool suite to explore the National Water Modle Reanalysis Version 2.  We read from a cloud-optimized [Zarr](https://zarr.readthedocs.io/en/stable/) dataset that is part of the [AWS Open Data Program](https://aws.amazon.com/opendata/), and we use a Dask cluster to parallelize computation and reading of data chunks.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import fsspec\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hvplot.pandas\n",
    "import hvplot.xarray\n",
    "import geoviews as gv\n",
    "from holoviews.operation.datashader import rasterize\n",
    "import cartopy.crs as ccrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start a Dask cluster\n",
    "This is not required, but speeds up computations.  Here we start a local cluster that just uses the cores available on the computer running the notebook server, but there are [many other ways to set up Dask clusters](https://docs.dask.org/en/latest/setup.html) that can scale larger than this. \n",
    "\n",
    "For example, if you use [Qhub](https://www.quansight.com/post/announcing-qhub) to install JuptyerHub with a Dask Gateway running on Kubernetes, you could start a cluster (with a specified environment and worker profile), scale it, and connect to it thusly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, progress\n",
    "\n",
    "from dask_gateway import Gateway\n",
    "gateway = Gateway()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = gateway.new_cluster()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#cluster.adapt(minimum=4, maximum=20);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Don't have a Dask Gateway or Kubecluster?  Use a local cluster!\n",
    "#from dask.distributed import Client\n",
    "#client = Client()\n",
    "#client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open Zarr datasets in Xarray using a mapper from fsspec.  We use `anon=True` for free-access public buckets like the AWS Open Data Program, and `requester_pays=True` for requester-pays public buckets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 's3://noaa-nwm-retro-v2-zarr-pds'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds = xr.open_zarr(fsspec.get_mapper(url, anon=True), consolidated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var='streamflow'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ds[var]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_nwm_field(da, label=None):\n",
    "    # Convert Xarray to Pandas dataframe so we can use hvplot.points for visualization\n",
    "    df = da.to_pandas().to_frame()\n",
    "    #The dataframe just has streamflow, so add longitude and latitude as columns\n",
    "    df = df.assign(latitude=ds['latitude'])\n",
    "    df = df.assign(longitude=ds['longitude'])\n",
    "    df.rename(columns={0: \"transport\"}, inplace=True)\n",
    "    p = df.hvplot.points('longitude', 'latitude', geo=True,\n",
    "                     c='transport', colorbar=True, size=14, label=label)\n",
    "    # We don't want to plot all the 2.7M points individually, so aggregate \n",
    "    # to 0.02 degree resolution and rasterize with datashader. \n",
    "    # Use a log scale for visualization since there is a large dynamic range in streamflow.\n",
    "    g = rasterize(p, aggregator='mean', x_sampling=0.02, y_sampling=0.02, width=500).opts(tools=['hover'], \n",
    "                aspect='equal', logz=True, cmap='viridis', clim=(1e-2, np.nan))\n",
    "    return (g * gv.tile_sources.OSM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read and plot data for all the stations at a specific time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "select_time = '2017-06-01 00:00:00'\n",
    "da = ds[var].sel(time=select_time)\n",
    "plot_nwm_field(da, label=f'{var}:{select_time}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read and plot data for entire time series at a specific location \n",
    "Just as an example we pick the location with the largest stream flow from the specific time above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "imax = da.argmax().values\n",
    "ds[var][:,imax].hvplot(grid=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute mean discharge during 2017 on all rivers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da= ds[var].sel(time=slice('2017-01-01 00:00','2017-12-31 23:00'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "var_mean = da.mean(dim='time').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_nwm_field(var_mean, 'Mean Streamflow: 2017')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close(); cluster.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}