{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0f551ac0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:12:34.491684Z",
     "start_time": "2022-05-25T08:12:34.085550Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "from tqdm import tqdm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "784a21fb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:12:34.988288Z",
     "start_time": "2022-05-25T08:12:34.983388Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "def plot_mean_and_CI_with_marker(time_steps, mean, lb, ub, color_mean=None, color_shading=None, marker=None, marker_size=12):\n",
    "    plt.fill_between(time_steps, ub, lb,\n",
    "                     color=color_shading, alpha=.2)\n",
    "    plt.plot(time_steps, mean, color_mean, marker=marker, markersize=marker_size, markevery=500)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8c68fd6",
   "metadata": {},
   "source": [
    "# cosine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8ac204d2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:06.637924Z",
     "start_time": "2022-05-25T08:30:05.543720Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "all_regrets_1 = []\n",
    "all_regrets_2 = []\n",
    "all_regrets_3 = []\n",
    "all_regrets_4 = []\n",
    "all_regrets_5 = []\n",
    "all_regrets_6 = []\n",
    "\n",
    "run_list = np.arange(3)\n",
    "for itr in run_list:\n",
    "    dataset = \"cosine\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_1.append(values_acc)\n",
    "\n",
    "    dataset = \"cosine\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_ts.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_2.append(values_acc)\n",
    "\n",
    "    dataset = \"cosine\"\n",
    "    N = 1\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_3.append(values_acc)\n",
    "    \n",
    "    dataset = \"cosine\"\n",
    "    N = 2\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_4.append(values_acc)\n",
    "\n",
    "    dataset = \"cosine\"\n",
    "    N = 5\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_5.append(values_acc)\n",
    "    \n",
    "    dataset = \"cosine\"\n",
    "    N = 10\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_6.append(values_acc)\n",
    "\n",
    "all_regrets_1_np = np.array(all_regrets_1)\n",
    "all_regrets_2_np = np.array(all_regrets_2)\n",
    "all_regrets_3_np = np.array(all_regrets_3)\n",
    "all_regrets_4_np = np.array(all_regrets_4)\n",
    "all_regrets_5_np = np.array(all_regrets_5)\n",
    "all_regrets_6_np = np.array(all_regrets_6)\n",
    "\n",
    "all_regrets_1_np_mean = np.mean(all_regrets_1_np, axis=0)\n",
    "all_regrets_2_np_mean = np.mean(all_regrets_2_np, axis=0)\n",
    "all_regrets_3_np_mean = np.mean(all_regrets_3_np, axis=0)\n",
    "all_regrets_4_np_mean = np.mean(all_regrets_4_np, axis=0)\n",
    "all_regrets_5_np_mean = np.mean(all_regrets_5_np, axis=0)\n",
    "all_regrets_6_np_mean = np.mean(all_regrets_6_np, axis=0)\n",
    "\n",
    "all_regrets_1_np_stderr = np.std(all_regrets_1_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_2_np_stderr = np.std(all_regrets_2_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_3_np_stderr = np.std(all_regrets_3_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_4_np_stderr = np.std(all_regrets_4_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_5_np_stderr = np.std(all_regrets_5_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_6_np_stderr = np.std(all_regrets_6_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "\n",
    "all_regrets_1_np_ub = all_regrets_1_np_mean + all_regrets_1_np_stderr\n",
    "all_regrets_1_np_lb = all_regrets_1_np_mean - all_regrets_1_np_stderr\n",
    "all_regrets_2_np_ub = all_regrets_2_np_mean + all_regrets_2_np_stderr\n",
    "all_regrets_2_np_lb = all_regrets_2_np_mean - all_regrets_2_np_stderr\n",
    "all_regrets_3_np_ub = all_regrets_3_np_mean + all_regrets_3_np_stderr\n",
    "all_regrets_3_np_lb = all_regrets_3_np_mean - all_regrets_3_np_stderr\n",
    "all_regrets_4_np_ub = all_regrets_4_np_mean + all_regrets_4_np_stderr\n",
    "all_regrets_4_np_lb = all_regrets_4_np_mean - all_regrets_4_np_stderr\n",
    "all_regrets_5_np_ub = all_regrets_5_np_mean + all_regrets_5_np_stderr\n",
    "all_regrets_5_np_lb = all_regrets_5_np_mean - all_regrets_5_np_stderr\n",
    "all_regrets_6_np_ub = all_regrets_6_np_mean + all_regrets_6_np_stderr\n",
    "all_regrets_6_np_lb = all_regrets_6_np_mean - all_regrets_6_np_stderr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df6c9dc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0559f29a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:20.380290Z",
     "start_time": "2022-05-25T08:30:20.028212Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85912 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "lw = 2.0\n",
    "plt.rc('font', size=22)\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.rcParams[\"legend.framealpha\"] = 0.5\n",
    "plt.rcParams[\"legend.loc\"] = \"best\"\n",
    "\n",
    "inds = np.arange(1, 100000)\n",
    "plot_len = 5000\n",
    "\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_1_np_mean[:plot_len], all_regrets_1_np_ub[:plot_len], all_regrets_1_np_lb[:plot_len], color_mean=\"b\", color_shading=\"b\", marker=\"v\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_2_np_mean[:plot_len], all_regrets_2_np_ub[:plot_len], all_regrets_2_np_lb[:plot_len], color_mean=\"r\", color_shading=\"r\", marker=\"D\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_3_np_mean[:plot_len], all_regrets_3_np_ub[:plot_len], all_regrets_3_np_lb[:plot_len], color_mean=\"g\", color_shading=\"g\", marker=\"s\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_4_np_mean[:plot_len], all_regrets_4_np_ub[:plot_len], all_regrets_4_np_lb[:plot_len], color_mean=\"c\", color_shading=\"c\", marker=\"o\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_5_np_mean[:plot_len], all_regrets_5_np_ub[:plot_len], all_regrets_5_np_lb[:plot_len], color_mean=\"y\", color_shading=\"y\", marker=\"x\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_6_np_mean[:plot_len], all_regrets_6_np_ub[:plot_len], all_regrets_6_np_lb[:plot_len], color_mean=\"m\", color_shading=\"m\", marker=\"^\", marker_size=12)\n",
    "\n",
    "plt.legend((r\"Neural UCB\", r\"Neural TS\", \\\n",
    "            r\"FN-UCB ($N=1$)\", r\"FN-UCB ($N=2$)\", r\"FN-UCB ($N=5$)\", r\"FN-UCB ($N=10$)\", \\\n",
    "           ), prop={'size':16})\n",
    "\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([0, 5000])\n",
    "axes.set_ylim([0, 130])\n",
    "\n",
    "plt.title(\"cosine, Fig. 1a\")\n",
    "plt.ylabel(\"Cumulative Regret\")\n",
    "plt.xlabel(\"Iterations\")\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0174a011",
   "metadata": {},
   "source": [
    "# square"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "525b933b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:23.592860Z",
     "start_time": "2022-05-25T08:30:23.273550Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "all_regrets_1 = []\n",
    "all_regrets_2 = []\n",
    "all_regrets_3 = []\n",
    "all_regrets_4 = []\n",
    "all_regrets_5 = []\n",
    "all_regrets_6 = []\n",
    "\n",
    "run_list = np.arange(3)\n",
    "for itr in run_list:\n",
    "    dataset = \"square\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_1.append(values_acc)\n",
    "\n",
    "    dataset = \"square\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_ts.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_2.append(values_acc)\n",
    "\n",
    "    dataset = \"square\"\n",
    "    N = 1\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_3.append(values_acc)\n",
    "    \n",
    "    dataset = \"square\"\n",
    "    N = 2\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_4.append(values_acc)\n",
    "\n",
    "    dataset = \"square\"\n",
    "    N = 5\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_5.append(values_acc)\n",
    "    \n",
    "    dataset = \"square\"\n",
    "    N = 10\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 0.1, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.1\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_6.append(values_acc)\n",
    "\n",
    "all_regrets_1_np = np.array(all_regrets_1)\n",
    "all_regrets_2_np = np.array(all_regrets_2)\n",
    "all_regrets_3_np = np.array(all_regrets_3)\n",
    "all_regrets_4_np = np.array(all_regrets_4)\n",
    "all_regrets_5_np = np.array(all_regrets_5)\n",
    "all_regrets_6_np = np.array(all_regrets_6)\n",
    "\n",
    "all_regrets_1_np_mean = np.mean(all_regrets_1_np, axis=0)\n",
    "all_regrets_2_np_mean = np.mean(all_regrets_2_np, axis=0)\n",
    "all_regrets_3_np_mean = np.mean(all_regrets_3_np, axis=0)\n",
    "all_regrets_4_np_mean = np.mean(all_regrets_4_np, axis=0)\n",
    "all_regrets_5_np_mean = np.mean(all_regrets_5_np, axis=0)\n",
    "all_regrets_6_np_mean = np.mean(all_regrets_6_np, axis=0)\n",
    "\n",
    "all_regrets_1_np_stderr = np.std(all_regrets_1_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_2_np_stderr = np.std(all_regrets_2_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_3_np_stderr = np.std(all_regrets_3_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_4_np_stderr = np.std(all_regrets_4_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_5_np_stderr = np.std(all_regrets_5_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_6_np_stderr = np.std(all_regrets_6_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "\n",
    "all_regrets_1_np_ub = all_regrets_1_np_mean + all_regrets_1_np_stderr\n",
    "all_regrets_1_np_lb = all_regrets_1_np_mean - all_regrets_1_np_stderr\n",
    "all_regrets_2_np_ub = all_regrets_2_np_mean + all_regrets_2_np_stderr\n",
    "all_regrets_2_np_lb = all_regrets_2_np_mean - all_regrets_2_np_stderr\n",
    "all_regrets_3_np_ub = all_regrets_3_np_mean + all_regrets_3_np_stderr\n",
    "all_regrets_3_np_lb = all_regrets_3_np_mean - all_regrets_3_np_stderr\n",
    "all_regrets_4_np_ub = all_regrets_4_np_mean + all_regrets_4_np_stderr\n",
    "all_regrets_4_np_lb = all_regrets_4_np_mean - all_regrets_4_np_stderr\n",
    "all_regrets_5_np_ub = all_regrets_5_np_mean + all_regrets_5_np_stderr\n",
    "all_regrets_5_np_lb = all_regrets_5_np_mean - all_regrets_5_np_stderr\n",
    "all_regrets_6_np_ub = all_regrets_6_np_mean + all_regrets_6_np_stderr\n",
    "all_regrets_6_np_lb = all_regrets_6_np_mean - all_regrets_6_np_stderr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a0e38a1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7bca62c1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:35.326024Z",
     "start_time": "2022-05-25T08:30:34.982715Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 89456 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "lw = 2.0\n",
    "plt.rc('font', size=22)\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.rcParams[\"legend.framealpha\"] = 0.5\n",
    "plt.rcParams[\"legend.loc\"] = \"best\"\n",
    "\n",
    "inds = np.arange(1, 100000)\n",
    "plot_len = 5000\n",
    "\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_1_np_mean[:plot_len], all_regrets_1_np_ub[:plot_len], all_regrets_1_np_lb[:plot_len], color_mean=\"b\", color_shading=\"b\", marker=\"v\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_2_np_mean[:plot_len], all_regrets_2_np_ub[:plot_len], all_regrets_2_np_lb[:plot_len], color_mean=\"r\", color_shading=\"r\", marker=\"D\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_3_np_mean[:plot_len], all_regrets_3_np_ub[:plot_len], all_regrets_3_np_lb[:plot_len], color_mean=\"g\", color_shading=\"g\", marker=\"s\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_4_np_mean[:plot_len], all_regrets_4_np_ub[:plot_len], all_regrets_4_np_lb[:plot_len], color_mean=\"c\", color_shading=\"c\", marker=\"o\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_5_np_mean[:plot_len], all_regrets_5_np_ub[:plot_len], all_regrets_5_np_lb[:plot_len], color_mean=\"y\", color_shading=\"y\", marker=\"x\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_6_np_mean[:plot_len], all_regrets_6_np_ub[:plot_len], all_regrets_6_np_lb[:plot_len], color_mean=\"m\", color_shading=\"m\", marker=\"^\", marker_size=12)\n",
    "\n",
    "plt.legend((r\"Neural UCB\", r\"Neural TS\", \\\n",
    "            r\"FN-UCB ($N=1$)\", r\"FN-UCB ($N=2$)\", r\"FN-UCB ($N=5$)\", r\"FN-UCB ($N=10$)\", \\\n",
    "           ), prop={'size':16})\n",
    "\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([0, 5000])\n",
    "# axes.set_ylim([0, 130])\n",
    "\n",
    "plt.title(\"square, Fig. 1b\")\n",
    "plt.ylabel(\"Cumulative Regret\")\n",
    "plt.xlabel(\"Iterations\")\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c188fd5d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e5c62969",
   "metadata": {},
   "source": [
    "# shuttle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b352466",
   "metadata": {},
   "source": [
    "## no diagonalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d43d1148",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:40.831444Z",
     "start_time": "2022-05-25T08:30:40.612163Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "all_regrets_1 = []\n",
    "all_regrets_2 = []\n",
    "all_regrets_3 = []\n",
    "all_regrets_4 = []\n",
    "all_regrets_5 = []\n",
    "\n",
    "run_list = np.arange(3)\n",
    "for itr in run_list:\n",
    "    dataset = \"shuttle\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_1.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_ts.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_2.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    N = 1\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_3.append(values_acc)\n",
    "    \n",
    "    dataset = \"shuttle\"\n",
    "    N = 2\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_4.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    N = 5\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_5.append(values_acc)\n",
    "    \n",
    "\n",
    "all_regrets_1_np = np.array(all_regrets_1)\n",
    "all_regrets_2_np = np.array(all_regrets_2)\n",
    "all_regrets_3_np = np.array(all_regrets_3)\n",
    "all_regrets_4_np = np.array(all_regrets_4)\n",
    "all_regrets_5_np = np.array(all_regrets_5)\n",
    "\n",
    "all_regrets_1_np_mean = np.mean(all_regrets_1_np, axis=0)\n",
    "all_regrets_2_np_mean = np.mean(all_regrets_2_np, axis=0)\n",
    "all_regrets_3_np_mean = np.mean(all_regrets_3_np, axis=0)\n",
    "all_regrets_4_np_mean = np.mean(all_regrets_4_np, axis=0)\n",
    "all_regrets_5_np_mean = np.mean(all_regrets_5_np, axis=0)\n",
    "\n",
    "all_regrets_1_np_stderr = np.std(all_regrets_1_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_2_np_stderr = np.std(all_regrets_2_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_3_np_stderr = np.std(all_regrets_3_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_4_np_stderr = np.std(all_regrets_4_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_5_np_stderr = np.std(all_regrets_5_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "\n",
    "all_regrets_1_np_ub = all_regrets_1_np_mean + all_regrets_1_np_stderr\n",
    "all_regrets_1_np_lb = all_regrets_1_np_mean - all_regrets_1_np_stderr\n",
    "all_regrets_2_np_ub = all_regrets_2_np_mean + all_regrets_2_np_stderr\n",
    "all_regrets_2_np_lb = all_regrets_2_np_mean - all_regrets_2_np_stderr\n",
    "all_regrets_3_np_ub = all_regrets_3_np_mean + all_regrets_3_np_stderr\n",
    "all_regrets_3_np_lb = all_regrets_3_np_mean - all_regrets_3_np_stderr\n",
    "all_regrets_4_np_ub = all_regrets_4_np_mean + all_regrets_4_np_stderr\n",
    "all_regrets_4_np_lb = all_regrets_4_np_mean - all_regrets_4_np_stderr\n",
    "all_regrets_5_np_ub = all_regrets_5_np_mean + all_regrets_5_np_stderr\n",
    "all_regrets_5_np_lb = all_regrets_5_np_mean - all_regrets_5_np_stderr\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de6fdb1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1b943414",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:30:55.454557Z",
     "start_time": "2022-05-25T08:30:55.186597Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 96172 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "lw = 2.0\n",
    "plt.rc('font', size=22)\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.rcParams[\"legend.framealpha\"] = 0.5\n",
    "plt.rcParams[\"legend.loc\"] = \"best\"\n",
    "\n",
    "inds = np.arange(1, 100000)\n",
    "plot_len = 5000\n",
    "\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_1_np_mean[:plot_len], all_regrets_1_np_ub[:plot_len], all_regrets_1_np_lb[:plot_len], color_mean=\"b\", color_shading=\"b\", marker=\"v\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_2_np_mean[:plot_len], all_regrets_2_np_ub[:plot_len], all_regrets_2_np_lb[:plot_len], color_mean=\"r\", color_shading=\"r\", marker=\"D\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_3_np_mean[:plot_len], all_regrets_3_np_ub[:plot_len], all_regrets_3_np_lb[:plot_len], color_mean=\"g\", color_shading=\"g\", marker=\"s\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_4_np_mean[:plot_len], all_regrets_4_np_ub[:plot_len], all_regrets_4_np_lb[:plot_len], color_mean=\"c\", color_shading=\"c\", marker=\"o\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_5_np_mean[:plot_len], all_regrets_5_np_ub[:plot_len], all_regrets_5_np_lb[:plot_len], color_mean=\"y\", color_shading=\"y\", marker=\"x\", marker_size=12)\n",
    "\n",
    "plt.legend((r\"Neural UCB\", r\"Neural TS\", \\\n",
    "            r\"FN-UCB ($N=1$)\", r\"FN-UCB ($N=2$)\", r\"FN-UCB ($N=5$)\", \\\n",
    "           ), prop={'size':16})\n",
    "\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([0, 5000])\n",
    "axes.set_ylim([0, 800])\n",
    "\n",
    "plt.title(\"shuttle, no diagonalization, Fig. 2a\")\n",
    "plt.ylabel(\"Cumulative Regret\")\n",
    "plt.xlabel(\"Iterations\")\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c3df2c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "db51d48d",
   "metadata": {},
   "source": [
    "## with diagonalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "20ba412f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:31:17.290259Z",
     "start_time": "2022-05-25T08:31:17.031988Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "all_regrets_1 = []\n",
    "all_regrets_2 = []\n",
    "all_regrets_3 = []\n",
    "all_regrets_4 = []\n",
    "all_regrets_5 = []\n",
    "\n",
    "run_list = np.arange(3)\n",
    "for itr in run_list:\n",
    "    dataset = \"shuttle\"\n",
    "    depth = 1\n",
    "    width = 50\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_diag.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_1.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    depth = 1\n",
    "    width = 50\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_diag.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_ts.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_2.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    N = 1\n",
    "    depth = 1\n",
    "    width = 50\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_diag.pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_3.append(values_acc)\n",
    "    \n",
    "    dataset = \"shuttle\"\n",
    "    N = 2\n",
    "    depth = 1\n",
    "    width = 50\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_diag.pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_4.append(values_acc)\n",
    "\n",
    "    dataset = \"shuttle\"\n",
    "    N = 5\n",
    "    depth = 1\n",
    "    width = 50\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_diag.pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_5.append(values_acc)\n",
    "    \n",
    "\n",
    "all_regrets_1_np = np.array(all_regrets_1)\n",
    "all_regrets_2_np = np.array(all_regrets_2)\n",
    "all_regrets_3_np = np.array(all_regrets_3)\n",
    "all_regrets_4_np = np.array(all_regrets_4)\n",
    "all_regrets_5_np = np.array(all_regrets_5)\n",
    "\n",
    "all_regrets_1_np_mean = np.mean(all_regrets_1_np, axis=0)\n",
    "all_regrets_2_np_mean = np.mean(all_regrets_2_np, axis=0)\n",
    "all_regrets_3_np_mean = np.mean(all_regrets_3_np, axis=0)\n",
    "all_regrets_4_np_mean = np.mean(all_regrets_4_np, axis=0)\n",
    "all_regrets_5_np_mean = np.mean(all_regrets_5_np, axis=0)\n",
    "\n",
    "all_regrets_1_np_stderr = np.std(all_regrets_1_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_2_np_stderr = np.std(all_regrets_2_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_3_np_stderr = np.std(all_regrets_3_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_4_np_stderr = np.std(all_regrets_4_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_5_np_stderr = np.std(all_regrets_5_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "\n",
    "all_regrets_1_np_ub = all_regrets_1_np_mean + all_regrets_1_np_stderr\n",
    "all_regrets_1_np_lb = all_regrets_1_np_mean - all_regrets_1_np_stderr\n",
    "all_regrets_2_np_ub = all_regrets_2_np_mean + all_regrets_2_np_stderr\n",
    "all_regrets_2_np_lb = all_regrets_2_np_mean - all_regrets_2_np_stderr\n",
    "all_regrets_3_np_ub = all_regrets_3_np_mean + all_regrets_3_np_stderr\n",
    "all_regrets_3_np_lb = all_regrets_3_np_mean - all_regrets_3_np_stderr\n",
    "all_regrets_4_np_ub = all_regrets_4_np_mean + all_regrets_4_np_stderr\n",
    "all_regrets_4_np_lb = all_regrets_4_np_mean - all_regrets_4_np_stderr\n",
    "all_regrets_5_np_ub = all_regrets_5_np_mean + all_regrets_5_np_stderr\n",
    "all_regrets_5_np_lb = all_regrets_5_np_mean - all_regrets_5_np_stderr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c79d3b4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b2a5fc1d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:31:30.774761Z",
     "start_time": "2022-05-25T08:31:30.538201Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88480 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "lw = 2.0\n",
    "plt.rc('font', size=22)\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.rcParams[\"legend.framealpha\"] = 0.5\n",
    "plt.rcParams[\"legend.loc\"] = \"best\"\n",
    "\n",
    "inds = np.arange(1, 100000)\n",
    "plot_len = 5000\n",
    "\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_1_np_mean[:plot_len], all_regrets_1_np_ub[:plot_len], all_regrets_1_np_lb[:plot_len], color_mean=\"b\", color_shading=\"b\", marker=\"v\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_2_np_mean[:plot_len], all_regrets_2_np_ub[:plot_len], all_regrets_2_np_lb[:plot_len], color_mean=\"r\", color_shading=\"r\", marker=\"D\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_3_np_mean[:plot_len], all_regrets_3_np_ub[:plot_len], all_regrets_3_np_lb[:plot_len], color_mean=\"g\", color_shading=\"g\", marker=\"s\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_4_np_mean[:plot_len], all_regrets_4_np_ub[:plot_len], all_regrets_4_np_lb[:plot_len], color_mean=\"c\", color_shading=\"c\", marker=\"o\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_5_np_mean[:plot_len], all_regrets_5_np_ub[:plot_len], all_regrets_5_np_lb[:plot_len], color_mean=\"y\", color_shading=\"y\", marker=\"x\", marker_size=12)\n",
    "\n",
    "plt.legend((r\"Neural UCB\", r\"Neural TS\", \\\n",
    "            r\"FN-UCB ($N=1$)\", r\"FN-UCB ($N=2$)\", r\"FN-UCB ($N=5$)\", \\\n",
    "           ), prop={'size':16})\n",
    "\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([0, 5000])\n",
    "axes.set_ylim([0, 305])\n",
    "\n",
    "plt.title(\"shuttle, with diagonalization, Fig. 2c\")\n",
    "plt.ylabel(\"Cumulative Regret\")\n",
    "plt.xlabel(\"Iterations\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f169961",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a6125e5a",
   "metadata": {},
   "source": [
    "# magic telescope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "84058ec4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:31:38.398413Z",
     "start_time": "2022-05-25T08:31:38.157117Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "all_regrets_1 = []\n",
    "all_regrets_2 = []\n",
    "all_regrets_3 = []\n",
    "all_regrets_4 = []\n",
    "all_regrets_5 = []\n",
    "\n",
    "run_list = np.arange(3)\n",
    "for itr in run_list:\n",
    "    dataset = \"MagicTelescope\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_1.append(values_acc)\n",
    "\n",
    "    dataset = \"MagicTelescope\"\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.01\n",
    "    local_training_iter = 30\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_train_steps_\" + str(local_training_iter) + \"_neural_ucb.pkl\"\n",
    "    log_file_name = log_file_name[:-4] + \"_ts.pkl\"\n",
    "    regrets = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(regrets)\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_2.append(values_acc)\n",
    "\n",
    "    dataset = \"MagicTelescope\"\n",
    "    N = 1\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_3.append(values_acc)\n",
    "    \n",
    "    dataset = \"MagicTelescope\"\n",
    "    N = 2\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_4.append(values_acc)\n",
    "\n",
    "    dataset = \"MagicTelescope\"\n",
    "    N = 5\n",
    "    depth = 1\n",
    "    width = 20\n",
    "    lam, nu = 10, 0.1\n",
    "    local_training_iter = 30\n",
    "    nu_2 = 0.01\n",
    "    log_file_name = \"results_fn_ucb/regret_iter_\" + str(itr) + \"_dataset_\" + dataset + \\\n",
    "                \"_N_\" + str(N) + \\\n",
    "                \"_depth_\" + str(depth) + \"_width_\" + str(width) + \"_lam_\" + str(lam) + \\\n",
    "                \"_nu_\" + str(nu) + \"_nu2_\" + str(nu_2) + \\\n",
    "                \"_train_steps_\" + str(local_training_iter) + \".pkl\"\n",
    "    res = pickle.load(open(log_file_name, \"rb\"))\n",
    "    regrets = np.array(res[\"regrets\"])\n",
    "    values = np.mean(regrets, axis=1)\n",
    "    values_acc = np.cumsum(values)\n",
    "    all_regrets_5.append(values_acc)\n",
    "    \n",
    "\n",
    "all_regrets_1_np = np.array(all_regrets_1)\n",
    "all_regrets_2_np = np.array(all_regrets_2)\n",
    "all_regrets_3_np = np.array(all_regrets_3)\n",
    "all_regrets_4_np = np.array(all_regrets_4)\n",
    "all_regrets_5_np = np.array(all_regrets_5)\n",
    "\n",
    "all_regrets_1_np_mean = np.mean(all_regrets_1_np, axis=0)\n",
    "all_regrets_2_np_mean = np.mean(all_regrets_2_np, axis=0)\n",
    "all_regrets_3_np_mean = np.mean(all_regrets_3_np, axis=0)\n",
    "all_regrets_4_np_mean = np.mean(all_regrets_4_np, axis=0)\n",
    "all_regrets_5_np_mean = np.mean(all_regrets_5_np, axis=0)\n",
    "\n",
    "all_regrets_1_np_stderr = np.std(all_regrets_1_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_2_np_stderr = np.std(all_regrets_2_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_3_np_stderr = np.std(all_regrets_3_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_4_np_stderr = np.std(all_regrets_4_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "all_regrets_5_np_stderr = np.std(all_regrets_5_np, axis=0) / (np.sqrt(len(run_list)))\n",
    "\n",
    "all_regrets_1_np_ub = all_regrets_1_np_mean + all_regrets_1_np_stderr\n",
    "all_regrets_1_np_lb = all_regrets_1_np_mean - all_regrets_1_np_stderr\n",
    "all_regrets_2_np_ub = all_regrets_2_np_mean + all_regrets_2_np_stderr\n",
    "all_regrets_2_np_lb = all_regrets_2_np_mean - all_regrets_2_np_stderr\n",
    "all_regrets_3_np_ub = all_regrets_3_np_mean + all_regrets_3_np_stderr\n",
    "all_regrets_3_np_lb = all_regrets_3_np_mean - all_regrets_3_np_stderr\n",
    "all_regrets_4_np_ub = all_regrets_4_np_mean + all_regrets_4_np_stderr\n",
    "all_regrets_4_np_lb = all_regrets_4_np_mean - all_regrets_4_np_stderr\n",
    "all_regrets_5_np_ub = all_regrets_5_np_mean + all_regrets_5_np_stderr\n",
    "all_regrets_5_np_lb = all_regrets_5_np_mean - all_regrets_5_np_stderr\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2e7498d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d46af53f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T08:31:39.621459Z",
     "start_time": "2022-05-25T08:31:39.407871Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 82136 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "lw = 2.0\n",
    "plt.rc('font', size=22)\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.rcParams[\"legend.framealpha\"] = 0.5\n",
    "plt.rcParams[\"legend.loc\"] = \"best\"\n",
    "\n",
    "inds = np.arange(1, 100000)\n",
    "plot_len = 5000\n",
    "\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_1_np_mean[:plot_len], all_regrets_1_np_ub[:plot_len], all_regrets_1_np_lb[:plot_len], color_mean=\"b\", color_shading=\"b\", marker=\"v\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_2_np_mean[:plot_len], all_regrets_2_np_ub[:plot_len], all_regrets_2_np_lb[:plot_len], color_mean=\"r\", color_shading=\"r\", marker=\"D\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_3_np_mean[:plot_len], all_regrets_3_np_ub[:plot_len], all_regrets_3_np_lb[:plot_len], color_mean=\"g\", color_shading=\"g\", marker=\"s\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_4_np_mean[:plot_len], all_regrets_4_np_ub[:plot_len], all_regrets_4_np_lb[:plot_len], color_mean=\"c\", color_shading=\"c\", marker=\"o\", marker_size=12)\n",
    "plot_mean_and_CI_with_marker(inds[:plot_len], all_regrets_5_np_mean[:plot_len], all_regrets_5_np_ub[:plot_len], all_regrets_5_np_lb[:plot_len], color_mean=\"y\", color_shading=\"y\", marker=\"x\", marker_size=12)\n",
    "\n",
    "plt.legend((r\"Neural UCB\", r\"Neural TS\", \\\n",
    "            r\"FN-UCB ($N=1$)\", r\"FN-UCB ($N=2$)\", r\"FN-UCB ($N=5$)\", \\\n",
    "           ), prop={'size':16})\n",
    "\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([0, 5000])\n",
    "axes.set_ylim([0, 955])\n",
    "\n",
    "plt.title(\"magic telescope, Fig. 2b\")\n",
    "plt.ylabel(\"Cumulative Regret\")\n",
    "plt.xlabel(\"Iterations\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e426dc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
