notebooks/domino_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random, os\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### User Variables/Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "root_dir = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/' # path to save toy data files (mnist, svhn, fmnist, kmnist, etc)\n",
    "img_size = 32\n",
    "trans1 = torchvision.transforms.Compose([  \n",
    "        torchvision.transforms.Resize((img_size,img_size)),\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Lambda(lambda x: x.repeat(3,1,1))\n",
    "    ])  # for mnist and fmnist we replicate channels\n",
    "trans2 = torchvision.transforms.Compose([\n",
    "        torchvision.transforms.Resize((img_size,img_size)),\n",
    "        torchvision.transforms.ToTensor()\n",
    "    ])  # for cifar10 we don't replicate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Setup Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Using downloaded and verified file: /jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/train_32x32.mat\n",
      "Files already downloaded and verified\n",
      "Using downloaded and verified file: /jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/test_32x32.mat\n"
     ]
    }
   ],
   "source": [
    "cifar10_train = torchvision.datasets.CIFAR10(root_dir, transform=trans2, download=True)\n",
    "fmnist_train = torchvision.datasets.FashionMNIST(root_dir, transform=trans1, download=True)\n",
    "mnist_train = torchvision.datasets.MNIST(root_dir, transform=trans1, download=True)\n",
    "kmnist_train = torchvision.datasets.KMNIST(root_dir, transform=trans1, download=True)\n",
    "svhn_train = torchvision.datasets.SVHN(root_dir, split='train', transform=trans2, download=True)\n",
    "\n",
    "cifar10_test = torchvision.datasets.CIFAR10(root_dir, train=False, transform=trans2, download=True)\n",
    "fmnist_test = torchvision.datasets.FashionMNIST(root_dir, train=False, transform=trans1, download=True)\n",
    "mnist_test = torchvision.datasets.MNIST(root_dir, train=False, transform=trans1, download=True)\n",
    "kmnist_test = torchvision.datasets.KMNIST(root_dir, train=False, transform=trans1, download=True)\n",
    "svhn_test = torchvision.datasets.SVHN(root_dir, split='test', transform=trans2, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cifar10_train[0][0].shape,fmnist_train[0][0].shape,mnist_train[0][0].shape,kmnist_train[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]),\n",
       " torch.Size([3, 32, 32]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cifar10_test[0][0].shape,fmnist_test[0][0].shape,mnist_test[0][0].shape,kmnist_test[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14faa3b75370>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAI4AAAD7CAYAAAC8Eqx6AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAd+ElEQVR4nO2de4zU13XHv2ee+4Rld3ksLBiwAWOMwQ51cLANxOBix2leSpVIteIqEomUVrbqqs1DitRKkVxZihpVSiWnceuqSVMrdhLbcdI4EXZtB2PAmJcB837DLrDvnffc/jHD755zvY/ZuzDs7JyPhDgz985vfjN75p577j3nXDLGQFHGSuhG34BSmajiKF6o4iheqOIoXqjiKF6o4ihejEtxiGgTER0ioiNE9M1rdVPKxId813GIKAzgQwAbAZwBsB3Al40xH1y721MmKpFxvPZuAEeMMccAgIh+BuAzAIZVnFAsZkJ1tcVHNOyF3RZiz4SYntfG46Jf05Sptq22xrmIvUYunw/kgcEB0W2gv8/2y2VFW57JIXaT5Nywgb3J/Ag/TOO8cLiew39TH31vYk9w2b04HzDcwYM/Hrzcc8kYM9193/EozhwAp9njMwA+PtILQnW1aLzvXgDOh3JvKh8Wj2tgH8cy9kMtX7hI9PvMnz4UyLcvWyrawrFoIPcmBgN5246tot+2rVsCuau7S7QlQrlArg/Zry7i6Gg6nwnkZDYj2vJM43IhOVPIYWg+qhz2O4g414hGo0PK5ChONmXvK5VKybaMbdv5n78+OdQ9jUdxhvrLf+RHQ0SbAWwGAHJHAaViGY/inAEwlz1uB3DO7WSMeQbAMwAQaZpq2PNOPytnnZ9HihkJw35ge44cEv2u9NgRYsO6taLt3ns+EcjNTdMCedO9nxT92ltaA/mFP/xatCU6zwZyLmvHh3xajhUUtSNkhP3qASDNzF/UGUqibDTKM3P6kZ/jCKN1Nmuvz7/jWEzeB8L2nsMx2ZQbyTYWGY9XtR3AIiJaQEQxAF8C8NI4rqdUEN4jjjEmS0R/BeB/AYQBPGuM2X/N7kyZ0IzHVMEY8yqAV6/RvSgVxLgUZzyMtH6UHcGAGu6Px2TH01c6AvmFV18WbcdPWufgkQ0PBvLSm24W/VYsuyuQI23Nou3t39nfyLH9ewO5L5eW98jmaOG4nFvEI3b+Y3LO3Eg8sJ8tb/KiX5a70iM463y+kzdyaYHPcUIxeY1IRHq1Q6FbDooXqjiKFzfMVLnwBUHjjL7CPQyzlVHXpWeLYd2ZhGh7a8e2QO44fz6QP7X2QdHvT9asCeRlC24Rbcsf/mwgv8SW61478L7ol2H+s8lLMxOO2q88FHJWjvlKNWvKOZ8zy0yXc3nEYta35i/LGbkQmckw80quqRpdLXTEUbxQxVG8UMVRvCjzHIeCuUzI2Zzju8g0govJ5z/upmDe2GfcbYs8e7v9J48F8oXnfyr6HTp/JpAfWn2naLvpQmcgxy5fDuSIM1fhG5nOFETuSpNspQib57HJSyTsuMdkXfx0fvhljXDYfuhcztlQNfaaqaTc5IxERw+10RFH8UIVR/Gi7O74VVPlxuNwDf7IQMme4K9y3XY+8OccI8FNVbTGDtOdeem2v/LGa4HcvfUt0bZh0JrCvoZe2zDVCRcJjfC18kCriPu7tdcnbsqdzxkNW1MVcd5rcNDGGvGVY3fHO8f/9M7tJtPOKvMQ6IijeKGKo3hRVlNFBERCBTORzTsbfGxoDjk2aPj9UKeBh9iS47WxhxkeNBaW70XMm0mdPSvaoilr4sKzmXfUKGOf83zzkhyPSHhV8v6JeUEykMvxEPPWlITD8k8Yi9uVY262co73lQf34OR3ZYxucirXCVUcxQtVHMWL8s5x8oRYouiOx+VbJzHCvIC51hRmOUuQO77c5wznZQAVn8nwXfWsMxcKsaCsOUa66iZSF8i9aArkiJHueCZn7ytH0rUNhfjnlFHiWTYPCbPVYsrKa/Drp5x5Et8dj4CtMA/KzxLOsrlWTi5dhMzo44mOOIoXqjiKF2U1VSFjUFscIvPOFmUsztNm5dDJvXPDg5hCcpgWG6fO7qIR8vD5uzXs+tOcKKlMmKUO83ypvLwGD9bKhKQ55THTkZy8f26eDMumdPZQEY7az5l00pSzSZa3xe4xXht3+tnruxuxaqqU64YqjuKFKo7iRVnnOCabRerSJQBAXXO9aON21nGykWW2P8vsfd4N9mZzl/zwsWCizZCca8XTNqipzokUy8bsnCQR54Fncq7C5x2DTl52bb116dN9SdFWH2PLECEru/NBEfTmbEek0+zbY252KCSXOCJsayIXdiIJstcgkIuIniWiDiLax55rJqLXiOhw8f9pI11DmXyUYqr+A8Am57lvAviDMWYRgD8UHytVxKimyhjzf0Q033n6MwDWFeXnALwO4O9Hu1ZDbR3W3LECALBz33uiLcTSecM1jaItX2Pb0iw9Nee4s7lhAr4A6Y5zU0WOM1rPzEyd45b2s28rVc92yp2Y4IYWe//zl8sU487uK4E88KGsWZQZtGYyGmfXjLm/b3vPYWflmxda4rv0mbQzAYhYV904K/UhN8Z5CHwnxzONMecBoPj/DM/rKBXKdfeqiGgzEe0goh3JVGL0FygVga9XdZGI2owx54moDUDHcB15Ra7maa0m21co1viJ5StFv3d3vBvIgwm5GhpvbbHXY+WjYmEnAIkNuVln1dcMU+0q7KSoNLCHOed31VNrv64sq0LhWEz05O1G6dy2maKtqdF6VfF++TkHTrGKXxnblnZX0vkGsRuglR26kmDeqYyRZtd3TVUO189UvQTgK0X5KwB+5XkdpUIpxR3/bwBbASwhojNE9FUATwHYSESHUahz/NT1vU1lolGKV/XlYZoeuMb3olQQZV05TqfTOH2qUJi0tXmKaFtyky0pcvTSRdHWccE+rp1uq4LGa+SObyhmP07CqZLFLXyWOedOLU5EWcd+JxD8Eiu3m2P5TOG4nBOc7bVTvsEDsiziwxsfDuRTXdJZ6LtgXxdjKbupvEzRTafsZ3OXHaLsngdZ8e+8Mxci5sZnnUCujFu9awh0r0rxQhVH8aLMKcAhIFJwpwf6pSmpYUWlFy5cINqaU3YzsLPbpt6anDQ0yYRdHa2JOS4mW1Hl+4KUlfcRZscBDLjBT22zAjk1aIfzfjdYq96+7rJzVsTWnTttv15Z8t/U2KUGStr7qonKmOY8M13plNwo5UvkMZYqnMo4FSn4Rq+znmDcyK4h0BFH8UIVR/FCFUfxoqxznEwuh44rPQCAGVOnirZoxNr3M2flWSKzFtxk5Vn23JGM4zUeOHYkkJPOwRz8jKooC/aOOMvrxHbEkw0y2Gz5ffcH8tQztiLXG6d3iX5plo8edoK8evrtnGdGiwxjqo/ZuVH/4RP2es6xQCF2/VhUzsPSadn3Km5Vr3SCze2cKIBQCeOJjjiKF6o4ihdlNVXxmlrcfOsdAICTx4+Ltq6EHcLTTpWsZOpEIN+1YkUgtzY0iH67L14IZOOcT1DfZFeqDTOLyDvuOMt76oE0A+d32d3rhavusO8bkS53/lQQZYuIs4LddsvsQI6RvL7psddpYMlk6f3yXK4oO+wi55RpCbFcqgQrEk5OJAGvSpZOyntEdnR/XEccxQtVHMWL8pqqWBzz5hdicKdMaxVtJ0+cCOTOizIWt+tSdyDvfs96ME2N0utJ9vQEcsYJ0OrrsddommODq+rrpLmIRq0Z63Mqg4XP2njhw4k3A/lEtEf0a663ZrH3Sr9o69hrayx//J51oi0TshW0LpyzNZV7rnSLfi2NbIOY5J+Qxz/HmElOGWmORJq1GxA3TDAYR0ccxQtVHMULVRzFi/LujhMFZyE1T5OrpnW1tYHc2izd7DOnreve19cdyP29faJffR2b80Tk/KQ3aecPZw7beUbTNLmCHU+woPaYnJ/cPtPujg+ySluD/XKO09QyPZCb6+X1D720PZDDx+S84/aP3RbIx3dal37w8iV5/VqWHuzkVRkT4Q/seznlXAwrzk2OS+9+d0OhI47ihSqO4kXZz3IY7vTf+nprZuoXyiMNa1hs8Um2kXmp47zoN3vunEAeSEgzxsNqIyw3K+rEFc9qbwvk1vkyfbd2qr3HTKM1VVkn8ulEl13BrknIchWNaXuNjp1HRds7p+zmbiJpN1EXtM0W/erZqu+AE4/M44ejNPyflx/XGIvJgLj8CMc+XUVHHMULVRzFC1UcxYsbdl6VO9fhjyMxOS+4ZfGtgVxfZ3Ovd++SQeILFy0O5A5n/nN5v81vCouqo/I+orX2+rcsXiTa+Nxif+fBQM7HndzuBnv/SeeAjfAUO8eJ9Up3PNll52Vx9l4RI/9M8ZCdk5DjOqcy7HhJ9tEizrZCiM3tEin5Pbp9h6KUFOC5RLSFiA4Q0X4ierz4vFblqmJKMVVZAE8aY5YCWA3gG0R0G7QqV1VTSu74eQBXiyj1EdEBAHPgUZUrFAqhrmhqkkmZD5RjZTicihwYZPlS8xcwV92JNzp9/lQgN0yRO+cpFqC8mLnZrbPkLv25w9bdp21/FG1rl1mTOb3XripPDcng5/aW9kA+m+wWbfULmgM5eVTGVt88z+aTUd5+PzOdzzLYZ1eSa+qkWY9G7b2kWBBZMivddv4qd/XZNd9DMabJcbGk250AtkGrclU1JSsOETUAeAHAE8aY3tH6s9cFFbkSg32jv0CpCEpSHCKKoqA0PzHGvFh8+mKxGhdGqspljHnGGLPKGLOqtq5xqC5KBTLqHIcK/vOPARwwxnyfNV2tyvUUSqzKFQoR4vH4kG2ptLXH3KUEgDBzDy932cGOu+kAEGK7xm9tfUO0cZezocEq8JQpcvf6WMLOXfo75Bykn0UVttdaeWHOyWXK27Is4Yj8intq7WerXSgd0Y6sDVZvjdt88Y/ffb/o13/JznHe379DtPEUqdooW+IIye80xY6I/shR3iW446Ws46wB8CiAvUT0fvG5b6OgMM8XK3SdAvDFEq6lTBJK8arewkfr91xFq3JVKWVdOQ6HI2huLrij/f1yomx4UFZI7tammRkLszMJurpkAFVbu3WDNzwki8Fvf+PtQO4bsCYhc04eET2QsKYw4VSm6h+wwWANzAWf6fysmlLWHHfPlKaws9264zvPnRBtIba+0BiygW0HD8l+H1tsA77W3yeXE95+9/eB3Ju0u/RR50yJCM8tc3KpyT0gawh0r0rxQhVH8aKspiqfz6O/eBpt8/QW0RZhm3pXnDwinr6aY0cOZp0RdbDfrrZOjcvrb3zgkUDe9pZdET5xQuZwdfdZc9Q+TV6j4TbrxR08uMdeg+V9AUA7+z0ucYK8piyaF8hnbp4n2rInbIWuWQm72ZpxKlBs27c3kFcsuU20rb/3IdvvPetZXrwkP2eUeXdOtjTyec2rUq4TqjiKF6o4ihflneMYg0SxutTFTrlDMX2GXUWtb6gTbR2dNo96cMCuhuacbfQQqxiaHZBt/Cjle9asC+R9e+W5WYeztjTIZaf8R8uqOwN562kbaH6hV85Bkt3W3W93dqVb77DBYa1pOf+Js2opt9Xa+RXVyjyzLrY8sX/fAdF2C5tDffI+uySx98B20W/fsXfZI7kbHovI6l1DoSOO4oUqjuJFmWOODUzRPU2n5WplZ4fduGttlauh7XPsivC5czaWOJGQlbsMKxgZds5g4jHNhp1rsPLOj4l+LS3WLOzb9a5o2/WBDfLqYXFobQvuEP3mt9jV4a7tvxdtV163K9hznFTnmY32c85tseY6FZdRBRF24EQiJAPiTp2ywWx52LYVd6wS/RpbbKmUne9LMzaYHsRo6IijeKGKo3ihiqN4cQPyqgJJPJ9hwVsXL3aKtilTrD2eNcvmdl+4cEH0y7NgpVzSOS6ZEeZHUDvnOM2YMz+Qb3bc8d+/tdXeL5tfLZvdJPrVt9g5WibiXL/DfrbZTkx4vMbO+8JR9v2E5SEg9UNXMgEA8CpsZ07b+WDKOT761tvt0kIsJnfw3972OkZDRxzFC1UcxYsbUOak4I6Tk8vDq2A6Jz+jn+1Y86DaqVObRb9Ejw0Gy+SlKeTBYNmsNQmhqOO2h2yA04z2+aJtMYtbPrjP7o7nUtJ9vXzRxio3QK4cz2QFuec7SwY9rDJWL9uyTrhFsNmxjg1R6dInU7avMdZuXemUgXP7d9uV71uWzhdt96+xgZ3P/8uzGAodcRQvVHEUL8q+coxiikkoJINgDTv1NucUaOYhsL0s9baxUa6oTp3aFMiDJIfmEDtakZuqrPNehpmuukbpbSxYvDSQo+yIoEvHDop+veftbuXMhDRjibj93BecMxPS9dbsxGfbQpX9fXKVPcrSdNwNybixHhj3uDIZ6VV1X7Gx1XvflycVu6ZrKHTEUbxQxVG8UMVRvCjrHCebzeJSMX21edp00RZlRyS7rjq31Ya57f39cv6Qi1hXt5EV3Hbh9j6Uk/OHFHOX886xyvEau2M9d8HCQJ5aI+cZR/bYXfTDaZnyHGdHObbk5NffyM7HmsW8+JltbaJf32V7GEnGOT6aH2JSy9z9EMl+YebiJ9Oybf8eOWcbilIqctUQ0btEtLtYkesfis9rRa4qphRTlQLwSWPMCgArAWwiotXQilxVTSm54wbAVR84Wvxn4FGRKxKJYsaMgom6clmW2JnSaG8lHpeueo7l+fBVUzfmOMNsWq/jftbUWFOYZ0vTxinrZZiZTDnXMCz9mBe4nOKsMC+vt1XDjs6TJvn48cOBfCEmNy+nD9p7SZ20FS/mOAlkc9h5W11d3aItkbA2LspOCOYx1wCQYZUlw87JioNp54khKLU+TrhYqaIDwGvGGK3IVeWUpDjGmJwxZiWAdgB3E9Htpb6BqMg1oBW5JgtjcseNMd0omKRN8KnIVa8VuSYLpVTkmg4gY4zpJqJaABsA/BM8KnLFwmHMLh7j3Fwj3eWTLAg9l5cK1siW/rkr7RSSQo7naTsRToPMbQ2z7YeQUzGLFbFCKCznBUlh++2bxyLys9Q228eLp9wl2rpb7UElWWcuEYnY9+vtsUsNYVwU/QxbJrhpwXx5/R4b9N/Xa7c+ws5hJ7GorWSazztqQPJzD0Up6zhtAJ4jojAKI9TzxphXiGgrtCJX1VKKV7UHhRK17vOXoRW5qpayrhzPmNGEx7/x+XFdwy10OF4+UjgxNPy0r5atRt93332B/J3vfEf0O8GOwv7BD34k2t57z6Ycu8sJlYTuVSleqOIoXpQ95niiwY90BIC1a9cG8ty5c0XbH/9oK3l97WtfC+Tly5eLfrfdZqtkHTwoNww7OuyqxcmTskpWJaEjjuKFKo7ihSqO4kXVz3FSKZn3xAPgH3nkEdH20EO2oueaNWuGfA0gXXz3GufP2xXyV155RbRdvChXiCcyOuIoXqjiKF5Uvaly840usSN93FXk9evXBzJfRR5pNXvZsmXi8WOPPRbI7hFML774YiC7lTgmGjriKF6o4iheqOIoXlT9HMeFbwPs3btXtG3cuDGQS92lb2iQZUhWr14dyHw+BQC7du0KZJ3jKJMSVRzFi7Kaqp07d17zQKxrzblztprW9u2ycDQ3YwsX2hTgsXwmHryVd0uPVRA64iheqOIoXqhX5TDATgjevXu3aPvlL38ZyJs3bw5k13PiwVruJir3lt544w3Rtm/fvrHf8A1CRxzFC1UcxQtVHMULneOMwJEjR8TjH/3I5kjxvKq77pJpvj/84Q8D+ejRo6Lt9OnTgXzo0CHR1tdXOUUZSh5xiqVOdhHRK8XHWpGrihmLqXocAD85VCtyVTElmSoiagfwKQDfA/A3xafHXJGr0nBXhHngVXu7PQbRTeXdsmVLIO/YsUO08cCxalg5/mcAfweIumdakauKKaXq6CMAOowxO33egFfk8nm9MjEpxVStAfBnRPQwgBoAU4jov1CsyGWMOT9aRS4AzwAAEZmh+iiVRyn1cb4F4FsAQETrAPytMeYviOhpjLEiV6XhBpPPnGmrffLjHt2tAh6glUzK4tPGPQuxQhnPAuBTADYS0WEAG4uPlSphTAuAxpjXUfCetCJXlaMrxyMwe/Zs8XjDhg1D9vvVr6SVvnLFnrUwWUyTi+5VKV6o4iheqKlyqKuzRwvxylqANFXcc/rNb34j+vX09Fynu5s46IijeKGKo3ihiqN4UfVznFrnCMZVq1YF8uc/L4t58x1xXoH08OHDop8boD4Z0RFH8UIVR/GiKk0V37xcuXKlaHv00UcD2TVVPAiLu+NNTU2iH8/NymblKcOTBR1xFC9UcRQvVHEUL6pyjsPnJF/4whdEG5/X8GAtAEizoxB5XtXnPvc50e+FF14IZF42BajsM6o4OuIoXqjiKF5Upani+VLNzc2izXWtObFYLJCXLFkSyE8//bTo193dHcgvv/yyaLt8+fJYbnXCoiOO4oUqjuJFVZqqrq6uQO7s7BRtPJ2FB3WNRMQ59P6ee+4J5HfeeUe0qalSqhpVHMULVRzFi6qc4/Ada+46A8D+/fsD2a3IxUuW8OMUn3zySdFv6dKlgTxt2uSsN1VqfZwTAPoA5ABkjTGriKgZwP8AmA/gBIA/N8Z0DXcNZXIxFlO13hiz0hhzNbZSK3JVMeMxVRVbkYtXwuJFrwHpPruuM3fjefzx17/+ddGPV7VwY5onC6WOOAbA74hoJxFdLSmuFbmqmFJHnDXGmHNENAPAa0R0sNQ3KCra5lE7KhVFSSOOMeZc8f8OAL8AcDeKFbkAYLSKXMaYVWxupEwCRh1xiKgeQMgY01eUHwTwjwBeQoVW5OKlRz788EPRxh+7QVc1NTWBPGvWrEB2txx4wNdkCdxyKcVUzQTwi2IoQgTAT40xvyWi7QCeJ6KvAjgF4IvX7zaViUYpNQCPAVgxxPNakauKqcqVY85YTEk0Gg1kXq3LLTLJU4Ina8kT3atSvFDFUbxQxVG8qPo5zljgc5lFixYFsuuO83M3w+GwaOMuvVs8u5LQEUfxQhVH8UJN1RjgO93Lly8P5FBI/v54INcDD8ilLh5E5h6tWEmmS0ccxQtVHMULNVVjgHtEixcvHrbf2rVrA3n16tWibeHChYHMTwsGgD179oz3FsuGjjiKF6o4iheqOIoXOscZAzzInRfB5tVIAbla7O6cf/rTnw7k48ePizae0zXRA8B0xFG8UMVRvFBTNQbOnj0byI8//nggr1+/XvTbtGlTIPOzIQBZAcw9urG+vj6Qe3t7x3ez1xkdcRQvVHEUL1RxFC90jjMGuAvOXWe3HNy2bdsC+bvf/a5oW7duXSDzUikA0NDQEMg6x1EmJao4ihdqqsYATx3mAVnueQ3z5s0LZDfIi6cHJxIJ0VZJRzKWNOIQURMR/ZyIDhLRASK6h4iaieg1Ijpc/H9y1ixThqRUU/UDAL81xtyKQjrwAWhFrqqmlGoVUwDcD+AxADDGpAGkiahiK3Jda9z0mAcffDCQudkCZJzxrl27RFtfX991uLvrQykjzkIAnQD+nYh2EdG/FcudaEWuKqYUxYkAuAvAvxpj7gQwgDGYJSLaTEQ7iGjH6L2VSqEUxTkD4Iwx5uqq1s9RUCStyFXFlFIf5wIRnSaiJcaYQyjUxPmg+K8iK3JdC7ib3dbWJtr4jjh3vwF5ftWWLVtEm9t3IlPqOs5fA/gJEcUAHAPwlyiMVlqRq0opSXGMMe8DGMrUaEWuKkVXjscAP5KRpwNz9xuQLvibb74p2ripOnr06LW+xbKhe1WKF6o4iheqOIoXOscZA3xrgR8C8sQTT4h+vDrpq6++Ktp4AFgloyOO4oUqjuIF8eCk6/5mRJ0ATgJoBXCpbG888ZnI38dNxpjp7pNlVZzgTYl26N6VpRK/DzVViheqOIoXN0pxnrlB7ztRqbjv44bMcZTKR02V4kVZFYeINhHRISI6QkRVlxVBRHOJaEsxxWg/ET1efL7iUo3KZqqIKAzgQwAbUQhH3Q7gy8aYD8pyAxOAYohtmzHmPSJqBLATwGdRyCC5Yox5qviDmmaMmdAZI+Ucce4GcMQYc6yYYvMzFA69rxqMMeeNMe8V5T4U8tPmoPA9PFfs9hwKyjShKafizAFwmj0+U3yuKiGi+QDuBLANFZhqVE7FoSGeq0qXjogaALwA4AljzMSuZzIM5VScMwDmssftAM4N03fSQkRRFJTmJ8aYF4tPl5RqNJEop+JsB7CIiBYUsyW+hMKh91UDFYKWfwzggDHm+6zpJRRSjIAKSTUq9+74wwD+GUAYwLPGmO+V7c0nAER0L4A3AewFcLXa9rdRmOc8D2AeiqlGxpgrN+QmS0RXjhUvdOVY8UIVR/FCFUfxQhVH8UIVR/FCFUfxQhVH8UIVR/Hi/wHW9RQbkkobegAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = kmnist_train[50][0]\n",
    "img[:,:5,:5] = 1.0\n",
    "test_img = torch.hstack((cifar10_train[50][0],img))\n",
    "plt.imshow(test_img.permute(1,2,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[100., 100., 100.],\n",
       "         [100., 100., 100.],\n",
       "         [100., 100., 100.]],\n",
       "\n",
       "        [[100., 100., 100.],\n",
       "         [100., 100., 100.],\n",
       "         [100., 100., 100.]],\n",
       "\n",
       "        [[100., 100., 100.],\n",
       "         [100., 100., 100.],\n",
       "         [100., 100., 100.]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kmnist_train[50][0][:,:3,:3]=1.0\n",
    "img = kmnist_train[50][0][:,:3,:3]\n",
    "img[:,:3,:3]=100\n",
    "img[:,:3,:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14faa3a99a60>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAI4AAAD7CAYAAAC8Eqx6AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAZdUlEQVR4nO2da4xV13XH/+ue+5g3MDM8BgYbsMFvYzvUieNHnMROsRPV+ZIqkSqlVSR/aatEadUm/dZKkSJVilL1QyWrceuqSVMrblQ3jZI4UezGLsWAwQYMBIx5DK8ZmPfMfd/VD3N89lrbF+aygcsMd/2k0axz9r7ncWfNXnvtvfbaxMwwjMsldb0fwFicmOIYQZjiGEGY4hhBmOIYQZjiGEFckeIQ0VYiOkRER4joG1froYyFD4WO4xBRBOC3AJ4EMARgB4AvMfO7V+/xjIVK+go++yCAI8x8FACI6IcAngFwUcUhIhttXHycZ+bl/skrMVVrAJwUx0PxOePG4ni9k1fS4lCdcx9qUYjoWQDPXsF9jAXIlSjOEIC14ngQwGm/EjM/B+A5wEzVjcSVmKodADYS0XoiygL4IoCXr85jGQud4BaHmStE9CcAfg4gAvA8M++/ak9mLGiC3fGgm5mpWozsYuYt/kkbOTaCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCMMUxgjDFMYIwxTGCmFdxiOh5Ihomon3iXC8RvUJEh+Pfy67tYxoLjUZanH8GsNU79w0Av2LmjQB+FR8bLcS8isPM/wNg1Dv9DIAXYvkFAJ+/uo9lLHRC+zgrmfkMAMS/V1y9RzIWA1eSWKkhLCPXjUloi3OOiAYAIP49fLGKzPwcM2+plyrDWLyEKs7LAL4cy18G8J9X53GMRQMzX/IHwL8BOAOgjLm8f18B0Ic5b+pw/Lt3vuvE12L7WXQ/O+v9LS0jlzEflpHLuHqY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEGY4hhBmOIYQZjiGEE0kpFrLRH9mogOENF+IvpqfN6ycrUwjbQ4FQB/xsx3APgYgD8mojthWblamkYycp1h5rdieQrAAQBrYFm5WprL6uMQ0ToA9wPYDsvK1dI0nJGLiLoAvATga8w8SUSNfs4yct2INJjXJgPg5wC+Ls4dAjAQywMADll+nBvyp25+nEa8KgLwPQAHmPk7osiycrUyDbQSj2BO894BsCf+eRoBWblw/f977OcqtTiWkcuYD8vIZVw9THGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwgGlk73kZEbxLR23FGrr+Oz1tGrhamkRanCOBTzLwZwH0AthLRx2AZuVqaRjJyMTNPx4eZ+IdhGblamob6OEQUEdEeAMMAXmFmy8jV4jSkOMxcZeb7AAwCeJCI7m70BkT0LBHtJKKdgc9oLEAuy6ti5nEArwLYCuAcEQ0AQPx7+CKfeY6Zt9RLlWEsXhrxqpYT0dJYbgfwBICDsIxcLU0jySMHALxARBHmFO1FZv4JEW0D8CIRfQXACQBfuIbPaSwwLCOXMR91M3I1nK72ahAt6UH3o48AAC6V7jZdi9RxG9xxtux0754NG1W9Z373qUS++6479L2zmUSezM8m8vad21S97dt+nchj42OqLJ+qJnJnyn116Tb9/KVaOZELlbIqq6Xce1dTuqdQRX38r0r+/6W9a2Qymbqy/y9bKbrnKhaLuqzsynb9y3/XfSabcjCCMMUxgmiqqZL4fSt5WPHa1SJqrp5Q9XeOHFL1RiecaXni8U+oskce+ngi9y5102pbH/mUqjfY15/IL/1KN9P5kVOJXK04w1IraSNDGWda08JcAECpWknkjGeDMsKM1WruneH3DC9h5isVd335HWez+jkQuWeOsrqo2kDSfGtxjCBMcYwgTHGMIBZMH0dSuYQ6c0p8Lqsrnhx1sx4v/fS/VNn7x48n8uee+Ewi33HzLare5rseSOT0QK8qe+MXP03ko/v3JvJUtaSfUfTRopzuW+TSrv/DVa9vpA7cu9W4pupVxHfHaKy/U+OKLhR9nFRWXyOd1sMh9bAWxwjCFMcI4rqZKh85ksxe66vcw8gdkO/Si1HU8XJelb2+c3siD585k8if/cRnVL3fefjhRL5r/a2q7J6nP5/IL4tx3lcO7FH1ysJ/5po2M1HGfeWplH5RFq66tEBV7z0rwnR5l0c263xr+bEq6xHsclmYV/JN1fxqYS2OEYQpjhGEKY4RRJP7OJT0ZVLerG5NGGS6hIsp+z/+bHKN3Rl/2qImbrf/+NFEPvviD1S9Q2eGEvmpj92vym4+O5LI2QsXEjnt9VXkDLjXBVHDEEy6lNKinyc6L+nIc4/Jufil2sWHNaLIvXS16s3Es7tmsaBnx9OZ+aNfrMUxgjDFMYJoujv+ganyA7mkBn+ooRQn5Kd8t102/FXPSEhTlWlzzfRITbvtP3ntlUQe3/a6Knti1pnCqa5JV7DEi+RKXeJrFe9Naf//1l2fpCn33jMTOVOV9u41O+uC1OTIsT/jXZV/eu9xCyVvlLkO1uIYQZjiGEE01VQRAenUnJmo1LwJPtE0pzwbdPH5UK9AfIzJ89rEYVkGjUX6XiS8meKpU6osU3QmLlotvKPunL6XnLwkzyNSXpV+fhJekA7k8jzEmjMlUaT/hNmcGzmWZqvqeV81SA9Of1fMNslpXCNMcYwgTHGMIJrbx6kRsvnYHc/pWxdwiX6BcK0pYnFWz/hKnzOq6QAq2ZORs+oVry+UEkFZa1i76pzuSORJLE3kNGt3vFx1z1Ul7dqmUvI9dZR4RfRDIjFaTBV9DXn9otdPkrPjaYgR5ln9LlFF9LWqeugixfO3Jw23OHGqk91E9JP42DJytTCXY6q+CuCAOLaMXC1MQ6aKiAYBfBbAtwB8PT79DIDHY/kFzKU/+ctLXSfFjPa4iax5U5TZnDBBXoyt9M5ZBjGldDOtJk692UVWsrigN4LdJq6/zIuSKkfueEaul6rpa8hgrXJKm1MZM52u6ueX5onFMlxvDhVRxr1noarNWKUg1m2JZ8y157x67vr+ROzVNFXfBfAX3j0sI1cL00h+nM8BGGbmXSE3kBm5quXS/B8wFgWNmKqHAfweET0NoA1ADxH9K+KMXMx8Zr6MXACeA4C2nqWW5uQGYV7FYeZvAvgmABDR4wD+nJn/gIj+FnOZuL6NBjNycaWC4vnzAICO3k5VJm2g52SjImx/Rdj7mh/sLfouXrdD30tNTei+Vq7kgpo6vEixStbpfT4nA8/0/4Psd8x667LbO51LX5oqqLLOrBiGSDnZ7w+qoDdvOqJUEt+ecLNTKT3EkRZTE9XIiySoXNtArm8DeJKIDgN4Mj42WoTLGgBk5lcx5z2BmS8A+PTVfyRjMdDUkeOu9g48fO9mAMCufW+pspRYzhu1dauyWpsrK4nlqVXPnZWHvqWSNaWpIs8Z7RRmpsNzS6fFt1XsFDPlXkxwV597/nX36CXGI+OjiTzz2+OqrDzrzGQmJ66Z9Q2De+bIG/mWGbrkLH255HUA0s5VZ2+kPuXHONfB5qqMIExxjCCaaqq4VkNlagYA8PF77lNlb+58M5Fn83o0NNff564h0kdlIy8ASTS5FW/Uly+S7Srylqh0icOq93810e6+rorIQuFZTEzU3HjV2oGVqmxpt/OqctP6PWdOiIxfZVdW8kfS5QSxH6BVqZ+CsuZlxiiJ6/umqgozVcY1whTHCMIUxwiiqX2cUqmEkydOAwD6e3tU2W03u5Qi750/p8qGz7rj9uUuK2iuTc/4prLudfJelixp4SvCOfdycSIjKk57geDn213AVlWsZ4pyuk9watLNvswe2K/Knn7y6UQ+MaaDq6bOus9lxZLdYk0v0S0V3bv5ww4Z8cyzszOJXPP6QiTc+IoXyFX2s3fVwVocIwhTHCOIJi8BTgHpOXd6ZlqbkjaRVHrDhvWqrLfoJgNHxt3SW65qQ1PIu9HRtqznYooRVTkvSBX9HJHYR2LGD34aWJXIxVnXnE/7wVqd7nMXhLkAgG27XHQKT+q9IrjNDTVQwT1XW0bHNNeE6SoV9USpHCLPiqXCxbKXkUJO9HrjCexHdtXBWhwjCFMcIwhTHCOIpvZxytUqhkcnAAArlixRZZm0s+9Dp06rslXrb3byqrXuep7XeODokUQueBtzVMU0Q0YEe6e94XUSM+KFLh1sds+jjyXykiGXkeu1k7tVvZJYjx55QV4T067Ps6JPryjqzLq+0fThY+563n5SKXH9bEb3w0olXfcD/Kxepbzo23lRAKkG2hNrcYwgTHGMIJpqqnJt7bjl9nsBAMfff1+VjeVdE17ysmQViscS+YHNmxO5v6tL1Xv73NlEZm9/gs6lbqSahVlEzXPHxbqnCWgzcGa3m73esOVed9+0drlrJ/YlctobwR64dXUiZ0lfnyfcdbrEYrLSfr0vV0ZsdlH10rSkxFqqvEgSTl4kgcxKVip4q08q8/vj1uIYQZjiGEE011Rlc7hp3VwMbs+yflV2/NixRB45p2Nxx86PJ/LbbzkPZmm39noKExOJXPYCtKYm3DWWrnHBVZ0d2lxkMs6MTXmZwaJTLl74cP43iXwsM6Hq9XY6szg5Oq3Khve6HMsffehxVVZOuQxaZ0+7nMoTo+OqXl+3mCAm/SeU8c9ZYZKLrM2RWmbtB8RdJBhMYi2OEYQpjhGEKY4RRHNnx4mSvZB6l+lR04729kTu79Vu9tBJ57pPTY0n8vTklKrX2SH6PGndP5ksuP7D0GHXz1i6TI9g5/IiqD2r+yd3r3Sz47Mi09bstO7jLO1bnsi9nfr6h17ekcjRUd3vuPsjdyby+7ucSz974by+frtYHuytq2JOywN3Ly+dC4vk3OS59P53V49G8+McAzCFuUC6CjNvIaJeAP8OYB2AYwB+n5nHLnYN48bickzVJ5n5PmbeEh9bRq4W5kpM1WVn5AIuvvtvZ6czM50b9JaGbSK2+LiYyDw/fEbVW712TSLP5LUZk2G1abE2K+PFFa8aHEjk/nV6+W77EveM5W5nqipe5NOxMTeC3ZbX6Sq6S+4aw7veU2X/d8JN7uYLbhJ1/cBqVa9TjPrOePHIMn44Qxf/88rtGrNZHRBXu8S2Tx/QaIvDAH5BRLuI6Nn4nGXkamEabXEeZubTRLQCwCtEdLDRG8SK9iwAdHVbYtIbhYZaHGY+Hf8eBvBjAA8izsgFAPNl5GLmLcy8pa29q14VYxEyb4tDRJ0AUsw8FcufAfA3AF7GZWbkiq8H4MN9HXmczup+wa2bbk/kzg639vrt3TpIfMPGTYk87PV/Lux365silXVUP0em3V3/1k0bVZnsW+wfcY1uLeet7e5yz1/wNtiIelwfJzup3fHCmOuX5cS90qz/TLmU65OQ5zoXy2J7SfFqaW9aISX6dvmi/h79uvVoxFStBPDj+A+eBvADZv4ZEe0A8CIRfQXACQBfaOBaxg1CIzkAjwLYXOe8ZeRqYZo6cpxKpdARm5pCQa8Hqoo0HF5GDsyK9VLr1gtX3Ys3OnnmRCJ39eiZ86IIUN4k3Oz+VXqW/vRh5+7T9v9VZZ+4y5nM5ZNuVHlJSgc/D/YNJvKpwrgq61zfm8iF93Rs9S03ufVkVHPfz0rvXWan3EhyW4c265mMe5aiCCIrVLTbLj/ljz775rseNldlBGGKYwRhimME0eQ+DiGXy9UtK5acPZYuJQBEwj28MObWjks3HQBSYtb49W2vqTLpcnZ1uaygPT169vpo3vVdpod1H2RaRBUOtjt5Q9Vby1RzaVmitP6KJ9rdu7Vv0AOiwxUXrN6fc+vFP/rgY6re9HnXx9mzf6cqk0uk2jNiiCOlv9Oi2CL6Q1t5N+COW4tjBGGKYwTRVFMVRWn09s65o9PTevaaZVBWSs/WloQZi8SeBGNjOoBqYNC5wU88tVWV7XjtjUSemnEmoXxabxE9k3emMO9lppqeccFgXcIFX+lNJi8tOnM8vlKbwpFB547vOn1MlaXE+EJ3ygW2HTyk631kkwv4+uSjejjhjTd/mciTBTdLn9FeO9JybZm3lpr8DbLqYC2OEYQpjhFEU01VrVbDdLwbbe/yPlWWFpN6o946Irl8tSq2HKx4LerstBttXZLT13/y059L5O2vuxHhY8f0Gq7xKWeOBpfpa3Td6by4gwffcdcQ674AYFD8P97mBXn1bLwpkYduuUmVVY65yNtVeTfZWvYyUGzftzeRN992pyr75CNPuXpvOc/y3Hn9nhnh3XmrpVGr2boq4xphimMEYYpjBNHcPg4z8nF2qXMjOmBw+Qo3itrZ1aHKhkfcOurZGTcaWvWm0VMiY2hlRpfJrZQfevjxRN63V++bdbjiUoNc8NJ/9G25P5G3nXSB5mcndR+kMO7c/UFvVrr/Xhcc1l/S/Z+cyJZyZ7vrX5EXOTkmhif27zugym4VfahPPeqGJPYe2KHq7Tv6pjjSs+HZtM7eVQ9rcYwgTHGMIJqcIJvBsXtaKunRypFhN3HX369HQwfXuBHh06ddLHE+rzN3sUgYGXl7MMmYZhb7Gtx3/0dUvb4+Zxb27X5Tle1+1wV5TYg4tIH196p66/rc6PDYjl+qstFX3Qj2Gm+p88pu955r+5y5Lub0VpNpseFEPqUD4k6ccMFsNbiyzfduUfW6+1yqlF17tBmbLc1iPqzFMYIwxTGCMMUxgmhyHwdwMUN6nLssgrfOnRtRZT09zh6vWuXWdp89e1bVq4lgpWrB2y5ZEMktqL19nFasWZfIt3ju+C9f3+aeV/Sv7lq9VNXr7HN9tHLau/6we7fVXkx4rs31+6KM+H4ivQlIZ/1MJgAAmYVt6KTrDxa97aNvv9sNLWSzegb/je2vYj6sxTGCMMUxgmi6qfrAHSdvLY/Mgunt/IxpMWMtg2qXLOlV9fITLhisXNOmUAaDVSrOJKQyntuecgFOKwbXqbJNIm754D43O14tavf1wjkXq9wFPXK8UiTkXucNGUyIzFiTYso67yfBFts6dmW0S18ourrMzm6NjujAuf1vu5HvW+9Yp8oee9its3zx759HPRpqcYhoKRH9iIgOEtEBInqIiHqJ6BUiOhz/tlQULUSjpurvAPyMmW/H3HLgA7CMXC1NI9kqegA8BuAPAYCZSwBKRBSQkYuBeIlJKqWDYFnselv1EjTLENhJsfS2u1uPqC5ZsjSRZ0k3zSmxtaI0VRXvXixMV0e39jbWb7ojkTNii6DzR3W6oMkzbrZyZV6bsXzOvfdZb8+EUqczO7nVLlHl9JQeZc+IZTr+hGSOnQcmPa5yWXtV46MutnrvHr1TsW+66tFIi7MBwAiAfyKi3UT0j3G6E8vI1cI0ojhpAA8A+Admvh/ADC7DLBHRs0S0k4h2zs5Mz/8BY1HQiOIMARhi5u3x8Y8wp0iXnZGro9Myct0oNJIf5ywRnSSi25j5EOZy4rwb/1xWRq5KpYLz8fLV3mXLVVlGbJHsu+rSVrNw26endf+hmnaubrdIuO0j7X2qqvsPReEu17xtlXNtbsZ67foNibykTfczjrzjZtEPl/SS55zYyrGvqr/+brE/1irhxa8cGFD1pi64zUjK3vbRchOTduHup0jXi4SLXyjpsv3vzJ/isdFxnD8F8H0iygI4CuCPMNdaWUauFqUhxWHmPQC21CmyjFwtSlNHjtPpDFasmDNRoxcmVVlPt3uUXE676lWxzkeOmvoxx2Vh0yY997OtzZnCmhiaZi+tFwszWfSuwWL5sUxw2eONMN/T6bKGvXeTNsnvv384kc9m9eTl8ln3LMXjLuPFGm8B2Rqx39bY2Lgqy+edjcuIHYJlzDUAlEVmycjbWXG25J2og81VGUGY4hhBmOIYQTS1j5ONIqyOt3HubdPu8nERhF6t6amEbjH0L11pL5EUqnKdthfhNCvc1khMP6S8jFkiiRVSke4XFJTtdzfPpvW7tPe64009D6iy8X63UUnF60uk0+5+kxNuqCHCOVWPxTDBzevX6etPuKD/qUk39RF5m51kMy6Taa3mqQHp966HtThGEKY4RhB0sf2jrsnNiEYAHAfQD+D8PNVbiYX8fdzMzMv9k01VnOSmRDvFTnstz2L8PsxUGUGY4hhBXC/Fee463Xehsui+j+vSxzEWP2aqjCCaqjhEtJWIDhHRESJquVURRLSWiH4dLzHaT0Rfjc8vuqVGTTNVRBQB+C2AJzEXjroDwJeY+d2mPMACIA6xHWDmt4ioG8AuAJ/H3AqSUWb+dvwPtYyZ593D/XrSzBbnQQBHmPlovMTmh5jb9L5lYOYzzPxWLE9hbn3aGsx9Dy/E1V7AnDItaJqpOGsAnBTHQ/G5loSI1gG4H8B2LMKlRs1UnHo7S7SkS0dEXQBeAvA1Zp6cr/5CpJmKMwRgrTgeBHD6InVvWIgogzml+T4z/0d8uqGlRguJZirODgAbiWh9vFrii5jb9L5loLmt6L4H4AAzf0cUvYy5JUZAg0uNrjfNnh1/GsB3AUQAnmfmbzXt5gsAInoEwG8A7IXb/PqvMNfPeRHATYiXGjHzaN2LLBBs5NgIwkaOjSBMcYwgTHGMIExxjCBMcYwgTHGMIExxjCBMcYwg/h86a2p61AEaswAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "test_img = torch.hstack((torch.zeros((3,32,32)),cifar10_train[50][0]))\n",
    "plt.imshow(test_img.permute(1,2,0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# 2-Class Domino Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Datasets and Subfolder Heirarchies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "num_train = 3500\n",
    "num_val = 500\n",
    "num_test = 800\n",
    "spur_corr = 1.0\n",
    "save_dir = root_dir = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_svhn_bot_kmnist/' # path to save domino data\n",
    "full_save_path = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_svhn_bot_kmnist_2class_ro_1p0.csv' \n",
    "top_dataset = {'name':'svhn', 'train_data':svhn_train, 'test_data':svhn_test}\n",
    "bot_dataset = {'name':'kmnist', 'train_data':kmnist_train, 'test_data':kmnist_test}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/jet/home/nmurali/asc170022p/nmurali/anaconda3/envs/pl/lib/python3.9/site-packages/torchvision/datasets/mnist.py:52: UserWarning: train_labels has been renamed targets\n",
      "  warnings.warn(\"train_labels has been renamed targets\")\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "8it [00:00, 74.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "16it [00:00, 49.91it/s]\u001b[A\u001b[A\n",
      "\n",
      "22it [00:02,  6.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "31it [00:02, 10.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "42it [00:02, 17.52it/s]\u001b[A\u001b[A\n",
      "\n",
      "56it [00:02, 28.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "70it [00:03, 38.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "98it [00:03, 70.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "119it [00:03, 91.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "147it [00:03, 122.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "166it [00:05, 30.18it/s] \u001b[A\u001b[A\n",
      "\n",
      "186it [00:06, 29.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "208it [00:06, 40.22it/s]\u001b[A\u001b[A\n",
      "\n",
      "222it [00:06, 46.41it/s]\u001b[A\u001b[A\n",
      "\n",
      "238it [00:06, 57.19it/s]\u001b[A\u001b[A\n",
      "\n",
      "259it [00:07, 27.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "282it [00:10, 15.84it/s]\u001b[A\u001b[A\n",
      "\n",
      "327it [00:10, 30.05it/s]\u001b[A\u001b[A\n",
      "\n",
      "365it [00:10, 45.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "390it [00:10, 57.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "415it [00:12, 37.45it/s]\u001b[A\u001b[A\n",
      "\n",
      "445it [00:12, 51.11it/s]\u001b[A\u001b[A\n",
      "\n",
      "494it [00:12, 82.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "540it [00:12, 117.23it/s]\u001b[A\u001b[A\n",
      "\n",
      "587it [00:13, 82.11it/s] \u001b[A\u001b[A\n",
      "\n",
      "636it [00:13, 114.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "688it [00:13, 155.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "727it [00:14, 103.89it/s]\u001b[A\u001b[A\n",
      "\n",
      "770it [00:14, 133.85it/s]\u001b[A\u001b[A\n",
      "\n",
      "824it [00:14, 178.25it/s]\u001b[A\u001b[A\n",
      "\n",
      "861it [00:14, 202.24it/s]\u001b[A\u001b[A\n",
      "\n",
      "897it [00:14, 167.86it/s]\u001b[A\u001b[A\n",
      "\n",
      "934it [00:15, 197.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "966it [00:15, 125.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "990it [00:18, 29.41it/s] \u001b[A\u001b[A\n",
      "\n",
      "1027it [00:18, 41.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "1063it [00:21, 26.25it/s]\u001b[A\u001b[A\n",
      "\n",
      "1079it [00:21, 27.25it/s]\u001b[A\u001b[A\n",
      "\n",
      "1127it [00:21, 44.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "1167it [00:21, 62.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "1224it [00:22, 90.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "1269it [00:22, 108.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "1327it [00:22, 152.79it/s]\u001b[A\u001b[A\n",
      "\n",
      "1360it [00:22, 150.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "1388it [00:23, 112.89it/s]\u001b[A\u001b[A\n",
      "\n",
      "1443it [00:23, 159.73it/s]\u001b[A\u001b[A\n",
      "\n",
      "1483it [00:23, 173.41it/s]\u001b[A\u001b[A\n",
      "\n",
      "1511it [00:25, 53.34it/s] \u001b[A\u001b[A\n",
      "\n",
      "1546it [00:25, 70.06it/s]\u001b[A\u001b[A\n",
      "\n",
      "1571it [00:26, 60.65it/s]\u001b[A\u001b[A\n",
      "\n",
      "1590it [00:26, 69.83it/s]\u001b[A\u001b[A\n",
      "\n",
      "1655it [00:26, 123.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "1708it [00:26, 171.25it/s]\u001b[A\u001b[A\n",
      "\n",
      "1747it [00:26, 140.04it/s]\u001b[A\u001b[A\n",
      "\n",
      "1805it [00:26, 194.31it/s]\u001b[A\u001b[A\n",
      "\n",
      "1843it [00:26, 216.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "1881it [00:30, 36.79it/s] \u001b[A\u001b[A\n",
      "\n",
      "1918it [00:30, 43.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "1939it [00:30, 48.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "1987it [00:31, 72.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "2044it [00:31, 109.74it/s]\u001b[A\u001b[A\n",
      "\n",
      "2083it [00:31, 136.63it/s]\u001b[A\u001b[A\n",
      "\n",
      "2125it [00:31, 168.49it/s]\u001b[A\u001b[A\n",
      "\n",
      "2162it [00:33, 51.32it/s] \u001b[A\u001b[A\n",
      "\n",
      "2204it [00:33, 69.52it/s]\u001b[A\u001b[A\n",
      "\n",
      "2247it [00:33, 93.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "2311it [00:33, 141.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "2354it [00:38, 28.68it/s] \u001b[A\u001b[A\n",
      "\n",
      "2384it [00:38, 35.48it/s]\u001b[A\u001b[A\n",
      "\n",
      "2439it [00:38, 53.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "2481it [00:38, 70.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "2528it [00:38, 96.14it/s]\u001b[A\u001b[A\n",
      "\n",
      "2573it [00:38, 125.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "2615it [00:40, 61.73it/s] \u001b[A\u001b[A\n",
      "\n",
      "2645it [00:41, 57.11it/s]\u001b[A\u001b[A\n",
      "\n",
      "2668it [00:41, 49.85it/s]\u001b[A\u001b[A\n",
      "\n",
      "2685it [00:42, 52.31it/s]\u001b[A\u001b[A\n",
      "\n",
      "2748it [00:42, 93.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "2791it [00:42, 111.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "2835it [00:42, 146.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "2879it [00:42, 184.32it/s]\u001b[A\u001b[A\n",
      "\n",
      "2921it [00:42, 211.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "2957it [00:46, 28.91it/s] \u001b[A\u001b[A\n",
      "\n",
      "2982it [00:47, 33.58it/s]\u001b[A\u001b[A\n",
      "\n",
      "3026it [00:49, 28.93it/s]\u001b[A\u001b[A\n",
      "\n",
      "3041it [00:50, 26.20it/s]\u001b[A\u001b[A\n",
      "\n",
      "3094it [00:50, 43.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "3131it [00:50, 58.74it/s]\u001b[A\u001b[A\n",
      "\n",
      "3185it [00:50, 88.93it/s]\u001b[A\u001b[A\n",
      "\n",
      "3232it [00:50, 120.19it/s]\u001b[A\u001b[A\n",
      "\n",
      "3271it [00:50, 148.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "3310it [00:50, 179.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "3357it [00:50, 224.16it/s]\u001b[A\u001b[A\n",
      "\n",
      "3409it [00:50, 273.34it/s]\u001b[A\u001b[A\n",
      "\n",
      "3452it [00:51, 115.44it/s]\u001b[A\u001b[A\n",
      "\n",
      "3498it [00:51, 147.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "3532it [00:54, 44.81it/s] \u001b[A\u001b[A\n",
      "\n",
      "3556it [00:54, 48.95it/s]\u001b[A\u001b[A\n",
      "\n",
      "3597it [00:54, 68.23it/s]\u001b[A\u001b[A\n",
      "\n",
      "3622it [00:54, 80.52it/s]\u001b[A\u001b[A\n",
      "\n",
      "3646it [00:57, 27.44it/s]\u001b[A\u001b[A\n",
      "\n",
      "3691it [00:58, 35.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "3737it [00:58, 53.01it/s]\u001b[A\u001b[A\n",
      "\n",
      "3760it [00:58, 62.33it/s]\u001b[A\u001b[A\n",
      "\n",
      "3788it [01:00, 34.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "3809it [01:03, 18.11it/s]\u001b[A\u001b[A\n",
      "\n",
      "3867it [01:03, 32.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "3916it [01:03, 49.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "3948it [01:04, 45.20it/s]\u001b[A\u001b[A\n",
      "\n",
      "4000it [01:04, 61.88it/s]\u001b[A\u001b[A\n",
      "\n",
      " 50%|█████     | 1/2 [01:04<01:04, 64.66s/it]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "9it [00:00, 49.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "21it [00:01, 15.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "59it [00:01, 53.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "114it [00:01, 117.65it/s]\u001b[A\u001b[A\n",
      "\n",
      "165it [00:01, 178.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "205it [00:01, 207.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "251it [00:01, 256.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "291it [00:01, 270.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "334it [00:02, 296.15it/s]\u001b[A\u001b[A\n",
      "\n",
      "371it [00:02, 291.31it/s]\u001b[A\u001b[A\n",
      "\n",
      "416it [00:02, 305.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "451it [00:02, 204.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "494it [00:02, 245.31it/s]\u001b[A\u001b[A\n",
      "\n",
      "545it [00:02, 286.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "580it [00:10, 16.23it/s] \u001b[A\u001b[A\n",
      "\n",
      "605it [00:13, 14.98it/s]\u001b[A\u001b[A\n",
      "\n",
      "650it [00:13, 22.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "687it [00:13, 30.84it/s]\u001b[A\u001b[A\n",
      "\n",
      "715it [00:13, 37.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "767it [00:13, 59.01it/s]\u001b[A\u001b[A\n",
      "\n",
      "812it [00:13, 76.37it/s]\u001b[A\u001b[A\n",
      "\n",
      "841it [00:14, 70.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "864it [00:14, 81.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "900it [00:14, 107.16it/s]\u001b[A\u001b[A\n",
      "\n",
      "947it [00:14, 147.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "994it [00:14, 193.65it/s]\u001b[A\u001b[A\n",
      "\n",
      "1031it [00:14, 222.95it/s]\u001b[A\u001b[A\n",
      "\n",
      "1067it [00:16, 65.46it/s] \u001b[A\u001b[A\n",
      "\n",
      "1093it [00:17, 40.60it/s]\u001b[A\u001b[A\n",
      "\n",
      "1112it [00:18, 45.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "1132it [00:18, 54.91it/s]\u001b[A\u001b[A\n",
      "\n",
      "1161it [00:18, 73.69it/s]\u001b[A\u001b[A\n",
      "\n",
      "1212it [00:18, 117.46it/s]\u001b[A\u001b[A\n",
      "\n",
      "1266it [00:18, 167.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "1300it [00:18, 186.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "1332it [00:18, 163.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "1383it [00:19, 128.33it/s]\u001b[A\u001b[A\n",
      "\n",
      "1428it [00:19, 131.58it/s]\u001b[A\u001b[A\n",
      "\n",
      "1453it [00:21, 55.19it/s] \u001b[A\u001b[A\n",
      "\n",
      "1488it [00:21, 72.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "1530it [00:21, 99.41it/s]\u001b[A\u001b[A\n",
      "\n",
      "1564it [00:21, 122.81it/s]\u001b[A\u001b[A\n",
      "\n",
      "1592it [00:29, 13.07it/s] \u001b[A\u001b[A\n",
      "\n",
      "1617it [00:29, 16.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "1665it [00:29, 27.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "1711it [00:30, 35.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "1744it [00:30, 46.30it/s]\u001b[A\u001b[A\n",
      "\n",
      "1780it [00:30, 62.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "1820it [00:30, 84.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "1853it [00:30, 104.20it/s]\u001b[A\u001b[A\n",
      "\n",
      "1884it [00:42,  9.64it/s] \u001b[A\u001b[A\n",
      "\n",
      "1902it [00:42, 11.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "1928it [00:42, 15.68it/s]\u001b[A\u001b[A\n",
      "\n",
      "1973it [00:42, 25.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "2024it [00:42, 40.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "2067it [00:42, 57.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "2119it [00:42, 84.26it/s]\u001b[A\u001b[A\n",
      "\n",
      "2165it [00:42, 112.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "2208it [00:43, 123.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "2245it [00:43, 148.52it/s]\u001b[A\u001b[A\n",
      "\n",
      "2280it [00:44, 60.64it/s] \u001b[A\u001b[A\n",
      "\n",
      "2305it [00:44, 70.81it/s]\u001b[A\u001b[A\n",
      "\n",
      "2329it [00:45, 54.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "2375it [00:45, 77.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "2422it [00:45, 109.93it/s]\u001b[A\u001b[A\n",
      "\n",
      "2463it [00:46, 141.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "2515it [00:46, 191.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "2553it [00:46, 219.97it/s]\u001b[A\u001b[A\n",
      "\n",
      "2605it [00:46, 265.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "2653it [00:46, 308.20it/s]\u001b[A\u001b[A\n",
      "\n",
      "2695it [00:47, 162.21it/s]\u001b[A\u001b[A\n",
      "\n",
      "2740it [00:50, 40.34it/s] \u001b[A\u001b[A\n",
      "\n",
      "2763it [00:50, 38.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "2817it [00:50, 59.44it/s]\u001b[A\u001b[A\n",
      "\n",
      "2872it [00:51, 86.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "2926it [00:51, 120.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "2968it [00:52, 77.76it/s] \u001b[A\u001b[A\n",
      "\n",
      "3003it [00:52, 95.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "3038it [00:52, 115.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "3069it [00:52, 136.38it/s]\u001b[A\u001b[A\n",
      "\n",
      "3123it [00:52, 190.31it/s]\u001b[A\u001b[A\n",
      "\n",
      "3185it [00:52, 259.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "3230it [00:52, 263.42it/s]\u001b[A\u001b[A\n",
      "\n",
      "3270it [00:53, 171.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "3306it [00:53, 190.09it/s]\u001b[A\u001b[A\n",
      "\n",
      "3364it [00:53, 252.45it/s]\u001b[A\u001b[A\n",
      "\n",
      "3408it [00:53, 278.94it/s]\u001b[A\u001b[A\n",
      "\n",
      "3446it [00:53, 283.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "3500it [00:53, 332.74it/s]\u001b[A\u001b[A\n",
      "\n",
      "3546it [00:54, 362.38it/s]\u001b[A\u001b[A\n",
      "\n",
      "3588it [00:54, 371.44it/s]\u001b[A\u001b[A\n",
      "\n",
      "3629it [00:54, 381.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "3672it [00:54, 393.21it/s]\u001b[A\u001b[A\n",
      "\n",
      "3714it [00:54, 345.04it/s]\u001b[A\u001b[A\n",
      "\n",
      "3763it [00:54, 370.45it/s]\u001b[A\u001b[A\n",
      "\n",
      "3818it [00:58, 35.73it/s] \u001b[A\u001b[A\n",
      "\n",
      "3847it [00:58, 43.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "3876it [00:59, 46.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "3910it [00:59, 60.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "3961it [00:59, 89.72it/s]\u001b[A\u001b[A\n",
      "\n",
      "4000it [01:02, 64.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "100%|██████████| 2/2 [02:06<00:00, 63.45s/it]\u001b[A\n",
      " 50%|█████     | 1/2 [02:06<02:06, 126.91s/it]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "30it [00:00, 298.65it/s]\u001b[A\u001b[A\n",
      "\n",
      "91it [00:00, 481.18it/s]\u001b[A\u001b[A\n",
      "\n",
      "145it [00:00, 160.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "193it [00:08, 16.31it/s] \u001b[A\u001b[A\n",
      "\n",
      "233it [00:08, 23.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "283it [00:08, 35.30it/s]\u001b[A\u001b[A\n",
      "\n",
      "335it [00:08, 51.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "371it [00:08, 65.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "406it [00:08, 65.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "432it [00:09, 76.84it/s]\u001b[A\u001b[A\n",
      "\n",
      "456it [00:09, 83.89it/s]\u001b[A\u001b[A\n",
      "\n",
      "501it [00:09, 119.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "536it [00:09, 147.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "566it [00:11, 39.06it/s] \u001b[A\u001b[A\n",
      "\n",
      "587it [00:12, 42.18it/s]\u001b[A\u001b[A\n",
      "\n",
      "634it [00:12, 66.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "682it [00:12, 97.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "714it [00:12, 118.49it/s]\u001b[A\u001b[A\n",
      "\n",
      "764it [00:12, 164.45it/s]\u001b[A\u001b[A\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "802it [00:12, 161.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "833it [00:13, 121.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "857it [00:16, 29.71it/s] \u001b[A\u001b[A\n",
      "\n",
      "912it [00:16, 48.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "962it [00:16, 64.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "1029it [00:16, 101.20it/s]\u001b[A\u001b[A\n",
      "\n",
      "1076it [00:16, 130.48it/s]\u001b[A\u001b[A\n",
      "\n",
      "1117it [00:17, 158.86it/s]\u001b[A\u001b[A\n",
      "\n",
      "1157it [00:17, 168.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "1196it [00:17, 198.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "1231it [00:17, 172.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "1260it [00:19, 59.76it/s] \u001b[A\u001b[A\n",
      "\n",
      "1298it [00:19, 79.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "1323it [00:19, 82.06it/s]\u001b[A\u001b[A\n",
      "\n",
      "1375it [00:19, 114.03it/s]\u001b[A\u001b[A\n",
      "\n",
      "1420it [00:19, 151.18it/s]\u001b[A\u001b[A\n",
      "\n",
      "1449it [00:19, 163.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "1505it [00:20, 225.63it/s]\u001b[A\u001b[A\n",
      "\n",
      "1559it [00:20, 279.13it/s]\u001b[A\u001b[A\n",
      "\n",
      "1599it [00:20, 165.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "1633it [00:20, 186.77it/s]\u001b[A\u001b[A\n",
      "\n",
      "1692it [00:20, 251.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "1731it [00:21, 211.72it/s]\u001b[A\u001b[A\n",
      "\n",
      "1770it [00:21, 216.70it/s]\u001b[A\u001b[A\n",
      "\n",
      "1821it [00:21, 268.57it/s]\u001b[A\u001b[A\n",
      "\n",
      "1857it [00:21, 280.13it/s]\u001b[A\u001b[A\n",
      "\n",
      "1908it [00:21, 330.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "1948it [00:21, 303.90it/s]\u001b[A\u001b[A\n",
      "\n",
      "1984it [00:24, 52.11it/s] \u001b[A\u001b[A\n",
      "\n",
      "2009it [00:24, 48.04it/s]\u001b[A\u001b[A\n",
      "\n",
      "2028it [00:25, 51.30it/s]\u001b[A\u001b[A\n",
      "\n",
      "2069it [00:25, 74.23it/s]\u001b[A\u001b[A\n",
      "\n",
      "2093it [00:25, 88.28it/s]\u001b[A\u001b[A\n",
      "\n",
      "2142it [00:25, 131.73it/s]\u001b[A\u001b[A\n",
      "\n",
      "2173it [00:25, 121.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "2218it [00:31, 19.46it/s] \u001b[A\u001b[A\n",
      "\n",
      "2236it [00:31, 22.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "2285it [00:31, 36.79it/s]\u001b[A\u001b[A\n",
      "\n",
      "2312it [00:31, 43.34it/s]\u001b[A\u001b[A\n",
      "\n",
      "2334it [00:32, 46.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "2373it [00:32, 67.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "2423it [00:32, 102.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "2475it [00:32, 145.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "2513it [00:32, 168.37it/s]\u001b[A\u001b[A\n",
      "\n",
      "2563it [00:32, 217.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "2617it [00:32, 274.57it/s]\u001b[A\u001b[A\n",
      "\n",
      "2661it [00:33, 285.79it/s]\u001b[A\u001b[A\n",
      "\n",
      "2712it [00:33, 332.91it/s]\u001b[A\u001b[A\n",
      "\n",
      "2756it [00:33, 272.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "2802it [00:33, 309.81it/s]\u001b[A\u001b[A\n",
      "\n",
      "2848it [00:33, 332.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "2888it [00:37, 36.26it/s] \u001b[A\u001b[A\n",
      "\n",
      "2916it [00:43, 13.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "2961it [00:44, 20.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "3008it [00:44, 28.83it/s]\u001b[A\u001b[A\n",
      "\n",
      "3032it [00:48, 15.73it/s]\u001b[A\u001b[A\n",
      "\n",
      "3068it [00:48, 21.85it/s]\u001b[A\u001b[A\n",
      "\n",
      "3115it [00:48, 32.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "3145it [00:49, 42.06it/s]\u001b[A\u001b[A\n",
      "\n",
      "3175it [00:49, 50.63it/s]\u001b[A\u001b[A\n",
      "\n",
      "3200it [00:51, 27.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "3219it [00:52, 24.81it/s]\u001b[A\u001b[A\n",
      "\n",
      "3232it [00:52, 28.21it/s]\u001b[A\u001b[A\n",
      "\n",
      "3244it [00:53, 25.80it/s]\u001b[A\u001b[A\n",
      "\n",
      "3288it [00:53, 47.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "3314it [00:53, 62.15it/s]\u001b[A\u001b[A\n",
      "\n",
      "3337it [01:00, 11.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "3366it [01:00, 16.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "3413it [01:00, 28.15it/s]\u001b[A\u001b[A\n",
      "\n",
      "3440it [01:00, 36.50it/s]\u001b[A\u001b[A\n",
      "\n",
      "3473it [01:00, 50.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "3503it [01:01, 49.09it/s]\u001b[A\u001b[A\n",
      "\n",
      "3540it [01:01, 69.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "3572it [01:01, 90.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "3599it [01:02, 62.26it/s]\u001b[A\u001b[A\n",
      "\n",
      "3620it [01:02, 63.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "3636it [01:02, 64.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "3650it [01:03, 34.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "3711it [01:04, 72.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "3750it [01:04, 99.79it/s]\u001b[A\u001b[A\n",
      "\n",
      "3780it [01:04, 94.49it/s]\u001b[A\u001b[A\n",
      "\n",
      "3803it [01:09, 16.32it/s]\u001b[A\u001b[A\n",
      "\n",
      "3820it [01:12, 12.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "3846it [01:12, 17.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "3881it [01:12, 27.03it/s]\u001b[A\u001b[A\n",
      "\n",
      "3899it [01:13, 27.42it/s]\u001b[A\u001b[A\n",
      "\n",
      "3958it [01:13, 52.05it/s]\u001b[A\u001b[A\n",
      "\n",
      "4000it [01:13, 54.26it/s]\u001b[A\u001b[A\n",
      "\n",
      " 50%|█████     | 1/2 [01:13<01:13, 73.73s/it]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "2it [00:00, 12.90it/s]\u001b[A\u001b[A\n",
      "\n",
      "12it [00:00, 36.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "23it [00:04,  4.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "76it [00:04, 22.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "98it [00:05, 21.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "117it [00:05, 24.75it/s]\u001b[A\u001b[A\n",
      "\n",
      "141it [00:07, 23.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "184it [00:07, 41.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "204it [00:07, 48.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "251it [00:07, 80.19it/s]\u001b[A\u001b[A\n",
      "\n",
      "286it [00:07, 97.22it/s]\u001b[A\u001b[A\n",
      "\n",
      "310it [00:07, 109.14it/s]\u001b[A\u001b[A\n",
      "\n",
      "338it [00:07, 127.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "360it [00:08, 140.24it/s]\u001b[A\u001b[A\n",
      "\n",
      "382it [00:08, 70.83it/s] \u001b[A\u001b[A\n",
      "\n",
      "426it [00:11, 30.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "446it [00:11, 34.33it/s]\u001b[A\u001b[A\n",
      "\n",
      "471it [00:11, 45.16it/s]\u001b[A\u001b[A\n",
      "\n",
      "508it [00:11, 66.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "530it [00:11, 78.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "570it [00:12, 111.41it/s]\u001b[A\u001b[A\n",
      "\n",
      "626it [00:12, 170.76it/s]\u001b[A\u001b[A\n",
      "\n",
      "660it [00:13, 54.04it/s] \u001b[A\u001b[A\n",
      "\n",
      "685it [00:14, 44.77it/s]\u001b[A\u001b[A\n",
      "\n",
      "713it [00:14, 56.78it/s]\u001b[A\u001b[A\n",
      "\n",
      "756it [00:15, 82.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "791it [00:15, 105.49it/s]\u001b[A\u001b[A\n",
      "\n",
      "830it [00:15, 137.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "863it [00:15, 159.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "893it [00:15, 175.05it/s]\u001b[A\u001b[A\n",
      "\n",
      "921it [00:15, 188.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "948it [00:15, 201.21it/s]\u001b[A\u001b[A\n",
      "\n",
      "992it [00:15, 253.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "1024it [00:16, 248.46it/s]\u001b[A\u001b[A\n",
      "\n",
      "1065it [00:16, 265.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "1095it [00:25, 11.83it/s] \u001b[A\u001b[A\n",
      "\n",
      "1116it [00:31,  7.70it/s]\u001b[A\u001b[A\n",
      "\n",
      "1159it [00:31, 12.25it/s]\u001b[A\u001b[A\n",
      "\n",
      "1188it [00:32, 16.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "1231it [00:32, 25.22it/s]\u001b[A\u001b[A\n",
      "\n",
      "1276it [00:32, 37.69it/s]\u001b[A\u001b[A\n",
      "\n",
      "1310it [00:35, 22.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "1335it [00:35, 28.07it/s]\u001b[A\u001b[A\n",
      "\n",
      "1359it [00:35, 34.38it/s]\u001b[A\u001b[A\n",
      "\n",
      "1392it [00:35, 47.84it/s]\u001b[A\u001b[A\n",
      "\n",
      "1432it [00:35, 69.40it/s]\u001b[A\u001b[A\n",
      "\n",
      "1461it [00:36, 81.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "1486it [00:39, 24.42it/s]\u001b[A\u001b[A\n",
      "\n",
      "1541it [00:39, 40.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "1588it [00:39, 57.45it/s]\u001b[A\u001b[A\n",
      "\n",
      "1612it [00:40, 51.48it/s]\u001b[A\u001b[A\n",
      "\n",
      "1654it [00:40, 72.97it/s]\u001b[A\u001b[A\n",
      "\n",
      "1679it [00:46, 14.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "1705it [00:47, 18.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "1745it [00:47, 28.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "1778it [00:47, 38.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "1812it [00:47, 52.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "1847it [00:47, 63.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "1894it [00:47, 93.13it/s]\u001b[A\u001b[A\n",
      "\n",
      "1940it [00:47, 120.66it/s]\u001b[A\u001b[A\n",
      "\n",
      "1979it [00:48, 150.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "2011it [00:48, 167.10it/s]\u001b[A\u001b[A\n",
      "\n",
      "2041it [00:48, 176.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "2082it [00:50, 50.24it/s] \u001b[A\u001b[A\n",
      "\n",
      "2102it [00:50, 55.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "2154it [00:50, 68.04it/s]\u001b[A\u001b[A\n",
      "\n",
      "2190it [00:51, 88.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "2211it [00:51, 85.80it/s]\u001b[A\u001b[A\n",
      "\n",
      "2228it [00:51, 83.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "2257it [00:51, 106.89it/s]\u001b[A\u001b[A\n",
      "\n",
      "2299it [00:51, 150.46it/s]\u001b[A\u001b[A\n",
      "\n",
      "2325it [00:56, 19.40it/s] \u001b[A\u001b[A\n",
      "\n",
      "2344it [00:58, 15.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "2387it [00:58, 25.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "2417it [00:58, 33.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "2473it [00:59, 57.42it/s]\u001b[A\u001b[A\n",
      "\n",
      "2508it [00:59, 74.63it/s]\u001b[A\u001b[A\n",
      "\n",
      "2542it [00:59, 93.77it/s]\u001b[A\u001b[A\n",
      "\n",
      "2584it [00:59, 126.01it/s]\u001b[A\u001b[A\n",
      "\n",
      "2619it [00:59, 130.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "2652it [00:59, 156.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "2687it [00:59, 184.77it/s]\u001b[A\u001b[A\n",
      "\n",
      "2725it [00:59, 219.59it/s]\u001b[A\u001b[A\n",
      "\n",
      "2765it [01:00, 256.03it/s]\u001b[A\u001b[A\n",
      "\n",
      "2800it [01:00, 195.03it/s]\u001b[A\u001b[A\n",
      "\n",
      "2837it [01:00, 227.03it/s]\u001b[A\u001b[A\n",
      "\n",
      "2868it [01:06, 18.11it/s] \u001b[A\u001b[A\n",
      "\n",
      "2890it [01:06, 22.21it/s]\u001b[A\u001b[A\n",
      "\n",
      "2916it [01:06, 28.84it/s]\u001b[A\u001b[A\n",
      "\n",
      "2935it [01:06, 34.54it/s]\u001b[A\u001b[A\n",
      "\n",
      "2983it [01:07, 58.14it/s]\u001b[A\u001b[A\n",
      "\n",
      "3035it [01:07, 90.64it/s]\u001b[A\u001b[A\n",
      "\n",
      "3069it [01:07, 100.49it/s]\u001b[A\u001b[A\n",
      "\n",
      "3125it [01:07, 147.95it/s]\u001b[A\u001b[A\n",
      "\n",
      "3162it [01:08, 118.42it/s]\u001b[A\u001b[A\n",
      "\n",
      "3212it [01:08, 160.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "3247it [01:08, 126.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "3274it [01:10, 46.68it/s] \u001b[A\u001b[A\n",
      "\n",
      "3293it [01:10, 44.98it/s]\u001b[A\u001b[A\n",
      "\n",
      "3327it [01:11, 61.97it/s]\u001b[A\u001b[A\n",
      "\n",
      "3375it [01:11, 93.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "3404it [01:11, 111.81it/s]\u001b[A\u001b[A\n",
      "\n",
      "3432it [01:11, 131.87it/s]\u001b[A\u001b[A\n",
      "\n",
      "3479it [01:11, 151.35it/s]\u001b[A\u001b[A\n",
      "\n",
      "3531it [01:11, 206.12it/s]\u001b[A\u001b[A\n",
      "\n",
      "3565it [01:11, 182.41it/s]\u001b[A\u001b[A\n",
      "\n",
      "3593it [01:15, 30.01it/s] \u001b[A\u001b[A\n",
      "\n",
      "3628it [01:15, 41.09it/s]\u001b[A\u001b[A\n",
      "\n",
      "3663it [01:15, 54.75it/s]\u001b[A\u001b[A\n",
      "\n",
      "3688it [01:15, 66.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "3716it [01:15, 75.92it/s]\u001b[A\u001b[A\n",
      "\n",
      "3737it [01:16, 74.86it/s]\u001b[A\u001b[A\n",
      "\n",
      "3762it [01:17, 35.26it/s]\u001b[A\u001b[A\n",
      "\n",
      "3786it [01:18, 45.14it/s]\u001b[A\u001b[A\n",
      "\n",
      "3811it [01:18, 59.06it/s]\u001b[A\u001b[A\n",
      "\n",
      "3848it [01:18, 86.61it/s]\u001b[A\u001b[A\n",
      "\n",
      "3896it [01:18, 130.82it/s]\u001b[A\u001b[A\n",
      "\n",
      "3926it [01:18, 153.58it/s]\u001b[A\u001b[A\n",
      "\n",
      "3956it [01:18, 118.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "4000it [01:19, 50.46it/s] \u001b[A\u001b[A\n",
      "\n",
      "100%|██████████| 2/2 [02:33<00:00, 76.51s/it]\u001b[A\n",
      "100%|██████████| 2/2 [04:39<00:00, 139.97s/it]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "45it [00:00, 448.65it/s]\u001b[A\u001b[A\n",
      "\n",
      "98it [00:00, 459.57it/s]\u001b[A\u001b[A\n",
      "\n",
      "149it [00:00, 481.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "198it [00:00, 177.86it/s]\u001b[A\u001b[A\n",
      "\n",
      "231it [00:00, 203.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "266it [00:07, 16.09it/s] \u001b[A\u001b[A\n",
      "\n",
      "311it [00:08, 24.24it/s]\u001b[A\u001b[A\n",
      "\n",
      "344it [00:08, 32.22it/s]\u001b[A\u001b[A\n",
      "\n",
      "380it [00:08, 44.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "412it [00:08, 57.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "444it [00:08, 72.09it/s]\u001b[A\u001b[A\n",
      "\n",
      "490it [00:08, 104.08it/s]\u001b[A\u001b[A\n",
      "\n",
      "534it [00:08, 139.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "572it [00:08, 155.23it/s]\u001b[A\u001b[A\n",
      "\n",
      "624it [00:09, 207.89it/s]\u001b[A\u001b[A\n",
      "\n",
      "667it [00:09, 245.60it/s]\u001b[A\u001b[A\n",
      "\n",
      "707it [00:10, 63.08it/s] \u001b[A\u001b[A\n",
      "\n",
      "800it [00:11, 71.68it/s]\u001b[A\u001b[A\n",
      "\n",
      " 50%|█████     | 1/2 [00:11<00:11, 11.16s/it]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "38it [00:00, 379.63it/s]\u001b[A\u001b[A\n",
      "\n",
      "76it [00:00, 127.46it/s]\u001b[A\u001b[A\n",
      "\n",
      "114it [00:00, 182.53it/s]\u001b[A\u001b[A\n",
      "\n",
      "150it [00:00, 224.74it/s]\u001b[A\u001b[A\n",
      "\n",
      "198it [00:00, 222.26it/s]\u001b[A\u001b[A\n",
      "\n",
      "238it [00:01, 260.69it/s]\u001b[A\u001b[A\n",
      "\n",
      "286it [00:01, 308.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "328it [00:01, 335.50it/s]\u001b[A\u001b[A\n",
      "\n",
      "372it [00:01, 361.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "415it [00:01, 380.40it/s]\u001b[A\u001b[A\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "456it [00:01, 266.90it/s]\u001b[A\u001b[A\n",
      "\n",
      "504it [00:01, 312.57it/s]\u001b[A\u001b[A\n",
      "\n",
      "542it [00:04, 43.01it/s] \u001b[A\u001b[A\n",
      "\n",
      "569it [00:09, 18.98it/s]\u001b[A\u001b[A\n",
      "\n",
      "610it [00:09, 26.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "640it [00:09, 34.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "685it [00:09, 50.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "714it [00:09, 57.90it/s]\u001b[A\u001b[A\n",
      "\n",
      "738it [00:16, 13.11it/s]\u001b[A\u001b[A\n",
      "\n",
      "800it [00:16, 47.91it/s]\u001b[A\u001b[A\n",
      "\n",
      "100%|██████████| 2/2 [00:27<00:00, 13.93s/it]\u001b[A\n",
      " 50%|█████     | 1/2 [00:27<00:27, 27.87s/it]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "58it [00:00, 562.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "115it [00:00, 325.51it/s]\u001b[A\u001b[A\n",
      "\n",
      "153it [00:00, 331.43it/s]\u001b[A\u001b[A\n",
      "\n",
      "190it [00:00, 342.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "227it [00:00, 338.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "263it [00:00, 336.60it/s]\u001b[A\u001b[A\n",
      "\n",
      "307it [00:01, 250.11it/s]\u001b[A\u001b[A\n",
      "\n",
      "336it [00:08, 15.92it/s] \u001b[A\u001b[A\n",
      "\n",
      "364it [00:08, 20.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "408it [00:08, 31.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "446it [00:08, 44.34it/s]\u001b[A\u001b[A\n",
      "\n",
      "497it [00:08, 67.17it/s]\u001b[A\u001b[A\n",
      "\n",
      "533it [00:08, 83.00it/s]\u001b[A\u001b[A\n",
      "\n",
      "565it [00:09, 67.27it/s]\u001b[A\u001b[A\n",
      "\n",
      "601it [00:09, 87.62it/s]\u001b[A\u001b[A\n",
      "\n",
      "651it [00:09, 125.09it/s]\u001b[A\u001b[A\n",
      "\n",
      "684it [00:09, 147.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "732it [00:11, 50.25it/s] \u001b[A\u001b[A\n",
      "\n",
      "756it [00:12, 44.19it/s]\u001b[A\u001b[A\n",
      "\n",
      "800it [00:12, 62.00it/s]\u001b[A\u001b[A\n",
      "\n",
      " 50%|█████     | 1/2 [00:12<00:12, 12.91s/it]\u001b[A\n",
      "\n",
      "0it [00:00, ?it/s]\u001b[A\u001b[A\n",
      "\n",
      "47it [00:00, 465.30it/s]\u001b[A\u001b[A\n",
      "\n",
      "94it [00:06, 12.72it/s] \u001b[A\u001b[A\n",
      "\n",
      "114it [00:06, 15.67it/s]\u001b[A\u001b[A\n",
      "\n",
      "127it [00:06, 18.32it/s]\u001b[A\u001b[A\n",
      "\n",
      "170it [00:07, 33.73it/s]\u001b[A\u001b[A\n",
      "\n",
      "226it [00:07, 60.38it/s]\u001b[A\u001b[A\n",
      "\n",
      "258it [00:07, 77.94it/s]\u001b[A\u001b[A\n",
      "\n",
      "289it [00:07, 82.99it/s]\u001b[A\u001b[A\n",
      "\n",
      "315it [00:07, 99.98it/s]\u001b[A\u001b[A\n",
      "\n",
      "340it [00:08, 96.96it/s]\u001b[A\u001b[A\n",
      "\n",
      "372it [00:08, 124.16it/s]\u001b[A\u001b[A\n",
      "\n",
      "397it [00:08, 138.86it/s]\u001b[A\u001b[A\n",
      "\n",
      "420it [00:11, 23.04it/s] \u001b[A\u001b[A\n",
      "\n",
      "443it [00:11, 28.55it/s]\u001b[A\u001b[A\n",
      "\n",
      "467it [00:19,  8.91it/s]\u001b[A\u001b[A\n",
      "\n",
      "491it [00:19, 12.32it/s]\u001b[A\u001b[A\n",
      "\n",
      "515it [00:19, 16.88it/s]\u001b[A\u001b[A\n",
      "\n",
      "529it [00:19, 20.15it/s]\u001b[A\u001b[A\n",
      "\n",
      "550it [00:19, 27.36it/s]\u001b[A\u001b[A\n",
      "\n",
      "585it [00:19, 43.60it/s]\u001b[A\u001b[A\n",
      "\n",
      "617it [00:20, 58.56it/s]\u001b[A\u001b[A\n",
      "\n",
      "654it [00:20, 84.68it/s]\u001b[A\u001b[A\n",
      "\n",
      "679it [00:27, 12.29it/s]\u001b[A\u001b[A\n",
      "\n",
      "700it [00:27, 15.46it/s]\u001b[A\u001b[A\n",
      "\n",
      "715it [00:31,  9.47it/s]\u001b[A\u001b[A\n",
      "\n",
      "727it [00:31, 11.39it/s]\u001b[A\u001b[A\n",
      "\n",
      "752it [00:31, 17.16it/s]\u001b[A\u001b[A\n",
      "\n",
      "800it [00:31, 25.02it/s]\u001b[A\u001b[A\n",
      "\n",
      "100%|██████████| 2/2 [00:44<00:00, 22.44s/it]\u001b[A\n",
      "100%|██████████| 2/2 [01:12<00:00, 36.38s/it]\n"
     ]
    }
   ],
   "source": [
    "num_imgs = {'train':num_train+num_val, 'test':num_test}  # number of imgs in subfolders\n",
    "\n",
    "# append labels to dataset dictionaries\n",
    "for dataset in [top_dataset,bot_dataset]:\n",
    "    if dataset['name']=='cifar10':\n",
    "        dataset['train_labels'] = torch.tensor(cifar10_train.targets)\n",
    "        dataset['test_labels'] = torch.tensor(cifar10_test.targets)\n",
    "    elif dataset['name']=='mnist':\n",
    "        dataset['train_labels'] = mnist_train.train_labels\n",
    "        dataset['test_labels'] = mnist_test.train_labels\n",
    "    elif dataset['name']=='kmnist' or dataset['name']=='kmnistPatch':\n",
    "        dataset['train_labels'] = kmnist_train.train_labels\n",
    "        dataset['test_labels'] = kmnist_test.train_labels\n",
    "    elif dataset['name']=='fmnist':\n",
    "        dataset['train_labels'] = fmnist_train.train_labels\n",
    "        dataset['test_labels'] = fmnist_test.train_labels\n",
    "    elif dataset['name']=='svhn':\n",
    "        dataset['train_labels'] = torch.tensor(svhn_train.labels)\n",
    "        dataset['test_labels'] = torch.tensor(svhn_test.labels)\n",
    "\n",
    "# populate the various train/test folders and domino label combinations\n",
    "for partition in ['train','test']:    \n",
    "    for bot_label in tqdm(range(2)):\n",
    "        for top_label in tqdm(range(2)):            \n",
    "            bot_inds = (bot_dataset[f'{partition}_labels']==bot_label).nonzero().squeeze()  # indices for bottom images\n",
    "            top_inds = (top_dataset[f'{partition}_labels']==top_label).nonzero().squeeze()  # indices for top images\n",
    "            Path(os.path.join(save_dir,f'{partition}/top{top_label}_bot{bot_label}/')).mkdir(parents=True, exist_ok=True)\n",
    "            \n",
    "            for file_idx,(bot_idx,top_idx) in tqdm(enumerate(zip(bot_inds[:num_imgs[partition]],top_inds[:num_imgs[partition]]))):\n",
    "                if bot_dataset['name'] == 'kmnistPatch' and bot_label==1:\n",
    "                    temp_img = bot_dataset[f'{partition}_data'][bot_idx][0]\n",
    "                    temp_img[:,:5,:5]=1.0\n",
    "                    img = torch.hstack((top_dataset[f'{partition}_data'][top_idx][0],temp_img))\n",
    "                else:\n",
    "                    img = torch.hstack((top_dataset[f'{partition}_data'][top_idx][0],bot_dataset[f'{partition}_data'][bot_idx][0]))\n",
    "                final_path = os.path.join(save_dir,f'{partition}/top{top_label}_bot{bot_label}/{file_idx}.jpg')\n",
    "                plt.imsave(final_path,np.array(img.permute(1,2,0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate CSV file for corresponding dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code assumes the heirarchy root_dir/{train/test}/top{0-9}_bot{0-9}/{0-num_train/0-num_test}.jpg\n",
    "# code currently only supports spurious correlation = 1.0\n",
    "\n",
    "val_inds = random.sample(range(num_train+num_val), num_val)\n",
    "all_inds = list(range(num_train+num_val))\n",
    "train_inds = [x for x in all_inds if x not in val_inds]\n",
    "csv_list = [] # all rows will be appended here\n",
    "csv_cols = ['path','top','bottom','val_train_test_split','test_nontest_split','all_zeros']\n",
    "\n",
    "for lbl_idx in range(2):\n",
    "    # add validation info to csv_list\n",
    "    for val_idx in val_inds:\n",
    "        path = os.path.join(root_dir,f'train/top{lbl_idx}_bot{lbl_idx}/{val_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 0\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add train info to csv_list\n",
    "    for tr_idx in train_inds:\n",
    "        path = os.path.join(root_dir,f'train/top{lbl_idx}_bot{lbl_idx}/{tr_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 1\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add test info to csv_list\n",
    "    for test_idx in range(num_test):\n",
    "        path = os.path.join(root_dir,f'test/top{lbl_idx}_bot{lbl_idx}/{test_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 2\n",
    "        test_nontest_split = 0\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "\n",
    "df = pd.DataFrame(csv_list)\n",
    "df.columns = csv_cols\n",
    "df.to_csv(full_save_path,index=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Top Blank Dataset Domino"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "num_train = 3500\n",
    "num_val = 500\n",
    "num_test = 800\n",
    "spur_corr = 1.0\n",
    "save_dir = root_dir = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_blank_bot_svhn/'\n",
    "full_save_path = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_blank_bot_svhn_2class_ro_1p0.csv'\n",
    "bot_dataset = {'name':'svhn', 'train_data':svhn_train, 'test_data':svhn_test}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "1it [00:00,  7.10it/s]\u001b[A\n",
      "7it [00:03,  2.18it/s]\u001b[A\n",
      "14it [00:03,  5.16it/s]\u001b[A\n",
      "24it [00:03, 10.78it/s]\u001b[A\n",
      "43it [00:03, 24.76it/s]\u001b[A\n",
      "73it [00:03, 52.35it/s]\u001b[A\n",
      "116it [00:03, 98.70it/s]\u001b[A\n",
      "141it [00:03, 118.27it/s]\u001b[A\n",
      "167it [00:03, 125.36it/s]\u001b[A\n",
      "214it [00:04, 140.65it/s]\u001b[A\n",
      "234it [00:04, 109.02it/s]\u001b[A\n",
      "270it [00:04, 144.88it/s]\u001b[A\n",
      "292it [00:07, 26.02it/s] \u001b[A\n",
      "337it [00:07, 42.31it/s]\u001b[A\n",
      "362it [00:07, 51.81it/s]\u001b[A\n",
      "404it [00:08, 76.70it/s]\u001b[A\n",
      "439it [00:08, 100.71it/s]\u001b[A\n",
      "469it [00:09, 47.25it/s] \u001b[A\n",
      "491it [00:10, 47.56it/s]\u001b[A\n",
      "513it [00:11, 36.99it/s]\u001b[A\n",
      "536it [00:11, 37.44it/s]\u001b[A\n",
      "561it [00:18,  9.96it/s]\u001b[A\n",
      "568it [00:18, 10.83it/s]\u001b[A\n",
      "616it [00:19, 21.39it/s]\u001b[A\n",
      "651it [00:19, 31.63it/s]\u001b[A\n",
      "683it [00:19, 43.92it/s]\u001b[A\n",
      "709it [00:22, 20.29it/s]\u001b[A\n",
      "748it [00:22, 30.89it/s]\u001b[A\n",
      "813it [00:22, 55.10it/s]\u001b[A\n",
      "864it [00:22, 79.22it/s]\u001b[A\n",
      "925it [00:22, 116.61it/s]\u001b[A\n",
      "972it [00:22, 142.13it/s]\u001b[A\n",
      "1014it [00:23, 170.44it/s]\u001b[A\n",
      "1055it [00:23, 195.12it/s]\u001b[A\n",
      "1093it [00:24, 99.42it/s] \u001b[A\n",
      "1121it [00:27, 27.37it/s]\u001b[A\n",
      "1165it [00:27, 39.58it/s]\u001b[A\n",
      "1201it [00:28, 52.54it/s]\u001b[A\n",
      "1254it [00:28, 78.03it/s]\u001b[A\n",
      "1303it [00:28, 106.85it/s]\u001b[A\n",
      "1341it [00:30, 45.67it/s] \u001b[A\n",
      "1368it [00:30, 50.53it/s]\u001b[A\n",
      "1411it [00:30, 70.85it/s]\u001b[A\n",
      "1445it [00:30, 89.80it/s]\u001b[A\n",
      "1474it [00:31, 107.73it/s]\u001b[A\n",
      "1513it [00:31, 111.30it/s]\u001b[A\n",
      "1537it [00:35, 20.71it/s] \u001b[A\n",
      "1554it [00:36, 23.09it/s]\u001b[A\n",
      "1589it [00:38, 19.85it/s]\u001b[A\n",
      "1643it [00:38, 34.15it/s]\u001b[A\n",
      "1685it [00:38, 46.72it/s]\u001b[A\n",
      "1720it [00:38, 59.78it/s]\u001b[A\n",
      "1745it [00:39, 70.75it/s]\u001b[A\n",
      "1768it [00:39, 83.67it/s]\u001b[A\n",
      "1800it [00:39, 106.00it/s]\u001b[A\n",
      "1855it [00:39, 162.76it/s]\u001b[A\n",
      "1892it [00:39, 194.47it/s]\u001b[A\n",
      "1928it [00:39, 222.01it/s]\u001b[A\n",
      "1972it [00:39, 208.55it/s]\u001b[A\n",
      "2017it [00:39, 252.86it/s]\u001b[A\n",
      "2051it [00:40, 213.44it/s]\u001b[A\n",
      "2080it [00:40, 211.39it/s]\u001b[A\n",
      "2118it [00:40, 244.87it/s]\u001b[A\n",
      "2148it [00:42, 54.19it/s] \u001b[A\n",
      "2170it [00:42, 54.66it/s]\u001b[A\n",
      "2187it [00:42, 58.30it/s]\u001b[A\n",
      "2222it [00:42, 81.80it/s]\u001b[A\n",
      "2265it [00:42, 117.62it/s]\u001b[A\n",
      "2316it [00:43, 168.95it/s]\u001b[A\n",
      "2349it [00:43, 130.33it/s]\u001b[A\n",
      "2374it [00:44, 94.09it/s] \u001b[A\n",
      "2393it [00:53,  9.24it/s]\u001b[A\n",
      "2407it [00:54,  9.77it/s]\u001b[A\n",
      "2426it [00:54, 12.74it/s]\u001b[A\n",
      "2463it [00:54, 21.02it/s]\u001b[A\n",
      "2518it [00:55, 38.03it/s]\u001b[A\n",
      "2548it [00:55, 44.51it/s]\u001b[A\n",
      "2571it [00:55, 52.31it/s]\u001b[A\n",
      "2629it [00:55, 87.82it/s]\u001b[A\n",
      "2659it [00:55, 102.94it/s]\u001b[A\n",
      "2688it [00:56, 110.43it/s]\u001b[A\n",
      "2712it [00:56, 121.13it/s]\u001b[A\n",
      "2734it [00:56, 132.46it/s]\u001b[A\n",
      "2760it [01:04, 10.54it/s] \u001b[A\n",
      "2784it [01:06, 10.53it/s]\u001b[A\n",
      "2795it [01:07, 11.61it/s]\u001b[A\n",
      "2836it [01:07, 20.46it/s]\u001b[A\n",
      "2875it [01:07, 30.86it/s]\u001b[A\n",
      "2932it [01:07, 52.89it/s]\u001b[A\n",
      "2962it [01:07, 60.92it/s]\u001b[A\n",
      "3006it [01:07, 86.42it/s]\u001b[A\n",
      "3064it [01:08, 118.17it/s]\u001b[A\n",
      "3094it [01:15, 16.67it/s] \u001b[A\n",
      "3115it [01:15, 19.88it/s]\u001b[A\n",
      "3142it [01:15, 25.90it/s]\u001b[A\n",
      "3163it [01:15, 29.73it/s]\u001b[A\n",
      "3204it [01:16, 45.83it/s]\u001b[A\n",
      "3228it [01:23, 11.35it/s]\u001b[A\n",
      "3276it [01:23, 18.84it/s]\u001b[A\n",
      "3302it [01:23, 24.15it/s]\u001b[A\n",
      "3327it [01:23, 30.85it/s]\u001b[A\n",
      "3373it [01:23, 48.79it/s]\u001b[A\n",
      "3411it [01:23, 67.17it/s]\u001b[A\n",
      "3448it [01:23, 89.64it/s]\u001b[A\n",
      "3485it [01:23, 116.50it/s]\u001b[A\n",
      "3520it [01:24, 132.81it/s]\u001b[A\n",
      "3551it [01:24, 138.12it/s]\u001b[A\n",
      "3577it [01:31, 13.31it/s] \u001b[A\n",
      "3622it [01:31, 20.96it/s]\u001b[A\n",
      "3654it [01:31, 28.16it/s]\u001b[A\n",
      "3703it [01:32, 43.66it/s]\u001b[A\n",
      "3737it [01:32, 56.62it/s]\u001b[A\n",
      "3769it [01:32, 71.30it/s]\u001b[A\n",
      "3799it [01:47,  7.06it/s]\u001b[A\n",
      "3801it [01:47,  6.91it/s]\u001b[A\n",
      "3829it [01:54,  5.56it/s]\u001b[A\n",
      "3872it [01:54,  9.49it/s]\u001b[A\n",
      "3901it [01:55, 12.51it/s]\u001b[A\n",
      "3920it [01:55, 15.29it/s]\u001b[A\n",
      "3949it [01:55, 21.68it/s]\u001b[A\n",
      "3967it [01:55, 25.97it/s]\u001b[A\n",
      "3982it [01:56, 29.62it/s]\u001b[A\n",
      "4000it [02:03, 32.31it/s]\u001b[A\n",
      " 50%|█████     | 1/2 [02:03<02:03, 123.81s/it]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "8it [00:00, 62.49it/s]\u001b[A\n",
      "32it [00:00, 107.12it/s]\u001b[A\n",
      "55it [00:00, 114.15it/s]\u001b[A\n",
      "71it [00:00, 73.39it/s] \u001b[A\n",
      "91it [00:11,  4.77it/s]\u001b[A\n",
      "129it [00:11,  9.56it/s]\u001b[A\n",
      "159it [00:11, 14.70it/s]\u001b[A\n",
      "180it [00:11, 19.20it/s]\u001b[A\n",
      "209it [00:11, 28.08it/s]\u001b[A\n",
      "255it [00:12, 47.58it/s]\u001b[A\n",
      "297it [00:12, 70.04it/s]\u001b[A\n",
      "343it [00:12, 101.10it/s]\u001b[A\n",
      "379it [00:12, 123.63it/s]\u001b[A\n",
      "421it [00:12, 160.02it/s]\u001b[A\n",
      "461it [00:12, 195.76it/s]\u001b[A\n",
      "498it [00:12, 221.04it/s]\u001b[A\n",
      "540it [00:20, 16.65it/s] \u001b[A\n",
      "565it [00:28,  8.29it/s]\u001b[A\n",
      "601it [00:29, 11.75it/s]\u001b[A\n",
      "641it [00:29, 17.18it/s]\u001b[A\n",
      "667it [00:30, 18.37it/s]\u001b[A\n",
      "686it [00:31, 17.37it/s]\u001b[A\n",
      "712it [00:31, 23.40it/s]\u001b[A\n",
      "730it [00:31, 28.49it/s]\u001b[A\n",
      "756it [00:32, 36.33it/s]\u001b[A\n",
      "801it [00:32, 60.12it/s]\u001b[A\n",
      "825it [00:32, 68.52it/s]\u001b[A\n",
      "860it [00:32, 94.12it/s]\u001b[A\n",
      "894it [00:32, 121.24it/s]\u001b[A\n",
      "941it [00:34, 52.40it/s] \u001b[A\n",
      "965it [00:35, 40.86it/s]\u001b[A\n",
      "989it [00:35, 42.79it/s]\u001b[A\n",
      "1006it [00:35, 49.53it/s]\u001b[A\n",
      "1019it [00:36, 52.10it/s]\u001b[A\n",
      "1065it [00:36, 89.31it/s]\u001b[A\n",
      "1112it [00:36, 133.23it/s]\u001b[A\n",
      "1141it [00:42, 16.13it/s] \u001b[A\n",
      "1184it [00:42, 24.63it/s]\u001b[A\n",
      "1212it [00:43, 23.83it/s]\u001b[A\n",
      "1232it [00:44, 28.66it/s]\u001b[A\n",
      "1250it [00:44, 34.13it/s]\u001b[A\n",
      "1318it [00:44, 68.45it/s]\u001b[A\n",
      "1364it [00:44, 96.08it/s]\u001b[A\n",
      "1402it [00:44, 121.49it/s]\u001b[A\n",
      "1445it [00:44, 156.82it/s]\u001b[A\n",
      "1483it [00:44, 185.05it/s]\u001b[A\n",
      "1537it [00:44, 242.35it/s]\u001b[A\n",
      "1579it [00:44, 246.46it/s]\u001b[A\n",
      "1616it [00:47, 44.38it/s] \u001b[A\n",
      "1643it [00:47, 54.07it/s]\u001b[A\n",
      "1670it [00:47, 66.72it/s]\u001b[A\n",
      "1698it [00:48, 82.26it/s]\u001b[A\n",
      "1739it [00:48, 114.39it/s]\u001b[A\n",
      "1781it [00:48, 151.83it/s]\u001b[A\n",
      "1816it [00:48, 181.17it/s]\u001b[A\n",
      "1851it [00:48, 210.54it/s]\u001b[A\n",
      "1886it [00:51, 34.95it/s] \u001b[A\n",
      "1922it [00:51, 48.03it/s]\u001b[A\n",
      "1950it [00:51, 59.10it/s]\u001b[A\n",
      "1995it [00:51, 79.61it/s]\u001b[A\n",
      "2019it [00:52, 87.05it/s]\u001b[A\n",
      "2040it [00:52, 90.03it/s]\u001b[A\n",
      "2085it [00:52, 131.61it/s]\u001b[A\n",
      "2111it [00:52, 149.35it/s]\u001b[A\n",
      "2137it [00:59, 12.71it/s] \u001b[A\n",
      "2161it [00:59, 16.60it/s]\u001b[A\n",
      "2195it [01:00, 24.54it/s]\u001b[A\n",
      "2245it [01:00, 40.66it/s]\u001b[A\n",
      "2276it [01:00, 50.99it/s]\u001b[A\n",
      "2315it [01:00, 70.46it/s]\u001b[A\n",
      "2376it [01:00, 111.93it/s]\u001b[A\n",
      "2422it [01:00, 146.51it/s]\u001b[A\n",
      "2464it [01:00, 180.34it/s]\u001b[A\n",
      "2505it [01:00, 208.75it/s]\u001b[A\n",
      "2544it [01:01, 231.27it/s]\u001b[A\n",
      "2585it [01:01, 265.22it/s]\u001b[A\n",
      "2623it [01:03, 54.91it/s] \u001b[A\n",
      "2662it [01:03, 72.86it/s]\u001b[A\n",
      "2692it [01:03, 88.32it/s]\u001b[A\n",
      "2733it [01:03, 117.42it/s]\u001b[A\n",
      "2804it [01:03, 186.38it/s]\u001b[A\n",
      "2853it [01:03, 229.59it/s]\u001b[A\n",
      "2906it [01:03, 279.97it/s]\u001b[A\n",
      "2954it [01:03, 314.90it/s]\u001b[A\n",
      "3001it [01:04, 246.33it/s]\u001b[A\n",
      "3039it [01:04, 264.31it/s]\u001b[A\n",
      "3078it [01:04, 284.57it/s]\u001b[A\n",
      "3115it [01:04, 301.19it/s]\u001b[A\n",
      "3152it [01:04, 315.17it/s]\u001b[A\n",
      "3209it [01:04, 332.81it/s]\u001b[A\n",
      "3255it [01:04, 362.83it/s]\u001b[A\n",
      "3295it [01:08, 37.37it/s] \u001b[A\n",
      "3348it [01:08, 54.73it/s]\u001b[A\n",
      "3399it [01:08, 76.53it/s]\u001b[A\n",
      "3439it [01:08, 96.62it/s]\u001b[A\n",
      "3478it [01:10, 54.87it/s]\u001b[A\n",
      "3506it [01:11, 42.25it/s]\u001b[A\n",
      "3532it [01:11, 51.79it/s]\u001b[A\n",
      "3577it [01:11, 75.48it/s]\u001b[A\n",
      "3629it [01:11, 109.73it/s]\u001b[A\n",
      "3664it [01:12, 111.74it/s]\u001b[A\n",
      "3692it [01:20, 13.23it/s] \u001b[A\n",
      "3737it [01:20, 19.91it/s]\u001b[A\n",
      "3770it [01:20, 26.53it/s]\u001b[A\n",
      "3813it [01:20, 38.50it/s]\u001b[A\n",
      "3856it [01:22, 30.79it/s]\u001b[A\n",
      "3880it [01:23, 32.31it/s]\u001b[A\n",
      "3904it [01:23, 40.10it/s]\u001b[A\n",
      "3940it [01:23, 56.23it/s]\u001b[A\n",
      "4000it [01:24, 47.54it/s]\u001b[A\n",
      "100%|██████████| 2/2 [03:27<00:00, 103.98s/it]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "42it [00:00, 419.45it/s]\u001b[A\n",
      "86it [00:00, 430.49it/s]\u001b[A\n",
      "130it [00:00, 280.35it/s]\u001b[A\n",
      "172it [00:00, 321.09it/s]\u001b[A\n",
      "209it [00:00, 303.96it/s]\u001b[A\n",
      "242it [00:00, 247.12it/s]\u001b[A\n",
      "270it [00:04, 28.52it/s] \u001b[A\n",
      "315it [00:04, 44.05it/s]\u001b[A\n",
      "372it [00:04, 70.30it/s]\u001b[A\n",
      "409it [00:04, 88.59it/s]\u001b[A\n",
      "448it [00:04, 113.80it/s]\u001b[A\n",
      "483it [00:04, 137.28it/s]\u001b[A\n",
      "532it [00:08, 32.70it/s] \u001b[A\n",
      "569it [00:08, 43.21it/s]\u001b[A\n",
      "622it [00:08, 64.57it/s]\u001b[A\n",
      "656it [00:13, 22.51it/s]\u001b[A\n",
      "680it [00:13, 23.84it/s]\u001b[A\n",
      "700it [00:14, 26.06it/s]\u001b[A\n",
      "724it [00:15, 26.31it/s]\u001b[A\n",
      "735it [00:17, 15.72it/s]\u001b[A\n",
      "748it [00:22,  8.37it/s]\u001b[A\n",
      "772it [00:23, 10.90it/s]\u001b[A\n",
      "800it [00:23, 34.33it/s]\u001b[A\n",
      " 50%|█████     | 1/2 [00:23<00:23, 23.31s/it]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "40it [00:00, 397.21it/s]\u001b[A\n",
      "80it [00:00, 386.88it/s]\u001b[A\n",
      "119it [00:00, 344.24it/s]\u001b[A\n",
      "173it [00:00, 406.69it/s]\u001b[A\n",
      "215it [00:00, 309.75it/s]\u001b[A\n",
      "250it [00:00, 310.26it/s]\u001b[A\n",
      "284it [00:01, 132.62it/s]\u001b[A\n",
      "309it [00:01, 101.90it/s]\u001b[A\n",
      "346it [00:01, 133.62it/s]\u001b[A\n",
      "384it [00:02, 167.68it/s]\u001b[A\n",
      "412it [00:03, 62.42it/s] \u001b[A\n",
      "453it [00:03, 87.73it/s]\u001b[A\n",
      "478it [00:03, 93.79it/s]\u001b[A\n",
      "500it [00:04, 46.76it/s]\u001b[A\n",
      "565it [00:04, 85.47it/s]\u001b[A\n",
      "595it [00:05, 62.99it/s]\u001b[A\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "617it [00:06, 70.04it/s]\u001b[A\n",
      "663it [00:07, 40.26it/s]\u001b[A\n",
      "686it [00:08, 46.68it/s]\u001b[A\n",
      "700it [00:08, 39.60it/s]\u001b[A\n",
      "735it [00:09, 48.62it/s]\u001b[A\n",
      "800it [00:09, 84.74it/s]\u001b[A\n",
      "100%|██████████| 2/2 [00:32<00:00, 16.38s/it]\n"
     ]
    }
   ],
   "source": [
    "num_imgs = {'train':num_train+num_val, 'test':num_test}  # number of imgs in subfolders\n",
    "\n",
    "# append labels to dataset dictionaries\n",
    "for dataset in [bot_dataset]:\n",
    "    if dataset['name']=='cifar10':\n",
    "        dataset['train_labels'] = torch.tensor(cifar10_train.targets)\n",
    "        dataset['test_labels'] = torch.tensor(cifar10_test.targets)\n",
    "    elif dataset['name']=='mnist':\n",
    "        dataset['train_labels'] = mnist_train.train_labels\n",
    "        dataset['test_labels'] = mnist_test.train_labels\n",
    "    elif dataset['name']=='kmnist' or dataset['name']=='kmnistPatch':\n",
    "        dataset['train_labels'] = kmnist_train.train_labels\n",
    "        dataset['test_labels'] = kmnist_test.train_labels\n",
    "    elif dataset['name']=='fmnist':\n",
    "        dataset['train_labels'] = fmnist_train.train_labels\n",
    "        dataset['test_labels'] = fmnist_test.train_labels\n",
    "    elif dataset['name']=='svhn':\n",
    "        dataset['train_labels'] = torch.tensor(svhn_train.labels)\n",
    "        dataset['test_labels'] = torch.tensor(svhn_test.labels)\n",
    "\n",
    "# populate the various train/test folders and domino label combinations\n",
    "for partition in ['train','test']:    \n",
    "    for bot_label in tqdm(range(2)):           \n",
    "        bot_inds = (bot_dataset[f'{partition}_labels']==bot_label).nonzero().squeeze()  # indices for bottom images\n",
    "        Path(os.path.join(save_dir,f'{partition}/{bot_label}/')).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        for file_idx,bot_idx in tqdm(enumerate(bot_inds[:num_imgs[partition]])):\n",
    "            if bot_dataset['name'] == 'kmnistPatch' and bot_label==1:\n",
    "                temp_img = bot_dataset[f'{partition}_data'][bot_idx][0]\n",
    "                temp_img[:,:5,:5]=1.0\n",
    "                img = torch.hstack((torch.zeros((3,32,32)),temp_img))\n",
    "            else:\n",
    "                img = torch.hstack((torch.zeros((3,32,32)),bot_dataset[f'{partition}_data'][bot_idx][0]))\n",
    "            final_path = os.path.join(save_dir,f'{partition}/{bot_label}/{file_idx}.jpg')\n",
    "            plt.imsave(final_path,np.array(img.permute(1,2,0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate CSV File for corresponding Blank Domino Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code assumes the heirarchy root_dir/{train/test}/top{0-9}_bot{0-9}/{0-num_train/0-num_test}.jpg\n",
    "# code currently only supports spurious correlation = 1.0\n",
    "\n",
    "val_inds = random.sample(range(num_train+num_val), num_val)\n",
    "all_inds = list(range(num_train+num_val))\n",
    "train_inds = [x for x in all_inds if x not in val_inds]\n",
    "csv_list = [] # all rows will be appended here\n",
    "csv_cols = ['path','top','bottom','val_train_test_split','test_nontest_split','all_zeros']\n",
    "\n",
    "for lbl_idx in range(2):\n",
    "    # add validation info to csv_list\n",
    "    for val_idx in val_inds:\n",
    "        path = os.path.join(root_dir,f'train/{lbl_idx}/{val_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 0\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add train info to csv_list\n",
    "    for tr_idx in train_inds:\n",
    "        path = os.path.join(root_dir,f'train/{lbl_idx}/{tr_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 1\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add test info to csv_list\n",
    "    for test_idx in range(num_test):\n",
    "        path = os.path.join(root_dir,f'test/{lbl_idx}/{test_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 2\n",
    "        test_nontest_split = 0\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "\n",
    "df = pd.DataFrame(csv_list)\n",
    "df.columns = csv_cols\n",
    "df.to_csv(full_save_path,index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Bottom Blank Dataset Domino"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "num_train = 3500\n",
    "num_val = 500\n",
    "num_test = 800\n",
    "spur_corr = 1.0\n",
    "save_dir = root_dir = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_svhn_bot_blank/'\n",
    "full_save_path = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_svhn_bot_blank_2class_ro_1p0.csv'\n",
    "top_dataset = {'name':'svhn', 'train_data':svhn_train, 'test_data':svhn_test}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "12it [00:00, 37.92it/s]\u001b[A\n",
      "35it [00:04,  7.55it/s]\u001b[A\n",
      "70it [00:04, 18.91it/s]\u001b[A\n",
      "87it [00:04, 25.42it/s]\u001b[A\n",
      "107it [00:06, 16.17it/s]\u001b[A\n",
      "130it [00:06, 22.25it/s]\u001b[A\n",
      "155it [00:07, 27.22it/s]\u001b[A\n",
      "188it [00:07, 43.07it/s]\u001b[A\n",
      "204it [00:07, 44.34it/s]\u001b[A\n",
      "226it [00:08, 48.30it/s]\u001b[A\n",
      "244it [00:08, 59.59it/s]\u001b[A\n",
      "278it [00:08, 90.15it/s]\u001b[A\n",
      "298it [00:08, 92.44it/s]\u001b[A\n",
      "322it [00:15, 10.36it/s]\u001b[A\n",
      "346it [00:15, 14.59it/s]\u001b[A\n",
      "384it [00:15, 23.93it/s]\u001b[A\n",
      "419it [00:15, 35.41it/s]\u001b[A\n",
      "451it [00:16, 46.13it/s]\u001b[A\n",
      "473it [00:16, 51.56it/s]\u001b[A\n",
      "491it [00:19, 20.93it/s]\u001b[A\n",
      "514it [00:19, 24.54it/s]\u001b[A\n",
      "562it [00:19, 40.95it/s]\u001b[A\n",
      "609it [00:20, 53.39it/s]\u001b[A\n",
      "646it [00:20, 72.02it/s]\u001b[A\n",
      "666it [00:26, 14.45it/s]\u001b[A\n",
      "680it [00:26, 16.32it/s]\u001b[A\n",
      "691it [00:27, 16.59it/s]\u001b[A\n",
      "727it [00:27, 27.50it/s]\u001b[A\n",
      "756it [00:27, 38.88it/s]\u001b[A\n",
      "799it [00:28, 46.65it/s]\u001b[A\n",
      "847it [00:28, 66.33it/s]\u001b[A\n",
      "875it [00:28, 81.72it/s]\u001b[A\n",
      "929it [00:28, 124.36it/s]\u001b[A\n",
      "962it [00:28, 147.09it/s]\u001b[A\n",
      "993it [00:36, 14.77it/s] \u001b[A\n",
      "1015it [00:36, 17.47it/s]\u001b[A\n",
      "1040it [00:36, 22.77it/s]\u001b[A\n",
      "1083it [00:37, 35.58it/s]\u001b[A\n",
      "1109it [00:38, 31.63it/s]\u001b[A\n",
      "1133it [00:39, 30.10it/s]\u001b[A\n",
      "1173it [00:39, 45.58it/s]\u001b[A\n",
      "1205it [00:39, 60.90it/s]\u001b[A\n",
      "1238it [00:39, 80.96it/s]\u001b[A\n",
      "1272it [00:39, 105.93it/s]\u001b[A\n",
      "1301it [00:39, 114.62it/s]\u001b[A\n",
      "1342it [00:39, 154.00it/s]\u001b[A\n",
      "1390it [00:39, 206.20it/s]\u001b[A\n",
      "1434it [00:39, 249.54it/s]\u001b[A\n",
      "1472it [00:40, 275.40it/s]\u001b[A\n",
      "1523it [00:40, 328.77it/s]\u001b[A\n",
      "1565it [00:47, 18.43it/s] \u001b[A\n",
      "1595it [00:47, 22.93it/s]\u001b[A\n",
      "1626it [00:47, 30.01it/s]\u001b[A\n",
      "1659it [00:48, 38.53it/s]\u001b[A\n",
      "1682it [00:48, 46.64it/s]\u001b[A\n",
      "1704it [00:48, 51.80it/s]\u001b[A\n",
      "1722it [00:51, 22.38it/s]\u001b[A\n",
      "1762it [00:51, 35.95it/s]\u001b[A\n",
      "1783it [00:51, 44.10it/s]\u001b[A\n",
      "1803it [00:51, 52.91it/s]\u001b[A\n",
      "1830it [00:51, 70.27it/s]\u001b[A\n",
      "1850it [00:52, 54.30it/s]\u001b[A\n",
      "1878it [00:52, 70.72it/s]\u001b[A\n",
      "1894it [00:52, 79.83it/s]\u001b[A\n",
      "1910it [00:53, 52.31it/s]\u001b[A\n",
      "1922it [00:53, 58.69it/s]\u001b[A\n",
      "1945it [00:55, 20.79it/s]\u001b[A\n",
      "1969it [00:55, 30.28it/s]\u001b[A\n",
      "2027it [00:55, 62.49it/s]\u001b[A\n",
      "2077it [00:55, 95.11it/s]\u001b[A\n",
      "2109it [00:56, 67.25it/s]\u001b[A\n",
      "2132it [00:56, 78.73it/s]\u001b[A\n",
      "2163it [00:56, 100.70it/s]\u001b[A\n",
      "2188it [00:57, 115.14it/s]\u001b[A\n",
      "2236it [00:57, 166.96it/s]\u001b[A\n",
      "2268it [00:57, 188.75it/s]\u001b[A\n",
      "2298it [00:57, 154.73it/s]\u001b[A\n",
      "2323it [00:59, 39.22it/s] \u001b[A\n",
      "2341it [00:59, 44.53it/s]\u001b[A\n",
      "2371it [00:59, 61.51it/s]\u001b[A\n",
      "2418it [01:00, 96.86it/s]\u001b[A\n",
      "2465it [01:00, 136.80it/s]\u001b[A\n",
      "2503it [01:00, 169.61it/s]\u001b[A\n",
      "2559it [01:00, 232.30it/s]\u001b[A\n",
      "2600it [01:01, 117.85it/s]\u001b[A\n",
      "2640it [01:01, 147.58it/s]\u001b[A\n",
      "2680it [01:01, 153.82it/s]\u001b[A\n",
      "2709it [01:04, 32.91it/s] \u001b[A\n",
      "2729it [01:04, 38.78it/s]\u001b[A\n",
      "2772it [01:04, 57.42it/s]\u001b[A\n",
      "2805it [01:04, 75.06it/s]\u001b[A\n",
      "2858it [01:05, 113.56it/s]\u001b[A\n",
      "2919it [01:05, 167.10it/s]\u001b[A\n",
      "2962it [01:05, 198.08it/s]\u001b[A\n",
      "3003it [01:12, 19.59it/s] \u001b[A\n",
      "3032it [01:12, 24.06it/s]\u001b[A\n",
      "3081it [01:12, 35.88it/s]\u001b[A\n",
      "3116it [01:13, 37.66it/s]\u001b[A\n",
      "3152it [01:13, 49.84it/s]\u001b[A\n",
      "3195it [01:13, 69.61it/s]\u001b[A\n",
      "3236it [01:13, 92.28it/s]\u001b[A\n",
      "3290it [01:13, 132.04it/s]\u001b[A\n",
      "3329it [01:16, 41.30it/s] \u001b[A\n",
      "3357it [01:16, 50.41it/s]\u001b[A\n",
      "3392it [01:16, 66.33it/s]\u001b[A\n",
      "3421it [01:17, 72.16it/s]\u001b[A\n",
      "3468it [01:17, 104.57it/s]\u001b[A\n",
      "3499it [01:21, 24.28it/s] \u001b[A\n",
      "3521it [01:23, 18.44it/s]\u001b[A\n",
      "3569it [01:23, 29.29it/s]\u001b[A\n",
      "3601it [01:23, 38.87it/s]\u001b[A\n",
      "3629it [01:23, 49.71it/s]\u001b[A\n",
      "3655it [01:23, 61.85it/s]\u001b[A\n",
      "3700it [01:24, 91.84it/s]\u001b[A\n",
      "3755it [01:24, 137.87it/s]\u001b[A\n",
      "3806it [01:24, 182.14it/s]\u001b[A\n",
      "3846it [01:24, 138.49it/s]\u001b[A\n",
      "3877it [01:24, 150.76it/s]\u001b[A\n",
      "3925it [01:25, 192.19it/s]\u001b[A\n",
      "3957it [01:25, 211.57it/s]\u001b[A\n",
      "4000it [01:25, 46.90it/s] \u001b[A\n",
      " 50%|█████     | 1/2 [01:25<01:25, 85.30s/it]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "40it [00:00, 397.80it/s]\u001b[A\n",
      "80it [00:07,  8.54it/s] \u001b[A\n",
      "113it [00:08, 14.19it/s]\u001b[A\n",
      "137it [00:10, 12.62it/s]\u001b[A\n",
      "177it [00:10, 20.99it/s]\u001b[A\n",
      "239it [00:10, 39.07it/s]\u001b[A\n",
      "287it [00:10, 57.28it/s]\u001b[A\n",
      "332it [00:10, 79.02it/s]\u001b[A\n",
      "373it [00:11, 102.58it/s]\u001b[A\n",
      "413it [00:11, 122.98it/s]\u001b[A\n",
      "455it [00:11, 156.38it/s]\u001b[A\n",
      "492it [00:19, 14.48it/s] \u001b[A\n",
      "526it [00:20, 19.42it/s]\u001b[A\n",
      "555it [00:21, 18.30it/s]\u001b[A\n",
      "592it [00:22, 25.85it/s]\u001b[A\n",
      "617it [00:22, 30.16it/s]\u001b[A\n",
      "637it [00:22, 36.44it/s]\u001b[A\n",
      "664it [00:22, 48.41it/s]\u001b[A\n",
      "691it [00:22, 63.19it/s]\u001b[A\n",
      "721it [00:22, 84.04it/s]\u001b[A\n",
      "746it [00:23, 93.36it/s]\u001b[A\n",
      "768it [00:23, 97.17it/s]\u001b[A\n",
      "806it [00:23, 134.50it/s]\u001b[A\n",
      "833it [00:23, 156.46it/s]\u001b[A\n",
      "858it [00:30, 11.76it/s] \u001b[A\n",
      "879it [00:30, 15.30it/s]\u001b[A\n",
      "934it [00:31, 28.63it/s]\u001b[A\n",
      "963it [00:35, 15.24it/s]\u001b[A\n",
      "1023it [00:35, 26.63it/s]\u001b[A\n",
      "1079it [00:35, 40.65it/s]\u001b[A\n",
      "1114it [00:35, 51.58it/s]\u001b[A\n",
      "1147it [00:35, 65.17it/s]\u001b[A\n",
      "1202it [00:35, 97.09it/s]\u001b[A\n",
      "1250it [00:36, 129.91it/s]\u001b[A\n",
      "1309it [00:36, 180.11it/s]\u001b[A\n",
      "1356it [00:38, 57.10it/s] \u001b[A\n",
      "1401it [00:38, 75.55it/s]\u001b[A\n",
      "1456it [00:38, 105.75it/s]\u001b[A\n",
      "1498it [00:38, 129.08it/s]\u001b[A\n",
      "1537it [00:38, 152.25it/s]\u001b[A\n",
      "1588it [00:38, 197.38it/s]\u001b[A\n",
      "1629it [00:39, 220.23it/s]\u001b[A\n",
      "1668it [00:39, 243.49it/s]\u001b[A\n",
      "1705it [00:39, 214.40it/s]\u001b[A\n",
      "1746it [00:39, 249.42it/s]\u001b[A\n",
      "1795it [00:39, 298.25it/s]\u001b[A\n",
      "1844it [00:39, 337.93it/s]\u001b[A\n",
      "1885it [00:39, 344.16it/s]\u001b[A\n",
      "1927it [00:39, 360.43it/s]\u001b[A\n",
      "1968it [00:40, 372.22it/s]\u001b[A\n",
      "2014it [00:40, 395.79it/s]\u001b[A\n",
      "2056it [00:47, 19.39it/s] \u001b[A\n",
      "2086it [00:47, 24.22it/s]\u001b[A\n",
      "2131it [00:47, 35.09it/s]\u001b[A\n",
      "2163it [00:47, 44.56it/s]\u001b[A\n",
      "2191it [00:47, 52.23it/s]\u001b[A\n",
      "2233it [00:48, 74.20it/s]\u001b[A\n",
      "2262it [00:51, 25.44it/s]\u001b[A\n",
      "2294it [00:51, 34.39it/s]\u001b[A\n",
      "2339it [00:51, 51.68it/s]\u001b[A\n",
      "2370it [00:52, 53.72it/s]\u001b[A\n",
      "2417it [00:59, 14.72it/s]\u001b[A\n",
      "2434it [00:59, 16.05it/s]\u001b[A\n",
      "2474it [01:00, 24.05it/s]\u001b[A\n",
      "2495it [01:02, 18.30it/s]\u001b[A\n",
      "2548it [01:02, 31.11it/s]\u001b[A\n",
      "2577it [01:02, 38.03it/s]\u001b[A\n",
      "2614it [01:02, 52.77it/s]\u001b[A\n",
      "2640it [01:02, 63.34it/s]\u001b[A\n",
      "2687it [01:03, 94.54it/s]\u001b[A\n",
      "2740it [01:03, 137.41it/s]\u001b[A\n",
      "2777it [01:03, 133.22it/s]\u001b[A\n",
      "2839it [01:03, 190.78it/s]\u001b[A\n",
      "2897it [01:03, 248.25it/s]\u001b[A\n",
      "2941it [01:03, 271.85it/s]\u001b[A\n",
      "2996it [01:04, 180.36it/s]\u001b[A\n",
      "3029it [01:04, 161.07it/s]\u001b[A\n",
      "3029it [01:21, 161.07it/s]\u001b[A\n",
      "3031it [02:14,  1.51it/s] \u001b[A\n",
      "3055it [02:15,  1.99it/s]\u001b[A\n",
      "3079it [02:15,  2.66it/s]\u001b[A\n",
      "3104it [02:15,  3.66it/s]\u001b[A\n",
      "3127it [02:15,  4.95it/s]\u001b[A\n",
      "3151it [02:17,  5.78it/s]\u001b[A\n",
      "3170it [02:19,  6.81it/s]\u001b[A\n",
      "3201it [02:19, 10.54it/s]\u001b[A\n",
      "3236it [02:19, 16.40it/s]\u001b[A\n",
      "3277it [02:19, 25.87it/s]\u001b[A\n",
      "3323it [02:19, 40.31it/s]\u001b[A\n",
      "3382it [02:19, 65.20it/s]\u001b[A\n",
      "3423it [02:30, 11.62it/s]\u001b[A\n",
      "3455it [02:30, 15.17it/s]\u001b[A\n",
      "3493it [02:30, 20.98it/s]\u001b[A\n",
      "3535it [02:30, 29.92it/s]\u001b[A\n",
      "3586it [02:30, 44.87it/s]\u001b[A\n",
      "3627it [02:30, 58.26it/s]\u001b[A\n",
      "3662it [02:31, 70.09it/s]\u001b[A\n",
      "3701it [02:31, 91.99it/s]\u001b[A\n",
      "3748it [02:31, 125.79it/s]\u001b[A\n",
      "3784it [02:31, 151.66it/s]\u001b[A\n",
      "3820it [02:31, 173.30it/s]\u001b[A\n",
      "3854it [02:34, 32.42it/s] \u001b[A\n",
      "3906it [02:35, 49.79it/s]\u001b[A\n",
      "3949it [02:35, 68.21it/s]\u001b[A\n",
      "4000it [02:35, 25.76it/s]\u001b[A\n",
      "100%|██████████| 2/2 [04:00<00:00, 120.30s/it]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "13it [00:03,  3.68it/s]\u001b[A\n",
      "31it [00:03, 10.43it/s]\u001b[A\n",
      "40it [00:03, 14.49it/s]\u001b[A\n",
      "84it [00:03, 43.35it/s]\u001b[A\n",
      "132it [00:04, 78.84it/s]\u001b[A\n",
      "160it [00:04, 100.53it/s]\u001b[A\n",
      "186it [00:07, 21.41it/s] \u001b[A\n",
      "213it [00:07, 29.48it/s]\u001b[A\n",
      "246it [00:07, 42.82it/s]\u001b[A\n",
      "271it [00:08, 53.47it/s]\u001b[A\n",
      "293it [00:08, 63.92it/s]\u001b[A\n",
      "313it [00:14, 10.73it/s]\u001b[A\n",
      "327it [00:14, 13.07it/s]\u001b[A\n",
      "356it [00:15, 20.19it/s]\u001b[A\n",
      "404it [00:15, 36.52it/s]\u001b[A\n",
      "434it [00:15, 49.22it/s]\u001b[A\n",
      "462it [00:15, 63.14it/s]\u001b[A\n",
      "489it [00:15, 67.00it/s]\u001b[A\n",
      "528it [00:15, 90.08it/s]\u001b[A\n",
      "582it [00:15, 138.49it/s]\u001b[A\n",
      "621it [00:16, 171.98it/s]\u001b[A\n",
      "660it [00:16, 206.83it/s]\u001b[A\n",
      "709it [00:16, 259.39it/s]\u001b[A\n",
      "749it [00:18, 63.86it/s] \u001b[A\n",
      "778it [00:19, 46.98it/s]\u001b[A\n",
      "800it [00:19, 41.27it/s]\u001b[A\n",
      " 50%|█████     | 1/2 [00:19<00:19, 19.39s/it]\n",
      "0it [00:00, ?it/s]\u001b[A\n",
      "35it [00:00, 349.83it/s]\u001b[A\n",
      "87it [00:00, 446.23it/s]\u001b[A\n",
      "132it [00:00, 369.55it/s]\u001b[A\n",
      "171it [00:07, 14.06it/s] \u001b[A\n",
      "207it [00:07, 20.38it/s]\u001b[A\n",
      "240it [00:08, 28.25it/s]\u001b[A\n",
      "274it [00:08, 38.84it/s]\u001b[A\n",
      "342it [00:08, 69.14it/s]\u001b[A\n",
      "384it [00:08, 89.86it/s]\u001b[A\n",
      "428it [00:08, 118.32it/s]\u001b[A\n",
      "481it [00:08, 160.08it/s]\u001b[A\n",
      "525it [00:08, 195.78it/s]\u001b[A\n",
      "569it [00:15, 19.26it/s] \u001b[A\n",
      "600it [00:23, 10.20it/s]\u001b[A\n",
      "625it [00:23, 12.73it/s]\u001b[A\n",
      "648it [00:24, 15.86it/s]\u001b[A\n",
      "669it [00:24, 18.63it/s]\u001b[A\n",
      "685it [00:29,  9.67it/s]\u001b[A\n",
      "697it [00:31,  9.01it/s]\u001b[A\n",
      "729it [00:31, 14.59it/s]\u001b[A\n",
      "749it [00:31, 19.18it/s]\u001b[A\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "769it [00:31, 25.40it/s]\u001b[A\n",
      "800it [00:31, 25.05it/s]\u001b[A\n",
      "100%|██████████| 2/2 [00:51<00:00, 25.66s/it]\n"
     ]
    }
   ],
   "source": [
    "num_imgs = {'train':num_train+num_val, 'test':num_test}  # number of imgs in subfolders\n",
    "\n",
    "# append labels to dataset dictionaries\n",
    "for dataset in [top_dataset]:\n",
    "    if dataset['name']=='cifar10':\n",
    "        dataset['train_labels'] = torch.tensor(cifar10_train.targets)\n",
    "        dataset['test_labels'] = torch.tensor(cifar10_test.targets)\n",
    "    elif dataset['name']=='mnist':\n",
    "        dataset['train_labels'] = mnist_train.train_labels\n",
    "        dataset['test_labels'] = mnist_test.train_labels\n",
    "    elif dataset['name']=='kmnist' or dataset['name']=='kmnistPatch':\n",
    "        dataset['train_labels'] = kmnist_train.train_labels\n",
    "        dataset['test_labels'] = kmnist_test.train_labels\n",
    "    elif dataset['name']=='fmnist':\n",
    "        dataset['train_labels'] = fmnist_train.train_labels\n",
    "        dataset['test_labels'] = fmnist_test.train_labels\n",
    "    elif dataset['name']=='svhn':\n",
    "        dataset['train_labels'] = torch.tensor(svhn_train.labels)\n",
    "        dataset['test_labels'] = torch.tensor(svhn_test.labels)\n",
    "\n",
    "# populate the various train/test folders and domino label combinations\n",
    "for partition in ['train','test']:    \n",
    "    for top_label in tqdm(range(2)):           \n",
    "        top_inds = (top_dataset[f'{partition}_labels']==top_label).nonzero().squeeze()  # indices for bottom images\n",
    "        Path(os.path.join(save_dir,f'{partition}/{top_label}/')).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        for file_idx,top_idx in tqdm(enumerate(top_inds[:num_imgs[partition]])):\n",
    "            if top_dataset['name'] == 'kmnistPatch' and top_label==1:\n",
    "                temp_img = top_dataset[f'{partition}_data'][top_idx][0]\n",
    "                temp_img[:,:5,:5]=1.0\n",
    "                img = torch.hstack((temp_img,torch.zeros((3,32,32))))\n",
    "            else:\n",
    "                img = torch.hstack((top_dataset[f'{partition}_data'][top_idx][0],torch.zeros((3,32,32))))\n",
    "            final_path = os.path.join(save_dir,f'{partition}/{top_label}/{file_idx}.jpg')\n",
    "            plt.imsave(final_path,np.array(img.permute(1,2,0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate CSV File for corresponding Blank Domino Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code assumes the heirarchy root_dir/{train/test}/top{0-9}_bot{0-9}/{0-num_train/0-num_test}.jpg\n",
    "# code currently only supports spurious correlation = 1.0\n",
    "\n",
    "val_inds = random.sample(range(num_train+num_val), num_val)\n",
    "all_inds = list(range(num_train+num_val))\n",
    "train_inds = [x for x in all_inds if x not in val_inds]\n",
    "csv_list = [] # all rows will be appended here\n",
    "csv_cols = ['path','top','bottom','val_train_test_split','test_nontest_split','all_zeros']\n",
    "\n",
    "for lbl_idx in range(2):\n",
    "    # add validation info to csv_list\n",
    "    for val_idx in val_inds:\n",
    "        path = os.path.join(root_dir,f'train/{lbl_idx}/{val_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 0\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add train info to csv_list\n",
    "    for tr_idx in train_inds:\n",
    "        path = os.path.join(root_dir,f'train/{lbl_idx}/{tr_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 1\n",
    "        test_nontest_split = 1\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "    # add test info to csv_list\n",
    "    for test_idx in range(num_test):\n",
    "        path = os.path.join(root_dir,f'test/{lbl_idx}/{test_idx}.jpg')\n",
    "        top = bottom = lbl_idx\n",
    "        val_train_test_split = 2\n",
    "        test_nontest_split = 0\n",
    "        all_zeros = 0\n",
    "        csv_list.append([path,top,bottom,val_train_test_split,test_nontest_split,all_zeros])\n",
    "\n",
    "df = pd.DataFrame(csv_list)\n",
    "df.columns = csv_cols\n",
    "df.to_csv(full_save_path,index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pl",
   "language": "python",
   "name": "pl"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}