rgn2_prediction.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# RGN2 Colab\n",
        "## Instructions\n",
        "1. Paste your protein sequence in the input field.\n",
        "2. Run the cells in the Colab individually with the play button on the left or via _Runtime_ > _Run all._\n",
        "3. The predicted protein structure will be downloaded after the final \"Refinement\" cell is executed."
      ],
      "metadata": {
        "id": "3SDCdyj7XneF"
      },
      "id": "3SDCdyj7XneF"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Download RGN2 and Install Dependencies\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import subprocess\n",
        "from IPython import get_ipython\n",
        "from IPython.utils import io\n",
        "\n",
        "WORKDIR = './rgn2'\n",
        "GIT_REPO = 'https://github.com/aqlaboratory/rgn2.git'\n",
        "ENV_CONFIG = os.path.join(WORKDIR, 'environment.yml')\n",
        "RGN2_PARAM_SOURCE_URL = 'https://huggingface.co/christinafl/rgn2'\n",
        "RGN2_PARAMS_DIR = os.path.join(WORKDIR, 'resources')\n",
        "RGN2_PARAM_RUN_DIR = os.path.join(RGN2_PARAMS_DIR, 'rgn2_runs')\n",
        "RGN2_RUN_DIR = os.path.join(WORKDIR, 'runs')\n",
        "\n",
        "AF2_GIT_REPO = 'https://github.com/deepmind/alphafold.git'\n",
        "AF2_SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2022-03-02.tar'\n",
        "AF2_PARAMS_DIR = './alphafold/data/params'\n",
        "AF2_PARAMS_PATH = os.path.join(AF2_PARAMS_DIR, os.path.basename(AF2_SOURCE_URL))\n",
        "\n",
        "REFINER_DIR = os.path.join(WORKDIR, 'ter2pdb')\n",
        "REFINER_PATH = os.path.join(REFINER_DIR, 'ModRefiner-l.zip')\n",
        "REFINER_URL = 'https://zhanggroup.org/ModRefiner/ModRefiner-l.zip'\n",
        "\n",
        "try:\n",
        "  with io.capture_output() as captured:\n",
        "    %cd '/content'\n",
        "\n",
        "    # Different conda envs necessary due to conflicting dependencies.\n",
        "    %shell rm -rf /opt/conda\n",
        "    %shell wget -q -P /tmp \\\n",
        "      https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \\\n",
        "        && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \\\n",
        "        && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "    \n",
        "    PATH=%env PATH\n",
        "    %env PATH=/opt/conda/bin:{PATH}\n",
        "\n",
        "    CONDA_INIT = 'source /opt/conda/etc/profile.d/conda.sh && conda init'\n",
        "    RGN2_ENV_INIT = f'{CONDA_INIT} && conda activate rgn2'\n",
        "\n",
        "    # Download RGN2.\n",
        "    %shell rm -rf {WORKDIR}\n",
        "    %shell git clone {GIT_REPO} {WORKDIR}\n",
        "    %shell {CONDA_INIT} && conda env create -f {ENV_CONFIG}\n",
        "\n",
        "    # Download AF2 for AF2Rank-based refinement.\n",
        "    AF2_ENV_INIT = f'{CONDA_INIT} && conda activate af2'\n",
        "\n",
        "    %shell rm -rf alphafold\n",
        "    %shell git clone --branch v2.2.4 {AF2_GIT_REPO} alphafold\n",
        "    %shell {CONDA_INIT} && conda create -y -q --name af2 python=3.7\n",
        "    %shell {AF2_ENV_INIT} && pip install -r ./alphafold/requirements.txt\n",
        "    %shell {AF2_ENV_INIT} && pip install --no-dependencies ./alphafold\n",
        "    %shell {AF2_ENV_INIT} && pip install --upgrade jax==0.3.17 \\\n",
        "      jaxlib==0.3.15+cuda11.cudnn805 \\\n",
        "      -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
        "    \n",
        "    %shell mkdir --parents \"{AF2_PARAMS_DIR}\"\n",
        "    %shell wget -O \"{AF2_PARAMS_PATH}\" \"{AF2_SOURCE_URL}\"\n",
        "    %shell tar --extract --verbose --file=\"{AF2_PARAMS_PATH}\" \\\n",
        "      --directory=\"{AF2_PARAMS_DIR}\" --preserve-permissions\n",
        "    %shell rm \"{AF2_PARAMS_PATH}\"\n",
        "\n",
        "    # Download AminoBERT/RGN2 weights.\n",
        "    %shell GIT_LFS_SKIP_SMUDGE=1 git clone \"{RGN2_PARAM_SOURCE_URL}\" \"{RGN2_PARAMS_DIR}\"\n",
        "    %shell cd {RGN2_PARAMS_DIR} && git lfs pull\n",
        "    %shell mv {RGN2_PARAM_RUN_DIR} {RGN2_RUN_DIR}\n",
        "\n",
        "    # Download Modrefiner to initialize all atoms from CA trace.\n",
        "    %shell wget -O {REFINER_PATH} {REFINER_URL}\n",
        "    %shell unzip -o {REFINER_PATH} -d {REFINER_DIR}\n",
        "    %shell rm {REFINER_PATH}\n",
        "\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise"
      ],
      "metadata": {
        "id": "_ds-a3EgRUsb",
        "cellView": "form"
      },
      "id": "_ds-a3EgRUsb",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Import Python Packages\n",
        "\n",
        "%cd '/content/rgn2'\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import re\n",
        "import hashlib\n",
        "import json\n",
        "import subprocess\n",
        "from pathlib import Path\n",
        "from IPython.utils import io\n",
        "from google.colab import files\n",
        "\n",
        "%reload_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "\n",
        "sys.path.append('/content/alphafold')\n",
        "from ter2pdb import ter2pdb"
      ],
      "metadata": {
        "id": "7GBilHivRmgP",
        "cellView": "form"
      },
      "id": "7GBilHivRmgP",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### Enter the amino acid sequence to fold ⬇️\n",
        "sequence = 'DEEEIQKAIEELLRKGVSEEEAAIIIVQRFNVAVVVVVQDERQGKHISEYIRRYIPEADVILFANLVVIKVETHELSTRVWEAAQKAY'  #@param {type:\"string\"}\n",
        "jobname = 'test' #@param {type:\"string\"}\n",
        "\n",
        "# Remove whitespace\n",
        "sequence = \"\".join(sequence.split()).upper()\n",
        "jobname = \"\".join(jobname.split())\n",
        "\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "seq_hash = hashlib.blake2b(sequence.encode(), digest_size=3).hexdigest()\n",
        "seq_id = f'{jobname}_{seq_hash}'\n",
        "\n",
        "MAX_SEQUENCE_LENGTH = 1023\n",
        "\n",
        "# Remove all whitespaces, tabs and end lines; upper-case\n",
        "sequence = sequence.translate(str.maketrans('', '', ' \\n\\t')).upper()\n",
        "aatypes = set('ACDEFGHIKLMNPQRSTVWY')  # 20 standard aatypes\n",
        "if not set(sequence).issubset(aatypes):\n",
        "  raise Exception(f'Input sequence contains non-amino acid letters: {set(sequence) - aatypes}. AlphaFold only supports 20 standard amino acids as inputs.')\n",
        "if len(sequence) > MAX_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too long: {len(sequence)} amino acids, while the maximum is {MAX_SEQUENCE_LENGTH}. Please use the full AlphaFold system for long sequences.')\n",
        "\n",
        "run_inputs = {'sequence': sequence, 'seq_id': seq_id}\n",
        "with open(\"run.json\", \"w\") as f:\n",
        "    json.dump(run_inputs, f)\n",
        "\n",
        "DATA_DIR = 'aminobert_output'\n",
        "RUN_DIR = 'runs/15106000'\n",
        "OUTPUT_DIR = 'output'\n",
        "REFINE_DIR = 'output/refine_model1'\n",
        "SEQ_PATH = os.path.join(DATA_DIR, f'{seq_id}.fa')\n",
        "TER_PATH = os.path.join(RUN_DIR, '1', 'outputsTesting', f'{seq_id}.tertiary')"
      ],
      "metadata": {
        "id": "_5YeBefARz0D",
        "cellView": "form"
      },
      "id": "_5YeBefARz0D",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d93c3e44",
      "metadata": {
        "id": "d93c3e44",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Generate AminoBERT Embeddings\n",
        "\n",
        "%%bash\n",
        "source /opt/conda/etc/profile.d/conda.sh && conda init\n",
        "conda activate rgn2\n",
        "python\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import json\n",
        "sys.path.append(os.path.join(os.getcwd(), 'aminobert'))\n",
        "\n",
        "import shutil\n",
        "from aminobert.prediction import aminobert_predict_sequence\n",
        "from data_processing.aminobert_postprocessing import aminobert_postprocess\n",
        "\n",
        "DATA_DIR = 'aminobert_output'\n",
        "DATASET_NAME = '1'\n",
        "PREPEND_M = True\n",
        "AMINOBERT_CHKPT_DIR = 'resources/aminobert_checkpoint/AminoBERT_runs_v2_uniparc_dataset_v2_5-1024_fresh_start_model.ckpt-1100000'\n",
        "\n",
        "with open(\"run.json\", \"r\") as f:\n",
        "    run_inputs = json.load(f)\n",
        "\n",
        "# Remove old data since AminoBERT combines entire directory to create dataset.\n",
        "if os.path.exists(DATA_DIR):\n",
        "  shutil.rmtree(DATA_DIR)\n",
        "os.makedirs(DATA_DIR)\n",
        "\n",
        "aminobert_predict_sequence(seq=run_inputs['sequence'], header=run_inputs['seq_id'],\n",
        "                           prepend_m=PREPEND_M, checkpoint=AMINOBERT_CHKPT_DIR,\n",
        "                           data_dir=DATA_DIR)\n",
        "aminobert_postprocess(data_dir=DATA_DIR, dataset_name=DATASET_NAME, prepend_m=PREPEND_M)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run RGN2\n",
        "#@markdown This step generates the raw RGN2-predicted C-alpha trace.\n",
        "\n",
        "rgn2_env_init = 'source /opt/conda/etc/profile.d/conda.sh && conda init && conda activate rgn2'\n",
        "try:\n",
        "  with io.capture_output() as captured:\n",
        "    cmd = (f\"python rgn/protling.py {os.path.join(RUN_DIR, 'configuration')} \"\n",
        "           f\"-p -e 'weighted_testing' -a -g 0\")\n",
        "    %shell {rgn2_env_init} && {cmd}\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "print('Prediction completed!')"
      ],
      "metadata": {
        "id": "IJkrRgF3R63Q",
        "cellView": "form"
      },
      "id": "IJkrRgF3R63Q",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Refinement and Structure Download\n",
        "#@markdown Once this cell has been executed, a PDB file with the refined\n",
        "#@markdown structure will be automatically downloaded to your computer.\n",
        "#@markdown **Note**: Notebook refinement pipeline is ~2x slower compared\n",
        "#@markdown to local execution.\n",
        "recycles = 1 #@param {type:\"integer\"}\n",
        "\n",
        "ter2pdb.run_ca_to_allatom(seq_path=SEQ_PATH, ter_path=TER_PATH,\n",
        "                          output_dir=OUTPUT_DIR, seq_id=seq_id)\n",
        "\n",
        "out_suffix = '_prediction'\n",
        "af2_env_init = 'source /opt/conda/etc/profile.d/conda.sh && conda init && conda activate af2'\n",
        "jax_env_vars = 'TF_FORCE_UNIFIED_MEMORY=1 XLA_PYTHON_CLIENT_MEM_FRACTION=2.0'\n",
        "cmd = (f\"{jax_env_vars} python ter2pdb/run_af2rank.py refine_model1 \"\n",
        "       f\"--target_list {seq_id} --af2_dir /content/alphafold/ \"\n",
        "       f\"--out_suffix {out_suffix} --seq_dir {Path(SEQ_PATH).parent} \"\n",
        "       f\"--pdb_dir {OUTPUT_DIR} --output_dir {OUTPUT_DIR} --deterministic \"\n",
        "       f\"--seq_replacement - --mask_sidechains_add_cb --model_num 1 \"\n",
        "       f\"--recycles {recycles}\")\n",
        "try:\n",
        "  with io.capture_output() as captured:\n",
        "    %shell {af2_env_init} && {cmd}\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "print('Refinement completed!')\n",
        "\n",
        "files.download(os.path.join(REFINE_DIR, f'{seq_id}{out_suffix}.pdb'))"
      ],
      "metadata": {
        "id": "I0A_TUORSukb",
        "cellView": "form"
      },
      "id": "I0A_TUORSukb",
      "execution_count": null,
      "outputs": []
    }
  ],
  "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.8"
    },
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}