PM_01_Frequency-Domain_Analysis_of_Large_Datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Frequency-Domain Analysis of Large Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Author(s)\n",
    "**Paige Martin and Ryan Abernathey**\n",
    "\n",
    "- Author1 = {\"name\": \"Paige Martin\", \"affiliation\": \"Australian National University/Lamont-Doherty Earth Observatory\", \"email\": \"paigemar@umich.edu\", \"orcid\": \"0000-0003-3538-633X\"}\n",
    "- Author2 = {\"name\": \"Ryan Abernathey\", \"affiliation\": \"Lamont-Doherty Earth Observatory\", \"email\": \"rpa@ldeo.columbia.edu\", \"orcid\": \"0000-0001-5999-4917\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Purpose\n",
    "\n",
    "Climate model datasets are typically stored as global snapshots, i.e. chunked in time rather than space. For many workflows, this chunking works well (e.g. computations across spatial domains at every point in time). However, this storage format can create serious challenges for processing long time series at each point in space, as is the case for frequency-domain analysis. For large datasets with frequent (e.g. daily) output, it is not feasible to process each spatial point as a single time series, even with the help of distributed computing such as Dask.\n",
    "\n",
    "This notebook provides an example scientific workflow for performing frequency-domain analysis on large datasets. Specifically, this notebook presents a workflow for computing the power spectrum of sea surface temperature in the [Community Earth System Model](https://www.cesm.ucar.edu) (CESM). While we carry out computations on a specific model, the main goal of this notebook is not to interpret scientific results from our computations, but rather to provide a working example of frequency-domain analysis on large datasets that others could follow.\n",
    "\n",
    "Because the goal of this notebook is to provide an example of a workflow that works on large datasets, we have chosen to use a large dataset (CESM) that is available on Pangeo Cloud. This notebook is therefore developed for a Jupyter Hub environment that can access data stored on the Pangeo Cloud (e.g. [PangeoBinder](https://binder.pangeo.io)).\n",
    "\n",
    "\n",
    "## Technical contributions\n",
    "\n",
    "- demonstrates how to quickly rechunk data, e.g. from chunks in time to chunks in space, using the package [Rechunker](https://rechunker.readthedocs.io/en/latest/)\n",
    "- demonstrates how to easily perform Fourier analysis using the package [xrft](https://xrft.readthedocs.io/en/latest/)\n",
    "- shows that all of these steps, with the use of [Xarray](http://xarray.pydata.org/en/stable/) and [Dask](https://dask.org), can be taken with large datasets\n",
    "\n",
    "\n",
    "## Methodology\n",
    "\n",
    "The notebook follows three main steps:\n",
    "\n",
    "1) Rechunk the data. We begin by rechunking the CESM sea surface temperature (SST) output from global, daily snapshots to chunks in space and 5-year chunks in time. This step is accomplished using the package `rechunker`. \n",
    "\n",
    "2) Fourier analysis. Next we compute the power spectrum of SST using the package `xrft`, which nicely integrates with `Xarray` and `Dask`. Within `xrft`, we are also able to easily apply detrending and windowing functions to our data, and also account for the fact that our data are real (with no imaginary components).\n",
    "\n",
    "3) Visualize the data. Last, we average over various frequency bands to show the spatial distribution of the SST power spectrum as global maps.\n",
    "\n",
    "Between each of these steps, we would typically write out the intermediate data. Specifically, we would write out the rechunked data, as well as the processed power spectra. Being able to write out data at intermediate steps is crucial to this workflow. However, due to the inability to write out data from Binder, we instead import previously rechunked data and use a spatial subset to compute and plot the power spectrum in this notebook. We still include all code necessary to run every step if a user wishes to run this notebook elsewhere that allows for data to be written out.\n",
    "\n",
    "\n",
    "## Results\n",
    "This notebook presents a feasible example for performing frequency-domain analysis on large datasets. Specifically, this notebook demonstrates how to quickly rechunk ~500GB of data from chunking in time to chunking in space. It also demonstrates how to pair the `xrft` library with the `rechunker` library to perform frequency-domain spectral analysis (here power spectra) and obtain interpretable results. We finish with a few sample plots to round out the workflow. This notebook is meant to serve as an example for others who wish to perform similar types of analysis.\n",
    "\n",
    "\n",
    "## Funding\n",
    "\n",
    "- Award1 = {\"agency\": \"Gordon and Betty Moore Foundation\", \"award_code\": \"\", \"award_URL\": \"https://www.moore.org\"}\n",
    "\n",
    "## Keywords\n",
    "\n",
    "keywords=[\"frequency-domain\", \"Pangeo\", \"spectral analysis\", \"rechunking\", \"cloud computing\"]\n",
    "\n",
    "## Citation\n",
    "\n",
    "Martin and Abernathey 2021. Frequency-Domain Analysis of Large Datasets. Accessed at https://github.com/paigem/EC2021_Martin_and_Abernathey.\n",
    "\n",
    "## Acknowledgements \n",
    "\n",
    "We thank the Pangeo community for developing and maintaining most of the packages used in this notebook. We also acknowledge Pangeo Cloud, which provides the computing power for this analysis. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "## Library import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reading in data\n",
    "import intake\n",
    "import gcsfs\n",
    "import zarr\n",
    "import os\n",
    "\n",
    "# Data manipulation\n",
    "import xarray as xr\n",
    "import dask.array as dsa\n",
    "from rechunker import rechunk\n",
    "import xrft\n",
    "\n",
    "# Distributed computing\n",
    "from dask_gateway import Gateway\n",
    "from dask.distributed import Client\n",
    "\n",
    "# Visualization\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data import, processing, and analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Open and rechunk the original data\n",
    "\n",
    "This dataset is stored on the [Pangeo Cloud](https://catalog.pangeo.io/) and contains daily output for 41 model years."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Access the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Access CESM POP2 control output from Pangeo Cloud data catalog\n",
    "cat = intake.open_catalog(\"https://raw.githubusercontent.com/pangeo-data/pangeo-datastore/master/intake-catalogs/ocean/CESM_POP.yaml\")\n",
    "item = cat['CESM_POP_hires_control']  # Specify CESM high resolution control run\n",
    "ds_orig = item.to_dask().reset_coords(drop=True)  # drop unneeded coordinates for efficiency\n",
    "ds_orig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are interested in working with SST.\n",
    "Let's make a quick plot to see what the data look like.\n",
    "Pre-coarsening makes the plot faster to load."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_orig.SST[0].coarsen(nlat=5, nlon=5).mean().plot(figsize=(18, 10))\n",
    "plt.title('SST snapshot ($^{\\circ}$C)',fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the chunk structure of the original data: contiguous in the spatial dimension and chunked in the time dimension.\n",
    "This is not optimal for frequency-domain spectral analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_orig.SST"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rechunk the data\n",
    "\n",
    "Rechunker allows us to transform the chunk structure of the dataset using Dask distributed computing in order to facilitate time-domain analysis of the dataset.\n",
    "Since we are only interested in specific variables, we rechunk a single variable (in this case SST) at a time.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**Note:** \n",
    "Rechunking requires the ability to write hundreds of GB to cloud storage.\n",
    "If you are running this notebook in a binder, your environment has read access to Google Cloud Storage, but not write access.\n",
    "(Providing anonymous write access would be a major security threat.)\n",
    "Therefore, below we display the code needed for rechunking in a markdown cell, rather than an executable cell.\n",
    "Instead, in Step 2 below we point to a rechunked dataset that has previously been written and then continue with the Fourier transform steps.\n",
    "</div>\n",
    "\n",
    "Users running on Pangeo Cloud should be able to run this code with no changes, since credentials are automatically populated.\n",
    "Users who have their own Google Cloud account could modify this code by providing valid credentials and a path to their own storage buckets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Open data as Zarr group\n",
    "gcs = gcsfs.GCSFileSystem(requester_pays=True) # Connect to Google Cloud Storage (GCS)\n",
    "mapper = gcs.get_mapper(item.urlpath) # gcs.get_mapper() is like getting the pathname of the data on GCS\n",
    "zgroup = zarr.open_consolidated(mapper) # Open the dataset in Zarr format\n",
    "\n",
    "# Select SST variable\n",
    "varname = 'SST' # name of variable you wish to rechunk (SST in this case)\n",
    "array = zgroup[varname] # select only the SST variable from the Zarr group\n",
    "\n",
    "# Write intermediary data to a temporary path. Items stored here get deleted every 7 days.\n",
    "scratch_path = os.environ['PANGEO_SCRATCH']\n",
    "\n",
    "# Define options needed by the `rechunk()` function of `rechunker`.\n",
    "max_mem = '1GB' # request memory for dask workers\n",
    "target_chunks = (365*5, 90, 180) # chunk in time (5 years) and space\n",
    "\n",
    "# Set paths to temporary and target storage locations and set names of output data\n",
    "tmp_path = f'{scratch_path}/CESM_POP_hires_control/{varname}_tmp.zarr'\n",
    "target_path = f'paigem-pangeo/CESM_POP_hires_control/{varname}_target.zarr'\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Delete arrays of the same name that already exist at those paths\n",
    "def clear_targets():\n",
    "    for path in tmp_path,target_path:\n",
    "        try:\n",
    "            gcs.rm(path + '/.zarray')\n",
    "        except FileNotFoundError:\n",
    "            pass\n",
    "clear_targets()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Create mappings to temporary and target storage buckets\n",
    "store_tmp = gcs.get_mapper(tmp_path)\n",
    "store_target = gcs.get_mapper(target_path)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Call the `rechunk()` function. Below it shows the type of the source, intermediate, and target data.\n",
    "r = rechunk(array, target_chunks, max_mem,\n",
    "            store_target, temp_store=store_tmp, executor='dask')\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Set up Dask cluster)\n",
    "\n",
    "Note that to successfully run `rechunk()` it is necessary to spin up a Dask cluster. To see how we do so on Pangeo Cloud using Dask Gateway, see the following section on computing the power spectrum."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Execute the `rechunk()` function\n",
    "\n",
    "Executing this function will perform the rechunking and save the rechunked SST as intermediate data in the storage bucket designated by `target_path` defined above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "r.execute(retries=10) # `retries=10` sets the number of times a Dask worker will retry if computing fails (default is 0)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Compute the power spectrum of SST\n",
    "\n",
    "The power spectrum is defined as follows, where a hat denotes a Fourier transform and the star denotes a complex conjugate.\n",
    "\n",
    "$$\n",
    "\\widehat{SST}^* \\widehat{SST}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we read in a previously rechunked SST field, which was obtained by following the same steps as shown above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read previously rechunked data from my GCS storage bucket\n",
    "gcs = gcsfs.GCSFileSystem(requester_pays=True)\n",
    "varname = 'SST' \n",
    "target_path = f'pangeo-paigem/CESM_POP_hires_control/{varname}_target.zarr'\n",
    "T_rechunked = dsa.from_zarr(gcs.get_mapper(target_path)) # load as dask file\n",
    "T_rechunked"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Work with a 10-year subset of the rechunked data\n",
    "\n",
    "Though we only show the workflow for one 10-year subset below, when running this workflow for scientific output we compute the power spectrum across several (in this case 10-year) windows. By averaging over all windows in our final plots, we are able to smooth out the naturally noisy Fourier Transform while also getting better statistics. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create 10-year subset over first 10 years of output\n",
    "yr1 = 0\n",
    "yr2 = 10\n",
    "\n",
    "T_rechunked = T_rechunked[365*yr1:365*yr2,:,:]\n",
    "\n",
    "# Convert from dask to xarray DataArray\n",
    "Txr_rechunked = xr.DataArray(T_rechunked,dims=['time','nlat','nlon'])\n",
    "Txr_rechunked"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define function to take power spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def take_power_spectrum(var,real_arg):\n",
    "    \n",
    "    var = var.chunk({'time':None}) # there can be no chunking in the time dimension\n",
    "    var_filled = var.fillna(0) # fill NaNs with zeros\n",
    "    \n",
    "    # Take power spectrum in the time domain, setting time to be a real dimension, with both a linear detrend and a windowing function\n",
    "    var_hat = xrft.power_spectrum(var_filled,dim='time',real=real_arg,detrend='linear',window=True)\n",
    "    \n",
    "    return var_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the power spectrum of SST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the\n",
    "T_calc_power_spectrum = take_power_spectrum(Txr_rechunked,'time')\n",
    "\n",
    "# Take only the real output, and immediately coarsen to 0.5 degree grid to reduce the memory usage\n",
    "T_power_spectrum = T_power_spectrum.real.coarsen(nlat=5, nlon=5).mean()\n",
    "T_power_spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Some Results\n",
    "\n",
    "So far, all of our calculations have been \"lazy\".\n",
    "No computation has actually happened yet.\n",
    "In order to load some data, we will subset the data to a size that can fit into our notebook's memory.\n",
    "We select a region in the North Atlantic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = dict(nlat=slice(290, 380), nlon=slice(50, 220))\n",
    "T_power_spectrum_NAtl = T_power_spectrum.isel(**region)\n",
    "T_power_spectrum_NAtl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Start Dask cluster\n",
    "\n",
    "Now that we are ready to compute, we will start a Dask cluster.\n",
    "If running this notebook on Pangeo binder, this Dask cluster will spin up on Pangeo's Cloud allocation.\n",
    "We use Dask Gateway here (as opposed to, e.g., Dask Kubernetes), as Dask Gateway provides flexibility with cluster scaling, and is the Dask library installed in the Pangeo Cloud infrastructure.\n",
    "This notebook can also be run with other Dask configurations.\n",
    "This Dask cluster sets each worker to use 8 GB of memory, and allows for up to 20 workers to spin up (using adaptive scaling). \n",
    "Note that this may take several minutes to spin up - once you see nonzero numbers below, you are ready to continue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_gateway import Gateway\n",
    "gateway = Gateway()\n",
    "options = gateway.cluster_options()\n",
    "options.worker_memory = 8 # assign each worker 8 GB of memory\n",
    "cluster = gateway.new_cluster(options)\n",
    "cluster.adapt(minimum=1, maximum=20) # use adaptive scaling to allow up to 20 workers\n",
    "cluster # print out cluster information - you can use this to view when the cluster spins up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the client\n",
    "from dask.distributed import Client\n",
    "client = Client(cluster)\n",
    "client # the link below can be used to access the Dask dashboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `load` triggers the cluster to scale up elastically and finish the computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_power_spectrum_NAtl.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Re-mask the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = (ds_orig.SST[0].notnull().astype(int).coarsen(nlon=5, nlat=5).mean() > 0.2).isel(**region).load()\n",
    "T_power_spectrum_NAtl = T_power_spectrum_NAtl.where(mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write to Zarr\n",
    "\n",
    "At this point in our workflow, to do the full globe, we would write out the power spectrum computation back to disk. However, we again include this cell as Markdown only since Binder does not have write access."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Set path to save power spectrum\n",
    "url = f'paigem-pangeo/CESM_POP_hires_control/SST_power_spectrum_yr{yr1}_{yr2}_earthcube_test.zarr'\n",
    "\n",
    "# Save to Zarr\n",
    "T_power_spectrum_reset.to_dataset(name='SST_power_spectrum').to_zarr(gcs.get_mapper(url)) # need to convert xarray DataArray to Dataset first\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Plot Results\n",
    "\n",
    "### Area Average Power Spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_power_spectrum_NAtl.mean(dim=('nlon', 'nlat')).plot(yscale='log', xscale='log')\n",
    "plt.title('Power spectrum of SST in North Atlantic region')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Maps of different frequency band averages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands = {\n",
    "    \"> 1 year\": slice(1/(3650), 1/400),\n",
    "    \"~ 1 year\": slice(1/370, 1/360),\n",
    "    \"100 days - 1 year\": slice(1/360, 1/100),\n",
    "    \"10 days - 100 days\": slice(1/100, 1/10),\n",
    "    \"< 10 days\": slice(1/10, None)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, band_slice in bands.items():\n",
    "    plt.figure()\n",
    "    T_power_spectrum_NAtl.sel(freq_time=band_slice).mean('freq_time').plot(robust=True)\n",
    "    plt.title(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:notebook] *",
   "language": "python",
   "name": "conda-env-notebook-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.8.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}