{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MOM6/CESM Ocean Model Analysis\n",
    "\n",
    "This notebook shows how to load and analyze ocean data from an out-of-the-box MOM6/CESM G-case simulation (coupled ocean ocean/sea ice). \n",
    "\n",
    "**NOTE**: MOM6/CESM is not ready to be used for research. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import holoviews as hv\n",
    "import datashader\n",
    "from holoviews.operation.datashader import regrid, shade, datashade\n",
    "\n",
    "hv.extension('bokeh', width=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create and Connect to Dask Distributed Cluster\n",
    "\n",
    "This will launch a cluster of virtual machines in the cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, progress\n",
    "from dask_gateway import Gateway\n",
    "\n",
    "gateway = Gateway()\n",
    "cluster = gateway.new_cluster()\n",
    "cluster.scale(10)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "👆 Don't forget to click this link to get the cluster dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load MOM6/CESM Data\n",
    "\n",
    "This data is stored in [xarray-zarr](http://xarray.pydata.org/en/latest/io.html#zarr) format in Google Cloud Storage.\n",
    "This format is optimized for parallel distributed reads from within the cloud environment.\n",
    "\n",
    "It may take up to a minute to initialize the dataset when you run this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from intake import open_catalog\n",
    "\n",
    "cat = open_catalog(\"https://raw.githubusercontent.com/pangeo-data/pangeo-datastore/master/intake-catalogs/master.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = cat.ocean.cesm_mom6_example.to_dask()\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize SST Data with Holoviews and Datashader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cells below show how to interactively explore the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_ds = hv.Dataset(ds['SST'], kdims=['time', 'geolon', 'geolat'])\n",
    "sst = sst_ds.to(hv.QuadMesh, kdims=[\"geolon\", \"geolat\"], dynamic=True)\n",
    "%opts RGB [width=900 height=600] \n",
    "datashade(sst, precompute=True, cmap=plt.cm.RdBu_r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize SSS Data with Holoviews and Datashader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sss_ds = hv.Dataset(ds['SSS'], kdims=['time', 'geolon', 'geolat'])\n",
    "sss = sst_ds.to(hv.QuadMesh, kdims=[\"geolon\", \"geolat\"], dynamic=True)\n",
    "%opts RGB [width=900 height=600] \n",
    "datashade(sss, precompute=True, cmap=plt.cm.Spectral_r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data reduction\n",
    "\n",
    "Here we make a data reduction by taking the time of SST and SSS. This demonstrates how the cluster distributes the reads from storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SST_mean = ds.SST.mean(dim=('time'))\n",
    "SST_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SSS_mean = ds.SSS.mean(dim=('time'))\n",
    "SSS_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time SST_mean.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot mean SST\n",
    "qm = hv.QuadMesh((ds.geolon.values, ds.geolat.values, SST_mean))\n",
    "datashade(qm, precompute=True, cmap=plt.cm.RdBu_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time SSS_mean.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot mean SSS\n",
    "qm = hv.QuadMesh((ds.geolon.values, ds.geolat.values, SSS_mean))\n",
    "datashade(qm, precompute=True, cmap=plt.cm.Spectral_r)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}