brenowitz-data-loading.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a toy xarray dataset. It has a few 3D and 2D variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (time: 4, x: 2, y: 5, z: 4)\n",
       "Dimensions without coordinates: time, x, y, z\n",
       "Data variables:\n",
       "    a        (time, z, y, x) float64 1.0 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 1.0\n",
       "    b        (time, y, x) float64 1.0 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 1.0"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "from itertools import product\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "class XRTimeSeries(Dataset):\n",
    "    \"\"\"A pytorch Dataset class for time series data in xarray format\n",
    "\n",
    "    This function assumes the data has dimensions ['time', 'z', 'y', 'x'], and\n",
    "    that the axes of the data arrays are all stored in that order.\n",
    "\n",
    "    An individual \"sample\" is the full time time series from a single\n",
    "    horizontal location. The time-varying variables in this sample will have\n",
    "    shape (time, z, 1, 1).\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    >>> ds = xr.open_dataset(\"in.nc\")\n",
    "    >>> dataset = XRTimeSeries(ds)\n",
    "    >>> dataset[0]\n",
    "\n",
    "    \"\"\"\n",
    "    dims = ['time', 'z', 'x', 'y']\n",
    "\n",
    "    def __init__(self, data, time_length=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        data : xr.DataArray\n",
    "            An input dataset. This dataset must contain at least some variables\n",
    "            with all of the dimensions ['time' , 'z', 'x', 'y'].\n",
    "        time_length : int, optional\n",
    "            The length of the time sequences to use, must evenly divide the\n",
    "            total number of time points.\n",
    "        \"\"\"\n",
    "        self.time_length = time_length or len(data.time)\n",
    "        self.data = data\n",
    "        self.numpy_data = {key: data[key].values for key in data.data_vars}\n",
    "        self.data_vars = set(data.data_vars)\n",
    "        self.dims = {key: data[key].dims for key in data.data_vars}\n",
    "        self.constants = {\n",
    "            key\n",
    "            for key in data.data_vars\n",
    "            if len({'x', 'y', 'time'} & set(data[key].dims)) == 0\n",
    "        }\n",
    "        self.setup_indices()\n",
    "\n",
    "    def setup_indices(self):\n",
    "        len_x = len(self.data['x'].values)\n",
    "        len_y = len(self.data['y'].values)\n",
    "        len_t = len(self.data['time'].values)\n",
    "\n",
    "        x_iter = range(0, len_x, 1)\n",
    "        y_iter = range(0, len_y, 1)\n",
    "        t_iter = range(0, len_t, self.time_length)\n",
    "        assert len_t % self.time_length == 0\n",
    "        self.indices = list(product(t_iter, y_iter, x_iter))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.indices)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        t, y, x = self.indices[i]\n",
    "        output_tensors = {}\n",
    "        for key in self.data_vars:\n",
    "            if key in self.constants:\n",
    "                continue\n",
    "\n",
    "            data_array = self.numpy_data[key]\n",
    "            if 'z' in self.dims[key]:\n",
    "                this_array_index = (slice(t, t + self.time_length),\n",
    "                                    slice(None), y, x)\n",
    "            else:\n",
    "                this_array_index = (slice(t, t + self.time_length), None, y, x)\n",
    "\n",
    "            sample = data_array[this_array_index][:, :, np.newaxis, np.newaxis]\n",
    "            output_tensors[key] = sample.astype(np.float32)\n",
    "        return output_tensors\n",
    "\n",
    "    @property\n",
    "    def time_dim(self):\n",
    "        return self.dims[0][0]\n",
    "\n",
    "    def torch_constants(self):\n",
    "        return {\n",
    "            key: torch.tensor(self.data[key].values, requires_grad=False)\n",
    "            .float()\n",
    "            for key in self.constants\n",
    "        }\n",
    "\n",
    "    @property\n",
    "    def scale(self):\n",
    "        std = self.std\n",
    "        return valmap(lambda x: x.max(), std)\n",
    "    \n",
    "\n",
    "def get_xarray_dataset():\n",
    "\n",
    "    dims_3d = ['time', 'z', 'y', 'x']\n",
    "    dims_2d = ['time', 'y', 'x']\n",
    "\n",
    "    data_3d = np.ones((4, 4, 5, 2))\n",
    "    data_2d = np.ones((4, 5, 2))\n",
    "\n",
    "    return xr.Dataset({\n",
    "        'a': (dims_3d, data_3d),\n",
    "        'b': (dims_2d, data_2d)\n",
    "    })\n",
    "\n",
    "ds = get_xarray_dataset()\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_dataset = XRTimeSeries(ds, time_length=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(torch_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The length of the dataset is $x\\dot y$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': array([[[[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]],\n",
       " \n",
       "         [[1.]]]], dtype=float32), 'b': array([[[[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]]],\n",
       " \n",
       " \n",
       "        [[[1.]]]], dtype=float32)}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = torch_dataset[0]\n",
    "sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two dimensional and three dimensional variables have broadcastable shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 1, 1, 1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample['b'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 4, 1, 1)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample['a'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have made the torch dataset object, we can pass it to pytorch's DataLoader class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(torch_dataset, batch_size=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape of b torch.Size([4, 4, 1, 1, 1])\n",
      "shape of b torch.Size([4, 4, 1, 1, 1])\n",
      "shape of b torch.Size([2, 4, 1, 1, 1])\n"
     ]
    }
   ],
   "source": [
    "for batch in train_loader:\n",
    "    print(\"shape of b\", batch['b'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first dimension becomes the \"batch\" dimension. The other dimensions are the physical dimensions (time, z, y, x). My [model classes](https://github.com/nbren12/uwnet/blob/047a63b70985b12e17013355ecd25c908681ab76/uwnet/modules.py) accept data in this format."
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}