{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d8af36ae-5645-40e5-9262-e70ca42534d8",
   "metadata": {},
   "source": [
    "## Corrupted Data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "08daafde-d04e-41e2-a51e-c136f79dbe34",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 42\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pruned Dataset Size 100\n",
      "Number of Clean Samples = 1000\n",
      "Number of corrupt Samples = 250\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from src_robust_estimate import robust_mean_estimate\n",
    "from pytorch_lightning import seed_everything\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sns.set(style=\"white\")\n",
    "sns.set_context(\"poster\")\n",
    "\n",
    "seed_everything(42)\n",
    "\n",
    "n_clean = 1000\n",
    "CORRUPTION_FRAC = 0.2\n",
    "plt_folder = 'plots/toy/corr=20/'\n",
    "\n",
    "n_corrupted = int(CORRUPTION_FRAC * n_clean / (1 - CORRUPTION_FRAC))\n",
    "\n",
    "CORESET_FRAC = 0.1\n",
    "CORESET_SIZE = int(CORESET_FRAC * n_clean)\n",
    "print(\"Pruned Dataset Size {}\".format(CORESET_SIZE))\n",
    "# --------------------\n",
    "# clean distribution\n",
    "# --------------------\n",
    "mean_true = torch.tensor(\n",
    "\t[0, 0],\n",
    "\tdtype=torch.float32\n",
    ")\n",
    "cov_true = torch.tensor([\n",
    "\t[1, -0.5],\n",
    "\t[-0.5, 0.5]\n",
    "], dtype=torch.float32\n",
    ")\n",
    "clean_dist = torch.distributions.MultivariateNormal(mean_true, cov_true).sample((n_clean,))\n",
    "print(\"Number of Clean Samples = {}\".format(len(clean_dist)))\n",
    "\n",
    "# --------------------\n",
    "# adv distribution \n",
    "# --------------------\n",
    "mean_adv = torch.tensor(\n",
    "\t[10, 6],\n",
    "\tdtype=torch.float32\n",
    ")\n",
    "cov_adv = torch.tensor([\n",
    "\t[1, 0.5],\n",
    "\t[0.5, 0.5]\n",
    "], dtype=torch.float32)\n",
    "adv_dist = torch.distributions.MultivariateNormal(mean_adv, cov_adv).sample((n_corrupted,))\n",
    "print(\"Number of corrupt Samples = {}\".format(len(adv_dist)))\n",
    "\n",
    "# --------------------\n",
    "# corrupted training distribution \n",
    "# --------------------\n",
    "p = torch.cat((clean_dist, adv_dist), dim=0)\n",
    "\n",
    "\n",
    "def plot_dist():\n",
    "\tsns.kdeplot(\n",
    "\t\tx=clean_dist[:, 0],\n",
    "\t\ty=clean_dist[:, 1],\n",
    "\t\tcmap='Blues',\n",
    "\t\talpha=0.8,\n",
    "\t\tthresh=0.001,\n",
    "\t\tfill=True,\n",
    "\t)\n",
    "\t\n",
    "\tsns.kdeplot(\n",
    "\t\tx=adv_dist[:, 0],\n",
    "\t\ty=adv_dist[:, 1],\n",
    "\t\tcmap='Reds',\n",
    "\t\talpha=0.7,\n",
    "\t\tthresh=0.001,\n",
    "\t\tfill=True,\n",
    "\t)\n",
    "    \n",
    "def plot_label():\n",
    "    plt.xlabel(r'${\\bf X0}$', fontsize=14)\n",
    "    plt.ylabel(r'${\\bf X1}$', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.legend(\n",
    "    \tloc='upper left',\n",
    "    \tframeon=True,\n",
    "    \tncol=3,\n",
    "    \tedgecolor='black',\n",
    "    \tframealpha=0.8,\n",
    "    \tmarkerscale=1,\n",
    "    \tfancybox=True,\n",
    "    \tprop={'weight': 'bold', 'size': 14}\n",
    "    )\n",
    "    plt.grid()\n",
    "    plt.tick_params(left = False, right = False , labelleft = False , labelbottom = False, bottom = False) \n",
    "    # sns.despine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eaffc104-c453-402e-a207-e39b1f422db8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "# # Generate uniformly distributed data over [0, pi]\n",
    "# uniform_data = np.random.uniform(0, np.pi, 10000)\n",
    "\n",
    "# # Transform the data using cosine\n",
    "# transformed_data = np.cos(uniform_data)\n",
    "\n",
    "# # Plot the histogram of the transformed data\n",
    "# plt.hist(transformed_data, bins=30, density=True, alpha=0.75)\n",
    "# plt.title('Density of Transformed Uniform Distribution')\n",
    "# plt.xlabel('Values')\n",
    "# plt.ylabel('Density')\n",
    "# plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04fdba9e-0861-44d8-9cba-80e776f4f28c",
   "metadata": {},
   "source": [
    "## Mean Estimate "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b660ba82-a7b3-48ed-a264-a44c6b1f414b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ran out of Max iter for GM - returning sub optimal answer\n"
     ]
    }
   ],
   "source": [
    "# Calculate Center \n",
    "\n",
    "true_moment = torch.mean(clean_dist, axis=0)\n",
    "noisy_moment = torch.mean(p, axis=0)\n",
    "geo_med = robust_mean_estimate(\n",
    "\tdata=p.numpy(),\n",
    "\testimator='geo_med'\n",
    ")\n",
    "co_med = robust_mean_estimate(\n",
    "\tdata=p.numpy(),\n",
    "\testimator='co_med'\n",
    ")\n",
    "\n",
    "\n",
    "def plot_moments():\n",
    "\tplt.scatter(\n",
    "\t\tx=true_moment[0],\n",
    "\t\ty=true_moment[1],\n",
    "\t\tcolor='g',\n",
    "\t\tedgecolors='black',  # Boundary color\n",
    "\t\tlinewidth=2,\n",
    "\t\t# label=r'$E(x \\sim p(x))$',\n",
    "\t\tlabel='mean (oracle)',\n",
    "\t\ts=400,\n",
    "\t\talpha=0.8,\n",
    "\t\tmarker='X'\n",
    "\t)\n",
    "\tplt.scatter(\n",
    "\t\tx=noisy_moment[0],\n",
    "\t\ty=noisy_moment[1],\n",
    "\t\tcolor='red',\n",
    "\t\tedgecolors='black',  # Boundary color\n",
    "\t\tlinewidth=2,\n",
    "\t\tlabel='mean (estimate)',\n",
    "\t\ts=400,\n",
    "\t\talpha=0.8,\n",
    "\t\tmarker='X'\n",
    "\t)\n",
    "\tplt.scatter(\n",
    "\t\tx=geo_med[0],\n",
    "\t\ty=geo_med[1],\n",
    "\t\tcolor='violet',\n",
    "\t\tedgecolors='black',  # Boundary color\n",
    "\t\tlinewidth=2,\n",
    "\t\tlabel=r'GM',\n",
    "\t\ts=400,\n",
    "\t\talpha=0.8,\n",
    "\t\tmarker='X'\n",
    "\t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ff9a990d-2a1e-40ba-bbac-743c9a246984",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 125776 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "# plt.figure(dpi=300)\n",
    "plot_dist()\n",
    "plot_moments()\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'dist')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2353f34c-48f5-4a27-bce3-dab75da2bdcf",
   "metadata": {},
   "source": [
    "## Pruning "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "171d6755-9b84-4b14-bb65-780f8ca4c3ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src_sampling import (\n",
    "\tuniform,\n",
    "\teasy,\n",
    "\thard,\n",
    "\tmoderate,\n",
    "\therding,\n",
    ")\n",
    "\n",
    "# Random \n",
    "p_uniform_ix = uniform(phi=p, coreset_size=CORESET_SIZE)\n",
    "p_filter_uniform_ix = uniform(phi=p, coreset_size=CORESET_SIZE, robust_filtering=True)\n",
    "\n",
    "# Easy \n",
    "p_easy_ix = easy(phi=p, coreset_size=CORESET_SIZE)\n",
    "p_filter_easy_ix = easy(phi=p, coreset_size=CORESET_SIZE, robust_filtering=True)\n",
    "p_center_easy_ix = easy(phi=p, coreset_size=CORESET_SIZE, target_moment='geo_med')\n",
    "\n",
    "# Hard\n",
    "p_hard_ix = hard(phi=p, coreset_size=CORESET_SIZE)\n",
    "p_filter_hard_ix = hard(phi=p, coreset_size=CORESET_SIZE, robust_filtering=True)\n",
    "p_center_hard_ix = hard(phi=p, coreset_size=CORESET_SIZE, target_moment='geo_med')\n",
    "\n",
    "# Moderate\n",
    "p_moderate_ix = moderate(phi=p, coreset_size=CORESET_SIZE)\n",
    "p_filter_moderate_ix = moderate(phi=p, coreset_size=CORESET_SIZE, robust_filtering=True)\n",
    "p_center_moderate_ix = moderate(phi=p, coreset_size=CORESET_SIZE, target_moment='geo_med')\n",
    "\n",
    "# herding \n",
    "p_herd_ix = herding(\n",
    "\tphi=p,\n",
    "\tcoreset_size=CORESET_SIZE,\n",
    "\tdist_measure='L2',\n",
    "\tinit='random'\n",
    ")\n",
    "p_filter_herd_ix = herding(\n",
    "\tphi=p,\n",
    "\tcoreset_size=CORESET_SIZE,\n",
    "\trobust_filtering=True,\n",
    "\tdist_measure='L2',\n",
    "\tinit='random'\n",
    ")\n",
    "p_center_herd_ix = herding(\n",
    "\tphi=p,\n",
    "\tcoreset_size=CORESET_SIZE,\n",
    "\ttarget_moment='geo_med',\n",
    "\tdist_measure='L2',\n",
    "\tinit='random'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "161f7dea-3671-4d14-9bb3-73b9d176f0e2",
   "metadata": {},
   "source": [
    "## Pruning : Uniform / Random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c2ffe23f-c25e-461f-88f3-1ed71f63be68",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 320696 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "p_uniform = p[p_uniform_ix]\n",
    "plot_dist()\n",
    "plt.scatter(\n",
    "\tx=p_uniform[:, 0],\n",
    "\ty=p_uniform[:, 1],\n",
    "\talpha=0.8,\n",
    "\tcolor='g',\n",
    "\tedgecolors='black',  # Boundary color\n",
    "\ts=200,\n",
    "\tlabel='Uniform',\n",
    "\tmarker='o'\n",
    ")\n",
    "\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'uniform')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c958b276-b3f8-47e4-a4dc-8ee7e971dc55",
   "metadata": {},
   "source": [
    "## Pruning : Easy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "28b6ef3b-5416-41b6-8bbb-4175b55cc318",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 137856 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "p_easy = p[p_easy_ix]\n",
    "plot_dist()\n",
    "plt.scatter(\n",
    "\tx=p_easy[:, 0],\n",
    "\ty=p_easy[:, 1],\n",
    "\talpha=0.8,\n",
    "\tcolor='g',\n",
    "\tedgecolors='black',  # Boundary color\n",
    "\ts=200,\n",
    "\tlabel='Easy',\n",
    "\tmarker='v'\n",
    ")\n",
    "\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'easy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c9afb50-d627-4a3c-ab60-3d7c0b85e296",
   "metadata": {},
   "source": [
    "# Pruning : Hard "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "77f8b953-b95a-4aa5-b058-aafc8143c0fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 167732 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "p_hard = p[p_hard_ix]\n",
    "plot_dist()\n",
    "plt.scatter(\n",
    "\tx=p_hard[:, 0],\n",
    "\ty=p_hard[:, 1],\n",
    "\talpha=0.8,\n",
    "\tcolor='g',\n",
    "\tedgecolors='black',  # Boundary color\n",
    "\ts=200,\n",
    "\tlabel='Hard',\n",
    "\tmarker='^'\n",
    ")\n",
    "\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'hard')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4d02e0b-5875-427c-aaf9-66816a3d73e0",
   "metadata": {},
   "source": [
    "## Pruning: Moderate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "92259a59-ad47-40ed-acd8-0b1856126473",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 141984 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "p_moderate = p[p_moderate_ix]\n",
    "plot_dist()\n",
    "plt.scatter(\n",
    "\tx=p_moderate[:, 0],\n",
    "\ty=p_moderate[:, 1],\n",
    "\talpha=0.8,\n",
    "\tcolor='g',\n",
    "\tedgecolors='black',  # Boundary color\n",
    "\ts=200,\n",
    "\tlabel='Moderate',\n",
    "\tmarker='<'\n",
    ")\n",
    "\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'moderate')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a12aa4f-f4ee-4e1a-ae48-dc4e58071590",
   "metadata": {},
   "source": [
    "## Pruning: Herding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9d28f7d7-438b-4351-bb71-298bf7652f23",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 291764 characters]",
      "text/plain": [
       "<Figure size 2400x1800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(8, 6), dpi=300)\n",
    "p_herd = p[p_herd_ix]\n",
    "plot_dist()\n",
    "plt.scatter(\n",
    "\tx=p_herd[:, 0],\n",
    "\ty=p_herd[:, 1],\n",
    "\talpha=0.8,\n",
    "\tcolor='g',\n",
    "\tedgecolors='black',  # Boundary color\n",
    "\ts=200,\n",
    "\tlabel='Herding',\n",
    "\tmarker='d'\n",
    ")\n",
    "plot_label()\n",
    "# plt.savefig(plt_folder + 'herding')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
