{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import sklearn.datasets\n",
    "from twomoons_model import TwoMoonsMLP, TwoMoonsSTE\n",
    "from optim import BiNNOptimizer as BayesBiNN\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_points_train=200\n",
    "num_points_test = 200\n",
    "X_orig, y_train = sklearn.datasets.make_moons(num_points_train, noise=0.1, random_state=0)\n",
    "X_test, y_test = sklearn.datasets.make_moons(num_points_test, noise=0.1, random_state=0)\n",
    "X_orig = torch.from_numpy(X_orig).float()\n",
    "y_train = torch.from_numpy(y_train).float().unsqueeze(1)\n",
    "X_test = torch.from_numpy(X_test).float()\n",
    "y_test = torch.from_numpy(y_test).float().unsqueeze(1)\n",
    "X_mean = X_orig.mean(dim=0)\n",
    "X_std = X_orig.std(dim=0)\n",
    "X_train = (X_orig-X_mean)/X_std\n",
    "X_test = (X_test-X_mean)/X_std\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training setting for BayesBiNN\n",
    "learning_rate = 1e-3\n",
    "learning_rate_decay = 0.1\n",
    "learning_rate_decay_epochs = [1500, 2500]\n",
    "beta = 0.99\n",
    "temperature = 1\n",
    "initialize_lambda = 15\n",
    "mc_test_samples = 10\n",
    "N = 5\n",
    "hidden_layers = [64,64]\n",
    "epochs = 3000\n",
    "criterion = nn.BCEWithLogitsLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3000/3000 [00:17<00:00, 174.66it/s]\n"
     ]
    }
   ],
   "source": [
    "# Training BayesBiNN Network\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "model_binn = TwoMoonsMLP(X_train.shape[1], hidden_layers, y_train.shape[1])\n",
    "\n",
    "optimizer_binn = BayesBiNN(model_binn, train_set_size=X_train.shape[0], lr=learning_rate, beta=beta, \n",
    "                 N=N, temperature=temperature, initialize_lambda=initialize_lambda)\n",
    "\n",
    "for epoch in tqdm(range(epochs)):\n",
    "    if epoch in learning_rate_decay_epochs:\n",
    "        for params in optimizer_binn.param_groups:\n",
    "            params['lr'] *=  learning_rate_decay\n",
    "    def closure():\n",
    "        optimizer_binn.zero_grad()\n",
    "        y_pred = model_binn.forward(X_train)\n",
    "        loss = criterion(y_pred, y_train)\n",
    "        return loss, y_pred\n",
    "    loss, y_pred = optimizer_binn.step(closure)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training STE Network \n",
    "learning_rate = 1e-1\n",
    "learning_rate_decay = 0.1\n",
    "learning_rate_decay_epochs = [1500, 2500]\n",
    "hidden_layers = [64,64]\n",
    "epochs = 3000\n",
    "criterion = nn.BCEWithLogitsLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3000/3000 [00:05<00:00, 558.76it/s]\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "# epochs_ste = 3000\n",
    "# lr_start = 1e-1\n",
    "# lr_decay = 0.1\n",
    "# lr_decay_epochs = [1500, 2500]\n",
    "\n",
    "model_ste = TwoMoonsSTE(X_train.shape[1], hidden_layers, y_train.shape[1])\n",
    "optimizer_ste = optim.Adam(model_ste.parameters(), lr=learning_rate)\n",
    "\n",
    "for i in tqdm(range(epochs)):\n",
    "    if i in learning_rate_decay_epochs:\n",
    "        for param_group in optimizer_ste.param_groups:\n",
    "            param_group['lr'] *=  learning_rate_decay\n",
    "    \n",
    "    optimizer_ste.zero_grad()\n",
    "    y_pred = model_ste.forward(X_train)\n",
    "    loss = criterion(y_pred, y_train)\n",
    "    loss.backward()\n",
    "    \n",
    "    for p in list(model_ste.parameters()):\n",
    "        if hasattr(p, 'org'):\n",
    "            p.data.copy_(p.org)\n",
    "\n",
    "    optimizer_ste.step()\n",
    "\n",
    "    for p in list(model_ste.parameters()):\n",
    "        if hasattr(p, 'org'):\n",
    "            p.org.copy_(p.data.clamp_(-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_grid(X, h=0.01, padding=0.5):\n",
    "    x_min, x_max = X[:, 0].min() - padding, X[:, 0].max() + padding\n",
    "    y_min, y_max = X[:, 1].min() - padding, X[:, 1].max() + padding\n",
    "    xx,yy=np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))\n",
    "    grid = torch.from_numpy(np.c_[xx.ravel(), yy.ravel()]).float()\n",
    "    return grid, xx, yy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "yf = y_train.flatten()\n",
    "grid_orig, xx, yy = setup_grid(X_orig, h=0.01, padding=0.3)\n",
    "grid = (grid_orig - X_mean) / X_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7012 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'BayesBiNN Mean')"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 177272 characters]\n",
      "text/plain": [
       "<Figure size 504x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.ticker as plticker\n",
    "\n",
    "fig, axes = plt.subplots(3,1,figsize=(7,15))\n",
    "\n",
    "loc_y = plticker.MultipleLocator(base=0.5)\n",
    "loc_x = plticker.MultipleLocator(base=1.0)\n",
    "\n",
    "axes[0].tick_params(axis='both', which='major', labelsize=16)\n",
    "axes[0].tick_params(axis='both', which='minor', labelsize=16)\n",
    "axes[1].tick_params(axis='both', which='major', labelsize=16)\n",
    "axes[1].tick_params(axis='both', which='minor', labelsize=16)\n",
    "axes[2].tick_params(axis='both', which='major', labelsize=16)\n",
    "axes[2].tick_params(axis='both', which='minor', labelsize=16)\n",
    "\n",
    "ms = 100\n",
    "\n",
    "Z = model_ste.predict(grid).reshape(xx.shape)\n",
    "c = axes[0].contourf(xx, yy, Z, alpha=0.7, antialiased=True, cmap='PuOr')\n",
    "axes[0].scatter(X_orig[yf==0, 0], X_orig[yf==0, 1], s=ms, lw=1.5, edgecolors='black', color='#EFBC38', marker='o')\n",
    "axes[0].scatter(X_orig[yf==1, 0], X_orig[yf==1, 1], s=ms, lw=1.5, edgecolors='black',color='#CE84EE',marker = 'p')\n",
    "axes[0].xaxis.set_major_locator(loc_x)\n",
    "axes[0].yaxis.set_major_locator(loc_y)\n",
    "axes[0].set_ylabel('STE Prediction', fontsize=18)\n",
    "\n",
    "Z = optimizer_binn.montecarlo_predictions(model_binn.forward,grid,epsilons=None)\n",
    "Z = torch.sigmoid(torch.stack(Z)).mean(dim=0).reshape(xx.shape).detach()\n",
    "c = axes[1].contourf(xx, yy, Z, alpha=0.7, antialiased=True, cmap='PuOr')\n",
    "axes[1].scatter(X_orig[yf==0, 0], X_orig[yf==0, 1], s=ms, edgecolors='black', color='#EFBC38',lw=1.5, marker = 'o')\n",
    "axes[1].scatter(X_orig[yf==1, 0], X_orig[yf==1, 1], s=ms, edgecolors='black', color='#CE84EE',lw=1.5, marker = 'p')\n",
    "axes[1].xaxis.set_major_locator(loc_x)\n",
    "axes[1].yaxis.set_major_locator(loc_y)\n",
    "axes[1].set_ylabel('BayesBiNN Mode', fontsize=18)\n",
    "\n",
    "torch.manual_seed(0)\n",
    "n_samples = 10\n",
    "raw_noises = []\n",
    "for mc_sample in range(n_samples):\n",
    "    raw_noises.append(torch.bernoulli(torch.sigmoid(2*optimizer_binn.state['lambda'])))\n",
    "Z = optimizer_binn.montecarlo_predictions(model_binn.forward, grid, epsilons=raw_noises)\n",
    "Z = torch.sigmoid(torch.stack(Z)).mean(dim=0).reshape(xx.shape).detach()\n",
    "c = axes[2].contourf(xx, yy, Z, alpha=0.7, antialiased=True, cmap='PuOr')\n",
    "axes[2].scatter(X_orig[yf==0, 0], X_orig[yf==0, 1], s=ms, edgecolors='black', color='#EFBC38',lw=1.5, marker = 'o')\n",
    "axes[2].scatter(X_orig[yf==1, 0], X_orig[yf==1, 1], s=ms, edgecolors='black', color='#CE84EE',lw=1.5, marker = 'p')\n",
    "axes[2].xaxis.set_major_locator(loc_x)\n",
    "axes[2].yaxis.set_major_locator(loc_y)\n",
    "axes[2].set_ylabel('BayesBiNN Mean', fontsize=18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7f5098fa38d0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52352 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "n_samples = 10\n",
    "raw_noises = []\n",
    "for mc_sample in range(n_samples):\n",
    "    raw_noises.append(torch.bernoulli(torch.sigmoid(2*optimizer_binn.state['lambda'])))\n",
    "Z = optimizer_binn.montecarlo_predictions(model_binn.forward, grid, epsilons=raw_noises)\n",
    "Z = torch.sigmoid(torch.stack(Z)).mean(dim=0).reshape(xx.shape).detach()\n",
    "\n",
    "c = plt.contourf(xx, yy, Z, alpha=0.7, antialiased=True, cmap='PuOr')\n",
    "plt.scatter(X_orig[yf==0, 0], X_orig[yf==0, 1], s=ms, edgecolors='black', color='#EFBC38',lw=1.5, marker = 'o')\n",
    "plt.scatter(X_orig[yf==1, 0], X_orig[yf==1, 1], s=ms, edgecolors='black', color='#CE84EE',lw=1.5, marker = 'p')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7f5098f4ef50>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 38444 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Z = model_ste.predict(grid).reshape(xx.shape)\n",
    "\n",
    "c = plt.contourf(xx, yy, Z, alpha=0.7, antialiased=True, cmap='PuOr')\n",
    "plt.scatter(X_orig[yf==0, 0], X_orig[yf==0, 1], s=ms, lw=1.5, edgecolors='black', marker='s', color='#EFBC38')\n",
    "plt.scatter(X_orig[yf==1, 0], X_orig[yf==1, 1], s=ms, lw=1.5, edgecolors='black',color='#CE84EE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
