{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from snelson_model import SnelsonMLP, adjust_learning_rate\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from optim import BiNNOptimizer as BayesBiNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def DataLoader(data_path):\n",
    "    def load_snelson(filename):\n",
    "        with open(data_path + '/{fn}'.format(fn=filename), 'r') as f:\n",
    "            data_temp = [float(i) for i in f.read().strip().split('\\n')]\n",
    "            return np.array(data_temp, dtype=np.float64).reshape(-1,1)\n",
    "    \n",
    "    def preprocess_snelson(X_train, y_train, X_test):\n",
    "        X_train_mean, X_train_std = X_train.mean(), X_train.std()\n",
    "        y_train_mean, y_train_std = y_train.mean(), y_train.std()\n",
    "\n",
    "        X_train = (X_train-X_train_mean)/X_train_std\n",
    "        y_train = (y_train-y_train_mean)/y_train_std\n",
    "        X_test = (X_test-X_train_mean)/X_train_std\n",
    "\n",
    "        return X_train, y_train, X_test\n",
    "\n",
    "    X_train = load_snelson('train_inputs')\n",
    "    y_train = load_snelson('train_outputs')\n",
    "    X_test = load_snelson('test_inputs')\n",
    "    perm = np.random.permutation(X_train.shape[0])\n",
    "    X_train = X_train[perm]\n",
    "    y_train = y_train[perm]\n",
    "\n",
    "    X_train, y_train, X_test = preprocess_snelson(X_train, y_train, X_test)\n",
    "    \n",
    "    return torch.from_numpy(X_train).float(), torch.from_numpy(y_train).float(), torch.from_numpy(X_test).float()\n",
    "\n",
    "\n",
    "X_train, y_train, X_test = DataLoader(data_path='./snelson data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training settings given in the paper\n",
    "hidden_layers = [64,64]\n",
    "epochs = 5000\n",
    "learning_rate_start = 1e-4\n",
    "learning_rate_end = 1e-5\n",
    "N = 1 \n",
    "temperature = 1\n",
    "beta = 0.99 \n",
    "\n",
    "# Learning Rate decay\n",
    "lr_decay = (learning_rate_end/learning_rate_start)**(1.0/epochs)\n",
    "\n",
    "# Loss function\n",
    "criterion = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "model = SnelsonMLP(X_train.shape[1], hidden_layers, y_train.shape[1], use_bn=True, learn_bn=False, only_last_bn=True)\n",
    "optimizer = BayesBiNN(model, train_set_size=X_train.shape[0], learning_rate=learning_rate_start, beta=beta, \n",
    "                 N=N, temperature=temperature)\n",
    "\n",
    "for i in range(epochs):\n",
    "    adjust_learning_rate(lr_decay, optimizer, i)\n",
    "    def closure():\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = model.forward(X_train)\n",
    "        loss = criterion(y_pred, y_train)\n",
    "        return loss, y_pred\n",
    "    loss, y_pred = optimizer.step(closure)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get posterior mean prediction\n",
    "\n",
    "y_pred_mean = optimizer.montecarlo_predictions(model.forward, X_test, epsilons=None)[0].detach()\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "n_samples = 10\n",
    "raw_noises = []\n",
    "for i in range(n_samples):        \n",
    "    raw_noises.append(torch.bernoulli(torch.sigmoid(2*optimizer.state['lambda'])))\n",
    "    \n",
    "y_preds = optimizer.montecarlo_predictions(model.forward, X_test, raw_noises)\n",
    "y_preds = torch.stack(y_preds).detach()\n",
    "\n",
    "y_preds_mean = y_preds.mean(dim=0)\n",
    "y_preds_std = y_preds.std(dim=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training Setting for STE\n",
    "epochs_ste = 5000\n",
    "hidden_layers_ste = [64,64]\n",
    "learning_rate_start_ste = 0.1\n",
    "\n",
    "# Loss Function STE\n",
    "criterion_ste = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "epochs_ste = 5000\n",
    "lr_start = 1e-1\n",
    "\n",
    "model_ste = SnelsonMLP(X_train.shape[1], [64,64], y_train.shape[1],\n",
    "            use_bn=True, learn_bn=False, only_last_bn=True)\n",
    "criterion_ste = nn.MSELoss()\n",
    "optimizer_ste = optim.Adam(model_ste.parameters(), lr=lr_start)\n",
    "\n",
    "for i in range(epochs_ste):\n",
    "    optimizer_ste.zero_grad()\n",
    "    y_pred = model_ste.forward(X_train)\n",
    "    loss = criterion_ste(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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model_ste.forward(X_test).detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 99648 characters]\n",
      "text/plain": [
       "<Figure size 720x468 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "color1 = '#D62038'\n",
    "color2 = '#20B2AA'\n",
    "\n",
    "plt.figure(figsize=(10,6.5))\n",
    "\n",
    "plt.scatter(X_train, y_train, s=20, c='black', marker='D', zorder=100)\n",
    "plt.fill_between(X_test.flatten(), \n",
    "                 (y_preds_mean-2*y_preds_std).flatten(), \n",
    "                 (y_preds_mean+2*y_preds_std).flatten(), \n",
    "                 alpha=0.2, color=color2)\n",
    "plt.plot(X_test, y_pred_mean, c=color2, ls='--', lw=4, label='BayesBiNN (mode)')\n",
    "plt.plot(X_test, y_preds_mean, c=color2, lw=4, label='BayesBiNN (mean)\\n(± 2 standard deviations)'.format(n_samples))\n",
    "plt.plot(X_test, y_pred, c=color1, lw=4, label='STE prediction')\n",
    "\n",
    "plt.yticks(fontsize=20, rotation=0)\n",
    "plt.xticks(fontsize=20, rotation=0)\n",
    "plt.xlim(X_test.min(), X_test.max())\n",
    "plt.ylim(-3.6, 3)\n",
    "plt.legend(fontsize=20, fancybox=False, loc='lower right')\n",
    "\n",
    "plt.tight_layout()\n",
    "# os.makedirs('./figs', exist_ok=True)\n",
    "# plt.savefig('./figs/bb_regression_snelson.pdf', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "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
}
