insulin_dosage_exp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure plotting\n",
    "%matplotlib inline\n",
    "\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "matplotlib.rcParams['figure.figsize'] = (10, 4)\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "matplotlib.rcParams['text.latex.preamble'] = r'\\usepackage{amsmath}'\n",
    "matplotlib.rcParams['font.size'] = 16\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import pickle\n",
    "from pathlib import Path\n",
    "from datetime import datetime\n",
    "from itertools import product\n",
    "\n",
    "import GPy\n",
    "import numpy as np\n",
    "RANDOM_SEED=123\n",
    "np.random.seed(RANDOM_SEED)\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "from utils.model import GPyModel\n",
    "\n",
    "from utils.utils import discretized_normal_distribution, sample_index_from_p\n",
    "from utils.utils import w_x_t, kappa_x_w, MMD, worst_context_distribution_DRO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Domain definition\n",
    "\n",
    "# Lower and upper limits\n",
    "ins_lim = [0, 30]\n",
    "carb_lim = [20, 80]\n",
    "fbg_lim = [100, 150]\n",
    "\n",
    "N_ins = ins_lim[1] - ins_lim[0] + 1  # Size of the action (insulin) set, we use integer values\n",
    "N_carb = 20  # Size of the context (carbohydrate) set\n",
    "N_fbg = 1  # Size of the fasting blood glucose set, we use a fixed value\n",
    "\n",
    "ins = np.linspace(ins_lim[0], ins_lim[1], N_ins)  # Vector of ins values\n",
    "carb = np.linspace(carb_lim[0], carb_lim[1], N_carb)  # Vector of carb values\n",
    "fbg = np.array([100])  # We fix the fasting blood glucose level to 100\n",
    "\n",
    "InsCarbFbg_pairs = np.array(list(product(ins, carb, fbg))).reshape(N_ins, N_carb, N_fbg, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_lengthscale = 10  # Lengthscale for M and GP kernel\n",
    "\n",
    "# Generation of matrix M for MMD\n",
    "Carb_kern = GPy.kern.Matern52(input_dim=1, lengthscale=obj_lengthscale)\n",
    "M = Carb_kern.K(carb.reshape(N_carb, 1))\n",
    "\n",
    "# plt.imshow(M, interpolation='bicubic')\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_w_true_and_ref(t=None):\n",
    "    \"\"\"Generates and returns true (w_true) and reference (w_ref) distributions according to the\n",
    "    experiment setup. If t is not None, distributional shifts converge to 0 with rate 1/log(t).\"\"\"\n",
    "\n",
    "    w_true = np.zeros((N_carb, 1))\n",
    "    w_true_ind = np.random.randint(N_carb)\n",
    "    \n",
    "    w_ref_ind_offset = (np.random.randint(13) - 6)  # Drift RV is U(-6, 6)\n",
    "    if t is not None:\n",
    "        w_ref_ind_offset = int(np.round(w_ref_ind_offset / np.log2(t)))\n",
    "    w_ref_ind = np.clip(w_true_ind + w_ref_ind_offset, a_min=0, a_max=N_carb-1)\n",
    "    \n",
    "    w_ref = discretized_normal_distribution(carb, carb[w_ref_ind], (1.5)**2)\n",
    "    w_true = discretized_normal_distribution(carb, carb[w_true_ind], (3)**2)\n",
    "\n",
    "    return w_true, w_ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simglucose_simulation(meal, fasting_bg, insulin):\n",
    "    \"\"\"\n",
    "    :param meal: float, amount of glucose intake\n",
    "    :param fasting_bg: float, fasting blood glucose of the patient\n",
    "    :param insulin: float, amount of insulin intake\n",
    "\n",
    "    Simulation results precomputed due to license restrictions on simulator usage.\n",
    "    \"\"\"\n",
    "    \n",
    "    result_dict_path = 'functions/simglucose_outputs.pkl'\n",
    "    if os.path.exists(result_dict_path):\n",
    "        with open(result_dict_path, 'rb') as f:\n",
    "            old_vals = pickle.load(f)\n",
    "    else:\n",
    "        old_vals = {}\n",
    "\n",
    "    old_bgl = old_vals.get((fasting_bg, meal, insulin))\n",
    "    if old_bgl is not None:\n",
    "        return old_bgl\n",
    "\n",
    "    raise RuntimeError(\"Given simulation configuration is not precomputed.\")\n",
    "\n",
    "    # new_bgl = simpatient(\"adolescent#001\", fasting_bg, meal, insulin, 0, 150)['bgl']\n",
    "    # old_vals[(fasting_bg, meal, insulin)] = new_bgl\n",
    "    # with open(result_dict_path, 'rb+') as f:\n",
    "    #     pickle.dump(old_vals, f)\n",
    "    \n",
    "    # return new_bgl\n",
    "\n",
    "\n",
    "def get_observation(meal, fasting_bg, insulin, noise_std=0):\n",
    "    \"\"\"Given a simulation configuration, returns the noisy observation.\"\"\"\n",
    "    bgl = run_simglucose_simulation(meal, fasting_bg, insulin)\n",
    "\n",
    "    bgl = bgl + np.random.randn(*bgl.shape) * noise_std\n",
    "\n",
    "    # Our pseudo-objective\n",
    "    return -abs(bgl - 112.5)\n",
    "\n",
    "print(run_simglucose_simulation(0, 80, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create objective function as an array. InsCarb_mat corresponds to f.\n",
    "InsCarb_mat = np.empty((N_carb, N_ins))\n",
    "for i in range(N_carb):\n",
    "    for j in range(N_ins):\n",
    "        InsCarb_mat[i, j] = get_observation(carb[i], fbg[0], ins[j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BO_loop_DRO(model, distributions, contexts, epsilon_coef, beta, T):\n",
    "    InsCarbFbg_t_indices = []\n",
    "    for t in range(T):\n",
    "        # True and reference distributions\n",
    "        w_true, w_ref = distributions[t]\n",
    "\n",
    "        # The variable epsilon here is the radius of the ambiguity set of DRBO\n",
    "        # true distance multiplied with a coefficient, e.g. 1/3\n",
    "        epsilon = MMD(M, w_true, w_ref) * epsilon_coef\n",
    "        print(\"Epsilon:\", epsilon)\n",
    "\n",
    "        # Fasting blood glucose is fixed.\n",
    "        fbg_index = contexts[t, 1]\n",
    "\n",
    "        # Modelling\n",
    "        mu, var = model.predict(InsCarbFbg_pairs[:, :, fbg_index, :].reshape(-1, model.input_dim))\n",
    "        InsCarb_pair_UCB = mu + beta*np.sqrt(var)\n",
    "        InsCarb_mat_UCB = InsCarb_pair_UCB.reshape(N_ins, N_carb)\n",
    "\n",
    "        # For each action(insulin), calculate worst expected reward in the ambiguity set\n",
    "        worst_expected_reward_x = np.zeros(N_ins)\n",
    "        for ins_ind in range(N_ins):\n",
    "            UCB_ins = InsCarb_mat_UCB[ins_ind]\n",
    "            worst_w_x = worst_context_distribution_DRO(M, w_ref, UCB_ins, epsilon)\n",
    "            worst_expected_reward_x[ins_ind] = np.dot(worst_w_x.T, UCB_ins)\n",
    "        \n",
    "        # Choose maximum of these worst expected rewards.\n",
    "        chosen_action_index = np.argmax(worst_expected_reward_x)\n",
    "        # Context is observed from precomputed context vector.\n",
    "        observed_carb_index = contexts[t, 0]\n",
    "        # Add selected action and observed context indices to the array to be returned\n",
    "        InsCarbFbg_t_indices.append([chosen_action_index, observed_carb_index, fbg_index])\n",
    "        \n",
    "        # Add new sample to model and update the model\n",
    "        InsCarbFbg_sample = InsCarbFbg_pairs[chosen_action_index, observed_carb_index, fbg_index].reshape(1, -1)\n",
    "        insulin_ts, carb_ts, fbg_ts = InsCarbFbg_sample[0]\n",
    "        Y_sample_n = np.array(\n",
    "            get_observation(carb_ts, fbg_ts, insulin_ts, noise_std=model.noise_std)\n",
    "        ).reshape(1, -1)\n",
    "        model.add_sample(InsCarbFbg_sample, Y_sample_n)\n",
    "        model.update()\n",
    "        \n",
    "        print(f\"Chosen point - insulin: {insulin_ts}, carb.: {carb_ts}, fbg: {fbg_ts}\")\n",
    "        print(\"Blood glucose(noisy) discrepancy:\", Y_sample_n[0][0])\n",
    "\n",
    "    return InsCarbFbg_t_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BO_loop_WRO(model, distributions, contexts, beta, T):\n",
    "    InsCarbFbg_t_indices = []\n",
    "    for t in range(T):\n",
    "        # Fasting blood glucose is fixed.\n",
    "        fbg_index = contexts[t, 1]\n",
    "\n",
    "        # Modelling\n",
    "        mu, var = model.predict(InsCarbFbg_pairs[:, :, fbg_index, :].reshape(-1, model.input_dim))\n",
    "        InsCarb_pair_UCB = mu + beta*np.sqrt(var)\n",
    "        InsCarb_mat_UCB = InsCarb_pair_UCB.reshape(N_ins, N_carb)\n",
    "\n",
    "        # Choose the action that maximizes the minimum over context.\n",
    "        chosen_action_index = np.argmax(np.min(InsCarb_mat_UCB, axis=1))\n",
    "        # Context is observed from precomputed context vector.\n",
    "        observed_carb_index = contexts[t, 0]\n",
    "        # Add selected action and observed context indices to the array to be returned\n",
    "        InsCarbFbg_t_indices.append([chosen_action_index, observed_carb_index, fbg_index])\n",
    "\n",
    "        # Add new sample to model and update the model\n",
    "        InsCarbFbg_sample = InsCarbFbg_pairs[chosen_action_index, observed_carb_index, fbg_index].reshape(1, -1)\n",
    "        insulin_ts, carb_ts, fbg_ts = InsCarbFbg_sample[0]\n",
    "        Y_sample_n = np.array(\n",
    "            get_observation(carb_ts, fbg_ts, insulin_ts, noise_std=model.noise_std)\n",
    "        ).reshape(1, -1)\n",
    "        model.add_sample(InsCarbFbg_sample, Y_sample_n)\n",
    "        model.update()\n",
    "\n",
    "        print(f\"Chosen point - insulin: {insulin_ts}, carb.: {carb_ts}, fbg: {fbg_ts}\")\n",
    "        print(\"Blood glucose(noisy) discrepancy:\", Y_sample_n[0][0])\n",
    "\n",
    "    return InsCarbFbg_t_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BO_loop_SO(model, distributions, contexts, beta, T):\n",
    "    InsCarbFbg_t_indices = []\n",
    "    for t in range(T):\n",
    "        # True and reference distributions\n",
    "        w_true, w_ref = distributions[t]\n",
    "\n",
    "        # Fasting blood glucose is fixed.\n",
    "        fbg_index = contexts[t, 1]\n",
    "\n",
    "        # Modelling\n",
    "        mu, var = model.predict(InsCarbFbg_pairs[:, :, fbg_index, :].reshape(-1, model.input_dim))\n",
    "        InsCarb_pair_UCB = mu + beta*np.sqrt(var)\n",
    "        InsCarb_mat_UCB = InsCarb_pair_UCB.reshape(N_ins, N_carb)\n",
    "\n",
    "        # Choose the action that maximizes the expected reward over reference distribution.\n",
    "        chosen_action_index = np.argmax(InsCarb_mat_UCB @ w_ref)\n",
    "        # Context is observed from precomputed context vector.\n",
    "        observed_carb_index = contexts[t, 0]\n",
    "        # Add selected action and observed context indices to the array to be returned\n",
    "        InsCarbFbg_t_indices.append([chosen_action_index, observed_carb_index, fbg_index])\n",
    "\n",
    "        # Add new sample to model and update the model\n",
    "        InsCarbFbg_sample = InsCarbFbg_pairs[chosen_action_index, observed_carb_index, fbg_index].reshape(1, -1)\n",
    "        insulin_ts, carb_ts, fbg_ts = InsCarbFbg_sample[0]\n",
    "        Y_sample_n = np.array(\n",
    "            get_observation(carb_ts, fbg_ts, insulin_ts, noise_std=model.noise_std)\n",
    "        ).reshape(1, -1)\n",
    "        model.add_sample(InsCarbFbg_sample, Y_sample_n)\n",
    "        model.update()\n",
    "        \n",
    "        print(f\"Chosen point - insulin: {insulin_ts}, carb.: {carb_ts}, fbg: {fbg_ts}\")\n",
    "        print(\"Blood glucose(noisy) discrepancy:\", Y_sample_n[0][0])\n",
    "\n",
    "    return InsCarbFbg_t_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BO_loop_RS(model, distributions, contexts, beta, T, taus):    \n",
    "    kappas = np.zeros(T)\n",
    "    InsCarbFbg_t_indices = []\n",
    "    for t in range(T):\n",
    "        # True and reference distributions\n",
    "        w_true, w_ref = distributions[t]\n",
    "        w_true = w_true.reshape(N_carb, 1)\n",
    "        w_ref = w_ref.reshape(N_carb, 1)\n",
    "        \n",
    "        tau = taus[t]\n",
    "\n",
    "        # Fasting blood glucose is fixed.\n",
    "        fbg_index = contexts[t, 1]\n",
    "            \n",
    "        # Modelling\n",
    "        mu, var = model.predict(InsCarbFbg_pairs[:, :, fbg_index, :].reshape(-1, model.input_dim))\n",
    "        InsCarb_pair_UCB = mu + beta*np.sqrt(var)\n",
    "        InsCarb_mat_UCB = InsCarb_pair_UCB.reshape(N_ins, N_carb)\n",
    "\n",
    "        # For each action(insulin), calculate kappa_hat_tau,t and kappa_tau,t\n",
    "        list_kappa_x = np.zeros(N_ins)\n",
    "        list_kappa_hat_x = np.zeros(N_ins)\n",
    "        for ins_ind in range(N_ins):\n",
    "            UCB_ins = InsCarb_mat_UCB[ins_ind]\n",
    "            w_bar_x_t = w_x_t(M, w_ref, UCB_ins, tau)\n",
    "            kappa_hat_x = kappa_x_w(M, w_ref, w_bar_x_t, UCB_ins, tau, clip=False)\n",
    "            list_kappa_hat_x[ins_ind] = kappa_hat_x\n",
    "            \n",
    "            f_x = InsCarb_mat[:, ins_ind].reshape(-1, 1)\n",
    "            w_dbar_x_t = w_x_t(M, w_ref, f_x, tau)\n",
    "            kappa_x = kappa_x_w(M, w_ref, w_dbar_x_t, f_x, tau, clip=False)\n",
    "            list_kappa_x[ins_ind] = kappa_x\n",
    "\n",
    "        # Choose action with minimum kappa_hat.\n",
    "        chosen_action_index = np.argmin(list_kappa_hat_x)\n",
    "        # Context is observed from precomputed context vector.\n",
    "        observed_carb_index = contexts[t, 0]\n",
    "        # Add selected action and observed context indices to the array to be returned\n",
    "        InsCarbFbg_t_indices.append([chosen_action_index, observed_carb_index, fbg_index])\n",
    "        \n",
    "        # Add new sample to model and update the model\n",
    "        InsCarbFbg_sample = InsCarbFbg_pairs[chosen_action_index, observed_carb_index, fbg_index].reshape(1, -1)\n",
    "        insulin_ts, carb_ts, fbg_ts = InsCarbFbg_sample[0]\n",
    "        Y_sample_n = np.array(\n",
    "            get_observation(carb_ts, fbg_ts, insulin_ts, noise_std=model.noise_std)\n",
    "        ).reshape(1, -1)\n",
    "        model.add_sample(InsCarbFbg_sample, Y_sample_n)\n",
    "        model.update()\n",
    "        \n",
    "        print(f\"Chosen point - insulin: {insulin_ts}, carb.: {carb_ts}, fbg: {fbg_ts}\", \" - Kappa:\", np.min(list_kappa_hat_x))\n",
    "        print(\"Blood glucose(noisy) discrepancy:\", Y_sample_n[0][0])\n",
    "\n",
    "        # Save minimum kappa for regret calculation.\n",
    "        min_kappa = np.min(list_kappa_x)\n",
    "        kappas[t] = min_kappa\n",
    "\n",
    "    return InsCarbFbg_t_indices, kappas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model simulation parameters: simulation count, timestep count and dro epsilons\n",
    "\n",
    "# Reset the randomness\n",
    "np.random.seed(RANDOM_SEED)\n",
    "\n",
    "sim_count = 50\n",
    "timestep_count = 200 + 1  # +1 for undersampled plotting (last tick in the plot)\n",
    "initial_sample_cnt = 10\n",
    "\n",
    "# Different ball radius coefficients for DRO\n",
    "DRO_epsilons = [1/3, 1, 3]\n",
    "RS_aspirations = [0.9]  # Currently not used\n",
    "\n",
    "# Time independent tau value\n",
    "fixed_tau = -10\n",
    "assert fixed_tau is None or (fixed_tau is not None and len(RS_aspirations)==1)\n",
    "\n",
    "# Modelling parameters\n",
    "input_dim = 3\n",
    "output_dim = 1\n",
    "noise_std = 1\n",
    "noise_var = np.square(noise_std)\n",
    "beta = 2\n",
    "\n",
    "ker = GPy.kern.Matern52(input_dim=input_dim, ARD=True, lengthscale=obj_lengthscale)\n",
    "\n",
    "CONVERGE_REF_TO_TRUE = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Model simulations\n",
    "\n",
    "# Initial samples for each simulation\n",
    "initial_samples_sims = np.empty((sim_count, initial_sample_cnt, input_dim+output_dim))\n",
    "# Time dependent true and reference distributions for each simulation\n",
    "distributions_w_true_and_ref_sims = np.empty((sim_count, timestep_count, 2, N_carb))\n",
    "# Chosen (action, context) pairs for each simulation of each DRO epsilon\n",
    "indices_DRO_sims_eps = np.empty((len(DRO_epsilons), sim_count, timestep_count, 3))\n",
    "# Chosen (action, context) pairs for each simulation of WRO\n",
    "indices_WRO_sims = np.empty((1, sim_count, timestep_count, 3))\n",
    "# Chosen (action, context) pairs for each simulation of SO\n",
    "indices_SO_sims = np.empty((1, sim_count, timestep_count, 3))\n",
    "# Chosen (action, context) pairs for each simulation of RS\n",
    "indices_RS_sims_asp = np.empty((len(RS_aspirations), sim_count, timestep_count, 3))\n",
    "# Kappa prime values for each simulation and timestep, we calculate these in RS loop\n",
    "kappas_sims_asp = np.empty((len(RS_aspirations), sim_count, timestep_count))\n",
    "# Tau values for each simulation and timestep, since we use a fixed tau this is constant\n",
    "taus_sims_asp = np.empty((len(RS_aspirations), sim_count, timestep_count))\n",
    "if fixed_tau is not None:\n",
    "    taus_sims_asp[:] = fixed_tau\n",
    "\n",
    "\n",
    "for sim in range(sim_count):\n",
    "    print(f\"\\n\\n------------------------ SIMULATION NUMBER: {sim:02d} ------------------------\\n\\n\")\n",
    "\n",
    "    # Randomly select initial samples\n",
    "    initial_ins_inds = np.random.randint(low=0, high=N_ins, size=initial_sample_cnt)\n",
    "    initial_carb_inds = np.random.randint(low=0, high=N_carb, size=initial_sample_cnt)\n",
    "    initial_fbgs_inds = np.full(initial_sample_cnt, 0)\n",
    "    initial_indices = np.stack([initial_ins_inds, initial_carb_inds, initial_fbgs_inds], axis=1)\n",
    "\n",
    "    initial_InsCarbFbg = np.empty((initial_sample_cnt, 3))\n",
    "    initial_Y = np.empty((initial_sample_cnt, 1))\n",
    "    for i in range(initial_sample_cnt):\n",
    "        initial_InsCarbFbg[i] = InsCarbFbg_pairs[\n",
    "            initial_indices[i][0], initial_indices[i][1], initial_indices[i][2]\n",
    "        ]\n",
    "        insulin_tmp, carb_tmp, fbg_tmp = initial_InsCarbFbg[i]\n",
    "        initial_Y[i, 0] = get_observation(carb_tmp, fbg_tmp, insulin_tmp)\n",
    "\n",
    "    initial_samples_sims[sim] = np.hstack((initial_InsCarbFbg, initial_Y))\n",
    "\n",
    "\n",
    "    # Create context distributions and sample context at each timestep for all algorithms.\n",
    "    # All algorithms run with same distributions and contexts for each simulation.\n",
    "    sim_contexts = np.empty((timestep_count, 2), dtype=int)\n",
    "    for t in range(timestep_count):\n",
    "        if CONVERGE_REF_TO_TRUE:\n",
    "            w_true, w_ref = create_w_true_and_ref(t=t+2)\n",
    "        else:\n",
    "            w_true, w_ref = create_w_true_and_ref()\n",
    "        distributions_w_true_and_ref_sims[sim, t, 0] = w_true.flatten()\n",
    "        distributions_w_true_and_ref_sims[sim, t, 1] = w_ref.flatten()\n",
    "        ts_carb_context = sample_index_from_p(w_true)\n",
    "        \n",
    "        sim_contexts[t] = np.array([ts_carb_context, 0])\n",
    "    \n",
    "        if fixed_tau is None:\n",
    "            pass\n",
    "        \n",
    "    # DRO simulations with different epsilons\n",
    "    for eps_i, eps in enumerate(DRO_epsilons):\n",
    "        print(f\"\\nDRO with eps={eps:.3f}\\n\")\n",
    "        model_DRO = GPyModel(input_dim, output_dim, noise_var, ker=ker)\n",
    "\n",
    "        model_DRO.add_sample(initial_InsCarbFbg, initial_Y)\n",
    "        model_DRO.update()\n",
    "        \n",
    "        chosen_indices_DRO = BO_loop_DRO(\n",
    "            model_DRO, distributions=distributions_w_true_and_ref_sims[sim], contexts=sim_contexts,\n",
    "            epsilon_coef=eps, beta=beta, T=timestep_count\n",
    "        )\n",
    "        indices_DRO_sims_eps[eps_i][sim] = np.array(chosen_indices_DRO).astype(int)\n",
    "    \n",
    "    # WRO simulation\n",
    "    print(f\"\\nWRO\\n\")\n",
    "    model_WRO = GPyModel(input_dim, output_dim, noise_var, ker=ker)\n",
    "\n",
    "    model_WRO.add_sample(initial_InsCarbFbg, initial_Y)\n",
    "    model_WRO.update()\n",
    "    \n",
    "    chosen_indices_WRO = BO_loop_WRO(\n",
    "        model_WRO, distributions=distributions_w_true_and_ref_sims[sim], contexts=sim_contexts, beta=beta, T=timestep_count\n",
    "    )\n",
    "    indices_WRO_sims[0, sim] = np.array(chosen_indices_WRO).astype(int)\n",
    "\n",
    "    # RS simulation\n",
    "    for asp_i, asp in enumerate(RS_aspirations):\n",
    "        print(f\"\\nRS with aspiration={asp:.3f}\\n\")\n",
    "        model_RS = GPyModel(input_dim, output_dim, noise_var, ker=ker)\n",
    "\n",
    "        model_RS.add_sample(initial_InsCarbFbg, initial_Y)\n",
    "        model_RS.update()\n",
    "        \n",
    "        chosen_indices_RS, kappas = BO_loop_RS(\n",
    "            model_RS, distributions=distributions_w_true_and_ref_sims[sim], contexts=sim_contexts,\n",
    "            beta=beta, T=timestep_count, taus=taus_sims_asp[asp_i, sim]\n",
    "        )\n",
    "        indices_RS_sims_asp[asp_i][sim] = np.array(chosen_indices_RS).astype(int)\n",
    "        kappas_sims_asp[asp_i][sim] = kappas\n",
    "\n",
    "    print(f\"\\nSO\\n\")\n",
    "    model_SO = GPyModel(input_dim, output_dim, noise_var, ker=ker)\n",
    "\n",
    "    model_SO.add_sample(initial_InsCarbFbg, initial_Y)\n",
    "    model_SO.update()\n",
    "\n",
    "    chosen_indices_SO = BO_loop_SO(\n",
    "        model_SO, distributions=distributions_w_true_and_ref_sims[sim], contexts=sim_contexts, beta=beta, T=timestep_count\n",
    "    )\n",
    "    indices_SO_sims[0, sim] = np.array(chosen_indices_SO).astype(int)\n",
    "\n",
    "print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write simulation results\n",
    "\n",
    "simulation_name = datetime.now().strftime(\"%m_%d_%Y-%H_%M_%S\")\n",
    "simulation_folder = os.path.join(\n",
    "    \"results\",\n",
    "    f\"simg_conv{int(CONVERGE_REF_TO_TRUE)}_ts_sim{sim_count}_ts{timestep_count}_tau-{fixed_tau}-std-{noise_std:.2f}-{simulation_name}\"\n",
    ")\n",
    "os.makedirs(simulation_folder, exist_ok=False)\n",
    "\n",
    "np.save(os.path.join(simulation_folder, \"initial_samples_sims\"), initial_samples_sims)\n",
    "np.save(os.path.join(simulation_folder, \"distributions_w_true_and_ref_sims\"), distributions_w_true_and_ref_sims)\n",
    "np.save(os.path.join(simulation_folder, \"indices_DRO_sims_eps\"), indices_DRO_sims_eps)\n",
    "np.save(os.path.join(simulation_folder, \"indices_WRO_sims\"), indices_WRO_sims)\n",
    "np.save(os.path.join(simulation_folder, \"indices_SO_sims\"), indices_SO_sims)\n",
    "np.save(os.path.join(simulation_folder, \"indices_RS_sims_asp\"), indices_RS_sims_asp)\n",
    "np.save(os.path.join(simulation_folder, \"kappas_sims_asp\"), kappas_sims_asp)\n",
    "np.save(os.path.join(simulation_folder, \"taus_sims_asp\"), taus_sims_asp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read simulation results\n",
    "\n",
    "# Change to True to read older experiments\n",
    "if False:\n",
    "    simulation_folder = \"\"\n",
    "    sim_count = 50  # Sim count of older experiment\n",
    "\n",
    "    initial_samples_sims = np.load(os.path.join(simulation_folder, \"initial_samples_sims.npy\"))\n",
    "    distributions_w_true_and_ref_sims = np.load(os.path.join(simulation_folder, \"distributions_w_true_and_ref_sims.npy\"))\n",
    "    indices_DRO_sims_eps = np.load(os.path.join(simulation_folder, \"indices_DRO_sims_eps.npy\")).astype(int)\n",
    "    indices_WRO_sims = np.load(os.path.join(simulation_folder, \"indices_WRO_sims.npy\")).astype(int)\n",
    "    indices_SO_sims = np.load(os.path.join(simulation_folder, \"indices_SO_sims.npy\")).astype(int)\n",
    "    indices_RS_sims_asp = np.load(os.path.join(simulation_folder, \"indices_RS_sims_asp.npy\")).astype(int)\n",
    "    kappas_sims_asp = np.load(os.path.join(simulation_folder, \"kappas_sims_asp.npy\"))\n",
    "    taus_sims_asp = np.load(os.path.join(simulation_folder, \"taus_sims_asp.npy\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot palettes\n",
    "colors = sns.color_palette(\"dark\")\n",
    "greens = sns.color_palette(\"BuGn\", 10)\n",
    "blues = sns.color_palette(\"PuBu\", 10)\n",
    "reds = sns.color_palette(\"YlOrRd\", 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Expected rewards of simulations and timesteps w.r.t. f\n",
    "expected_reward_true_sims_ts = distributions_w_true_and_ref_sims[:, :, 0, :] @ InsCarb_mat\n",
    "\n",
    "# MMD distances of simulations and timesteps\n",
    "mmd_sims_ts = np.empty((sim_count, timestep_count))\n",
    "for sim_i in range(sim_count):\n",
    "    for ts_i in range(timestep_count):\n",
    "        mmd_sims_ts[sim_i, ts_i] = MMD(M, *distributions_w_true_and_ref_sims[sim_i, ts_i])\n",
    "\n",
    "\n",
    "# Calculate rewards\n",
    "\n",
    "# DRO rewards\n",
    "true_rewards_of_actions_DRO_eps = np.empty((len(DRO_epsilons), sim_count, timestep_count))\n",
    "for e_i in range(len(DRO_epsilons)):\n",
    "    true_rewards_of_actions_DRO_eps[e_i] = expected_reward_true_sims_ts[\n",
    "        np.arange(sim_count)[:, None],\n",
    "        np.arange(timestep_count),\n",
    "        indices_DRO_sims_eps[e_i, :, :, 0].astype(int)\n",
    "    ]\n",
    "\n",
    "# WRO rewards\n",
    "true_rewards_of_actions_WRO = expected_reward_true_sims_ts[\n",
    "    np.arange(sim_count)[:, None],\n",
    "    np.arange(timestep_count),\n",
    "    indices_WRO_sims[0, :, :, 0].astype(int)\n",
    "]\n",
    "\n",
    "# SO rewards\n",
    "true_rewards_of_actions_SO = expected_reward_true_sims_ts[\n",
    "    np.arange(sim_count)[:, None],\n",
    "    np.arange(timestep_count),\n",
    "    indices_SO_sims[0, :, :, 0].astype(int)\n",
    "]\n",
    "\n",
    "# RS rewards\n",
    "true_rewards_of_actions_RS_asp = np.empty((len(RS_aspirations), sim_count, timestep_count))\n",
    "for a_i in range(len(RS_aspirations)):\n",
    "    true_rewards_of_actions_RS_asp[a_i] = expected_reward_true_sims_ts[\n",
    "        np.arange(sim_count)[:, None],\n",
    "        np.arange(timestep_count),\n",
    "        indices_RS_sims_asp[a_i, :, :, 0].astype(int)\n",
    "    ]\n",
    "\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = (10, 4)\n",
    "fig, ax = plt.subplots(nrows=len(RS_aspirations), ncols=2)\n",
    "\n",
    "\n",
    "# REGRET PLOTS\n",
    "\n",
    "\n",
    "# Label, color, linestyle\n",
    "lcls = [\n",
    "    (r\"DRBO: $3\\epsilon$\", colors[6], \":\",),\n",
    "    (r\"DRBO: $\\epsilon$\", blues[-2], \":\",),\n",
    "    (r\"DRBO: $\\epsilon/3$\", colors[5], \":\",),\n",
    "    (\"WRBO\", colors[8], \"-.\",),\n",
    "    (\"SO\", greens[-2], \"--\",),\n",
    "    (\"RoBOS\", reds[-1], \"-\",),\n",
    "]\n",
    "\n",
    "\n",
    "def plot_regret_std_bar(ax, regret, ind):\n",
    "    std_div = 2\n",
    "    subsample_rat = 25\n",
    "    ealph = 0.6\n",
    "    elinew = 0.3\n",
    "    c_size = 4\n",
    "\n",
    "    label, color, linestyle = lcls[ind]\n",
    "\n",
    "    x = np.arange(timestep_count)[::subsample_rat]\n",
    "    y = np.cumsum(regret.mean(axis=0))[::subsample_rat]\n",
    "    y_err = (np.cumsum(regret, axis=1).std(axis=0) / std_div)[::subsample_rat]\n",
    "\n",
    "    ax.errorbar(\n",
    "        x, y, y_err, label=label, c=color,\n",
    "        fmt=linestyle, ecolor=(*color, ealph), elinewidth=elinew, capsize=c_size\n",
    "    )\n",
    "\n",
    "# Plot regret values with std regions as error bars\n",
    "plot_func = plot_regret_std_bar\n",
    "\n",
    "\n",
    "# Robust satisficing regret plot\n",
    "for a_i in range(len(RS_aspirations)):\n",
    "    tmp_ax = ax[0] if len(RS_aspirations) == 1 else ax[a_i, 0]\n",
    "    tmp_ax.set_xlabel(\"Round\")\n",
    "    tmp_ax.set_ylabel(r\"$R_t^{\\textit{rs}}$\")\n",
    "\n",
    "    taus_sims = taus_sims_asp[a_i]\n",
    "    kappas_sims = kappas_sims_asp[a_i]\n",
    "\n",
    "    # DROs\n",
    "    for e_i in range(len(DRO_epsilons)):\n",
    "        rs_regret = np.maximum(0, taus_sims - kappas_sims*mmd_sims_ts - true_rewards_of_actions_DRO_eps[e_i])\n",
    "        plot_func(tmp_ax, rs_regret, e_i)\n",
    "    \n",
    "    # WRO\n",
    "    rs_regret = np.maximum(0, taus_sims - kappas_sims*mmd_sims_ts - true_rewards_of_actions_WRO)\n",
    "    plot_func(tmp_ax, rs_regret, e_i+1)\n",
    "\n",
    "    # SO\n",
    "    rs_regret = np.maximum(0, taus_sims - kappas_sims*mmd_sims_ts - true_rewards_of_actions_SO)\n",
    "    plot_func(tmp_ax, rs_regret, e_i+2)\n",
    "    \n",
    "    # RS\n",
    "    rs_regret = np.maximum(0, taus_sims - kappas_sims*mmd_sims_ts - true_rewards_of_actions_RS_asp[a_i])\n",
    "    plot_func(tmp_ax, rs_regret, e_i+3)\n",
    "\n",
    "    leg = tmp_ax.legend(prop={'size': 12}, frameon=True)\n",
    "\n",
    "\n",
    "# Lenient regret plot\n",
    "for a_i in range(len(RS_aspirations)):\n",
    "    tmp_ax = ax[1] if len(RS_aspirations) == 1 else ax[a_i, 1]\n",
    "    tmp_ax.set_xlabel(\"Round\")\n",
    "    tmp_ax.set_ylabel(r\"$R_t^{\\textit{l}}$\")\n",
    "\n",
    "    taus_sims = taus_sims_asp[a_i]\n",
    "    \n",
    "    # DROs\n",
    "    for e_i in range(len(DRO_epsilons)):\n",
    "        lenient_regret = np.maximum(0, taus_sims - true_rewards_of_actions_DRO_eps[e_i])\n",
    "        plot_func(tmp_ax, lenient_regret, e_i)\n",
    "\n",
    "    # WRO\n",
    "    lenient_regret = np.maximum(0, taus_sims - true_rewards_of_actions_WRO)\n",
    "    plot_func(tmp_ax, lenient_regret, e_i+1)\n",
    "\n",
    "    # SO\n",
    "    lenient_regret = np.maximum(0, taus_sims - true_rewards_of_actions_SO)\n",
    "    plot_func(tmp_ax, lenient_regret, e_i+2)\n",
    "    \n",
    "    # RS\n",
    "    lenient_regret = np.maximum(0, taus_sims - true_rewards_of_actions_RS_asp[a_i])\n",
    "    plot_func(tmp_ax, lenient_regret, e_i+3)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(os.path.join(Path(simulation_folder), Path(simulation_folder).name + '.pdf'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation_folders = [\n",
    "    \"\", # None convergent experiment path\n",
    "    \"\", # Convergent experiment path\n",
    "]\n",
    "\n",
    "\n",
    "def plot_reward_std_bar(ax, reward, ind):\n",
    "    std_div = 2\n",
    "    subsample_rat = 25\n",
    "    ealph = 0.6\n",
    "    elinew = 0.3\n",
    "    c_size = 4\n",
    "\n",
    "    label, color, linestyle = lcls[ind]\n",
    "\n",
    "    x = np.arange(timestep_count)[::subsample_rat]\n",
    "    y = np.cumsum(reward.mean(axis=0))[::subsample_rat]\n",
    "    y_err = (np.cumsum(reward, axis=1).std(axis=0) / std_div)[::subsample_rat]\n",
    "\n",
    "    ax.errorbar(\n",
    "        x, y, y_err, label=label, c=color,\n",
    "        fmt=linestyle, ecolor=(*color, ealph), elinewidth=elinew, capsize=c_size\n",
    "    )\n",
    "\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = (10, 4)\n",
    "fig, ax = plt.subplots(nrows=len(RS_aspirations), ncols=2)\n",
    "\n",
    "for plt_i in range(2):\n",
    "    simulation_folder = simulation_folders[plt_i]\n",
    "\n",
    "    initial_samples_sims = np.load(os.path.join(simulation_folder, \"initial_samples_sims.npy\"))\n",
    "    distributions_w_true_and_ref_sims = np.load(os.path.join(simulation_folder, \"distributions_w_true_and_ref_sims.npy\"))\n",
    "    indices_DRO_sims_eps = np.load(os.path.join(simulation_folder, \"indices_DRO_sims_eps.npy\")).astype(int)\n",
    "    indices_WRO_sims = np.load(os.path.join(simulation_folder, \"indices_WRO_sims.npy\")).astype(int)\n",
    "    indices_SO_sims = np.load(os.path.join(simulation_folder, \"indices_SO_sims.npy\")).astype(int)\n",
    "    indices_RS_sims_asp = np.load(os.path.join(simulation_folder, \"indices_RS_sims_asp.npy\")).astype(int)\n",
    "    kappas_sims_asp = np.load(os.path.join(simulation_folder, \"kappas_sims_asp.npy\"))\n",
    "    taus_sims_asp = np.load(os.path.join(simulation_folder, \"taus_sims_asp.npy\"))\n",
    "\n",
    "    # Expected rewards of simulations and timesteps w.r.t. f\n",
    "    expected_reward_true_sims_ts = distributions_w_true_and_ref_sims[:, :, 0, :] @ InsCarb_mat\n",
    "\n",
    "    # MMD distances of simulations and timesteps\n",
    "    mmd_sims_ts = np.empty((sim_count, timestep_count))\n",
    "    for sim_i in range(sim_count):\n",
    "        for ts_i in range(timestep_count):\n",
    "            mmd_sims_ts[sim_i, ts_i] = MMD(M, *distributions_w_true_and_ref_sims[sim_i, ts_i])\n",
    "\n",
    "\n",
    "    # Calculate rewards\n",
    "\n",
    "    # DRO rewards\n",
    "    true_rewards_of_actions_DRO_eps = np.empty((len(DRO_epsilons), sim_count, timestep_count))\n",
    "    for e_i in range(len(DRO_epsilons)):\n",
    "        true_rewards_of_actions_DRO_eps[e_i] = expected_reward_true_sims_ts[\n",
    "            np.arange(sim_count)[:, None],\n",
    "            np.arange(timestep_count),\n",
    "            indices_DRO_sims_eps[e_i, :, :, 0].astype(int)\n",
    "        ]\n",
    "\n",
    "    # WRO rewards\n",
    "    true_rewards_of_actions_WRO = expected_reward_true_sims_ts[\n",
    "        np.arange(sim_count)[:, None],\n",
    "        np.arange(timestep_count),\n",
    "        indices_WRO_sims[0, :, :, 0].astype(int)\n",
    "    ]\n",
    "\n",
    "    # SO rewards\n",
    "    true_rewards_of_actions_SO = expected_reward_true_sims_ts[\n",
    "        np.arange(sim_count)[:, None],\n",
    "        np.arange(timestep_count),\n",
    "        indices_SO_sims[0, :, :, 0].astype(int)\n",
    "    ]\n",
    "\n",
    "    # RS rewards\n",
    "    true_rewards_of_actions_RS_asp = np.empty((len(RS_aspirations), sim_count, timestep_count))\n",
    "    for a_i in range(len(RS_aspirations)):\n",
    "        true_rewards_of_actions_RS_asp[a_i] = expected_reward_true_sims_ts[\n",
    "            np.arange(sim_count)[:, None],\n",
    "            np.arange(timestep_count),\n",
    "            indices_RS_sims_asp[a_i, :, :, 0].astype(int)\n",
    "        ]\n",
    "\n",
    "\n",
    "    # REGRET PLOTS\n",
    "\n",
    "\n",
    "    # Label, color, linestyle\n",
    "    lcls = [\n",
    "        (r\"DRBO: $3\\epsilon$\", colors[6], \":\",),\n",
    "        (r\"DRBO: $\\epsilon$\", blues[-2], \":\",),\n",
    "        (r\"DRBO: $\\epsilon/3$\", colors[5], \":\",),\n",
    "        (\"WRBO\", colors[8], \"-.\",),\n",
    "        (\"SO\", greens[-2], \"--\",),\n",
    "        (\"RoBOS\", reds[-1], \"-\",),\n",
    "    ]\n",
    "\n",
    "    # Plot std regions as error bars or colored regions\n",
    "    plot_func = plot_reward_std_bar\n",
    "\n",
    "    # Reward plot\n",
    "    tmp_ax = ax[plt_i]\n",
    "    tmp_ax.set_xlabel(\"Round\")\n",
    "    tmp_ax.set_ylabel(r\"Cumulative Reward\")\n",
    "\n",
    "    taus_sims = taus_sims_asp[a_i]\n",
    "    kappas_sims = kappas_sims_asp[a_i]\n",
    "\n",
    "    # DROs\n",
    "    for e_i in range(len(DRO_epsilons)):\n",
    "        reward = true_rewards_of_actions_DRO_eps[e_i]\n",
    "        plot_func(tmp_ax, reward, e_i)\n",
    "\n",
    "    # WRO\n",
    "    reward = true_rewards_of_actions_WRO\n",
    "    plot_func(tmp_ax, reward, e_i+1)\n",
    "\n",
    "    # SO\n",
    "    reward = true_rewards_of_actions_SO\n",
    "    plot_func(tmp_ax, reward, e_i+1)\n",
    "    \n",
    "    # RS\n",
    "    reward = true_rewards_of_actions_RS_asp[a_i]\n",
    "    plot_func(tmp_ax, reward, -1)\n",
    "\n",
    "    leg = tmp_ax.legend(prop={'size': 12}, frameon=True)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "for simulation_folder in simulation_folders:\n",
    "    plt.savefig(os.path.join(Path(simulation_folder), Path(simulation_folder).name + '_both_exp_rewards.pdf'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}