examples/ODBO_for_different_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ODBO as a pipeline for different datasets\n",
    "To show the generality of our ODBO method, we combine the notebooks for three different datasets into this one notebook. GB1_2016 has combintorially measured all the possible mutations on the 20^4 space, which is different from the three datasets listed here. We have separate notebooks for GB1_2016.\n",
    "To use this notebook, please change the global control varibles. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import pandas as pd\n",
    "import odbo\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global control varibles\n",
    "This section describe the parameters used for three different datasets. We also have a detailed walkthrough for method comparisons using GB1_2016 dataset in the seperated notebooks. Since the experimental setups are different for GB1_2016 and the other datasets, this notebook can still work for GB1_2016, but not ideal. We recommend to use the [GB1_2016_ODBO.ipynb](./GB1_2016_ODBO.ipynb) instead.\n",
    "We note that we didn't do extensive hyperparameter tuning, so other values can work and might even work better.\n",
    "Values used in the data collection is listed in the comments. Due to different randomness for different device, the results might differ slightly from the results we obtain using local computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset must be'GB1_2014', 'Ube4b_2013', 'avGFP_2016'\n",
    "dataset_name = 'GB1_2014'\n",
    "# Experiment settings \n",
    "random_seed = 9 #Random seed for the trial\n",
    "search_iter = 100 #Number of new observations, GB1_2014=100, Ube4b_2013=50, avGFP_2016=50\n",
    "# Initialization method protocol\n",
    "update_method='correlate'#find round 0 experiments to initiate BO. For the datasets with few changes in the sequences, 'correlate' mode is recommended. \n",
    "allow_abundance=False #If we allow the top scoring experiments to take abundance of a mutation in different sites into account.\n",
    "# Featurization settings\n",
    "method=['Avg','Max','Avg','Max'] #switching order for feature spaces to overcome local maxima in one certain representation\n",
    "#GB1_2014 and Ube4b_2013 using ['Avg','Max','Avg','Max'], avGFP_2016 using ['Avg','Avg','Avg','Max']\n",
    "mode='hybrid' #Feature computing mode. \n",
    "# Adaptive search space predicted by XGBOD model (Prescreening step)\n",
    "threshold_ratio = 0.9 #Use the value of top (1-threshold_ratio) of training measurements as threshold. GB1_2014=0.9, Ube4b=0.95, avGFP_2016=0.8\n",
    "cMat_plot = False #Plot the confusion matrix to check the accuracy of search space prescreening or not\n",
    "# BO method settings (Optimization step)\n",
    "BO_method = 'BO' #Must be 'ODBO_BO' or 'ODBO_TuRBO' or 'BO' or 'TuRBO'\n",
    "gp_method='gp_regression' #Must be 'gp_regression' or 'robust_regression'\n",
    "tr_length = [6.4] #Trust region length, used in the TuRBO, GB1_2014=[6.4], Ube4b_2013=[3.2], avGFP_2016=[12.8]. \n",
    "batch_size = 1 #Number of new oberservations provided by BO. We found 1 is the most cost-effective experimentally\n",
    "failure_tolerance =20 #Number of failure iterations to change TR length in TuRBO, GB1_2014=20, Ube4b=10, avGFP_2016=10\n",
    "data_augmentation = True #If do data augmentation. For double mutations, A1B,C2D==C2D,A1B, i.e.[1,2,0.1,0.2]=[2,1,0.2,0.1], we augment their symetric features to both training and test set\n",
    "# GB1_2014=True, Ube4b=False, avGFP_2016=False\n",
    "save_files = True #Save files or not"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data initalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(random_seed)\n",
    "\n",
    "#For the dataset with very small differences b/w each measurements, we use an exponential scaling to make BO easier\n",
    "if dataset_name == 'GB1_2014':\n",
    "    data_test = pd.read_csv('../datasets/GB1_2014_536085.csv', sep=',')\n",
    "    name_pre, Y_test = np.array(data_test['AACombo']), np.array(data_test['score'])\n",
    "    Y_test = 2**Y_test \n",
    "elif dataset_name == 'Ube4b_2013':\n",
    "    data_test = pd.read_csv('../datasets/Ube4b_2013_98299.csv', sep=',')\n",
    "    name_pre, Y_test = np.array(data_test['AACombo']), np.array(data_test['Log2Eratio'])\n",
    "elif dataset_name == 'avGFP_2016':\n",
    "    data_test = pd.read_csv('../datasets/avGFP_2016_54025.csv', sep=',')\n",
    "    name_pre, Y_test = np.array(data_test['AACombo']), np.array(data_test['medianBrightness'])\n",
    "    Y_test = 2**Y_test\n",
    "del data_test\n",
    "\n",
    "#Load the preselected indices using a certain shuffling order. Control Round 0 experiments to be the same for different trials\n",
    "if os.path.isfile('sele_indices_{}.npy'.format(dataset_name)) == True:\n",
    "    sele_indices = np.load('sele_indices_{}.npy'.format(dataset_name))\n",
    "    shuffle_order = np.load('shuffle_order_{}.npy'.format(dataset_name))\n",
    "    name_pre[1:], Y_test[1:] = name_pre[shuffle_order[1:]], Y_test[shuffle_order[1:]]\n",
    "    name = odbo.utils.code_to_array(name_pre)    \n",
    "else:\n",
    "    # If there is no preselected indices, we initialize by using the settings in the Globall variable\n",
    "    shuffle_order = np.arange(len(Y_test))\n",
    "    np.random.shuffle(shuffle_order[1:])\n",
    "    np.save('shuffle_order_{}.npy'.format(dataset_name), shuffle_order)\n",
    "    name_pre[1:], Y_test[1:] = name_pre[shuffle_order[1:]], Y_test[shuffle_order[1:]]\n",
    "    name = odbo.utils.code_to_array(name_pre)\n",
    "    sele_indices = odbo.initialization.initial_design(name, least_occurance=np.ones(name.shape[1]),allow_abundance=allow_abundance,update_method=update_method,verbose=True)\n",
    "    np.save('sele_indices_{}.npy'.format(dataset_name), sele_indices)\n",
    "name_sele, Y_train = name[sele_indices, :], Y_test[sele_indices]\n",
    "ids_keep = np.delete(range(len(Y_test)), sele_indices)\n",
    "name, Y_test = name[ids_keep, :], Y_test[ids_keep]\n",
    "reduction = np.random.choice(np.arange(len(Y_test)), int(0.9*len(Y_test)), replace=False)\n",
    "name, Y_test = name[reduction, :], Y_test[reduction]\n",
    "print('Selected initial experiments no. is ', len(Y_train))\n",
    "print('Select max Y: ', Y_train.max(), 'True max Y:', Y_test.max())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Featurization and find the adaptive search space model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Featurization\n",
    "feature_model = odbo.featurization.FewFeatureTransform(raw_vars=name_sele, Y=Y_train, method=method[0], mode=mode)\n",
    "X_test = feature_model.transform(name)\n",
    "X_train = feature_model.transform(name_sele)\n",
    "\n",
    "if BO_method == 'ODBO_BO' or BO_method == 'ODBO_TuRBO':\n",
    "    # Find the threshold for outliers or inliers \n",
    "    threshold = Y_train[np.argsort(Y_train)[int(threshold_ratio*len(Y_train))]]\n",
    "    print('Lowest measurement value to be considered as inlier for BO: ', threshold)\n",
    "    labels_train = odbo.prescreening.sp_label(X_train, Y_train, thres=threshold)\n",
    "    # Find the XGBOD adaptive search space model\n",
    "    pre_model = odbo.prescreening.XGBOD(eval_metric = 'error', random_state = random_seed)\n",
    "    pre_model.fit(X_train, labels_train)\n",
    "    # Predict the entire search space to get the adapt search space\n",
    "    labels_test = odbo.prescreening.sp_label(X_test, Y_test, thres=threshold)\n",
    "    pred_test_labels = pre_model.predict(X_test)\n",
    "    sele_id_test = list(np.where(pred_test_labels == 0)[0])\n",
    "    print(\"Adapt space size, Entire space size: \", len(sele_id_test), name.shape[0])\n",
    "    # Plot the confusion matrix to check the accuracy of search space prescreening\n",
    "    if cMat_plot:\n",
    "        out_outlier, in_outlier, out_inlier, in_inlier = odbo.plot.plot_cm(labels_test, pred_test_labels, Y_test)\n",
    "        print(\"Correct ratio: {0:.3%}\".format((len(out_outlier)+len(in_inlier))/len(labels_test)))\n",
    "        print(\"FN ratio: {0:.3%}\".format(len(out_inlier)/len(labels_test)))\n",
    "        print(\"FP ratio: {0:.3%}\".format(len(in_outlier)/len(labels_test)))\n",
    "else:\n",
    "    sele_id_test = np.arange(len(Y_test))\n",
    "print(\"Adapt space size, Entire space size: \", len(sele_id_test), name.shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run optimizations on the datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Creat search data\n",
    "X_train_sele, Y_train_sele = torch.tensor(X_train), torch.tensor(Y_train.reshape(len(Y_train),1))\n",
    "search_name_sele, name_sele_temp = name[sele_id_test, :], name_sele\n",
    "X_test_sele, Y_test_sele = torch.tensor(X_test[sele_id_test, :]), torch.tensor(Y_test[sele_id_test].reshape(len(sele_id_test),1))\n",
    "\n",
    "## Run BO experiment with robust regression or directly gp\n",
    "l, failure_count = 0, 0 #l is the current iteration, failure_count controls when to switch to another featurization\n",
    "if BO_method == 'ODBO_TuRBO' or BO_method == 'TuRBO':\n",
    "    state = odbo.turbo.TurboState(dim=X_train_sele.shape[1], batch_size=batch_size, length=tr_length, n_trust_regions=len(tr_length), failure_tolerance = failure_tolerance)\n",
    "    state.best_value = Y_train_sele.max()\n",
    "\n",
    "while l < search_iter:\n",
    "    print(\"Iter: \", l, \"Current Max: \", Y_train_sele.max().detach().numpy(), \"Test max: \", Y_test_sele.max().detach().numpy())\n",
    "    top_ids = np.argsort(Y_train_sele.numpy().ravel())[-40-l:]#To better regress the high meaurement value region, we only pick the top 40+l experiments to regress the surrogate model\n",
    "    sele_feat = []#There might be some identical features across the current selected training experiments\n",
    "    for i in range(X_train_sele.shape[1]):\n",
    "        if (X_train_sele[top_ids,i]-X_train_sele[0,i]).any() !=0:\n",
    "            sele_feat.append(i)\n",
    "    X_sele = X_train_sele[:,sele_feat]\n",
    "    X_test_sele=X_test_sele[:, sele_feat]\n",
    "    if data_augmentation:\n",
    "        X_test_sele = torch.cat([X_test_sele, X_test_sele[:, [1,0,3,2]]])\n",
    "        X_sele = torch.cat([X_sele[top_ids, :], X_sele[top_ids, :][:, [1,0,3,2]]])\n",
    "        Y_sele = torch.cat([Y_train_sele[top_ids], Y_train_sele[top_ids]])\n",
    "    else:\n",
    "        X_sele, Y_sele=X_sele[top_ids, :], Y_train_sele[top_ids]\n",
    "    # Run optimization using different methods    \n",
    "    if BO_method == 'ODBO_BO' or BO_method == 'BO':\n",
    "        X_next, acq_value, next_exp_id = odbo.bo_design(X=X_sele, Y=Y_sele, X_pending=X_test_sele, gp_method=gp_method, batch_size=batch_size)\n",
    "        next_exp_id = np.mod(next_exp_id, len(Y_test_sele))\n",
    "    elif BO_method == 'ODBO_TuRBO' or BO_method == 'TuRBO':\n",
    "        X_next, acq_value, raw_next_exp_id = odbo.turbo_design(state=state, X=X_sele, Y=Y_sele, X_pending=X_test_sele, n_trust_regions=len(tr_length), batch_size=batch_size, gp_method=gp_method)\n",
    "        Y_next_m = torch.zeros((len(tr_length), batch_size, 1), device=Y_train_sele.device, dtype=Y_train_sele.dtype)\n",
    "        raw_next_exp_id = np.mod(raw_next_exp_id, len(Y_test_sele))\n",
    "        next_exp_id = []  \n",
    "        for i in range(batch_size):\n",
    "            next_exp_id_m = raw_next_exp_id[:, i]\n",
    "            Y_next_m[:, i, 0], idtoadd = Y_test_sele[next_exp_id_m].reshape(len(tr_length)), next_exp_id_m[np.argmax(Y_test_sele[next_exp_id_m])]\n",
    "            next_exp_id.append(idtoadd)\n",
    "\n",
    "    Y_train_sele = torch.cat([Y_train_sele, Y_test_sele[next_exp_id]])\n",
    "    ids_keep = list(np.delete(range(Y_test_sele.shape[0]), next_exp_id))\n",
    "    Y_test_sele = Y_test_sele[ids_keep]\n",
    "    name_sele_temp = np.concatenate((name_sele_temp, search_name_sele[next_exp_id]))\n",
    "    search_name_sele = search_name_sele[ids_keep]\n",
    "    print(\"Newly added value: \", Y_train_sele[-batch_size:].detach().numpy(), ''.join(name_sele_temp[-1, :]))\n",
    "    if BO_method == 'ODBO_TuRBO'or BO_method == 'TuRBO':\n",
    "        # Update the TuRBO state with the newly added Y values\n",
    "        state = odbo.turbo.update_state(state=state, Y_next=Y_next_m)\n",
    "\n",
    "    # Switch different representations if one representation fails in \n",
    "    if Y_train_sele[-batch_size:].detach().numpy().max() > Y_train_sele[:-batch_size].max():\n",
    "        failure_count = 0\n",
    "        feature_model = odbo.featurization.FewFeatureTransform(raw_vars=name_sele_temp, \n",
    "                                                                Y=Y_train_sele.detach().numpy(), method=method[1], mode=mode)\n",
    "    else:\n",
    "        failure_count = failure_count + 1\n",
    "        if failure_count >= 3:\n",
    "            feature_model = odbo.featurization.FewFeatureTransform(raw_vars=name_sele_temp, \n",
    "                                                                    Y=Y_train_sele.detach().numpy(), method=method[2], mode=mode)\n",
    "        else:\n",
    "            feature_model = odbo.featurization.FewFeatureTransform(raw_vars=name_sele_temp, \n",
    "                                                                    Y=Y_train_sele.detach().numpy(), method=method[3], mode=mode)\n",
    "    del X_test_sele, X_sele, Y_sele\n",
    "    X_test_sele= torch.tensor(feature_model.transform(search_name_sele))\n",
    "    X_train_sele = torch.tensor(feature_model.transform(name_sele_temp))\n",
    "    del feature_model\n",
    "    l = l + 1\n",
    "\n",
    "# Save the BO results. Note we save all the observations including the Round 0 ones\n",
    "if dataset_name == 'GB1_2014' or dataset_name == 'avGFP_2016':\n",
    "    Y_train_sele = np.log2(Y_train_sele)\n",
    "\n",
    "if save_files:\n",
    "    if gp_method == 'robust_regression':\n",
    "        np.save('results/{}/{}_{}_RobustGP_batch{}_{}.npy'.format(dataset_name, dataset_name, BO_method, batch_size, random_seed), Y_train_sele)\n",
    "    elif gp_method == 'gp_regression':\n",
    "        np.save('results/{}/{}_{}_GP_batch{}_{}.npy'.format(dataset_name, dataset_name, BO_method, batch_size, random_seed), Y_train_sele)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run random selection on the datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sele_Y, Y_train_sele = list(np.random.choice(Y_test, search_iter, replace = False)), list(Y_train.copy())\n",
    "Y_train_sele.extend(sele_Y)\n",
    "print('Max Y', max(sele_Y))\n",
    "if dataset_name == 'GB1_2014' or dataset_name == 'avGFP_2016':\n",
    "    Y_train_sele = np.log2(Y_train_sele)\n",
    "if save_files:\n",
    "    np.save('results/{}/{}_random_{}.npy'.format(dataset_name, dataset_name, random_seed), Y_train_sele)"
   ]
  },
  {
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}