{
 "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_kubernetes import KubeCluster\n",
    "cluster = KubeCluster(n_workers=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": [
    "# Load with Cloud object storage\n",
    "import gcsfs\n",
    "gcsmap = gcsfs.mapping.GCSMap('pangeo-data/MOM6.CESM/')\n",
    "ds = xr.open_zarr(gcsmap, decode_cf=True, decode_times=False)\n",
    "\n",
    "# Print dataset\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize SST Data with Holoviews and Datashader\n",
    "\n",
    "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'])\n",
    "sst = sst_ds.to(hv.QuadMesh, kdims=[\"xh\", \"yh\"], dynamic=True)\n",
    "%opts RGB [width=1000 height=600] \n",
    "datashade(sst, precompute=True, cmap=plt.cm.jet)"
   ]
  },
  {
   "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['SST'])\n",
    "sss = sss_ds.to(hv.QuadMesh, kdims=[\"xh\", \"yh\"], dynamic=True)\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": [
    "fig, ax = plt.subplots(figsize=(16,8))\n",
    "SST_mean.plot.contourf(levels=np.arange(-2,30))\n",
    "plt.title('Mean SST')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time SSS_mean.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(16,8))\n",
    "SSS_mean.plot.contourf(levels=np.linspace(31,38,20),cmap=plt.cm.jet)\n",
    "plt.title('Mean SSS')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}