{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp data.tsloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TimeSeriesLoader\n",
    "> Data Loader for Time Series data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import warnings\n",
    "from collections.abc import Mapping\n",
    "from typing import Dict, List, Optional, Union\n",
    "\n",
    "import numpy as np\n",
    "import torch as t\n",
    "from fastcore.foundation import patch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nixtlats.data.tsdataset import TimeSeriesDataset, WindowsDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inherited `DataLoader` from `pytorch` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TimeSeriesLoader(DataLoader):\n",
    "\n",
    "    def __init__(self, dataset: Union[TimeSeriesDataset, WindowsDataset], \n",
    "                 eq_batch_size: bool = False, \n",
    "                 **kwargs) -> 'TimeSeriesLoader':\n",
    "        \"\"\"Wraps the pytorch `DataLoader` with a special collate function \n",
    "        for the `TimeSeriesDataset` ouputs.\n",
    "        \n",
    "        The TimeSeriesDataset constructs all the trainable windows \n",
    "        of `batch_size` series. The number of windows can be greater \n",
    "        or smaller than the `batch_size`. For this reason, \n",
    "        an additional boolean parameter, `eq_batch_size` is included \n",
    "        that if `True` samples `batch_size` windows randomly, \n",
    "        while `False` returns all windows.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        dataset: TimeSeriesDataset\n",
    "            Stored time series.\n",
    "        eq_batch_size: bool\n",
    "            If `True` samples `batch_size` windows randomly,\n",
    "            while `False` or `batch_size=None` returns all windows.\n",
    "        \"\"\"\n",
    "        if 'collate_fn' in kwargs.keys():\n",
    "            warnings.warn(\n",
    "                'This class wraps the pytorch `DataLoader` with a '\n",
    "                'special collate function. If you want to use yours '\n",
    "                'simply use `DataLoader`. Removing collate_fn'\n",
    "            )\n",
    "            kwargs.pop('collate_fn')\n",
    "            \n",
    "        kwargs_ = {**kwargs, **dict(collate_fn=self._collate_fn)}\n",
    "        DataLoader.__init__(self, dataset=dataset, **kwargs_)\n",
    "        self.eq_batch_size = eq_batch_size\n",
    "        self.w_idxs: Optional[np.ndarray] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _check_batch_size(self: TimeSeriesLoader, batch: t.Tensor):\n",
    "    complete_batch = batch\n",
    "    if self.w_idxs is not None:\n",
    "        complete_batch = batch[self.w_idxs]\n",
    "    \n",
    "    return complete_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _collate_fn(self: TimeSeriesLoader, batch: Union[List, Dict[str, t.Tensor], t.Tensor]):\n",
    "    \"\"\"Special collate fn for the `TimeSeriesDataset`.\n",
    "    \n",
    "    Notes\n",
    "    -----\n",
    "    [1] Adapted from https://github.com/pytorch/pytorch/blob/master/torch/utils/data/_utils/collate.py.\n",
    "    \"\"\"\n",
    "    elem = batch[0]\n",
    "    if len(batch) == 1:\n",
    "        return {key: self._check_batch_size(elem[key]) for key in elem}\n",
    "    \n",
    "    elem_type = type(elem)\n",
    "    \n",
    "    if isinstance(elem, t.Tensor):\n",
    "        out = None\n",
    "        if t.utils.data.get_worker_info() is not None:\n",
    "            # If we're in a background process, concatenate directly into a\n",
    "            # shared memory tensor to avoid an extra copy\n",
    "            numel = sum([x.numel() for x in batch])\n",
    "            storage = elem.storage()._new_shared(numel)\n",
    "            out = elem.new(storage)\n",
    "        complete_batch = t.cat(batch, out=out)\n",
    "        return self._check_batch_size(complete_batch)\n",
    "    \n",
    "    elif isinstance(elem, Mapping):\n",
    "        n_windows = elem['Y'].size(0)\n",
    "        if self.eq_batch_size and self.batch_size is not None:\n",
    "            self.w_idxs = np.random.choice(n_windows, size=self.batch_size, \n",
    "                                           replace=(n_windows < self.batch_size))\n",
    "        return {key: self.collate_fn([d[key] for d in batch]) for key in elem}\n",
    "\n",
    "    raise TypeError(f'Unknown {elem_type}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tests WindowsDataset and TimeSeriesDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.utils import create_synthetic_tsdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_eq_batch_size(dataset, batch_size, loader_class):\n",
    "    # Check returns batch_size tensors\n",
    "    loader = loader_class(dataset=dataset, batch_size=batch_size, eq_batch_size=True)\n",
    "    sizes = [batch['Y'].size(0) == batch_size for batch in loader]\n",
    "    \n",
    "    assert all(sizes), 'Unexpected batch sizes.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_eq_batch_size_order(dataset, batch_size, loader_class):\n",
    "    #This test only works for TimeSeriesDataset class\n",
    "    loader = loader_class(dataset=dataset, batch_size=batch_size, eq_batch_size=True)\n",
    "    for batch in loader:\n",
    "        idxs = batch['idxs']\n",
    "        dataset_batch = dataset[idxs.numpy().tolist()]\n",
    "        for key in batch.keys():\n",
    "            assert t.equal(batch[key], dataset_batch[key]), (\n",
    "                f'Batch and dataset batch differ, key {key}'\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete timeseries dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df, X_df, S_df = create_synthetic_tsdata(sort=True)\n",
    "dataset = TimeSeriesDataset(S_df=S_df, Y_df=Y_df, X_df=X_df,\n",
    "                            input_size=5,\n",
    "                            output_size=2)\n",
    "dataloader = TimeSeriesLoader(dataset=dataset, batch_size=12, \n",
    "                              eq_batch_size=False, shuffle=True)\n",
    "\n",
    "for batch in dataloader:\n",
    "    batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size(dataset, 32, TimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size_order(dataset, 32, TimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Windowed timeseries dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df, X_df, S_df = create_synthetic_tsdata(sort=True)\n",
    "dataset = WindowsDataset(S_df=S_df, Y_df=Y_df, X_df=X_df,\n",
    "                         input_size=5,\n",
    "                         output_size=2,\n",
    "                         sample_freq=1)\n",
    "dataloader = TimeSeriesLoader(dataset=dataset, batch_size=12, \n",
    "                              eq_batch_size=False, shuffle=True)\n",
    "\n",
    "for batch in dataloader:\n",
    "    batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size(dataset, 32, TimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Faster implemention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FastTimeSeriesLoader:\n",
    "    \"\"\"\n",
    "    A DataLoader-like object for a set of tensors that can be much faster than\n",
    "    TensorDataset + DataLoader because dataloader grabs individual indices of\n",
    "    the dataset and calls cat (slow).\n",
    "    Source: https://discuss.pytorch.org/t/dataloader-much-slower-than-manual-batching/27014/6\n",
    "    \n",
    "    Notes\n",
    "    -----\n",
    "    [1] Adapted from https://github.com/hcarlens/pytorch-tabular/blob/master/fast_tensor_data_loader.py.\n",
    "    \"\"\"\n",
    "    def __init__(self, dataset: TimeSeriesDataset, batch_size: int = 32, \n",
    "                 eq_batch_size: bool = False,\n",
    "                 shuffle: bool = False) -> 'FastTimeSeriesLoader':\n",
    "        \"\"\"Initialize a FastTimeSeriesLoader.\n",
    "        \n",
    "        The TimeSeriesDataset constructs all the trainable windows \n",
    "        of `batch_size` series. The number of windows can be greater \n",
    "        or smaller than the `batch_size`. For this reason, \n",
    "        an additional boolean parameter, `eq_batch_size` is included \n",
    "        that if `True` samples `batch_size` windows randomly, \n",
    "        while `False` returns all windows.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        dataset: TimeSeriesDataset\n",
    "            Stored time series.\n",
    "        batch_size: int\n",
    "            Batch size to load.\n",
    "        shuffle: bool \n",
    "            If `True`, shuffle the data *in-place* whenever an\n",
    "            iterator is created out of this object.\n",
    "        \"\"\"\n",
    "        self.dataset = dataset\n",
    "        self.dataset_len = len(dataset)\n",
    "        self.batch_size = batch_size\n",
    "        self.eq_batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.idxs = np.arange(self.dataset_len)\n",
    "\n",
    "        # Calculate # batches\n",
    "        n_batches, remainder = divmod(self.dataset_len, self.batch_size)\n",
    "        if remainder > 0:\n",
    "            n_batches += 1\n",
    "        self.n_batches = n_batches\n",
    "        self.w_idxs: Optional[np.ndarray] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __iter__(self: FastTimeSeriesLoader):\n",
    "    if self.shuffle:\n",
    "        self.idxs = np.random.permutation(self.dataset_len)\n",
    "\n",
    "    self.i = 0\n",
    "    return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _check_batch_size(self: FastTimeSeriesLoader, batch: t.Tensor):\n",
    "    complete_batch = batch\n",
    "    if self.w_idxs is not None:\n",
    "        complete_batch = batch[self.w_idxs]\n",
    "    return complete_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __next__(self: FastTimeSeriesLoader):\n",
    "    if self.i >= self.dataset_len:\n",
    "        raise StopIteration\n",
    "    idxs = self.idxs[self.i:(self.i + self.batch_size)].tolist()\n",
    "    batch = self.dataset[idxs]\n",
    "    self.i += self.batch_size\n",
    "    \n",
    "    n_windows = batch['Y'].size(0)\n",
    "    if self.eq_batch_size and self.batch_size is not None:\n",
    "        self.w_idxs = np.random.choice(n_windows, size=self.batch_size, \n",
    "                                     replace=(n_windows < self.batch_size))\n",
    "    \n",
    "    return {key: self._check_batch_size(batch[key]) for key in batch}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __len__(self: FastTimeSeriesLoader):\n",
    "    return self.n_batches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tests WindowsDataset and TimeSeriesDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete timeseries dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df, X_df, S_df = create_synthetic_tsdata(sort=True)\n",
    "dataset = TimeSeriesDataset(S_df=S_df, Y_df=Y_df, X_df=X_df,\n",
    "                            input_size=5,\n",
    "                            output_size=2)\n",
    "dataloader = FastTimeSeriesLoader(dataset=dataset, batch_size=12, \n",
    "                                  eq_batch_size=False, shuffle=True)\n",
    "\n",
    "for batch in dataloader:\n",
    "    batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size(dataset, 32, FastTimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size_order(dataset, 32, FastTimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Windowed timeseries dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df, X_df, S_df = create_synthetic_tsdata(sort=True)\n",
    "dataset = WindowsDataset(S_df=S_df, Y_df=Y_df, X_df=X_df,\n",
    "                         input_size=5,\n",
    "                         output_size=2,\n",
    "                         sample_freq=1)\n",
    "dataloader = FastTimeSeriesLoader(dataset=dataset, batch_size=12, \n",
    "                                  eq_batch_size=False, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in dataloader:\n",
    "    batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq_batch_size(dataset, 32, FastTimeSeriesLoader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = TimeSeriesLoader(dataset=dataset, batch_size=12, shuffle=True)\n",
    "fast_dataloader = FastTimeSeriesLoader(dataset=dataset, batch_size=12, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n 50 -r 3  [batch for batch in dataloader]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n 50 -r 3 [batch for batch in fast_dataloader]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
