notebooks/inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using trained models to make predictions for new variants\n",
    "\n",
    "This notebook shows how to use trained models to make predictions for new variants.\n",
    "\n",
    "Prerequisites\n",
    "- A trained model. You can use the pre-trained models we provide in the `pub/trained_models` directory, train your own models similar to ours using the arguments in the `pub/regression_args` directory, or train your own models using your preferred arguments.\n",
    "\n",
    "Main steps\n",
    "- Encode variants with a combination of one-hot and AAindex encoding.\n",
    "- Use the trained models to get predictions for those variants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "# reload modules before executing code in order to make development and debugging easier\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this jupyter notebook is running inside of the \"notebooks\" directory\n",
    "# for relative paths to work properly, we need to set the current working directory to the root of the project\n",
    "# for imports to work properly, we need to add the code folder to the system path\n",
    "import os\n",
    "from os.path import abspath, join, isdir\n",
    "import sys\n",
    "if not isdir(\"notebooks\"):\n",
    "    # if there's a \"notebooks\" directory in the cwd, we've already set the cwd so no need to do it again\n",
    "    os.chdir(\"..\")\n",
    "module_path = abspath(\"code\")\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import constants\n",
    "import utils\n",
    "import encode as enc\n",
    "import inference as inf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encode variants\n",
    "Using a simple example with a few avGFP variants. For more detailed information on how to encode variants, check out the data encoding notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 237, 40)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variants = [\"Y64C,E170V\", \"I126T,N210H\", \"E15V,D17G,I169F\", \"A108G\"]\n",
    "# specifying \"ds_name\" only works if the dataset is defined in constants.py\n",
    "# alternatively, you can specify the wild-type sequence and offset\n",
    "# see the encoding notebook for details\n",
    "encoded_variants = enc.encode(encoding=\"one_hot,aa_index\", variants=variants, ds_name=\"avgfp\")\n",
    "encoded_variants.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use a pre-trained model to get predictions for these variants\n",
    "The saved models consist of three files (meta, index, data). This is due to how TensorFlow saves checkpoints. A single saved model will have the same prefix for all three files. Using the avGFP linear regression model as an example. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from pub/trained_models/avgfp/avgfp_lr\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([2.687976 , 3.7444162, 2.885042 , 3.2154722], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr_prefix = \"pub/trained_models/avgfp/avgfp_lr\"  # just the prefix, no file extension\n",
    "lr_predictions = inf.run_inference(encoded_data=encoded_variants, ckpt_prefix_fn=lr_prefix)\n",
    "lr_predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predictions for lots of variants (batches)\n",
    "If you want to make predictions for >64 variants a time, the script will automatically break the input into batches of size 64. You can change the batch size by using the `batch_size` argument to `inf.run_inference()`. A progress bar will show how progress through the batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from pub/trained_models/avgfp/avgfp_lr\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 88.57it/s]\n"
     ]
    }
   ],
   "source": [
    "lots_of_variants = enc.encode(encoding=\"one_hot,aa_index\", variants=[\"Y64C,E170V\"] * 200, ds_name=\"avgfp\")\n",
    "preds = inf.run_inference(encoded_data=lots_of_variants, ckpt_prefix_fn=lr_prefix, batch_size=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Repeated predictions (single session)\n",
    "If you need to run inference many times in a loop, the code above is inefficient since it restores the TensorFlow model on each call to `inf.run_inference()`. \n",
    "You can create a single TensorFlow session to use in the loop instead. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from pub/trained_models/avgfp/avgfp_lr\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([2.687976 , 3.7444162, 2.885042 , 3.2154722], dtype=float32)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([2.687976 , 3.7444162, 2.885042 , 3.2154722], dtype=float32)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([2.687976 , 3.7444162, 2.885042 , 3.2154722], dtype=float32)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# open the session\n",
    "lr_sess = inf.restore_sess(lr_prefix)\n",
    "\n",
    "# run inference many times in a loop\n",
    "for i in range(3):\n",
    "    display(inf.run_inference(encoded_data=encoded_variants, sess=lr_sess))\n",
    "\n",
    "# close the session when you're done\n",
    "lr_sess.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can have multiple sessions open at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from pub/trained_models/avgfp/avgfp_lr\n",
      "INFO:tensorflow:Restoring parameters from pub/trained_models/avgfp/avgfp_cnn\n",
      "LR: [2.687976  3.7444162 2.885042  3.2154722]\n",
      "CNN: [1.4410211 3.8864908 2.870633  3.6873856]\n",
      "LR: [2.687976  3.7444162 2.885042  3.2154722]\n",
      "CNN: [1.4410211 3.8864908 2.870633  3.6873856]\n",
      "LR: [2.687976  3.7444162 2.885042  3.2154722]\n",
      "CNN: [1.4410211 3.8864908 2.870633  3.6873856]\n"
     ]
    }
   ],
   "source": [
    "lr_sess = inf.restore_sess(lr_prefix)\n",
    "cnn_prefix = \"pub/trained_models/avgfp/avgfp_cnn\"\n",
    "cnn_sess = inf.restore_sess(cnn_prefix)\n",
    "for i in range(3):\n",
    "    print(\"LR:\", inf.run_inference(encoded_data=encoded_variants, sess=lr_sess))\n",
    "    print(\"CNN:\", inf.run_inference(encoded_data=encoded_variants, sess=cnn_sess))\n",
    "lr_sess.close()\n",
    "cnn_sess.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}