dask.ipynb

xarray.ipynb

xarray_dask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from dask.distributed import Client\n",
    "from dask.delayed import delayed\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors\n",
    "\n",
    "import cartopy.crs as ccrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Client('localhost:8786')\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls /scratch/jp492/gfdl_data/dry_earth/run*/daily.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "ds = xr.open_mfdataset('/scratch/jp492/gfdl_data/dry_earth/run*/daily.nc', decode_times=False, chunks={'time': 30})\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(12,4))\n",
    "\n",
    "# sel selects by coordinate value, isel selects by coordinate index.  Same rules apply as with normal python indexing\n",
    "ds.t_surf.isel(time=-1).plot.contourf(levels=22)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# take a slice time mean and apply a projection\n",
    "fig, ax = plt.subplots(1, figsize=(12,4), subplot_kw={'projection': ccrs.Robinson()})\n",
    "ax.coastlines(color='grey')\n",
    "\n",
    "ds.t_surf.sel(time=slice(200, 210)).mean('time').plot.contourf(levels=22, transform=ccrs.PlateCarree())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The nasty white line through the middle is because the model is periodic at (0, 360ºE), but the output doesn't have this.  We can easily fix by copying the values at 0 longitude to 360 longitude and concatenating onto the DataArray.  This function will do this for any arbitrary periodic axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def periodic(field, dim):\n",
    "    edge = field.isel(**{dim: 0})\n",
    "    ddim = field[dim].diff(dim)[0]\n",
    "    edge.coords[dim] = field[dim][-1] + ddim\n",
    "    return xr.concat([field, edge], dim=dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# take a slice time mean and apply a projection\n",
    "fig, ax = plt.subplots(1, figsize=(12,4), subplot_kw={'projection': ccrs.Robinson()})\n",
    "ax.coastlines(color='grey')\n",
    "\n",
    "ds.t_surf.sel(time=slice(200, 210)).mean('time').pipe(periodic, 'lon').plot.contourf(levels=22, transform=ccrs.PlateCarree())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(12,4), subplot_kw={'projection': ccrs.Robinson()})\n",
    "ax.coastlines(color='grey')\n",
    "ds.precipitation.sel(time=300, method='nearest').pipe(periodic, 'lon').plot.contourf(cmap=plt.cm.gray, levels=22, transform=ccrs.PlateCarree())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def surf_mean(field, *dims):\n",
    "    dA = ds.lat.diff('lat')*ds.lon.diff('lon')*np.cos(np.deg2rad(ds.lat))\n",
    "    return (field*dA).sum(('lat', 'lon'))/dA.sum(('lat', 'lon'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.t_surf.pipe(surf_mean).plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## persist\n",
    "\n",
    "If we have a very large data contraction, we can perform it on the cluster and then store the output using `client.persist`.  The object will then be available for further manipulation on the cluster, without having to recalculate or resend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_precip = c.persist(ds.precipitation.mean('lon'))\n",
    "mean_precip.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "mean_precip.plot.pcolormesh(x='time', y='lat', cmap=plt.cm.bone, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.to_netcdf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_bucket = c.persist(ds.bucket_depth.mean('lon'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wet_dry = matplotlib.colors.LinearSegmentedColormap.from_list('RdWtBl', ((224/255., 199/255., 152/255.), (24/255, 143/255, 224/255)))\n",
    "mean_bucket.plot(x='time', y='lat', cmap=wet_dry)"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}