ecodatascience_using_xarray_dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code is adapted from [this tutorial](https://github.com/planet-os/notebooks/blob/master/aws/era5-s3-via-boto.ipynb) on accessing ERA5 data stored on a public S3 bucket as a part of Amazon Web Services (AWS) and [this tutorial](https://github.com/pangeo-data/pangeo/blob/master/notebooks/newmann_ensemble_meteorology.ipynb) that analyzes a multi-model ensemble using dask and xarray.\n",
    "\n",
    "\n",
    "Using [ERA5](https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5) data downloaded from a public S3 bucket on AWS, we will utilize [xarray](http://xarray.pydata.org/en/stable/) and [dask](https://docs.dask.org/en/latest/) to plot and analyze variables.\n",
    "\n",
    "ERA5 data is hourly at a gridded resolution of 30 km. ERA5 combines vast amounts of historical observations into global estimates using advanced modelling and data assimilation systems to provide estimates of large number of atmospheric, land and oceanic climate variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Load python libraries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import boto3\n",
    "import botocore\n",
    "import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "import os.path\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "# import cartopy.crs as ccrs\n",
    "# from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER\n",
    "# import matplotlib.ticker as mticker\n",
    "# from IPython.display import Image, display\n",
    "from dask.distributed import Client, progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Connect to Dask Distributed Cluster**\n",
    "\n",
    "Set the `memory_limit` parameter in `Client()` if dask doesn't auto detect your memory limit accurately later on in the notebook. You will know this is occurring if processes start to get killed due o memory limit errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "client = Client(processes=True, n_workers=3)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set up access to S3 bucket using `boto3` and a low-level client**\n",
    "\n",
    "Rather than setting up access key and ID, we will use a low-level client to request data anonymously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "era5_bucket = 'era5-pds'\n",
    "# No AWS keys required\n",
    "client = boto3.client('s3', config=botocore.client.Config(signature_version=botocore.UNSIGNED))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's download main.nc file for the month indicated and use xarray to inspect the metadata relating to the data files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date = datetime.date(2017,1,1) # update to desired date\n",
    "prefix = date.strftime('%Y/%m/')\n",
    "\n",
    "metadata_file = 'main.nc'\n",
    "metadata_key = prefix + metadata_file\n",
    "client.download_file(era5_bucket, metadata_key, metadata_file)\n",
    "ds_meta = xr.open_dataset('main.nc', decode_times=False)\n",
    "ds_meta.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Download variables of interest**\n",
    "\n",
    "Note that each variable for one month is approximately 1 GB, so this takes a minute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select variable(s) of interest\n",
    "# var1 = 'precipitation_amount_1hour_Accumulation'\n",
    "var2 = 'air_pressure_at_mean_sea_level'\n",
    "var1 = 'air_temperature_at_2_metres'\n",
    "\n",
    "var_list = [var1, var2]\n",
    "\n",
    "for var in var_list:\n",
    "    # file path patterns for remote S3 objects and corresponding local file\n",
    "    s3_data_ptrn = '{year}/{month}/data/{var}.nc'\n",
    "    data_file_ptrn = '{year}{month}_{var}.nc'\n",
    "\n",
    "    year = date.strftime('%Y')\n",
    "    month = date.strftime('%m')\n",
    "    s3_data_key = s3_data_ptrn.format(year=year, month=month, var=var)\n",
    "    data_file = data_file_ptrn.format(year=year, month=month, var=var)\n",
    "\n",
    "    if not os.path.isfile(data_file): # check if file already exists\n",
    "        print(\"Downloading %s from S3...\" % s3_data_key)\n",
    "        client.download_file(era5_bucket, s3_data_key, data_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Open multiple data files into a single xarray dataset object**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Open multiple files as a single dataset\n",
    "file_pattern = '{year}{month}*.nc'\n",
    "filename_pattern = file_pattern.format(year=year, month=month)\n",
    "\n",
    "ds = xr.open_mfdataset(filename_pattern, engine='netcdf4', concat_dim='var')\n",
    "print('ds size in GB {:0.2f}\\n'.format(ds.nbytes / 1e9))\n",
    "ds.info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ds.info` output above shows us that there are four dimensions to the data: lat, lon, and time0; and two data variables: air_temperature_at_2_metres, and air_pressure_at_mean_sea_level."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Spatially subset to Western US**\n",
    "\n",
    "It is important near the beginning of your analysis to subset as much as possible. That way, when we start computing, we aren't taking up computational space with something we aren't interested in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.sel(lat=slice(50,30), lon=slice(360-140,360-105))\n",
    "print('ds size in GB {:0.2f}\\n'.format(ds.nbytes / 1e9))\n",
    "ds.info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See how much smaller that is? Now we will do some unit conversions so that our plots make a little more sense.\n",
    "\n",
    "**Convert units to something more familiar**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert mean sea level pressure units to 'hPa'\n",
    "ds['air_pressure_at_mean_sea_level'] = ds.air_pressure_at_mean_sea_level/100.0\n",
    "ds.air_pressure_at_mean_sea_level.attrs['units'] = 'hPa'\n",
    "\n",
    "# and surface temperature units to 'C'\n",
    "ds['air_temperature_at_2_metres'] = ds.air_temperature_at_2_metres - 273.15\n",
    "ds.air_temperature_at_2_metres.attrs['units'] = 'C'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's do some quick analysis - for example taking the average January surface temperature.\n",
    "\n",
    "**Compute the mean temperature across the time axis**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculates the monthly mean along the time dimension\n",
    "da_temp_mean = ds['air_temperature_at_2_metres'].mean(dim='time0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expressions above didn’t actually compute anything. They just build the dask task graph. To do the computations, we call the `compute` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time da_temp_mean = da_temp_mean.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Plot Average Surface Temperature**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da_temp_mean.plot(figsize=(10, 6))\n",
    "plt.title('January 2017 Mean Temperature')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of the time spent in the last calculation was loading data from disk. After we were done with this data, Dask threw it away to free up memory. If we plan to reuse the same dataset many times then we may want to `persist` it in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "air_temp = ds['air_temperature_at_2_metres'].persist()\n",
    "air_temp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the air_temperature DataArray is resident in memory on our workers. We can repeat our computation from last time much more quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time temp_mean = air_temp.mean(dim='time0').compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can also modify the computation and try something new. Keeping data in memory allows to iterate quickly, which is the whole point of this exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time temp_std = air_temp.std(dim='time0').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_std.plot(figsize=(10, 6))\n",
    "plt.title('January 2017 Temperature Standard Deviation')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's select various cities and plot their January temperature over time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# location coordinates\n",
    "locs = [\n",
    "    {'name': 'santa_barbara', 'lon': -119.6982, 'lat': 34.4208},\n",
    "    {'name': 'colorado_springs', 'lon': -104.8214, 'lat': 38.8339},\n",
    "    {'name': 'honolulu', 'lon': -157.835938, 'lat': 21.290014},\n",
    "    {'name': 'seattle', 'lon': -122.3321, 'lat': 47.6062},\n",
    "]\n",
    "\n",
    "# convert westward longitudes to degrees east\n",
    "for l in locs:\n",
    "    if l['lon'] < 0:\n",
    "        l['lon'] = 360 + l['lon']\n",
    "locs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_locs = xr.Dataset()\n",
    "air_temp_ds = air_temp.to_dataset()\n",
    "\n",
    "# interate through the locations and create a dataset\n",
    "# containing the temperature values for each location\n",
    "for l in locs:\n",
    "    name = l['name']\n",
    "    lon = l['lon']\n",
    "    lat = l['lat']\n",
    "    var_name = name\n",
    "\n",
    "    ds2 = air_temp_ds.sel(lon=lon, lat=lat, method='nearest')\n",
    "\n",
    "    lon_attr = '%s_lon' % name\n",
    "    lat_attr = '%s_lat' % name\n",
    "\n",
    "    ds2.attrs[lon_attr] = ds2.lon.values.tolist()\n",
    "    ds2.attrs[lat_attr] = ds2.lat.values.tolist()\n",
    "#     ds2 = ds2.drop(('lat', 'lon'))\n",
    "    ds2 = ds2.rename({var1 : var_name}).drop(('lat', 'lon'))\n",
    "\n",
    "    ds_locs = xr.merge([ds_locs, ds2])\n",
    "\n",
    "ds_locs.data_vars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Do some descriptive statistics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_f = ds_locs.to_dataframe()\n",
    "df_f.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Plot temperature across time for multiple cities**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# readability please\n",
    "plt.rcParams.update({'font.size': 16})\n",
    "\n",
    "ax = df_f.plot(figsize=(18, 10), title=\"ERA5 Air Temperature at 2 Meters\", grid=1)\n",
    "ax.set(xlabel='Date', ylabel='Air Temperature (deg C)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:dask-env]",
   "language": "python",
   "name": "conda-env-dask-env-py"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}