CASP14-Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AWS-RoseTTAFold: Bulk Job Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I. Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to analyze multiple protein simultaneously, in this case a subset of the CASP14 target set."
   ]
  },
  {
   "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",
    "from IPython.display import display\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## III. Input Protein Sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download and process CASP14 sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://predictioncenter.org/download_area/CASP14/sequences/casp14.seq.txt\" -O \"data/casp14.fa\"\n",
    "!sed '137,138d' \"data/casp14.fa\" > \"data/casp14_dedup.fa\" # Remove duplicate entry for T1085\n",
    "\n",
    "casp14_iterator = SeqIO.parse(\"data/casp14_dedup.fa\", \"fasta\")\n",
    "casp14_df = pd.DataFrame(\n",
    "    (\n",
    "        (record.id, record.description, len(record), record.seq)\n",
    "        for record in casp14_iterator\n",
    "    ),\n",
    "    columns=[\"id\", \"description\", \"length\", \"seq\"],\n",
    ").sort_values(by=\"length\")\n",
    "!rm data/casp14*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Display information about CASP14 proteins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pd.option_context(\"display.max_rows\", None):\n",
    "    display(casp14_df.loc[:, (\"id\", \"description\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot distribution of the protein lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "plt.hist(casp14_df.length, bins=50)\n",
    "plt.ylabel(\"Sample count\")\n",
    "plt.xlabel(\"Residue count\")\n",
    "plt.title(\"CASP-14 Protein Length Distribution\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the names of the AWS Batch resources deployed in your account."
   ]
  },
  {
   "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": [
    "Submit analysis jobs for a subset of CASP14 proteins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "protein_count = 84  # Change this to analyze a smaller number of CASP14 targets\n",
    "job_name_list = []\n",
    "\n",
    "for row in casp14_df[:protein_count].itertuples(index=False):\n",
    "    record = SeqRecord(row.seq, id=row.id, description=row.description)\n",
    "    print(f\"Protein sequence for analysis is \\n{record.description}\")\n",
    "    sequence_length = len(record.seq)\n",
    "    print(f\"Sequence length is {sequence_length}\")\n",
    "\n",
    "    if sequence_length < 400:\n",
    "        prep_cpu = 8\n",
    "        prep_mem = 32\n",
    "        predict_cpu = 4\n",
    "        predict_mem = 16\n",
    "        predict_gpu = True\n",
    "        predict_job_definition = gpu_predict_job_def\n",
    "        predict_queue = gpu_queue\n",
    "    else:\n",
    "        prep_cpu = 8\n",
    "        prep_mem = 64\n",
    "        predict_cpu = 4\n",
    "        predict_mem = 32\n",
    "        predict_gpu = False\n",
    "        predict_job_definition = cpu_predict_job_def\n",
    "        predict_queue = cpu_queue\n",
    "\n",
    "    job_name = rfutils.create_job_name(record.id)\n",
    "    print(f\"Automatically-generated job name is: {job_name}\")\n",
    "    job_name_list.append(job_name)\n",
    "    input_uri = rfutils.upload_fasta_to_s3(record, bucket, job_name)\n",
    "    two_step_response = rfutils.submit_2_step_job(\n",
    "        bucket=bucket,\n",
    "        job_name=job_name,\n",
    "        data_prep_input_file=\"input.fa\",\n",
    "        data_prep_job_definition=cpu_data_prep_job_def,\n",
    "        data_prep_queue=cpu_queue,\n",
    "        data_prep_cpu=prep_cpu,\n",
    "        data_prep_mem=prep_mem,\n",
    "        predict_job_definition=predict_job_definition,\n",
    "        predict_queue=predict_queue,\n",
    "        predict_cpu=predict_cpu,\n",
    "        predict_mem=predict_mem,\n",
    "        predict_gpu=predict_gpu,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IV. 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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jobs = []\n",
    "for job_name in job_name_list:\n",
    "    metrics = rfutils.get_rf_job_metrics(job_name, bucket, region)\n",
    "    row = [\n",
    "        job_name,\n",
    "        metrics[\"DATA_PREP\"][\"JOB_ID\"],\n",
    "        metrics[\"DATA_PREP\"][\"CPU\"],\n",
    "        metrics[\"DATA_PREP\"][\"MEM\"],\n",
    "        metrics[\"DATA_PREP\"][\"LENGTH\"],\n",
    "        metrics[\"DATA_PREP\"][\"MSA_COUNT\"],\n",
    "        metrics[\"DATA_PREP\"][\"TEMPLATE_COUNT\"],\n",
    "        metrics[\"DATA_PREP\"][\"MSA_DURATION\"],\n",
    "        metrics[\"DATA_PREP\"][\"SS_DURATION\"],\n",
    "        metrics[\"DATA_PREP\"][\"TEMPLATE_DURATION\"],\n",
    "        metrics[\"DATA_PREP\"][\"TOTAL_DATA_PREP_DURATION\"],\n",
    "        metrics[\"PREDICT\"][\"JOB_ID\"],\n",
    "        metrics[\"PREDICT\"][\"CPU\"],\n",
    "        metrics[\"PREDICT\"][\"MEM\"],\n",
    "        metrics[\"PREDICT\"][\"TOTAL_PREDICT_DURATION\"],\n",
    "    ]\n",
    "    jobs.append(row)\n",
    "metrics_df = pd.DataFrame(\n",
    "    jobs,\n",
    "    columns=[\n",
    "        \"jobName\",\n",
    "        \"dataPrepJobID\",\n",
    "        \"dataPrepCPU\",\n",
    "        \"dataPrepMEM\",\n",
    "        \"sequenceLength\",\n",
    "        \"MSACount\",\n",
    "        \"templateCount\",\n",
    "        \"MSADuration\",\n",
    "        \"SSDuration\",\n",
    "        \"templateDuration\",\n",
    "        \"dataPrepDuration\",\n",
    "        \"predictJobId\",\n",
    "        \"predictCPU\",\n",
    "        \"predictMEM\",\n",
    "        \"predictDuration\",\n",
    "    ],\n",
    ")\n",
    "metrics_df.sort_values(by=[\"dataPrepCPU\", \"dataPrepMEM\", \"predictCPU\", \"predictMEM\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_df.to_csv(\"results.csv\")"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "interpreter": {
   "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
  },
  "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
}