xBOUT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xBOUT\n",
    "\n",
    "## Readable and scalable BOUT++ data analysis in python\n",
    "\n",
    "Tom Nicholas, John Omotani & Rhys Doyle\n",
    "\n",
    "(thomas.nicholas@york.ac.uk)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Outline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "- Multidimensional data analysis\n",
    "- What is xarray?\n",
    "- numpy vs xarray - syntax comparison \n",
    "- xarray basic features\n",
    "- xarray + dask\n",
    "- xBOUT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Multidimensional data analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Most plasma physicists have some similar data analysis requirements:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "- Relatively large datasets (up to 100's GBs)\n",
    "- Multidimensional\n",
    "- Warped-grid\n",
    "- Several plasma quantities (density, temperature, ...)\n",
    "- Lots of metadata to attach (simulation input file, shot number ...)\n",
    "- Visualise multiple dimensions easily\n",
    "- Apply mathematical operations over many dimensions easily and clearly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# What is xarray?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "xarray is an open-source python library which aims to provide Pandas-like labelling, visualisation & analysis functionality for N-dimensional data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Developed by atmospheric physicists, who have similar data analysis needs to us.\n",
    "\n",
    "They also have relatively large, multidimensional, warped-grid, fluid turbulence datasets to visualise and analyse."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Sponsored by NumFocus, who also support NumPy, MatPlotLib, Pandas, Jupyter, IPython, Julia...\n",
    "\n",
    "Large open-source project: github says used by 1.5k other packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### xarray basic features: \n",
    "- Labelled multidimensional data\n",
    "- Clear syntax for operations\n",
    "- Lazy loading into memory\n",
    "- Plotting convenience"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xarray: Labelled multidimensional data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "xarray wraps numpy arrays with their `dims` as `DataArray`s.\n",
    "\n",
    "Can also select data via special variables called `coords`.\n",
    "\n",
    "Coordinates can be multidimensional\n",
    "(e.g. for mapping Orthogonal toroidal coordinates -> field-aligned coordinates)\n",
    "\n",
    "Multiple `DataArrays` are stored in same `Dataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "For example, imagine we had some output from an atmospheric fluid simulation..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "<img src=\"dataset-diagram.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "<xarray.Dataset>\n",
    "Dimensions:        (t: 8, x: 8, y: 8)\n",
    "Coordinates:\n",
    "  * t              (t) int64 0 1 2 3 4 5 6 7 8\n",
    "  * x              (x) int64 0 1 2 3 4 5 6 7 8\n",
    "  * y              (y) int64 0 1 2 3 4 5 6 7 8\n",
    "    latitude       (x, y) float32 numpy.array(8, 8)\n",
    "    longitude      (x, y) float32 numpy.array(8, 8)\n",
    "Data variables:\n",
    "    temperature    (t, x, y) float32 numpy.array(8, 8, 8)\n",
    "    precipitation  (t, x, y) float32 numpy.array(8, 8, 8)\n",
    "    \n",
    "Attributes:\n",
    "    reference_time 123.0\n",
    "    input          {'solver': {'mxstep': 100000000, 'type': 'cvode'}, 'timeste..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Also carries around a dictionary of \"attributes\". I use that to store units, normalisations etc, as well as to store my entire input file for my simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# numpy vs xarray: Clearer syntax for typical operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Problem: We have some data $n(t,x,y,z)$, and we want to find the maximum over time of the spatially-averaged density at the separatrix. \n",
    "\n",
    "i.e. find $\\text{max}(<n(t,x=\\text{separatrix})>)$, where $<...>$ is an average over $y$ & $z$: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Bare numpy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "max_separatrix_density = np.max(np.mean(n[:, sep_x, ...], axis=(2,3)), axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "xarray:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "max_separatrix_density = ds['n'].isel(x=sep_x).mean(dim=('y', 'z')).max(dim='t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "The xarray code is **clearer**, more **general**, contains **less magic** numbers, and the order of operations applied reads **left-to-right**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Copycats in other languages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Similar libraries implemented in at least 2 other languages, who credit inspiration to `xarray`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "C++ - [xframe](https://github.com/QuantStack/xframe)\n",
    "\n",
    "Julia - [AxisArrays.jl](https://github.com/JuliaArrays/AxisArrays.jl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xarray: Lazy loading into memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "xarray uses the netCDF format in the backend.\n",
    "\n",
    "Lazily loads data values - never waste RAM on unneeded values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "\n",
    "# Open a 100GB file\n",
    "ds = xr.open_dataset('BOUT_data.nc')\n",
    "\n",
    "# Select a 1GB subset of the data\n",
    "data = ds.isel(y=0)\n",
    "\n",
    "# Data is only loaded into memory here, when we actually need it\n",
    "result = some_maths(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xarray: Plotting convenience"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "xarray provides plotting functions which wrap matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "data_slice = ds['phi'].isel(y=34, t=-1)\n",
    "\n",
    "data_slice.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"phi_2D_y=34.png\" alt=\"Drawing\" style=\"width: 600px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "These plotting functions automatically use an appropriate type of plot for the dimension of the data (1D, 2D etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "data = ds['T'].mean(dim=('t', 'y', 'z'))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "data.plot.line(ax=ax)\n",
    "\n",
    "plot_separatrix(data, sep_position, ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"T_profile.png\" alt=\"Drawing\" style=\"width: 600px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xarray + dask: Memory chunking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "If you also install the dask library (literally that's it, no other compiling or anything required), xarray will provide the option to load data into memory in chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('example-data.nc', chunks={'time': 10})\n",
    "\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "<xarray.Dataset>\n",
    "Dimensions:      (latitude: 180, longitude: 360, time: 365)\n",
    "Coordinates:\n",
    "  * time         (time) datetime64[ns] 2015-01-01 2015-01-02 2015-01-03 ...\n",
    "  * longitude    (longitude) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...\n",
    "  * latitude     (latitude) float64 89.5 88.5 87.5 86.5 85.5 84.5 83.5 82.5 ...\n",
    "Data variables:\n",
    "    temperature  (time, latitude, longitude) float64 dask.array<shape=(365, 180, 360), chunksize=(10, 180, 360)>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now if you apply any xarray or any \"embarrassingly parallel\" numpy function to this dataset then it will compute the result only on one chunk at a time, and combine the results at the end.\n",
    "\n",
    "This is very useful when you have \"medium data\": data larger than your system's RAM but not \"big data\". "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xarray + dask: Parallel analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "dask can also automatically parallelize the operation of any xarray function, and most numpy and scipy functions, using the ``apply_ufunc`` helper function: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "result = xr.apply_ufunc(some_numpy_analysis_fn, ds, dask='parallelized', output=[float])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Chunking and parallelization through dask integration should allow you to easily scale up whatever analysis you were doing with numpy to work on datasets that are 100's of GBs in size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (x: 4)>\n",
       "array([1, 0, 1, 4])\n",
       "Dimensions without coordinates: x"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "\n",
    "squared_error = lambda x, y: (x - y) ** 2\n",
    "\n",
    "arr1 = xr.DataArray([0, 1, 2, 3], dims='x')\n",
    "\n",
    "xr.apply_ufunc(squared_error, arr1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def vector_norm(x, dim, ord=None):\n",
    "    return xr.apply_ufunc(np.linalg.norm, x,\n",
    "                          input_core_dims=[[dim]],\n",
    "                          kwargs={'ord': ord, 'axis': -1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray ()>\n",
       "array(3.741657)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vector_norm(arr1, dim='x')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# How does dask work?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"collections-schedulers.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Dask works by:\n",
    "    \n",
    "- Labelling the various operations you want to perform, using either dask objects like dask.arrays or encoding general functions using dask.delayed\n",
    "\n",
    "- Instead of evaluating these operations, it organises them into a Task Graph for later evaluation\n",
    "\n",
    "- Evaluates them using one of a set of Schedulers, which can perform in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "def inc(x):\n",
    "    return x + 1\n",
    "\n",
    "def add(x, y):\n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dask\n",
    "\n",
    "x = dask.delayed(inc)(1)\n",
    "y = dask.delayed(inc)(2)\n",
    "z = dask.delayed(add)(x, y)\n",
    "z.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "z.vizualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"./inc-add.svg\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# xBOUT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "### Loading data\n",
    "\n",
    "xBOUT allows you to load all your BOUT++ output into an xarray dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "(parallel output of BOUT meant this wasn't simple: required significant [upstream additions](https://github.com/pydata/xarray/pull/2616) to xarray)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "from xbout import open_boutdataset\n",
    "\n",
    "ds = open_boutdataset('./BOUT.dmp.*.nc', inputfilepath='./BOUT.inp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "Read in:\n",
    "<xbout.BoutDataset>\n",
    "Contains:\n",
    "<xarray.Dataset>\n",
    "Dimensions:   (t: 51, x: 260, y: 1, z: 256)\n",
    "Dimensions without coordinates: t, x, y, z\n",
    "Data variables:\n",
    "    ncalls_i  (t) int32 dask.array<chunksize=(51,), meta=np.ndarray>\n",
    "    t_array   (t) float64 dask.array<chunksize=(51,), meta=np.ndarray>\n",
    "    dx        (x, y) float64 dask.array<chunksize=(66, 1), meta=np.ndarray>\n",
    "    g11       (x, y) float64 dask.array<chunksize=(66, 1), meta=np.ndarray>\n",
    "    phi       (t, x, y, z) float64 dask.array<chunksize=(51, 66, 1, 256), meta=np.ndarray>\n",
    "    n         (t, x, y, z) float64 dask.array<chunksize=(51, 66, 1, 256), meta=np.ndarray>\n",
    "    omega     (t, x, y, z) float64 dask.array<chunksize=(51, 66, 1, 256), meta=np.ndarray>\n",
    "Attributes:\n",
    "    metadata:  {'MXG': 2, 'zperiod': 1, 'MZSUB': 256, 'NXPE': 4, 'ZMAX': 1.0,...\n",
    "    options:   <configparser.ConfigParser object at 0x7f99e7c81860>\n",
    "Metadata:\n",
    "{   'BOUT_VERSION': 4.3,\n",
    "    'MXG': 2,\n",
    "    'MXSUB': 64,\n",
    "    'NXPE': 4,}\n",
    "Options:\n",
    "<configparser.ConfigParser object at 0x7f99e7c81860>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Replacement for old collect function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from xbout import collect\n",
    "\n",
    "# API matches old boutdata.collect\n",
    "n = collect('n', path='./', prefix='BOUT.dmp')\n",
    "\n",
    "# Returns numpy array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Intended to help smooth transition from `numpy` -> `xarray` for users "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Accessors for BOUT++ specific functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "\"accessors\" provided for attaching your own methods to xarray objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Allows for attaching methods without subclassing xarray objects (which would be complex)\n",
    "ds.bout.do_something_bout_specific()\n",
    "\n",
    "# Also clearly separates the namespaces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Plotting tokamak geometries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ds['psi'].bout.contourf(x='R', y='Z')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"images/xbout_contourf.png\" alt=\"Drawing\" style=\"width: 300px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Animated plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "We wrote methods to animate 1D and 2D plots (by wrapping the animatplot library)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "ds['T'].bout.animate2D(animate_over='time', x='radial', y='binormal', sep_x=sep_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "![tempgif](T_over_time.gif \"segment\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Can even do multiple variables together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ds.bout.animate_list(['n', 'T', 'vort', 'U', 'T', 'phi'], poloidal_plot=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "![tempgif](xy-at-0.gif \"segment\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### calc module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Accessors would be a good place to build up a libary of analysis functions useful to wider BOUT++ community"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Subclassing accessors for new physics modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from xarray import register_dataset_accessor\n",
    "from xbout.boutdataset import BoutDatasetAccessor\n",
    "\n",
    "@register_dataset_accessor('storm')\n",
    "class StormAccessor(BoutAccessor):\n",
    "    def __init__(self, ds_object):\n",
    "        super().__init__(ds_object)\n",
    "\n",
    "    def special_method(self):\n",
    "        print(\"Do something only STORM users would want to do\")\n",
    "\n",
    "ds.storm.special_method()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "Do something only STORM users would want to do"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Registering new geometries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Default is to leave `'x', 'y', 'z', 't'` dimensions alone.\n",
    "\n",
    "Can also choose `geometry='toroidal'`, which adds `R` & `Z` coordinates, and allows poloidal plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "But you can use `register_geometry` to define new geometries, then apply them when loading data\n",
    "\n",
    "e.g. `geometry='stellarator'` ??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "@register_geometry('stellarator')\n",
    "def create_stellarator_coordinate(ds):\n",
    "    # Do some funky 3D maths here\n",
    "    ...\n",
    "\n",
    "# Will apply your stellarator coordinate logic internally\n",
    "ds = open_boutdataset(datapath, geometry='stellarator')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Example of xarray in use: Analysing parameter scan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Grouping multiple datasets into one xarray object makes it easier to analyse parameter scans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Specify folders containing parameter scan\n",
    "diffusions = ['1e-4', '2e-4', '1e-3', '2e-3', '1e-2', '2e-2', '1e-1', '2e-1']\n",
    "\n",
    "# Load everything \n",
    "file_globs = [file_root + f\"mu_n_{diff}/data/BOUT.dmp.*.nc\" for diff in diffusions]\n",
    "datasets = [open_stormdataset(file_glob) for file_glob in file_globs]\n",
    "\n",
    "# Join it all up\n",
    "diffusion = xr.DataArray(name='diff', data=diffusions, dims=['diff'])\n",
    "data = xr.concat(datasets, dim=diffusion)\n",
    "\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "<xarray.Dataset>\n",
    "Dimensions:     (binormal: 512, diff: 8, radial: 480, time: 201)\n",
    "Coordinates:\n",
    "  * time        (time) float64 0.005548 0.005555 0.005562 ... 0.006923 0.00693\n",
    "  * binormal    (binormal) float64 0.0 0.000682 0.001364 ... 0.3478 0.3485\n",
    "  * radial      (radial) float64 0.0 0.0006821 0.001364 ... 0.3254 0.326 0.3267\n",
    "  * diff        (diff) <U4 '1e-4' '2e-4' '1e-3' '2e-3' ... '2e-2' '1e-1' '2e-1'\n",
    "Data variables:\n",
    "    t_array     (diff, time) float64 dask.array<chunksize=(1, 201), meta=np.ndarray>\n",
    "    dx          (diff, radial) float64 dask.array<chunksize=(1, 10), meta=np.ndarray>\n",
    "    dy          (diff, radial) float64 dask.array<chunksize=(1, 10), meta=np.ndarray>\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# Do my power spectrum analysis...\n",
    "ps = xrft.power_spectrum(data['n'].isel(radial=slice(100,200)), dim=['binormal'], detrend='constant')\n",
    "\n",
    "# Add a descriptive coordinate for Fourier-transformed axis\n",
    "k_z_rho_s = ps.coords['freq_binormal'] * data.attrs['params']['rho_s']\n",
    "ps = ps.assign_coords(k_z_rho_s=k_z_rho_s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Plot absolute power spectrum on log-log against k_perp * rho_s\n",
    "abs_ps = np.log10(np.abs(ps))\n",
    "abs_ps = abs_ps.mean(dim=['radial', 'time'])\n",
    "\n",
    "# Plotting across whole scan\n",
    "abs_ps.name = 'log10(abs(power spectrum))'\n",
    "abs_ps.plot(x='k_z_rho_s', xscale='log', hue='diff')   \n",
    "plt.title('density binormal power spectrum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "<img src=\"images/power_spectrum.png\" alt=\"Drawing\" style=\"width: 500px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Future Bonus 1: Plans for array duck typing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Currently xarray wraps either `numpy.ndarray` or `dask.array` objects\n",
    "\n",
    "Possible because they have almost same API, so xarray can select elements in the same way etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "But what if we extended this idea so xarray could wrap any array-like object with the same API? (so-called \"duck-typing\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "There are lots of numpy-like arrays in python ecosystem which do clever things:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "GPU arrays with `cupy.ndarray`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "sparse arrays with `scipy.sparse`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Units with pint, which creates `astropy.units.Quantity` arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$[1,~2,~3] \\; \\mathrm{m}$"
      ],
      "text/plain": [
       "<Quantity [1., 2., 3.] m>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from astropy import units as u\n",
    "import numpy as np\n",
    "\n",
    "np.array([1., 2., 3.]) * u.meter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Plan is to support wrapping any of these with xarray, so you can get benefits of both simultaneously!\n",
    "\n",
    "e.g. A labelled, unit-aware, GPU-parallelised array in python!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Support for this is very close - now in master but not officially rolled out.\n",
    "\n",
    "Open-source so get involved!\n",
    "\n",
    "Start reading about units support on Github [here](https://github.com/pydata/xarray/issues/525), and sparse arrays [here](https://github.com/pydata/xarray/issues/3213)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Future Bonus 2: Data analysis software stack using xarray + dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "[Pangeo](https://pangeo.io/index.html) - \"A community platform for Big Data geoscience\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"images/pangeo_tech_1.png\" alt=\"Drawing\" style=\"width: 600px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "(The guy in charge of this project (an [Oceanography professor](https://rabernat.github.io/) at Columbia) has just advertised a full-time software position to support xarray, so will be even more solidly supported in future.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "- xarray is great\n",
    "\n",
    "- dask is powerful\n",
    "\n",
    "- We've made these both accessible to all BOUT++ users\n",
    "\n",
    "- Solved some common plotting problems\n",
    "\n",
    "- Structured to be easy to subclass for your pet physics module\n",
    "\n",
    "- Potential for more additions in future"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Resources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Blog post introducing xarray:\n",
    "http://stephanhoyer.com/2015/06/11/xray-dask-out-of-core-labeled-arrays/\n",
    "\n",
    "\n",
    "xarray GitHub:\n",
    "https://github.com/pydata/xarray/\n",
    "\n",
    "\n",
    "xarray documentation:\n",
    "http://xarray.pydata.org/en/stable/\n",
    "\n",
    "\n",
    "xarray documentation on dask integration:\n",
    "http://xarray.pydata.org/en/stable/dask.html\n",
    "\n",
    "\n",
    "Other useful blogs/tutorials:\n",
    "http://meteo.unican.es/work/xarray_seminar/xArray_seminar.html\n",
    "https://rabernat.github.io/research_computing/xarray.html\n",
    "\n",
    "\n",
    "Useful page from the dask documentation explaining the general idea:\n",
    "http://docs.dask.org/en/latest/delayed.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Secret bonus: xgcm for staggered grids and complex topologies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Interesting work going on in the xgcm package \n",
    "https://github.com/xgcm/xgcm\n",
    "\n",
    "xgcm (Xarray for Global Circulation Models) aims to provide objects which encode complex grids for use with xarray."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Can encode and perform operations on staggered grids:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"./grid2d_hv.svg\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "and encode complex topologies by storing connections between different cartesian grids:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "<img src=\"cubed_sphere.jpeg\" alt=\"Drawing\" style=\"width: 400px;\">"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}