llc4320.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLC4320 Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import intake\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import holoviews as hv\n",
    "from holoviews.operation.datashader import regrid\n",
    "hv.extension('bokeh')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "Data is stored in [zarr](http://zarr.readthedocs.io) format on Google Cloud Storage.\n",
    "This format is optimized for fast parallel access in the cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords = intake.cat.LLC4320_grid.to_dask().reset_coords()\n",
    "sst = intake.cat.LLC4320_SST.to_dask()\n",
    "u = intake.cat.LLC4320_SSU.to_dask()\n",
    "v = intake.cat.LLC4320_SSV.to_dask()\n",
    "ds = xr.merge([sst, u, v])\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Dataset Total Size: {ds.nbytes / 1e12:3.1f} TB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data is on a lat-lon-cap grid with 13 faces. To simplify things a bit, we just select one face from the North Atlantic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_natl = ds.sel(face=10)\n",
    "coords_natl = coords.sel(face=10)\n",
    "ds_natl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Launch Dask Cluster\n",
    "\n",
    "This allows us to parallelize calculations across cloud computing nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_kubernetes import KubeCluster\n",
    "from dask.distributed import Client\n",
    "cluster = KubeCluster()\n",
    "cluster.adapt(minimum=1, maximum=20)\n",
    "client = Client(cluster)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sea Surface Temperature\n",
    "\n",
    "An interactive visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocean_mask = coords_natl.hFacC.reset_coords(drop=True)>0\n",
    "sst_natl = (ds_natl.SST.where(ocean_mask)\n",
    "            .rename('SST'))\n",
    "hv_image = hv.Dataset(sst_natl).to(hv.Image, kdims=['j', 'i'], dynamic=True)\n",
    "\n",
    "%output holomap='scrubber' fps=1\n",
    "%opts Image [width=800 height=500 invert_yaxis=True colorbar=True bgcolor='gray'] (cmap='RdBu_r')\n",
    "regrid(hv_image, precompute=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use XGCM to Perform Calculus\n",
    "\n",
    "[XGCM](https://xgcm.readthedocs.io/en/latest/) is a python packge for working with the datasets produced by numerical General Circulation Models (GCMs) and similar gridded datasets that are amenable to finite volume analysis. In these datasets, different variables are located at different positions with respect to a volume or area element (e.g. cell center, cell face, etc.) xgcm solves the problem of how to interpolate and difference these variables from one position to another."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgcm\n",
    "grid = xgcm.Grid(coords_natl.drop(['k', 'k_p1']), periodic=False)\n",
    "grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vorticity\n",
    "\n",
    "An interactive map of relative vorticity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeta = (-grid.diff(ds_natl.U * coords_natl.dxC, 'Y', boundary='extend')\n",
    "        +grid.diff(ds_natl.V * coords_natl.dyC, 'X', boundary='extend'))/coords_natl.rAz\n",
    "zeta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vort_image = hv.Dataset(zeta.rename('vort')).to(hv.Image, kdims=['j_g', 'i_g'], dynamic=True)\n",
    "\n",
    "%output holomap='scrubber' fps=0.25\n",
    "%opts Image [width=800 height=500 invert_yaxis=True colorbar=True bgcolor='gray' logz=False] (cmap='RdBu_r')\n",
    "regrid(vort_image, precompute=True).redim.range(vort=(-1e-4, 1e-4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinetic Energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke = 0.5 * ( grid.interp(ds_natl.U, 'X', boundary='fill')**2 +\n",
    "              grid.interp(ds_natl.V, 'Y', boundary='fill')**2\n",
    "            ).where(ocean_mask).rename('EKE')\n",
    "eke"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke_image = hv.Dataset(np.log10(eke)).to(hv.Image, kdims=['j', 'i'], dynamic=True)\n",
    "\n",
    "%output holomap='scrubber' fps=0.25\n",
    "%opts Image [width=800 height=500 invert_yaxis=True colorbar=True bgcolor='gray' logz=False] (cmap='magma')\n",
    "rg = regrid(eke_image, precompute=True).redim.range(EKE=(-4, 0))\n",
    "rg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Daily-Averaged Kinetic Energy\n",
    "\n",
    "The data is resampled on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_natl_daily = ds_natl.resample(time='D').mean()\n",
    "ds_natl_daily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke_daily = 0.5 * (grid.interp(ds_natl_daily.U, 'X', boundary='fill')**2 +\n",
    "                   grid.interp(ds_natl_daily.V, 'Y', boundary='fill')**2\n",
    "                  ).where(ocean_mask).rename('EKE')\n",
    "eke_daily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke_image = hv.Dataset(np.log10(eke_daily)).to(hv.Image, kdims=['j', 'i'], dynamic=True)\n",
    "\n",
    "%output holomap='scrubber' fps=0.25\n",
    "%opts Image [width=800 height=500 invert_yaxis=True colorbar=True bgcolor='gray' logz=False] (cmap='magma')\n",
    "regrid(eke_image, precompute=True).redim.range(EKE=(-4, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Try to reproduce Peter Cornillon's Gulf Stream Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xc = coords_natl.XC[:, 0].load()\n",
    "yc = coords_natl.YC[0].load()\n",
    "eke_daily.coords['yc'] = yc.reset_coords(drop=True)\n",
    "eke_daily.coords['xc'] = xc.reset_coords(drop=True)\n",
    "eke_daily_swap = eke_daily.swap_dims({'i': 'yc', 'j': 'xc'})\n",
    "eke_daily_swap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke_gs = eke_daily_swap.sel(yc=slice(40, 30)).sel(xc=[-70, -68, -66], method='nearest')\n",
    "eke_gs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# takes a very long time on a small cluster\n",
    "eke_gs.isel(time=slice(0,30)).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eke_gs.transpose().plot(col='xc', figsize=(12, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}