rasp-data-loading.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# xarray use case: Neural network training\n",
    "\n",
    "\n",
    "**tl;dr**\n",
    "\n",
    "1. This notebook is an example of reading from a climate model netCDF file to train a neural network. Neural networks (for use in parameterization research) require random columns of several stacked variables at a time. \n",
    "\n",
    "2. Experiments in this notebook show:\n",
    "    1. Reading from raw climate model output files is super slow (1s per batch... need speeds on the order of ms)\n",
    "    2. open_mfdataset is half as fast as opening the same dataset with open_dataset\n",
    "    3. Pure h5py is much faster than reading the same dataset using xarray (even using the h5 backend)\n",
    "\n",
    "3. Currently, I revert to preformatting the dataset (flatten time, lat, lon). This gets the reading speed down to milliseconds per batch.\n",
    "\n",
    "**Conclusions**\n",
    "\n",
    "Reading straight from the raw netCDF files (with all dimensions intact) is handy and might be necessary for later applications (using continuous time slices or lat-lon regions for RNNs or CNNs).\n",
    "\n",
    "However, at the moment this is many orders of magnitude too slow. Preprocessing seems required.\n",
    "\n",
    "What would be a good way of speeding this up without too extensive post processing?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.11.2'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load an example dataset\n",
    "\n",
    "I uploaded a sample dataset here: http://doi.org/10.5281/zenodo.2559313\n",
    "\n",
    "The files are around 1GB large. Let's download it.\n",
    "\n",
    "NOTE: I have all my data on an SSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify this path!\n",
    "DATADIR = '/local/S.Rasp/tmp/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2019-02-07 13:08:52--  https://zenodo.org/record/2559183/files/sample_SPCAM_1.nc\n",
      "Resolving zenodo.org (zenodo.org)... 137.138.76.77\n",
      "Connecting to zenodo.org (zenodo.org)|137.138.76.77|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 923498891 (881M) [application/octet-stream]\n",
      "Saving to: ‘/local/S.Rasp/tmp/sample_SPCAM_1.nc’\n",
      "\n",
      "sample_SPCAM_1.nc   100%[===================>] 880.72M  6.59MB/s    in 1m 49s  \n",
      "\n",
      "2019-02-07 13:10:42 (8.09 MB/s) - ‘/local/S.Rasp/tmp/sample_SPCAM_1.nc’ saved [923498891/923498891]\n",
      "\n",
      "--2019-02-07 13:10:42--  https://zenodo.org/record/2559183/files/sample_SPCAM_2.nc\n",
      "Resolving zenodo.org (zenodo.org)... 137.138.76.77\n",
      "Connecting to zenodo.org (zenodo.org)|137.138.76.77|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 923498891 (881M) [application/octet-stream]\n",
      "Saving to: ‘/local/S.Rasp/tmp/sample_SPCAM_2.nc’\n",
      "\n",
      "sample_SPCAM_2.nc   100%[===================>] 880.72M  24.5MB/s    in 87s     \n",
      "\n",
      "2019-02-07 13:12:09 (10.1 MB/s) - ‘/local/S.Rasp/tmp/sample_SPCAM_2.nc’ saved [923498891/923498891]\n",
      "\n",
      "--2019-02-07 13:12:09--  https://zenodo.org/record/2559183/files/sample_SPCAM_concat.nc\n",
      "Resolving zenodo.org (zenodo.org)... 137.138.76.77\n",
      "Connecting to zenodo.org (zenodo.org)|137.138.76.77|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1846816429 (1.7G) [application/octet-stream]\n",
      "Saving to: ‘/local/S.Rasp/tmp/sample_SPCAM_concat.nc’\n",
      "\n",
      "sample_SPCAM_concat 100%[===================>]   1.72G  3.80MB/s    in 3m 28s  \n",
      "\n",
      "2019-02-07 13:15:38 (8.46 MB/s) - ‘/local/S.Rasp/tmp/sample_SPCAM_concat.nc’ saved [1846816429/1846816429]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget -P $DATADIR https://zenodo.org/record/2559313/files/sample_SPCAM_1.nc\n",
    "!wget -P $DATADIR https://zenodo.org/record/2559313/files/sample_SPCAM_2.nc\n",
    "!wget -P $DATADIR https://zenodo.org/record/2559313/files/sample_SPCAM_concat.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 S.Rasp ls-craig 881M Feb  7 13:00 /local/S.Rasp/tmp//sample_SPCAM_1.nc\r\n",
      "-rw-r--r-- 1 S.Rasp ls-craig 881M Feb  7 13:00 /local/S.Rasp/tmp//sample_SPCAM_2.nc\r\n",
      "-rw-r--r-- 1 S.Rasp ls-craig 1.8G Feb  7 13:00 /local/S.Rasp/tmp//sample_SPCAM_concat.nc\r\n"
     ]
    }
   ],
   "source": [
    "!ls -lh $DATADIR/sample_SPCAM*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The files are typical climate model output files. `sample_SPCAM_1.nc` and `sample_SPCAM_2.nc` are two contiguous output files. `sample_SPCAM_concat.nc` is the concatenated version of the two files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 56 ms, sys: 0 ns, total: 56 ms\n",
      "Wall time: 54.7 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ds = xr.open_mfdataset(DATADIR + 'sample_SPCAM_1.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:       (crm_x: 32, crm_y: 1, crm_z: 28, ilev: 31, isccp_prs: 7, isccp_prstau: 49, isccp_tau: 7, lat: 64, lev: 30, lon: 128, tbnd: 2, time: 48)\n",
       "Coordinates:\n",
       "  * lat           (lat) float64 -87.86 -85.1 -82.31 -79.53 ... 82.31 85.1 87.86\n",
       "  * lon           (lon) float64 0.0 2.812 5.625 8.438 ... 351.6 354.4 357.2\n",
       "  * crm_x         (crm_x) float64 0.0 4.0 8.0 12.0 ... 112.0 116.0 120.0 124.0\n",
       "  * crm_y         (crm_y) float64 0.0\n",
       "  * crm_z         (crm_z) float64 992.6 976.3 957.5 936.2 ... 38.27 24.61 14.36\n",
       "  * lev           (lev) float64 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6\n",
       "  * ilev          (ilev) float64 2.255 5.032 10.16 18.56 ... 967.5 985.1 1e+03\n",
       "  * isccp_prs     (isccp_prs) float64 90.0 245.0 375.0 500.0 620.0 740.0 900.0\n",
       "  * isccp_tau     (isccp_tau) float64 0.15 0.8 2.45 6.5 16.2 41.5 219.5\n",
       "  * isccp_prstau  (isccp_prstau) float64 90.0 90.0 90.0 ... 900.0 900.0 900.2\n",
       "  * time          (time) object 0000-01-01 00:00:00 ... 0000-01-01 23:29:59\n",
       "Dimensions without coordinates: tbnd\n",
       "Data variables:\n",
       "    P0            float64 ...\n",
       "    time_bnds     (time, tbnd) object dask.array<shape=(48, 2), chunksize=(48, 2)>\n",
       "    date_written  (time) |S8 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    time_written  (time) |S8 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    ntrm          int32 ...\n",
       "    ntrn          int32 ...\n",
       "    ntrk          int32 ...\n",
       "    ndbase        int32 ...\n",
       "    nsbase        int32 ...\n",
       "    nbdate        int32 ...\n",
       "    nbsec         int32 ...\n",
       "    mdt           int32 ...\n",
       "    nlon          (lat) int32 dask.array<shape=(64,), chunksize=(64,)>\n",
       "    wnummax       (lat) int32 dask.array<shape=(64,), chunksize=(64,)>\n",
       "    hyai          (ilev) float64 dask.array<shape=(31,), chunksize=(31,)>\n",
       "    hybi          (ilev) float64 dask.array<shape=(31,), chunksize=(31,)>\n",
       "    hyam          (lev) float64 dask.array<shape=(30,), chunksize=(30,)>\n",
       "    hybm          (lev) float64 dask.array<shape=(30,), chunksize=(30,)>\n",
       "    gw            (lat) float64 dask.array<shape=(64,), chunksize=(64,)>\n",
       "    ndcur         (time) int32 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    nscur         (time) int32 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    date          (time) int32 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    datesec       (time) int32 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    nsteph        (time) int32 dask.array<shape=(48,), chunksize=(48,)>\n",
       "    DTV           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    DTVKE         (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    FLNS          (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    FLNT          (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    FLUT          (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    FSNS          (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    FSNT          (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    LHFLX         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PHCLDICE      (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    PHCLDLIQ      (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    PHQ           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    PRECC         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECL         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECSC        (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECSL        (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECSTEN      (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECT         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PRECTEND      (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    PS            (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    QAP           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    QCAP          (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    QIAP          (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    QRL           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    QRS           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    SHFLX         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    SOLIN         (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    SPDQ          (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    SPDT          (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    T             (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    TAP           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    TPHYSTND      (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    TS            (time, lat, lon) float32 dask.array<shape=(48, 64, 128), chunksize=(48, 64, 128)>\n",
       "    UAP           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    VAP           (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    VD01          (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "    VPHYSTND      (time, lev, lat, lon) float32 dask.array<shape=(48, 30, 64, 128), chunksize=(48, 30, 64, 128)>\n",
       "Attributes:\n",
       "    Conventions:  CF-1.0\n",
       "    source:       CAM\n",
       "    case:         AndKua_aqua_SPCAM3.0_sp_fbp32\n",
       "    title:        \n",
       "    logname:      tg847872\n",
       "    host:         \n",
       "    Version:      $Name:  $\n",
       "    revision_Id:  $Id: history.F90,v 1.26.2.38 2003/12/15 18:52:35 hender Exp $"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random columns for machine learning parameterizations\n",
    "\n",
    "For the work on ML parameterizations that a few of us are doing now, we would like to work one column at a time. One simple example would be predicting the temperature and humidity tendencies (TPHYSTND and PHQ) from the temperature and humidity profiles (TAP and QAP). \n",
    "\n",
    "This means we would like to give the neural network a stacked vector containing the inputs (2 x 30 levels) and ask it to predict the outputs (also 2 x 30 levels).\n",
    "\n",
    "In NN training, we usually train on a batch of data at a time. Batches typically have a few hundred samples (columns in our case). It is really important that the samples in a batch are not correlated but rather represent a random sample of the entire dataset.\n",
    "\n",
    "To achieve this we will write a data generator that loads the batches by randomly selecting along the time, lat and lon dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataGenerator(object):\n",
    "    \"\"\"\n",
    "    Data generator that randomly (if shuffle = True) picks columns from the dataset and returns them in \n",
    "    batches. For each column the input variables and output variables will be stacked.\n",
    "    \"\"\"\n",
    "    def __init__(self, fn_or_ds, batch_size=128, input_vars=['TAP', 'QAP'], output_vars=['TPHYSTND', 'PHQ'], \n",
    "                 shuffle=True, engine='netcdf4'):\n",
    "        self.ds = xr.open_mfdataset(fn_or_ds, engine=engine) if type(fn_or_ds) is str else fn_or_ds\n",
    "        self.batch_size = batch_size\n",
    "        self.input_vars = input_vars\n",
    "        self.output_vars = output_vars\n",
    "        self.ntime, self.nlat, self.nlon = self.ds.time.size, self.ds.lat.size, self.ds.lon.size\n",
    "        self.ntot = self.ntime * self.nlat * self.ntime\n",
    "        self.n_batches = self.ntot // batch_size\n",
    "        self.indices = np.arange(self.ntot)\n",
    "        if shuffle:\n",
    "            self.indices = np.random.permutation(self.indices)\n",
    "    def __getitem__(self, index):\n",
    "        time_indices, lat_indices, lon_indices = np.unravel_index(\n",
    "            self.indices[index*self.batch_size:(index+1)*self.batch_size], (self.ntime, self.nlat, self.nlon)\n",
    "        )\n",
    "        \n",
    "        X, Y = [], []\n",
    "        for itime, ilat, ilon in zip(time_indices, lat_indices, lon_indices):\n",
    "            X.append(\n",
    "                np.concatenate(\n",
    "                    [self.ds[v].isel(time=itime, lat=ilat, lon=ilon).values for v in self.input_vars]\n",
    "                )\n",
    "            )\n",
    "            Y.append(\n",
    "                np.concatenate(\n",
    "                    [self.ds[v].isel(time=itime, lat=ilat, lon=ilon).values for v in self.output_vars]\n",
    "                )\n",
    "            )\n",
    "\n",
    "        return np.array(X), np.array(Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-file dataset\n",
    "\n",
    "Let's start by using the split dataset `sample_SPCAM_1.nc` and `sample_SPCAM_2.nc`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGenerator(DATADIR + 'sample_SPCAM_[1-2].nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is how we get one batch of inputs and corresponding outputs\n",
    "x, y = gen[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((128, 60), (128, 60))"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A little test function to check the timing.\n",
    "def test(g, n):\n",
    "    for i in range(n):\n",
    "        x, y = g[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13.3 s, sys: 1.34 s, total: 14.6 s\n",
      "Wall time: 14.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12.5 s, sys: 1.28 s, total: 13.8 s\n",
      "Wall time: 13.5 s\n"
     ]
    }
   ],
   "source": [
    "# does shuffling make a big difference\n",
    "gen = DataGenerator(DATADIR + 'sample_SPCAM_[1-2].nc', shuffle=True)\n",
    "%time test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So it takes more than one second to read one batch. This is way too slow to train a neural network in a reasonable amount of time. Shuffling doesn't seem to be a huge problem, but even without shuffling I am probably accessing the data in a different order than saved on disc. \n",
    "\n",
    "Let's check what actually takes that long."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The line_profiler extension is already loaded. To reload it, use:\n",
      "  %reload_ext line_profiler\n"
     ]
    }
   ],
   "source": [
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "%lprun -f gen.__getitem__ test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Output:\n",
    "\n",
    "```\n",
    "Timer unit: 1e-06 s\n",
    "\n",
    "Total time: 24.5229 s\n",
    "File: <ipython-input-57-78b9d254df3b>\n",
    "Function: __getitem__ at line 18\n",
    "\n",
    "Line #      Hits         Time  Per Hit   % Time  Line Contents\n",
    "==============================================================\n",
    "    18                                               def __getitem__(self, index):\n",
    "    19        10         17.0      1.7      0.0          time_indices, lat_indices, lon_indices = np.unravel_index(\n",
    "    20        10        267.0     26.7      0.0              self.indices[index*self.batch_size:(index+1)*self.batch_size], (self.ntime, self.nlat, self.nlon)\n",
    "    21                                                   )\n",
    "    22                                                   \n",
    "    23        10         10.0      1.0      0.0          X, Y = [], []\n",
    "    24      1290       4642.0      3.6      0.0          for itime, ilat, ilon in zip(time_indices, lat_indices, lon_indices):\n",
    "    25      1280       1399.0      1.1      0.0              X.append(\n",
    "    26      1280       1721.0      1.3      0.0                  np.concatenate(\n",
    "    27      1280   12256070.0   9575.1     50.0                      [self.ds[v].isel(time=itime, lat=ilat, lon=ilon).values for v in self.input_vars]\n",
    "    28                                                           )\n",
    "    29                                                       )\n",
    "    30      1280       2393.0      1.9      0.0              Y.append(\n",
    "    31      1280       1750.0      1.4      0.0                  np.concatenate(\n",
    "    32      1280   12253415.0   9573.0     50.0                      [self.ds[v].isel(time=itime, lat=ilat, lon=ilon).values for v in self.output_vars]\n",
    "    33                                                           )\n",
    "    34                                                       )\n",
    "    35                                           \n",
    "    36        10       1218.0    121.8      0.0          return np.array(X), np.array(Y)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the concatenated dataset\n",
    "\n",
    "Let's see whether it makes a difference to use the pre-concatenated dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.93 s, sys: 984 ms, total: 6.91 s\n",
      "Wall time: 6.91 s\n"
     ]
    }
   ],
   "source": [
    "ds = xr.open_dataset(f'{DATADIR}sample_SPCAM_concat.nc')\n",
    "gen = DataGenerator(ds, shuffle=True)\n",
    "%time test(gen, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 11.5 s, sys: 1.25 s, total: 12.8 s\n",
      "Wall time: 12.5 s\n"
     ]
    }
   ],
   "source": [
    "ds = xr.open_mfdataset(f'{DATADIR}sample_SPCAM_concat.nc')\n",
    "gen = DataGenerator(ds, shuffle=True)\n",
    "%time test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So yes, it approximately halves the time but only if the single dataset is NOT opened with `open_mfdataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With h5py engine\n",
    "\n",
    "Let's see whether using the h5py backend makes a difference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5netcdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(f'{DATADIR}sample_SPCAM_concat.nc', engine='h5netcdf')\n",
    "gen = DataGenerator(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.97 s, sys: 972 ms, total: 7.94 s\n",
      "Wall time: 7.8 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Doesn't seem to speed it up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using plain h5py\n",
    "\n",
    "Let's write a version of the data generator that uses plain h5py for data loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataGeneratorH5(object):\n",
    "    def __init__(self, fn, batch_size=128, input_vars=['TAP', 'QAP'], output_vars=['TPHYSTND', 'PHQ'], shuffle=True):\n",
    "        self.ds = xr.open_dataset(fn)\n",
    "        self.batch_size = batch_size\n",
    "        self.input_vars = input_vars\n",
    "        self.output_vars = output_vars\n",
    "        self.ntime, self.nlat, self.nlon = self.ds.time.size, self.ds.lat.size, self.ds.lon.size\n",
    "        self.ntot = self.ntime * self.nlat * self.ntime\n",
    "        self.n_batches = self.ntot // batch_size\n",
    "        self.indices = np.arange(self.ntot)\n",
    "        if shuffle:\n",
    "            self.indices = np.random.permutation(self.indices)\n",
    "            \n",
    "        # Close xarray dataset and open h5py object\n",
    "        self.ds.close()\n",
    "        self.ds = h5py.File(fn, 'r')\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        time_indices, lat_indices, lon_indices = np.unravel_index(\n",
    "            self.indices[index*self.batch_size:(index+1)*self.batch_size], (self.ntime, self.nlat, self.nlon)\n",
    "        )\n",
    "        \n",
    "        X, Y = [], []\n",
    "        for itime, ilat, ilon in zip(time_indices, lat_indices, lon_indices):\n",
    "            X.append(\n",
    "                np.concatenate(\n",
    "                    [self.ds[v][itime, :, ilat, ilon] for v in self.input_vars]\n",
    "                )\n",
    "            )\n",
    "            Y.append(\n",
    "                np.concatenate(\n",
    "                    [self.ds[v][itime, :, ilat, ilon] for v in self.output_vars]\n",
    "                )\n",
    "            )\n",
    "\n",
    "        return np.array(X), np.array(Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGeneratorH5(f'{DATADIR}sample_SPCAM_concat.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.78 s, sys: 860 ms, total: 2.64 s\n",
      "Wall time: 2.61 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.ds.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So this is significantly faster than xarray."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use in a simple neural network\n",
    "\n",
    "How would we actually use this data generator for network training...\n",
    "\n",
    "Note that this neural network will not actually learn much because we didn't normalize the input data. But we only care about computational performance here, right?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import Sequential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.1.6-tf'"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.keras.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Sequential([\n",
    "    Dense(128, input_shape=(60,), activation='relu'),\n",
    "    Dense(60),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 128)               7808      \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 60)                7740      \n",
      "=================================================================\n",
      "Total params: 15,548\n",
      "Trainable params: 15,548\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile('adam', 'mse')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the xarray version using the concatenated dataset\n",
    "ds = xr.open_dataset(f'{DATADIR}sample_SPCAM_concat.nc')\n",
    "gen = DataGenerator(ds, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1\n",
      "  37/4608 [..............................] - ETA: 1:04:11 - loss: 1733.6299"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-101-401f100b439b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit_generator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgen\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn_batches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/miniconda3/envs/tf/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit_generator\u001b[0;34m(self, generator, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch)\u001b[0m\n\u001b[1;32m   2175\u001b[0m         \u001b[0muse_multiprocessing\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0muse_multiprocessing\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2176\u001b[0m         \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mshuffle\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2177\u001b[0;31m         initial_epoch=initial_epoch)\n\u001b[0m\u001b[1;32m   2178\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2179\u001b[0m   def evaluate_generator(self,\n",
      "\u001b[0;32m~/miniconda3/envs/tf/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_generator.py\u001b[0m in \u001b[0;36mfit_generator\u001b[0;34m(model, generator, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch)\u001b[0m\n\u001b[1;32m    145\u001b[0m       \u001b[0mbatch_index\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    146\u001b[0m       \u001b[0;32mwhile\u001b[0m \u001b[0msteps_done\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m         \u001b[0mgenerator_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_generator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgenerator_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'__len__'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/tf/lib/python3.6/site-packages/tensorflow/python/keras/utils/data_utils.py\u001b[0m in \u001b[0;36mget\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    823\u001b[0m           \u001b[0;32mraise\u001b[0m \u001b[0mStopIteration\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    824\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 825\u001b[0;31m           \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwait_time\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    826\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    827\u001b[0m     \u001b[0;31m# Make sure to rethrow the first exception in the queue, if any\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "model.fit_generator(iter(gen), steps_per_epoch=gen.n_batches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So as you can see, it would take around 1 hour to go through one epoch (i.e. the entire dataset once). This is crazy slow since we only used 2 days of data. The full dataset contains a year of data..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-processing the dataset\n",
    "\n",
    "What I have resorted to to solve this issue is to prestack the data, preshuffle the data and save it to disc conveniently.\n",
    "\n",
    "These files contain the exactly same information for the input (features) and output (targets) variables required.\n",
    "\n",
    "The files only have two dimensions: sample, which is the shuffled, flattened time, lat and lon dimensions and lev which is the stacked vertical coordinate.\n",
    "\n",
    "The preprocessing for these two files only takes a few seconds but for an entire year of data, the preprocessing alone can take around an hour.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2019-02-07 15:42:32--  https://zenodo.org/record/2559313/files/preproc_features.nc\n",
      "Resolving zenodo.org (zenodo.org)... 137.138.76.77\n",
      "Connecting to zenodo.org (zenodo.org)|137.138.76.77|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 205465847 (196M) [application/octet-stream]\n",
      "Saving to: ‘/local/S.Rasp/tmp/preproc_features.nc.2’\n",
      "\n",
      "preproc_features.nc 100%[===================>] 195.95M  10.8MB/s    in 15s     \n",
      "\n",
      "2019-02-07 15:42:48 (13.0 MB/s) - ‘/local/S.Rasp/tmp/preproc_features.nc.2’ saved [205465847/205465847]\n",
      "\n",
      "--2019-02-07 15:42:48--  https://zenodo.org/record/2559313/files/preproc_targets.nc\n",
      "Resolving zenodo.org (zenodo.org)... 137.138.76.77\n",
      "Connecting to zenodo.org (zenodo.org)|137.138.76.77|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 205465846 (196M) [application/octet-stream]\n",
      "Saving to: ‘/local/S.Rasp/tmp/preproc_targets.nc.1’\n",
      "\n",
      "preproc_targets.nc. 100%[===================>] 195.95M  9.98MB/s    in 9.5s    \n",
      "\n",
      "2019-02-07 15:42:58 (20.6 MB/s) - ‘/local/S.Rasp/tmp/preproc_targets.nc.1’ saved [205465846/205465846]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget -P $DATADIR https://zenodo.org/record/2559313/files/preproc_features.nc\n",
    "!wget -P $DATADIR https://zenodo.org/record/2559313/files/preproc_targets.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 S.Rasp ls-craig 196M Feb  7 13:57 /local/S.Rasp/tmp//preproc_features.nc\r\n",
      "-rw-r--r-- 1 S.Rasp ls-craig 196M Feb  7 13:57 /local/S.Rasp/tmp//preproc_targets.nc\r\n"
     ]
    }
   ],
   "source": [
    "!ls -lh $DATADIR/preproc*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(f'{DATADIR}preproc_features.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:        (feature_lev: 60, sample: 778240)\n",
       "Coordinates:\n",
       "  * feature_lev    (feature_lev) int64 0 1 2 3 4 5 6 7 ... 53 54 55 56 57 58 59\n",
       "    time           (sample) int64 ...\n",
       "    lat            (sample) float64 ...\n",
       "    lon            (sample) float64 ...\n",
       "    feature_names  (feature_lev) object ...\n",
       "Dimensions without coordinates: sample\n",
       "Data variables:\n",
       "    features       (sample, feature_lev) float32 ...\n",
       "Attributes:\n",
       "    log:      \\n    Time: 2019-02-07T13:57:24\\n\\n    Executed command:\\n\\n   ..."
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write a new data generator\n",
    "class DataGeneratorPreproc(object):\n",
    "    \"\"\"\n",
    "    Data generator that randomly (if shuffle = True) picks columns from the dataset and returns them in \n",
    "    batches. For each column the input variables and output variables will be stacked.\n",
    "    \"\"\"\n",
    "    def __init__(self, feature_fn, target_fn, batch_size=128, shuffle=True, engine='netcdf4'):\n",
    "        self.feature_ds = xr.open_dataset(feature_fn, engine=engine)\n",
    "        self.target_ds = xr.open_dataset(target_fn, engine=engine)\n",
    "        self.batch_size = batch_size\n",
    "        self.ntot = self.feature_ds.sample.size\n",
    "        self.n_batches = self.ntot // batch_size\n",
    "        self.indices = np.arange(self.ntot)\n",
    "        if shuffle:\n",
    "            self.indices = np.random.permutation(self.indices)\n",
    "    def __getitem__(self, index):\n",
    "        batch_indices = self.indices[index*self.batch_size:(index+1)*self.batch_size]\n",
    "        \n",
    "        X = self.feature_ds.features.isel(sample=batch_indices)\n",
    "        Y = self.target_ds.targets.isel(sample=batch_indices)\n",
    "\n",
    "        return X, Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGeneratorPreproc(f'{DATADIR}preproc_features.nc', f'{DATADIR}preproc_targets.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = gen[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((128, 60), (128, 60))"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 84 ms, sys: 0 ns, total: 84 ms\n",
      "Wall time: 81.6 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGeneratorPreproc(f'{DATADIR}preproc_features.nc', f'{DATADIR}preproc_targets.nc', shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 84 ms, sys: 0 ns, total: 84 ms\n",
      "Wall time: 83.9 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.feature_ds.close(); gen.target_ds.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGeneratorPreproc(f'{DATADIR}preproc_features.nc', f'{DATADIR}preproc_targets.nc', engine='h5netcdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 84 ms, sys: 0 ns, total: 84 ms\n",
      "Wall time: 80.6 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So these are the sort of times that are required for training a neural network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pure h5py version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataGeneratorPreprocH5(object):\n",
    "    \"\"\"\n",
    "    Data generator that randomly (if shuffle = True) picks columns from the dataset and returns them in \n",
    "    batches. For each column the input variables and output variables will be stacked.\n",
    "    \"\"\"\n",
    "    def __init__(self, feature_fn, target_fn, batch_size=128):\n",
    "        self.feature_ds = xr.open_dataset(feature_fn)\n",
    "        self.target_ds = xr.open_dataset(target_fn)\n",
    "        self.batch_size = batch_size\n",
    "        self.ntot = self.feature_ds.sample.size\n",
    "        self.n_batches = self.ntot // batch_size\n",
    "        \n",
    "        # Close xarray dataset and open h5py object\n",
    "        self.feature_ds.close()\n",
    "        self.feature_ds = h5py.File(feature_fn, 'r')\n",
    "        self.target_ds.close()\n",
    "        self.target_ds = h5py.File(target_fn, 'r')\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        X = self.feature_ds['features'][index*self.batch_size:(index+1)*self.batch_size, :]\n",
    "        Y = self.target_ds['targets'][index*self.batch_size:(index+1)*self.batch_size, :]\n",
    "\n",
    "        return X, Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.feature_ds.close(); gen.target_ds.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = DataGeneratorPreprocH5(f'{DATADIR}preproc_features.nc', f'{DATADIR}preproc_targets.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8 ms, sys: 0 ns, total: 8 ms\n",
      "Wall time: 6.61 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "test(gen, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So again, the pure h5py version is an order of magnitude faster than the xarray version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End"
   ]
  },
  {
   "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.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": false,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}