notebooks/dataset_dist.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "import re\n",
    "from os.path import join\n",
    "from IPython.display import display\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['ps.fonttype'] = 42\n",
    "matplotlib.rc('text', usetex=True)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import joblib\n",
    "from joblib import Parallel, delayed\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.utils.data as data_utils\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "import joblib\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from lolip.models.torch_utils.archs import *\n",
    "from lolip.variables import auto_var\n",
    "from params import *\n",
    "from utils import params_to_dataframe\n",
    "\n",
    "fontsize=24"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nearest_oppo_dist(X, y, tstX, tsty, norm, n_jobs=10):\n",
    "    if len(X.shape) > 2:\n",
    "        X = X.reshape(len(X), -1)\n",
    "        tstX = tstX.reshape(len(tstX), -1)\n",
    "    p = norm\n",
    "\n",
    "    def helper(yi):\n",
    "        return NearestNeighbors(n_neighbors=1,\n",
    "                metric='minkowski', p=p, n_jobs=12).fit(X[y != yi])\n",
    "    nns = Parallel(n_jobs=n_jobs)(delayed(helper)(yi) for yi in np.unique(y))\n",
    "    ret = np.zeros(len(X))\n",
    "    tst_ret = np.zeros(len(tstX))\n",
    "    for yi in np.unique(y):\n",
    "        dist, _ = nns[yi].kneighbors(X[y==yi], n_neighbors=1)\n",
    "        ret[np.where(y==yi)[0]] = dist[:, 0]\n",
    "        \n",
    "        dist, _ = nns[yi].kneighbors(tstX[tsty==yi], n_neighbors=1)\n",
    "        tst_ret[np.where(tsty==yi)[0]] = dist[:, 0]\n",
    "\n",
    "    return nns, ret, tst_ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_var.set_variable_value(\"random_seed\", 0)\n",
    "trnX, trny, tstX, tsty = auto_var.get_var_with_argument(\"dataset\", \"mnist\")\n",
    "nns_linf, mnist_dists_linf, tst_mnist_dists_linf = get_nearest_oppo_dist(trnX, trny, tstX, tsty, np.inf, n_jobs=1)\n",
    "#data = {\n",
    "#    \"nns_linf\": nns_linf,\n",
    "#    \"dists_linf\": mnist_dists_linf,\n",
    "#    \"tst_dists_linf\": tst_mnist_dists_linf,\n",
    "#}\n",
    "#joblib.dump(data, \"./dist_data/mnist_data.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mnist_dists_linf.min(), mnist_dists_linf.mean())\n",
    "print(tst_mnist_dists_linf.min(), tst_mnist_dists_linf.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"MNIST Train-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "\n",
    "data = mnist_dists_linf\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()-0.27, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.1\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.2, 0.85, '2$\\epsilon$={:.1f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/mnist_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/mnist_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"MNIST Test-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "data = tst_mnist_dists_linf\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()-0.27, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.1\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.2, 0.85, '2$r$={:.1f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/mnist_test_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/mnist_test_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_var.set_variable_value(\"random_seed\", 0)\n",
    "trnX, trny, tstX, tsty = auto_var.get_var_with_argument(\"dataset\", \"cifar10\")\n",
    "nns_linf, cifar10_dists_linf, tst_cifar10_dists_linf = get_nearest_oppo_dist(trnX, trny, tstX, tsty, np.inf, n_jobs=1)\n",
    "#data = {\n",
    "#    \"nns_linf\": nns_linf,\n",
    "#    \"dists_linf\": cifar10_dists_linf,\n",
    "#    \"tst_dists_linf\": tst_cifar10_dists_linf,\n",
    "#}\n",
    "#joblib.dump(data, \"./dist_data/cifar10_data.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cifar10_dists_linf.min(), cifar10_dists_linf.mean())\n",
    "print(tst_cifar10_dists_linf.min(), tst_cifar10_dists_linf.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"CIFAR-10 Train-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "\n",
    "data = cifar10_dists_linf\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.031\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.11, 0.65, '2$\\epsilon$={:.3f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/cifar10_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/cifar10_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"CIFAR-10 Test-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "\n",
    "data = tst_cifar10_dists_linf\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.031\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.11, 0.65, '2$\\epsilon$={:.3f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/cifar10_test_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/cifar10_test_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_var.set_variable_value(\"random_seed\", 0)\n",
    "trnX, trny, tstX, tsty = auto_var.get_var_with_argument(\"dataset\", \"svhn\")\n",
    "#np.testing.assert_array_equal(trnX[25235], trnX[25236])\n",
    "#trnX = np.delete(trnX, 25235, axis=0) # image 25235 and 25236 are exactly the same\n",
    "#trny = np.delete(trny, 25235, axis=0)\n",
    "#trnX = np.delete(trnX, 65042, axis=0) # incorrectly labeled (image with idx 65043 and 11933 are same image but different label...)\n",
    "#trny = np.delete(trny, 65042, axis=0)\n",
    "nns_linf, svhn_dists_linf, tst_svhn_dists_linf = get_nearest_oppo_dist(trnX, trny, tstX, tsty, np.inf, n_jobs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_svhn_dists_linf = svhn_dists_linf[svhn_dists_linf >= np.partition(svhn_dists_linf, 5)[:5].max()]\n",
    "print(temp_svhn_dists_linf.min(), temp_svhn_dists_linf.mean())\n",
    "print(tst_svhn_dists_linf.min(), tst_svhn_dists_linf.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.partition(svhn_dists_linf, 5)[:5].max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"SVHN Train-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "#plt.hist(svhn_dists_linf, range=(0, 1), weights=np.ones(len(svhn_dists_linf)) / len(svhn_dists_linf))\n",
    "#plt.hist(tst_svhn_dists_linf, range=(0, 1), weights=np.ones(len(tst_svhn_dists_linf)) / len(tst_svhn_dists_linf))\n",
    "\n",
    "data = svhn_dists_linf[svhn_dists_linf >= np.partition(svhn_dists_linf, 5)[:5].max()]\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.031\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.11, 0.65, '2$\\epsilon$={:.3f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/svhn_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/svhn_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"SVHN Test-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "\n",
    "data = tst_svhn_dists_linf\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.031\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.11, 0.65, '2$\\epsilon$={:.3f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/svhn_test_linf_hist.png\", format=\"png\")\n",
    "plt.savefig(\"./figs/dataset_dist/svhn_test_linf_hist.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(trnX[11933])\n",
    "print(trny[11933])\n",
    "plt.savefig(\"svhn_11933.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(trnX[65043])\n",
    "print(trny[65043])\n",
    "plt.savefig(\"svhn_11933_close(65043).pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(trnX[25235])\n",
    "plt.savefig(\"svhn_25235.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "import torch\n",
    "from sklearn.metrics import pairwise_distances\n",
    "import numpy as np\n",
    "import joblib\n",
    "trn_ds = ImageFolder(\"../data/RestrictedImgNet/train\",\n",
    "    transform=transforms.Compose([\n",
    "        transforms.Resize(72),\n",
    "        transforms.CenterCrop(64),\n",
    "        transforms.ToTensor(),\n",
    "    ]))\n",
    "tst_ds = ImageFolder(\"../data/RestrictedImgNet/val\",\n",
    "    transform=transforms.Compose([\n",
    "        transforms.Resize(72),\n",
    "        transforms.CenterCrop(64),\n",
    "        transforms.ToTensor(),\n",
    "    ]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst_dists = joblib.load(\"../restricted_tst_linf.pkl\")\n",
    "np_tst_dist = tst_dists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "tst_loader = torch.utils.data.DataLoader(tst_ds, batch_size=batch_size, shuffle=False, num_workers=16)\n",
    "\n",
    "for i, (xi, yi) in enumerate(tst_loader):\n",
    "    yi = yi.numpy()\n",
    "    end = min(len(tst_ds), (i+1)*batch_size)\n",
    "    for j in range(i*batch_size, end):\n",
    "        np_tst_dist[j, yi[j-i*batch_size]] = np.inf\n",
    "print(np_tst_dist.min(), np_tst_dist.min(1).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"Restricted ImageNet Test-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "#plt.hist(svhn_dists_linf, range=(0, 1), weights=np.ones(len(svhn_dists_linf)) / len(svhn_dists_linf))\n",
    "#plt.hist(tst_svhn_dists_linf, range=(0, 1), weights=np.ones(len(tst_svhn_dists_linf)) / len(tst_svhn_dists_linf))\n",
    "\n",
    "data = np_tst_dist.min(1)\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.005\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.06, 0.65, '2$\\epsilon$={:.2f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/resImgNet_test_linf_hist.pdf\", format=\"pdf\")\n",
    "plt.savefig(\"./figs/dataset_dist/resImgNet_test_linf_hist.png\", format=\"png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_and_save_figure(ds, fig_ids, file_prefix=\"figs/dataset_dist/resImgNet_\"):\n",
    "    for i in fig_ids:\n",
    "        print(f\"id: {i}, label: {ds[i][1]}\")\n",
    "        plt.imshow(ds[i][0].numpy().transpose(1, 2, 0))\n",
    "        plt.savefig(file_prefix + f\"{i}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_and_save_figure(trn_ds, fig_ids=[27782, 28340, 31382, 51605, 89607, 93380], file_prefix=\"figs/dataset_dist/resImgNet_\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(trn_ds[51605][1])\n",
    "plt.imshow(trn_ds[51605][0].numpy().transpose(1, 2, 0))\n",
    "plt.savefig(\"figs/dataset_dist/resImgNet_51605.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(trn_ds[28340][0].numpy().transpose(1, 2, 0))\n",
    "plt.savefig(\"figs/dataset_dist/resImgNet_51605.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dists = joblib.load(\"./dist_data/restricted_trn_linf.pkl\")\n",
    "np_dist = dists\n",
    "batch_size = 256\n",
    "loader = torch.utils.data.DataLoader(trn_ds, batch_size=batch_size, shuffle=False, num_workers=16)\n",
    "\n",
    "for i, (xi, yi) in enumerate(loader):\n",
    "    yi = yi.numpy()\n",
    "    end = min(len(trn_ds), (i+1)*batch_size)\n",
    "    for j in range(i*batch_size, end):\n",
    "        np_dist[j, yi[j-i*batch_size]] = np.inf\n",
    "        \n",
    "np_dist = np_dist.min(1)\n",
    "np_dist_f = np_dist[np_dist >= np.partition(np_dist, 6)[6]]\n",
    "print(np_dist_f.min(), np_dist_f.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.5, 4))\n",
    "plt.title(\"Restricted ImageNet Train-Train ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.ylabel(\"Percentage of data\", fontsize=fontsize)\n",
    "plt.xlabel(\"Distance ($\\ell_\\infty$)\", fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(fontsize=fontsize)\n",
    "\n",
    "data = np_dist_f\n",
    "plt.hist(data, range=(0, 1), bins=20, weights=np.ones(len(data)) / len(data))\n",
    "\n",
    "plt.axvline(data.min(), color='k', linestyle='dashed', linewidth=1)\n",
    "#plt.text(-0.0, 0.35, 'min: {:.2f}'.format(np.partition(np_dist.min(1), 6)[7]), fontsize=fontsize)\n",
    "plt.text(data.min()+0.02, 0.85, 'min: {:.2f}'.format(data.min()), fontsize=fontsize)\n",
    "r = 0.005\n",
    "plt.axvline(2*r, color='r', linestyle='dashed', linewidth=1)\n",
    "plt.text(2*r-0.06, 0.65, '2$\\epsilonr$={:.2f}'.format(2*r), fontsize=fontsize, color='r')\n",
    "plt.ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"./figs/dataset_dist/resImgNet_linf_hist.pdf\", format=\"pdf\")\n",
    "plt.savefig(\"./figs/dataset_dist/resImgNet_linf_hist.png\", format=\"png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joblib.dump(tst_dists.numpy(), \"./dist_data/restricted_tst_linf.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_classes = 9\n",
    "trn_dists = torch.ones((len(trn_ds), n_classes)).float()\n",
    "batch_size = 192\n",
    "trn_loader = torch.utils.data.DataLoader(trn_ds, batch_size=batch_size, shuffle=False, num_workers=12)\n",
    "tst_loader = torch.utils.data.DataLoader(trn_ds, batch_size=batch_size, shuffle=False, num_workers=12)\n",
    "\n",
    "for x, y in tqdm(trn_loader):\n",
    "    y = y.numpy()\n",
    "    x = x.flatten(1).cuda()\n",
    "    for i, (xi, yi) in enumerate(tst_loader):\n",
    "        yi = yi.numpy()\n",
    "        xi = xi.flatten(1).cuda()\n",
    "        \n",
    "        cov = torch.norm(x.repeat((len(xi), 1)) - xi.repeat_interleave(len(x), dim=0), p=np.inf, dim=1)\n",
    "        cov = cov.view(len(xi), len(x)).cpu()\n",
    "        \n",
    "        #cov = pairwise_distances(xi, x, metric='minkowski', n_jobs=16, p=np.inf)\n",
    "        for j in range(n_classes):\n",
    "            ty = (y == j)\n",
    "            if ty.sum() >= 1:\n",
    "                min_dist = cov[:, ty].min(axis=1)[0]\n",
    "                trn_dists[i*batch_size: (i+1)*batch_size, j] = np.minimum(\n",
    "                    trn_dists[i*batch_size: (i+1)*batch_size, j],\n",
    "                    min_dist\n",
    "                )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}