AWS-RoseTTAFold.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AWS-RoseTTAFold"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I. Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook runs the [RoseTTAFold](https://www.ipd.uw.edu/2021/07/rosettafold-accurate-protein-structure-prediction-accessible-to-all/) algorithm developed by Minkyung Baek et al. and described in [M. Baek et al., Science \n",
    "10.1126/science.abj8754 2021](https://www.ipd.uw.edu/wp-content/uploads/2021/07/Baek_etal_Science2021_RoseTTAFold.pdf) on AWS."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/RF_workflow.png\" alt=\"RoseTTAFold Workflow\" width=\"800px\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The AWS workflow depends on a Batch compute environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/AWS-RoseTTAFold-arch.png\" alt=\"AWS-RoseTTAFold Architecture\" width=\"800px\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## II. Environment setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Install dependencies\n",
    "%pip install -q -q -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Import helper functions at rfutils/rfutils.py\n",
    "from rfutils import rfutils\n",
    "\n",
    "## Load additional dependencies\n",
    "from Bio import SeqIO\n",
    "from Bio.Seq import Seq\n",
    "from Bio.SeqRecord import SeqRecord\n",
    "import boto3\n",
    "import glob\n",
    "import json\n",
    "import pandas as pd\n",
    "import sagemaker\n",
    "\n",
    "pd.set_option(\"max_colwidth\", None)\n",
    "\n",
    "# Get service clients\n",
    "session = boto3.session.Session()\n",
    "sm_session = sagemaker.session.Session()\n",
    "region = session.region_name\n",
    "role = sagemaker.get_execution_role()\n",
    "s3 = boto3.client(\"s3\", region_name=region)\n",
    "account_id = boto3.client(\"sts\").get_caller_identity().get(\"Account\")\n",
    "\n",
    "bucket = sm_session.default_bucket()\n",
    "print(f\"S3 bucket name is {bucket}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## III. Input Protein Sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter a protein sequence manually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = SeqRecord(\n",
    "    Seq(\"MKQHKAMIVALIVICITAVVAALVTRKDLCEVHIRTGQTEVAVF\"),\n",
    "    id=\"YP_025292.1\",\n",
    "    name=\"HokC\",\n",
    "    description=\"toxic membrane protein, small\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or provide the path to a fasta file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = SeqIO.read(\"data/T1078.fa\", \"fasta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Protein sequence for analysis is \\n{seq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IV. Submit RoseTTAFold Jobs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Job Name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = rfutils.create_job_name(seq.id)\n",
    "print(f\"Automatically-generated job name is: {job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload fasta file to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_uri = rfutils.upload_fasta_to_s3(seq, bucket, job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit jobs to AWS Batch queues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select the job definitions and Batch queues for your job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_resources = rfutils.get_rosettafold_batch_resources(region=region)\n",
    "\n",
    "cpu_queue = batch_resources[\"CPUJobQueue\"][0]\n",
    "gpu_queue = batch_resources[\"GPUJobQueue\"][0]\n",
    "cpu_data_prep_job_def = batch_resources[\"CPUDataPrepJobDefinition\"][0]\n",
    "cpu_predict_job_def = batch_resources[\"CPUPredictJobDefinition\"][0]\n",
    "gpu_predict_job_def = batch_resources[\"GPUPredictJobDefinition\"][0]\n",
    "\n",
    "batch_resources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because our test sequence is small (less than 400 residues) we will run the prediction step on a GPU to decrease the job duration from hours to minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "two_step_response = rfutils.submit_2_step_job(\n",
    "    bucket=bucket,\n",
    "    job_name=job_name,\n",
    "    data_prep_job_definition=cpu_data_prep_job_def,\n",
    "    data_prep_queue=cpu_queue,\n",
    "    data_prep_cpu=8,\n",
    "    data_prep_mem=32,\n",
    "    predict_job_definition=gpu_predict_job_def,  # Change this to the cpu_predict_job_def for large proteins\n",
    "    predict_queue=gpu_queue,  # Change this to the cpu_queue for large proteins\n",
    "    predict_cpu=4,\n",
    "    predict_mem=16,\n",
    "    predict_gpu=True,  # Change this to False for large proteins\n",
    ")\n",
    "data_prep_jobId = two_step_response[0][\"jobId\"]\n",
    "predict_jobId = two_step_response[1][\"jobId\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## V. Check Status of Data Prep and Prediction Jobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rfutils.get_rf_job_info(\n",
    "    cpu_queue,\n",
    "    gpu_queue,\n",
    "    hrs_in_past=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VI. View Data Prep Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pause while the data prep job starts up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rfutils.wait_for_job_start(data_prep_jobId)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get logs for data prep job (Run this multiple times to see how the job progresses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_prep_logStreamName = rfutils.get_batch_job_info(data_prep_jobId)[\"logStreamName\"]\n",
    "rfutils.get_batch_logs(data_prep_logStreamName).tail(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve and Display Multiple Sequence Alignment (MSA) Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rfutils.display_msa(data_prep_jobId, bucket)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VII. View Prediction Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pause while the predict job starts up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rfutils.wait_for_job_start(predict_jobId)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get logs for prediction job (Run this multiple times to see how the job progresses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_prep_logStreamName = rfutils.get_batch_job_info(data_prep_jobId)[\"logStreamName\"]\n",
    "rfutils.get_batch_logs(data_prep_logStreamName).tail(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VIII. View Job Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = rfutils.get_rf_job_metrics(job_name, bucket, region)\n",
    "\n",
    "print(f'Number of sequences in MSA: {metrics[\"DATA_PREP\"][\"MSA_COUNT\"]}')\n",
    "print(f'Number of templates: {metrics[\"DATA_PREP\"][\"TEMPLATE_COUNT\"]}')\n",
    "print(f'MSA duration (sec): {metrics[\"DATA_PREP\"][\"MSA_DURATION\"]}')\n",
    "print(f'SS duration (sec): {metrics[\"DATA_PREP\"][\"SS_DURATION\"]}')\n",
    "print(f'Template search duration (sec): {metrics[\"DATA_PREP\"][\"TEMPLATE_DURATION\"]}')\n",
    "print(\n",
    "    f'Total data prep duration (sec): {metrics[\"DATA_PREP\"][\"TOTAL_DATA_PREP_DURATION\"]}'\n",
    ")\n",
    "print(f'Total predict duration (sec): {metrics[\"PREDICT\"][\"TOTAL_PREDICT_DURATION\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IX. Retrieve and Display Predicted Structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rfutils.display_structure(predict_jobId, bucket, vmin=0.5, vmax=0.9)"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "interpreter": {
   "hash": "8ad3a54da4d511af1a5c2549d8f1b22d83bfd1079fb699a3f5552b91d143b102"
  },
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}