{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:34:35.588620Z",
     "start_time": "2020-11-20T05:34:35.582780Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data\n",
    "\n",
    "> Tools to help construct datasets, which may be related to loading, processing, or encoding data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:38:10.862265Z",
     "start_time": "2021-08-07T20:38:09.126444Z"
    }
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:38:10.900421Z",
     "start_time": "2021-08-07T20:38:10.864875Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:38:13.138688Z",
     "start_time": "2021-08-07T20:38:10.905320Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "import boto3\n",
    "from collections import deque\n",
    "from functools import partial\n",
    "import mmh3\n",
    "import matplotlib.pyplot as plt\n",
    "from multiprocessing import Pool\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm.auto import tqdm\n",
    "import warnings\n",
    "\n",
    "from htools import auto_repr, save, load, func_name, ifnone, tolist, \\\n",
    "    BasicPipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:38:13.173539Z",
     "start_time": "2021-08-07T20:38:13.140956Z"
    }
   },
   "outputs": [],
   "source": [
    "# Only needed for testing.\n",
    "from collections import Counter\n",
    "from itertools import chain\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from htools import eprint, assert_raises"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:34:40.997460Z",
     "start_time": "2020-11-20T05:34:40.962702Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def probabilistic_hash_item(x, n_buckets, mode=int, n_hashes=3):\n",
    "    \"\"\"Slightly hacky way to probabilistically hash an integer by\n",
    "    first converting it to a string.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: int\n",
    "        The integer or string to hash.\n",
    "    n_buckets: int\n",
    "        The number of buckets that items will be mapped to. Typically \n",
    "        this would occur outside the hashing function, but since \n",
    "        the intended use case is so narrow here it makes sense to me \n",
    "        to include it here.\n",
    "    mode: type\n",
    "        The type of input you want to hash. This is user-provided to prevent\n",
    "        accidents where we pass in a different item than intended and hash \n",
    "        the wrong thing. One of (int, str). When using this inside a\n",
    "        BloomEmbedding layer, this must be `int` because there are no\n",
    "        string tensors. When used inside a dataset or as a one-time\n",
    "        pre-processing step, you can choose either as long as you\n",
    "        pass in the appropriate inputs.\n",
    "    n_hashes: int\n",
    "        The number of times to hash x, each time with a different seed.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[int]: A list of integers with length `n_hashes`, where each integer\n",
    "        is in [0, n_buckets).\n",
    "    \"\"\"\n",
    "    # Check type to ensure we don't accidentally hash Tensor(5) instead of 5.\n",
    "    assert isinstance(x, mode), f'Input `x` must have type {mode}.'\n",
    "    return [mmh3.hash(str(x), i, signed=False) % n_buckets \n",
    "            for i in range(n_hashes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:34:41.938165Z",
     "start_time": "2020-11-20T05:34:41.898962Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def probabilistic_hash_tensor(x_r2, n_buckets, n_hashes=3, pad_idx=0):\n",
    "    \"\"\"Hash a rank 2 LongTensor.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x_r2: torch.LongTensor\n",
    "        Rank 2 tensor of integers. Shape: (bs, seq_len)\n",
    "    n_buckets: int\n",
    "        Number of buckets to hash items into (i.e. the number of \n",
    "        rows in the embedding matrix). Typically a moderately large\n",
    "        prime number, like 251 or 997.\n",
    "    n_hashes: int\n",
    "        Number of hashes to take for each input index. This determines\n",
    "        the number of rows of the embedding matrix that will be summed\n",
    "        to get the representation for each word. Typically 2-5.\n",
    "    pad_idx: int or None\n",
    "        If you want to pad sequences with vectors of zeros, pass in an\n",
    "        integer (same as the `padding_idx` argument to nn.Embedding).\n",
    "        If None, no padding index will be used. The sequences must be\n",
    "        padded before passing them into this function.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.LongTensor: Tensor of indices where each row corresponds\n",
    "        to one of the input indices. Shape: (bs, seq_len, n_hashes)\n",
    "    \"\"\"\n",
    "    return torch.tensor(\n",
    "        [[probabilistic_hash_item(x.item(), n_buckets, int, n_hashes) \n",
    "          if x != pad_idx else [pad_idx]*n_hashes for x in row]\n",
    "         for row in x_r2]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sents = [\n",
    "    'I walked to the store so I hope it is not closed.',\n",
    "    'The theater is closed today and the sky is grey.',\n",
    "    'His dog is brown while hers is grey.'\n",
    "]\n",
    "labels = [0, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data(Dataset):\n",
    "    \n",
    "    def __init__(self, sentences, labels, seq_len):\n",
    "        x = [s.split(' ') for s in sentences]\n",
    "        self.w2i = self.make_w2i(x)\n",
    "        self.seq_len = seq_len\n",
    "        self.x = self.encode(x)\n",
    "        self.y = torch.tensor(labels)\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "    \n",
    "    def make_w2i(self, tok_rows):\n",
    "        return {k: i for i, (k, v) in \n",
    "                enumerate(Counter(chain(*tok_rows)).most_common(), 1)}\n",
    "    \n",
    "    def encode(self, tok_rows):\n",
    "        enc = np.zeros((len(tok_rows), self.seq_len), dtype=int)\n",
    "        for i, row in enumerate(tok_rows):\n",
    "            trunc = [self.w2i.get(w, 0) for w in row[:self.seq_len]]\n",
    "            enc[i, :len(trunc)] = trunc\n",
    "        return torch.tensor(enc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We construct a toy dataset with a vocabulary of size 23. In reality, you might wish to lowercase text or use a better tokenizer, but this is sufficient for the purposes of demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = Data(sents, labels, 10)\n",
    "len(ds.w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]), tensor([0, 1, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dl = DataLoader(ds, batch_size=3)\n",
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 10])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We hash each word index 4 times, as specified by the `n_hashes` parameter in `probabilistic_hash_tensor`. Notice that we only use 7 buckets, meaning the embedding matrix will have 7 rows rather than 23 (not counting a padding row)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x shape: torch.Size([3, 10])\n",
      "x_hashed shape: torch.Size([3, 10, 4])\n"
     ]
    }
   ],
   "source": [
    "x_hashed = probabilistic_hash_tensor(x, n_buckets=7, n_hashes=4)\n",
    "print('x shape:', x.shape)\n",
    "print('x_hashed shape:', x_hashed.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, each row of 4 numbers encodes a single word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2, 0, 2, 2],\n",
       "         [1, 6, 2, 4],\n",
       "         [2, 0, 1, 4],\n",
       "         [5, 2, 6, 5],\n",
       "         [1, 5, 1, 3],\n",
       "         [0, 4, 0, 0],\n",
       "         [2, 0, 2, 2],\n",
       "         [2, 4, 0, 2],\n",
       "         [3, 4, 4, 6],\n",
       "         [5, 0, 3, 6]],\n",
       "\n",
       "        [[5, 4, 4, 2],\n",
       "         [5, 3, 3, 1],\n",
       "         [5, 0, 3, 6],\n",
       "         [2, 1, 1, 1],\n",
       "         [2, 4, 4, 4],\n",
       "         [2, 4, 2, 6],\n",
       "         [5, 2, 6, 5],\n",
       "         [3, 5, 0, 0],\n",
       "         [5, 0, 3, 6],\n",
       "         [1, 5, 2, 6]],\n",
       "\n",
       "        [[5, 5, 3, 4],\n",
       "         [4, 5, 5, 1],\n",
       "         [5, 0, 3, 6],\n",
       "         [6, 2, 0, 6],\n",
       "         [4, 2, 6, 1],\n",
       "         [3, 6, 1, 6],\n",
       "         [5, 0, 3, 6],\n",
       "         [1, 5, 2, 6],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_hashed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See how each word is mapped to a list of 4 indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I [2, 0, 2, 2]\n",
      "walked [1, 6, 2, 4]\n",
      "to [2, 0, 1, 4]\n",
      "the [5, 2, 6, 5]\n",
      "store [1, 5, 1, 3]\n",
      "so [0, 4, 0, 0]\n",
      "I [2, 0, 2, 2]\n",
      "hope [2, 4, 0, 2]\n",
      "it [3, 4, 4, 6]\n",
      "is [5, 0, 3, 6]\n"
     ]
    }
   ],
   "source": [
    "for word, i in zip(sents[0].split(' '), x[0]):\n",
    "    print(word, probabilistic_hash_item(i.item(), 7, int, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that hashing the words directly is also possible, but the resulting hashes will be **different** than if hashing after encoding words as integers. This is fine as long as you are consistent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0: ('I', [0, 5, 5])\n",
      " 1: ('walked', [2, 5, 1])\n",
      " 2: ('to', [10, 4, 6])\n",
      " 3: ('the', [4, 1, 4])\n",
      " 4: ('store', [4, 6, 3])\n",
      " 5: ('so', [7, 8, 8])\n",
      " 6: ('I', [0, 5, 5])\n",
      " 7: ('hope', [1, 2, 7])\n",
      " 8: ('it', [3, 9, 0])\n",
      " 9: ('is', [3, 1, 3])\n",
      "10: ('not', [6, 10, 4])\n",
      "11: ('closed.', [3, 6, 10])\n",
      "\n",
      " 0: ('The', [1, 6, 9])\n",
      " 1: ('theater', [8, 10, 2])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('closed', [5, 5, 0])\n",
      " 4: ('today', [3, 10, 8])\n",
      " 5: ('and', [7, 2, 4])\n",
      " 6: ('the', [4, 1, 4])\n",
      " 7: ('sky', [1, 2, 9])\n",
      " 8: ('is', [3, 1, 3])\n",
      " 9: ('grey.', [7, 6, 7])\n",
      "\n",
      " 0: ('His', [0, 10, 3])\n",
      " 1: ('dog', [8, 6, 6])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('brown', [9, 8, 9])\n",
      " 4: ('while', [9, 2, 8])\n",
      " 5: ('hers', [0, 5, 4])\n",
      " 6: ('is', [3, 1, 3])\n",
      " 7: ('grey.', [7, 6, 7])\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for row in [s.split(' ') for s in sents]:\n",
    "    eprint(list(zip(row, (probabilistic_hash_item(word, 11, str) for word in row))))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we show that we can obtain unique representations for >99.9% of words in a vocabulary of 30,000 words with a far smaller embedding matrix. The number of buckets is the number of rows in the embedding matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unique_combos(tups):\n",
    "    return len(set(tuple(sorted(x)) for x in tups))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hash_all_idx(vocab_size, n_buckets, n_hashes):\n",
    "    return [probabilistic_hash_item(i, n_buckets, int, n_hashes) \n",
    "            for i in range(vocab_size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Buckets: 127 \n",
      "Hashes: 5 \n",
      "Unique combos: 29998 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 251 \n",
      "Hashes: 4 \n",
      "Unique combos: 29996 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 997 \n",
      "Hashes: 3 \n",
      "Unique combos: 29997 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 5003 \n",
      "Hashes: 2 \n",
      "Unique combos: 29969 \n",
      "% unique: 0.999\n"
     ]
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "buckets2hashes = {127: 5,\n",
    "                  251: 4,\n",
    "                  997: 3,\n",
    "                  5_003: 2}\n",
    "for b, h in buckets2hashes.items():\n",
    "    tups = hash_all_idx(vocab_size, b,  h)\n",
    "    unique = unique_combos(tups)\n",
    "    print('\\n\\nBuckets:', b, '\\nHashes:', h, '\\nUnique combos:', unique,\n",
    "          '\\n% unique:', round(unique/30_000, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:34:47.133190Z",
     "start_time": "2020-11-20T05:34:47.092417Z"
    }
   },
   "outputs": [],
   "source": [
    "@auto_repr\n",
    "class LazyDataset(Dataset):\n",
    "    \"\"\"Lazily load batches from an enormous dataframe that can't fit into \n",
    "    memory.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, df_path, length, shuffle, chunksize=1_000, \n",
    "                 c=2, classes=('neg', 'pos'), **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        df_path: str\n",
    "            File path of dataframe to load.\n",
    "        length: int\n",
    "            Number of rows of data to use. This is required so that we don't \n",
    "            have to go through the whole file and count the number of lines,\n",
    "            which can be enormous with a big dataset. It also makes it easy to\n",
    "            work with a subset (the data should already be shuffled, so \n",
    "            choosing the top n rows is fine).\n",
    "        shuffle: bool\n",
    "            If True, shuffle the data in each chunk. Note that if batch size\n",
    "            is close to chunk size, this will have minimal effect. If possible,\n",
    "            the training set should therefore load as large a chunk as \n",
    "            possible if we want to shuffle the data. Shuffling is unnecessary \n",
    "            for the validation set.\n",
    "        chunksize: int\n",
    "            Number of rows of df to load at a time. This should usually \n",
    "            be significantly larger than the batch size in order to retain\n",
    "            some randomness in the batches.\n",
    "        c: int\n",
    "            Number of classes. Used if training with FastAI.\n",
    "        classes: iterable\n",
    "            List of tuple of class names. Used if training with FastAI.\n",
    "        kwargs: any\n",
    "            Additional keyword arguments to pass to `read_csv`, eg. \n",
    "            compression='gzip'.\n",
    "        \"\"\"\n",
    "        if length < chunksize:\n",
    "            warnings.warn('Total # of rows < 1 full chunk. LazyDataset may '\n",
    "                          'not be necessary.')\n",
    "\n",
    "        self.length = length\n",
    "        self.shuffle = shuffle\n",
    "        self.chunksize = chunksize\n",
    "        self.df_path = df_path\n",
    "        self.df = None\n",
    "        self.chunk = None\n",
    "        self.chunk_idx = None\n",
    "        self.df_kwargs = kwargs\n",
    "        \n",
    "        # Additional attributes required by FastAI. \n",
    "        # c: Number of classes in model.\n",
    "        self.c = c\n",
    "        self.classes = list(classes)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Because not all indices are loaded at once, we must do shuffling\n",
    "        in the dataset rather than the dataloader (e.g. if the loader randomly\n",
    "        samples index 5000 but we have indices 0-500 loaded, it will be\n",
    "        unavailable).\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        idx: int\n",
    "            Retrieve item i in dataset.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        tuple[np.array]: x array, y array\n",
    "        \"\"\"\n",
    "        # Load next chunk of data if necessary. Must specify nrows, otherwise\n",
    "        # we will chunk through the whole file.\n",
    "        if not self.chunk_idx:\n",
    "            while True:\n",
    "                try:\n",
    "                    self.chunk = self.df.get_chunk()\n",
    "                    break\n",
    "                except (AttributeError, StopIteration):\n",
    "                    self.df = pd.read_csv(self.df_path, engine='python',\n",
    "                                          chunksize=self.chunksize,\n",
    "                                          nrows=len(self),\n",
    "                                          **self.df_kwargs)\n",
    "\n",
    "            self.chunk_idx = self.chunk.index.values\n",
    "            if self.shuffle: np.random.shuffle(self.chunk_idx)\n",
    "            self.chunk_idx = deque(self.chunk_idx)\n",
    "            \n",
    "        *x, y = self.chunk.loc[self.chunk_idx.popleft()].values\n",
    "        return np.array(x), y.astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:34:56.900958Z",
     "start_time": "2020-11-20T05:34:56.863957Z"
    }
   },
   "outputs": [],
   "source": [
    "class ImageMixer:\n",
    "    \"\"\"The transformation that powers MixupDS.\n",
    "\n",
    "    Inspired by the \"Visual-Spatial - Entangled Figures\" task here:\n",
    "    http://www.happy-neuron.com/brain-games/visual-spatial/entangled-figures\n",
    "\n",
    "    The key idea: when a human plays the happy neuron task, it is much easier\n",
    "    when viewing an entanglement of objects they recognize (e.g. bicycle, leaf,\n",
    "    etc.) than of random scribbles. I noticed that for the harder levels, my\n",
    "    strategy was to try to quickly identify a couple distinctive features, then\n",
    "    search for them in the individual images once they appeared. This quick\n",
    "    feature extraction seems very close to what we want to achieve during the\n",
    "    pre-training step.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n=3, a=5, b=8, dist=None, **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n: int\n",
    "            Number of images to use as inputs. With the current implementation,\n",
    "            the constructed image will use exactly 2 of these. The rest will\n",
    "            be negatives (zero weight).\n",
    "        a: int\n",
    "            Parameter in beta distribution.\n",
    "        b: int\n",
    "            Parameter in beta distribution.\n",
    "        dist: torch.distribution\n",
    "            This can be anything with a \"sample()\" method that generates a\n",
    "            random value between 0 and 1. By default, we use a Beta\n",
    "            distribution. If one is passed in, a and b are ignored.\n",
    "        kwargs: any\n",
    "            Makes it easier to use in `get_databunch` function. Extra kwargs\n",
    "            are ignored.\n",
    "        \"\"\"\n",
    "        assert n >= 2, 'n must be >=2 so we can combine images.'\n",
    "\n",
    "        self.dist = dist or torch.distributions.beta.Beta(a, b)\n",
    "        self.n = n\n",
    "\n",
    "    def transform(self, *images):\n",
    "        \"\"\"Create linear combination of images.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        images: torch.tensors\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        tuple(torch.tensor): First item is (n_channel, h, w*self.n), meaning\n",
    "        images are horizontally stacked. The first of these is the new image.\n",
    "        The second item is the rank 1 tensor of weights used to generate the\n",
    "        combination. These will serve as labels in our self-supervised task.\n",
    "        \"\"\"\n",
    "        w = self._generate_weights()\n",
    "        return (self._combine_images(images, w), *images), w\n",
    "\n",
    "    def _generate_weights(self):\n",
    "        \"\"\"\n",
    "        Returns\n",
    "        -------\n",
    "        weights: torch.Tensor\n",
    "            Vector of length self.n. Exactly 2 of these values are\n",
    "            nonzero and they sum to 1. This will be used to compute a linear\n",
    "            combination of a row of images.\n",
    "        \"\"\"\n",
    "        weights = np.zeros(self.n)\n",
    "        p = self.dist.sample()\n",
    "        indices = np.random.choice(self.n, size=2, replace=False)\n",
    "        weights[indices] = p, 1 - p\n",
    "        return torch.tensor(weights, dtype=torch.float)\n",
    "\n",
    "    def _combine_images(self, images, weights):\n",
    "        \"\"\"Create linear combination of multiple images.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        images: torch.Tensor\n",
    "        weights: torch.Tensor\n",
    "            Vector with 1 value for each image. Exactly 2 of these values are\n",
    "            nonzero and they sum to 1. I.e. if we have 3 images a, b, and c,\n",
    "            weights would look something like [0, .3, .7].\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.tensor: Shape (channels, height, width), same as each of the\n",
    "            input images.\n",
    "        \"\"\"\n",
    "        images = torch.stack(images, dim=0)\n",
    "        # 3 new dimensions correspond to (c, h, w), NOT self.n.\n",
    "        return (weights[:, None, None, None] * images).sum(0).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:35:00.758992Z",
     "start_time": "2020-11-20T05:35:00.716514Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def plot_images(images, titles=None, nrows=None, figsize=None, \n",
    "                tight_layout=True, title_colors=None):\n",
    "    nrows = nrows or int(np.ceil(np.sqrt(len(images))))\n",
    "    titles = ifnone(titles, [None] * len(images))\n",
    "    colors = ifnone(title_colors, ['black'] * len(images))\n",
    "    figsize = figsize or (nrows*2, nrows*2)\n",
    "    fig, ax = plt.subplots(nrows, nrows, figsize=figsize)\n",
    "    for axi, img, title, color in zip(ax.flatten(), images, titles, colors):\n",
    "        axi.imshow(img.permute(1, 2, 0))\n",
    "        axi.set_title(title, color=color)\n",
    "        axi.set_axis_off()\n",
    "    if tight_layout: plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-18T04:41:02.901675Z",
     "start_time": "2020-10-18T04:41:02.835076Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class RandomTransform:\n",
    "    \"\"\"Wrap a function to create a data transform that occurs with some\n",
    "    probability p.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, func, p=.5):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        func: function\n",
    "            Transforms an input tensor (x) in some way. X will be the first\n",
    "            argument passed in but the function can accept additional\n",
    "            arguments.\n",
    "        p: float\n",
    "            Between 0 and 1, determines the likelihood that func will be\n",
    "            applied to any input x.\n",
    "        \"\"\"\n",
    "        self.func = func\n",
    "        self.p = p\n",
    "\n",
    "    def __call__(self, x, *args, **kwargs):\n",
    "        \"\"\"Additional args and kwargs are forwarded to self.func.\"\"\"\n",
    "        if np.random.uniform() < self.p: x = self.func(x, *args, **kwargs)\n",
    "        return x\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'{type(self).__name__}({func_name(self.func)}, p={self.p})'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:54:30.765826Z",
     "start_time": "2020-11-20T05:54:30.703787Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class RandomPipeline(BasicPipeline):\n",
    "    \"\"\"Create a pipeline of callables that are applied in sequence, each with\n",
    "    some random probability p (this can be the same or different for each \n",
    "    step). This is useful for on-the-fly data augmentation (think in the\n",
    "    __getitem__ method of a torch Dataset).\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, *transforms, p=.5):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        transforms: callable\n",
    "            Functions or callable classes that accept a single argument (use\n",
    "            functools.partial if necessary). They will be applied in the order\n",
    "            you pass them in.\n",
    "        p: float or Iterable[float]\n",
    "            Probability that each transform will be applied. If a single \n",
    "            float, each transform will have the same probability. If a list,\n",
    "            its length msut match the number of transforms passed in: p[0] \n",
    "            will be assigned to transforms[0], p[1] to transforms[1], and so \n",
    "            on.\n",
    "        \"\"\"\n",
    "        p = tolist(p, transforms, error_message='p must be a float or a list '\n",
    "                   'with one float for each transform.')\n",
    "        if any(n <= 0 or n > 1 for n in p):\n",
    "            raise ValueError('p must be in range (0, 1]. I.E. you can choose '\n",
    "                             'to always apply a transform, but if you never '\n",
    "                             'want to apply it there\\'s no need to include '\n",
    "                             'it in the pipeline.')\n",
    "        \n",
    "        super().__init__(*[RandomTransform(t, p_) \n",
    "                           for t, p_ in zip(transforms, p)])\n",
    "        \n",
    "    @classmethod\n",
    "    def from_dict(cls, t2p):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        t2p: dict[callable, float]\n",
    "            Maps transform to its corresponding probability.\n",
    "            \n",
    "        Examples\n",
    "        --------\n",
    "        transforms = {times_3: .33,\n",
    "                      to_string: 1.0,\n",
    "                      dashed_join: .67,\n",
    "                      to_upper: .95}\n",
    "        pipeline = RandomPipeline.from_dict(transforms)\n",
    "        \"\"\"\n",
    "        return cls(*t2p.keys(), p=t2p.values())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we define a few toy functions to demonstrate how we can alter an input string. In practice, we'd use more useful transformations like the ones defined in `incendio.nlp`, e.g. `ParaphraseTransform`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_upper(t):\n",
    "    return t.upper()\n",
    "\n",
    "def times_3(t):\n",
    "    return t * 3\n",
    "\n",
    "def join(t, sep='---'):\n",
    "    return sep.join(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T05:12:18.232844Z",
     "start_time": "2020-11-19T05:12:18.226752Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomPipeline(\n",
       "\tRandomTransform(to_upper, p=0.5),\n",
       "\tRandomTransform(times_3, p=0.5),\n",
       "\tRandomTransform(join, p=0.5)\n",
       ")"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = 'dog'\n",
    "pipeline = RandomPipeline(to_upper, times_3, join)\n",
    "pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T05:12:18.900985Z",
     "start_time": "2020-11-19T05:12:18.894277Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d---o---g\n",
      "dog\n",
      "D---O---G---D---O---G---D---O---G\n",
      "D---O---G---D---O---G---D---O---G\n",
      "DOGDOGDOG\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    print(pipeline(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T04:59:15.309764Z",
     "start_time": "2020-11-19T04:59:15.304369Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D---O---G---D---O---G---D---O---G\n",
      "D---O---G\n",
      "D---O---G---D---O---G---D---O---G\n",
      "D---O---G\n",
      "D---O---G\n"
     ]
    }
   ],
   "source": [
    "pipeline = RandomPipeline(to_upper, times_3, join, p=[1., .4, 1])\n",
    "\n",
    "for i in range(5):\n",
    "    print(pipeline(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T04:59:44.549340Z",
     "start_time": "2020-11-19T04:59:44.535027Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got ValueError(p must be in range (0, 1]. I.E. you can choose to always apply a transform, but if you never want to apply it there's no need to include it in the pipeline.).\n"
     ]
    }
   ],
   "source": [
    "with assert_raises(ValueError):\n",
    "    pipeline = RandomPipeline(join, p=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T05:00:25.514140Z",
     "start_time": "2020-11-19T05:00:25.507963Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got AssertionError(p must be a float or a list with one float for each transform.).\n"
     ]
    }
   ],
   "source": [
    "with assert_raises(AssertionError):\n",
    "    pipeline = RandomPipeline(to_upper, times_3, join, p=[.2, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-19T04:59:18.046357Z",
     "start_time": "2020-11-19T04:59:18.037502Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D---O---G\n",
      "D---O---G\n",
      "D---O---G---D---O---G---D---O---G\n",
      "DOG\n",
      "D---O---G\n"
     ]
    }
   ],
   "source": [
    "transforms = {times_3: .33,\n",
    "              join: .67,\n",
    "              to_upper: .95}\n",
    "pipeline = RandomPipeline.from_dict(transforms)\n",
    "\n",
    "for i in range(5):\n",
    "    print(pipeline(text))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataLoaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:37:44.160242Z",
     "start_time": "2021-08-07T20:37:44.143931Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def dataloader_subset(dl, batches=1):\n",
    "    \"\"\"Like torch's dataset Subset but for dataloaders. This lets us easily\n",
    "    do things like fit on a single batch (or small number of batches) by\n",
    "    essentially copying an existing dataloader's options but using a smaller\n",
    "    dataset. A few limitations: custom samplers and batch_samplers are not\n",
    "    supported for reasons explained the error message within this function.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    dl: torch.utils.data.DataLoader\n",
    "    batches: int\n",
    "        Number of batches to use in the new dataloader.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    DataLoader: Should be identical to the input dataloader except it operates\n",
    "    one a subset of that dataset.\n",
    "    \n",
    "    Examples\n",
    "    --------\n",
    "    train_dataset = MyDataset()\n",
    "    train_dataloader = DataLoader(train_dataset, shuffle=True, drop_last=)\n",
    "    \"\"\"\n",
    "    ds_sub = torch.utils.data.Subset(dl.dataset, \n",
    "                                     range(dl.batch_size * batches))\n",
    "    \n",
    "    # Check with type rather than isinstance since subclasses of these types\n",
    "    # should raise an error.\n",
    "    custom_sampler = False\n",
    "    if type(dl.sampler) == torch.utils.data.sampler.SequentialSampler:\n",
    "        shuffle = False\n",
    "    elif type(dl.sampler) == torch.utils.data.sampler.RandomSampler:\n",
    "        shuffle = True\n",
    "    else:\n",
    "        custom_sampler = True\n",
    "        \n",
    "    if custom_sampler or (type(dl.batch_sampler) \n",
    "                          != torch.utils.data.sampler.BatchSampler):\n",
    "        raise RuntimeError(\n",
    "            'dataloader_subset does not support custom samplers or '\n",
    "            'batch_samplers. Sampler constructors can require non-standard '\n",
    "            'args so we can\\'t reliably create a new instance automatically, '\n",
    "            'and they often require a dataset as an argument so we can\\'t '\n",
    "            'simply copy the existing instance - that uses the full dataset '\n",
    "            'whereas we want a subset.'\n",
    "        )\n",
    "        \n",
    "    # Exclude shuffle since this isn't stored as a dataloader attribute.\n",
    "    # It can be inferred from samplers if the user specified a shuffle value\n",
    "    # rather than passing in non-default sampler objects.\n",
    "    defaults = {'batch_size', 'num_workers', 'collate_fn', 'pin_memory',\n",
    "                'drop_last', 'timeout', 'worker_init_fn', 'prefetch_factor',\n",
    "                'persistent_workers'}\n",
    "    kwargs = {attr: getattr(dl, attr) for attr in defaults}\n",
    "    return DataLoader(ds_sub, shuffle=shuffle, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## File Handling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class BotoUploader:\n",
    "    \"\"\"Uploads files to S3. Built as a public alternative to Accio. Note to \n",
    "    self: the interfaces are not identical so be careful to know which you're\n",
    "    using.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, bucket, verbose=True):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        bucket: str\n",
    "            Name of s3 bucket to upload to. For a single project, I generally\n",
    "            stick to a single bucket so we can usually keep this fixed. We can\n",
    "            always change the attribute later if necessary.\n",
    "        verbose: bool\n",
    "            If True, print message when downloading each file.\n",
    "        \"\"\"\n",
    "        self.s3 = boto3.resource('s3')\n",
    "        self.bucket = bucket\n",
    "        self.verbose = verbose\n",
    "        \n",
    "    def upload_file(self, path, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Upload a single local file. By default, its path in S3 will be the \n",
    "        same as its local path. Usually, this means your S3 bucket will have a \n",
    "        single directory called \"data\" which corresponds exactly to your local \n",
    "        \"data\" directory.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str or Path\n",
    "            Local path to the file to upload. This must be relative to the \n",
    "            project root directory (e.g. \"data/models/v1/model.gz\").\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}.\n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. All four combinations of retain_tree \n",
    "            (True/False) and s3_dir (empty/non-empty) are supported.\n",
    "        \"\"\"\n",
    "        path = str(path)\n",
    "        s3_path = self._convert_local_path(path, s3_dir, retain_tree)\n",
    "        if self.verbose: print(f'Uploading {path} -> {s3_path}.')\n",
    "        self.s3.meta.client.upload_file(path, self.bucket, s3_path)\n",
    "        \n",
    "    def upload_files(self, paths, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Upload multiple files. Currently does not support parallelized\n",
    "        uploads: need to figure out a workaround because self.s3 cannot be\n",
    "        pickled.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        paths: Iterable[str or Path]\n",
    "            Sequence of file paths to upload.\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}.\n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. All four combinations of retain_tree \n",
    "            (True/False) and s3_dir (empty/non-empty) are supported.\n",
    "        \"\"\"\n",
    "        for p in paths:\n",
    "            self.upload_file(p, s3_dir=s3_dir, retain_tree=retain_tree)\n",
    "            \n",
    "    def upload_folder(self, dirname, s3_dir, retain_tree=True, recurse=True,\n",
    "                      keep_fn=None):\n",
    "        \"\"\"Upload all files in a directory.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        dirname: str or Path\n",
    "            Directory to upload.\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}. \n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. When uploading recursively, we require\n",
    "            that the file tree is retained.\n",
    "        recurse: bool\n",
    "            If True, upload all files in subdirectories as well.\n",
    "        keep_fn: None or callable\n",
    "            If provided, this should be a function that accepts a filename as\n",
    "            input and returns a boolean specifying whether to include it in the\n",
    "            upload or not. For example:\n",
    "            \n",
    "            lambda x: os.path.splitext(x)[-1] != '.pkl'\n",
    "            \n",
    "            keeps all files except those with an '.pkl' extension. \n",
    "        \"\"\"\n",
    "        if recurse and not retain_tree:\n",
    "            raise ValueError('retain_tree must be True when uploading '\n",
    "                             'recursively.')\n",
    "        pat = os.path.join(str(dirname), '**' if recurse else '*')\n",
    "        # glob's recursive option only has an effect when using '**'.\n",
    "        paths = (o for o in glob(pat, recursive=True) if os.path.isfile(o))\n",
    "        if keep_fn: paths = filter(keep_fn, paths)\n",
    "        self.upload_files(paths, s3_dir, retain_tree)\n",
    "            \n",
    "    def _convert_local_path(self, path, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Convert local path to s3 path. See public methods for parameter\n",
    "        documentation.\n",
    "        \"\"\"\n",
    "        path = path if retain_tree else os.path.basename(path)\n",
    "        return os.path.join(s3_dir, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "up = BotoUploader('gg-datascience')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No S3 prefix, Yes retain file tree:\n",
      "data/v1/history.csv\n",
      "\n",
      "Yes S3 prefix, Yes retain file tree:\n",
      "hmamin/data/v1/history.csv\n",
      "\n",
      "Yes S3 prefix, No retain file tree:\n",
      "hmamin/history.csv\n",
      "\n",
      "No S3 prefix, No retain file tree:\n",
      "history.csv\n"
     ]
    }
   ],
   "source": [
    "ft = up._convert_local_path('data/v1/history.csv')\n",
    "tt = up._convert_local_path('data/v1/history.csv', 'hmamin')\n",
    "tf = up._convert_local_path('data/v1/history.csv', 'hmamin', retain_tree=False)\n",
    "ff = up._convert_local_path('data/v1/history.csv', retain_tree=False)\n",
    "\n",
    "print('No S3 prefix, Yes retain file tree:\\n' + ft)\n",
    "print('\\nYes S3 prefix, Yes retain file tree:\\n' + tt)\n",
    "print('\\nYes S3 prefix, No retain file tree:\\n' + tf)\n",
    "print('\\nNo S3 prefix, No retain file tree:\\n' + ff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:38:16.260214Z",
     "start_time": "2021-08-07T20:38:16.221250Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def plot_images(images, titles=None, nrows=None, figsize=None, \n",
    "                tight_layout=True, title_colors=None):\n",
    "    \"\"\"Plot a grid of images.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    images: Iterable\n",
    "        List of tensors/arrays to plot.\n",
    "    titles: Iterable[str] or None\n",
    "        Title for each subplot. Must be same length and order as `images`.\n",
    "    nrows: int or None\n",
    "        If provided, this manually sets the number of rows in the grid.\n",
    "    figsize: tuple[int] or None\n",
    "        Determines size of plot. By default, we double the number of rows and\n",
    "        columns, respectively, to get dimensions.\n",
    "    tight_layout: bool\n",
    "        Often helps matplotlib formatting when we have many images.\n",
    "    title_colors: Iterable[str] or None\n",
    "        If provided, this should have the same length as `titles` and \n",
    "        `images`. It will be used to determine the color of each title (see\n",
    "        PredictionExaminer in incendio.core for an example).\n",
    "    \"\"\"\n",
    "    nrows = nrows or int(np.ceil(np.sqrt(len(images))))\n",
    "    titles = ifnone(titles, [None] * len(images))\n",
    "    colors = ifnone(title_colors, ['black'] * len(images))\n",
    "    figsize = figsize or (nrows*2, nrows*2)\n",
    "    fig, ax = plt.subplots(nrows, nrows, figsize=figsize)\n",
    "    for axi, img, title, color in zip(ax.flatten(), images, titles, colors):\n",
    "        axi.imshow(np.transpose(img, (1, 2, 0)))\n",
    "        axi.set_title(title, color=color)\n",
    "        axi.set_axis_off()\n",
    "    if tight_layout: plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we demonstrate plotting a list of images stored as numpy arrays in a grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:51:51.542963Z",
     "start_time": "2020-11-20T05:51:50.981840Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38608 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "images = [np.clip(5 * np.random.uniform(size=(3, 16, 16)) / i, 0, 1)\n",
    "          for i in range(1, 17)]\n",
    "plot_images(images, titles=[f'Image {i}' for i in range(16)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a similar example showcasing the title color functionality. Notice we can use tensors here too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-20T05:52:38.217337Z",
     "start_time": "2020-11-20T05:52:37.586483Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39808 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "images = [torch.clamp(5 * torch.rand(3, 16, 16) / i, 0, 1)\n",
    "          for i in range(1, 17)]\n",
    "plot_images(images, titles=[f'Image {i}' for i in range(16)], \n",
    "            title_colors=['red' if i%2 == 0 else 'green' for i in range(16)])"
   ]
  }
 ],
 "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.7.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "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": 4
}
