gapTrick.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "n8phw4MCjJQH"
      },
      "outputs": [],
      "source": [
        "#@title gapTrick - AlphaFold2 with multimeric templates\n",
        "\n",
        "#@markdown This notebok uses code from https://github.com/gchojnowski/gapTrick and is based on third-party software, including\n",
        "#@markdown -  AlphaFold2 [code](https://github.com/google-deepmind/alphafold)\n",
        "#@markdown -  MMseqs2 API [code](https://github.com/soedinglab/MMseqs2)\n",
        "#@markdown -  ... and many others. See [acknowledgements](https://github.com/deepmind/alphafold/#acknowledgements) in AlphaFold's readme for more details.\n",
        "\n",
        "#@markdown If you found it uselful remember to cite\n",
        "#@markdown  - Chojnowski bioRxiv (2025) [01.31.635911](https://doi.org/10.1101/2025.01.31.635911)\n",
        "#@markdown  - Jumper et al Nature (2021) [596, 583–589](https://www.nature.com/articles/s41586-021-03819-2)\n",
        "#@markdown  - Mirdita et al Bioinformatics (2021) [btab184](https://doi.org/10.1093/bioinformatics/btab184)\n",
        "\n",
        "\n",
        "#@markdown **Select options below and hit `Runtime` -> `Run all` to start uploading protein sequences and template**\n",
        "\n",
        "#@markdown **A few tips on how to use this notebook**\n",
        "#@markdown - Use TPU v2-8 runtime (most RAM)\n",
        "#@markdown - Try disabling relaxation if you run out of memory anyway\n",
        "\n",
        "\n",
        "from google.colab import files\n",
        "import os, re\n",
        "import hashlib\n",
        "import random\n",
        "from pathlib import Path\n",
        "\n",
        "from datetime import datetime\n",
        "d=datetime.now()\n",
        "\n",
        "def add_hash(x,y):\n",
        "  return x+\"_\"+hashlib.sha1(y.encode()).hexdigest()[:7]\n",
        "chain_ids = ''\n",
        "\n",
        "n=d.strftime('%H%M%d%m%Y')\n",
        "\n",
        "#query_sequence = 'PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK' #@param {type:\"string\"}\n",
        "jobname = \"ja pobrusze a ty poczywaj\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown - Upload a template protein model (PDB or mmCIF)\n",
        "use_templates = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown - Run AMBER relaxation on top-scored prediction\n",
        "relax_top_model = True #@param {type:\"boolean\"}\n",
        "\n",
        "##@markdown - Force matching following template chains to consecutive input sequences (default: greedy automation)\n",
        "#chain_ids = '' # @param {type:\"string\", placeholder:\"A,B,C\"}\n",
        "\n",
        "#@markdown - Remove everything except models and figures from output (MSAs, AF2 pkl etc can be pretty large)\n",
        "clean_outputs = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown - Close session after downloading results\n",
        "disconnet_when_complete = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown - You can reduce MSA depth here\n",
        "max_msa_depth = 5000 #@param [10, 100, 1000, 2000, 3000, 4000, 5000, 10000] {type:\"raw\"}\n",
        "\n",
        "jobname_prefix = re.sub(r'\\s', '_', jobname.strip())\n",
        "jobname_prefix = re.sub(r'\\W+', '', jobname_prefix)\n",
        "\n",
        "print(\"Uploading FASTA\")\n",
        "uploaded_fasta = files.upload()\n",
        "fasta_fn, query_sequence = list(uploaded_fasta.items())[0]\n",
        "query_sequence = query_sequence.decode('ascii')\n",
        "jobfolder=add_hash(jobname_prefix, query_sequence)\n",
        "\n",
        "\n",
        "if os.path.exists(jobfolder):\n",
        "  n=1\n",
        "  while os.path.exists(f\"{jobfolder}_{n}\"): n += 1\n",
        "  jobfolder = f\"{jobfolder}_{n}\"\n",
        "\n",
        "# create a job directory\n",
        "os.makedirs(jobfolder, exist_ok=True)\n",
        "os.rename(fasta_fn,os.path.join(jobfolder,fasta_fn))\n",
        "fasta_fn = os.path.join('/content',jobfolder,fasta_fn)\n",
        "\n",
        "\n",
        "template_files=[]\n",
        "if use_templates:\n",
        "  print(\"Uploading template in PDB/mmCIF\")\n",
        "  uploaded_template = files.upload()\n",
        "  use_templates = True\n",
        "  fn = list(uploaded_template.keys())[0]\n",
        "  template_fn=os.path.join('/','content',jobfolder,fn)\n",
        "  os.rename(fn,template_fn)\n",
        "  template_files.append(template_fn)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NA9JHdIglArX",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# Set environment variables before running any other code.\n",
        "import os\n",
        "from sys import version_info\n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "\n",
        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '4.0'\n",
        "\n",
        "\n",
        "#@title 1. Prepare colab environment\n",
        "\n",
        "from IPython.utils import io\n",
        "import os\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "if os.path.exists('/opt/conda'):\n",
        "  print('Using existing conda environment')\n",
        "else:\n",
        "  try:\n",
        "    with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      with io.capture_output() as captured:\n",
        "        # Uninstall default Colab version of TF.\n",
        "        #%shell pip uninstall -y tensorflow keras\n",
        "        #%shell sudo apt install --quiet --yes hmmer\n",
        "        #pbar.update(6)\n",
        "\n",
        "        # Install py3dmol.\n",
        "        %shell pip install py3dmol\n",
        "        pbar.update(10)\n",
        "\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 -bfp /opt/conda \\\n",
        "        #    && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "\n",
        "        %shell wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh\n",
        "        %shell bash Miniforge3-Linux-x86_64.sh -bfp /opt/conda\n",
        "        PATH=%env PATH\n",
        "        %env PATH=/opt/conda/bin:{PATH}\n",
        "\n",
        "        pbar.update(10)\n",
        "\n",
        "        #%shell mamba config --set auto_update_conda false\n",
        "        %shell mamba install -y -c conda-forge -c bioconda \\\n",
        "                          python='{python_version}' \\\n",
        "                          matplotlib \\\n",
        "                          kalign2=2.04 \\\n",
        "                          hhsuite=3.3.0 \\\n",
        "                          openmm=8.0.0 \\\n",
        "                          pdbfixer\n",
        "\n",
        "        #%shell conda install -qy conda==24.11.1 \\\n",
        "        #    && conda install -qy -c conda-forge -c bioconda\\\n",
        "        #      python={python_version} \\\n",
        "        #      openmm=8.0.0 \\\n",
        "        #      matplotlib \\\n",
        "        #      kalign2 \\\n",
        "        #      hhsuite \\\n",
        "        #      pyopenssl==22.0.0 \\\n",
        "        #      pdbfixer\n",
        "        pbar.update(70)\n",
        "\n",
        "        %shell git clone --recursive https://github.com/gchojnowski/gapTrick.git\n",
        "        pbar.update(10)\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "M5odFfx3macA"
      },
      "outputs": [],
      "source": [
        "#@title 2. Download AlphaFold\n",
        "\n",
        "from sys import version_info\n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar'\n",
        "PARAMS_DIR = '/content/alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\n",
        "\n",
        "\n",
        "if os.path.exists('/content/alphafold'):\n",
        "  print(\"Using an existing AF2 instalation\")\n",
        "else:\n",
        "  try:\n",
        "    with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      with io.capture_output() as captured:\n",
        "\n",
        "        %shell git clone --branch main {GIT_REPO} /content/alphafold\n",
        "        pbar.update(10)\n",
        "        # Install AlphaFodl2 and its dependencies.\n",
        "        %shell pip3 install -r /content/alphafold/requirements.txt\n",
        "        %shell pip3 install --no-dependencies /content/alphafold\n",
        "        pbar.update(30)\n",
        "\n",
        "        # Download stereo_chemical_props.txt\n",
        "        %shell wget -q -P /content https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "        %shell mkdir -p /opt/conda/lib/python{python_version}/site-packages/alphafold/common/\n",
        "        %shell cp -f /content/stereo_chemical_props.txt /opt/conda/lib/python{python_version}/site-packages/alphafold/common/\n",
        "\n",
        "        # Download weights\n",
        "        %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "        %shell wget -O \"{PARAMS_PATH}\" \"{SOURCE_URL}\"\n",
        "        pbar.update(50)\n",
        "\n",
        "        %shell tar --extract --verbose --file=\"{PARAMS_PATH}\" --directory=\"{PARAMS_DIR}\" --preserve-permissions\n",
        "        %shell rm \"{PARAMS_PATH}\"\n",
        "        pbar.update(10)\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise\n",
        "\n",
        "import sys\n",
        "sys.path.append(f'/opt/conda/lib/python{python_version}/site-packages')\n",
        "sys.path.append('/content/alphafold')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pOUdqcfwnkzb",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title 3. Run gapTrick\n",
        "\n",
        "jobpath=Path('/', 'content', jobfolder)\n",
        "msapath=Path('/', 'content', jobfolder, 'msa')\n",
        "os.makedirs(msapath, exist_ok=True)\n",
        "\n",
        "from sys import version_info\n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "import sys,os\n",
        "import glob\n",
        "import datetime\n",
        "import matplotlib.pyplot as plt\n",
        "plt.close()\n",
        "\n",
        "import logging\n",
        "for handler in logging.root.handlers[:]:\n",
        "    logging.root.removeHandler(handler)\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "logging.basicConfig(level=logging.INFO, format=\"%(message)s\",\\\n",
        "      handlers=[logging.FileHandler(Path(jobpath, 'logfile.txt')),\n",
        "                logging.StreamHandler(sys.stdout)])\n",
        "\n",
        "sys.path.append(f'/opt/conda/lib/python{python_version}/site-packages')\n",
        "sys.path.append('/content/alphafold')\n",
        "sys.path.insert(1, '/content/gapTrick')\n",
        "\n",
        "from gapTrick.__main__ import *\n",
        "import gapTrick.version\n",
        "\n",
        "from Bio import SeqIO\n",
        "\n",
        "logger.info(f\"## gapTrick version {version.__version__} (colab)\\n\\n\")\n",
        "logger.info(f\"Started at {datetime.now().strftime('%H:%M:%S %d/%m/%Y')}\\n\")\n",
        "\n",
        "msas=[]\n",
        "local_msa_dict = {}\n",
        "with open(fasta_fn) as ifile:\n",
        "  for record in SeqIO.parse(ifile, \"fasta\"):\n",
        "    a3m_fname=local_msa_dict.get(record.seq, None)\n",
        "\n",
        "    if a3m_fname:\n",
        "      logger.info(\"Found existing MSA\")\n",
        "    else:\n",
        "      a3m_fname = os.path.join(msapath, f\"{len(local_msa_dict):04d}.a3m\")\n",
        "      query_mmseqs2(record.seq, a3m_fname, user_agent='colab')\n",
        "      local_msa_dict[record.seq]=a3m_fname\n",
        "\n",
        "    logger.info(f\"{record.id}: {a3m_fname}\")\n",
        "    logger.info('')\n",
        "    msas.append(a3m_fname)\n",
        "\n",
        "\n",
        "runme(msa_filenames     = msas,\n",
        "      query_cardinality = [1]*len(msas),\n",
        "      query_trim        = [[0,9999]]*len(msas),\n",
        "      jobname           = jobpath,\n",
        "      do_relax          = relax_top_model,\n",
        "      data_dir          = Path('/', 'content', 'alphafold', 'data'),\n",
        "      max_seq           = max_msa_depth,\n",
        "      pbty_cutoff       = 0.8,\n",
        "      chain_ids         = chain_ids if chain_ids else None,\n",
        "      template_fn_list  = template_files if use_templates else [])\n",
        "\n",
        "\n",
        "logger.info(f\"\\n\\nNormal termination at {datetime.now().strftime('%H:%M:%S %d/%m/%Y')}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iv_KG_SQHV0X"
      },
      "outputs": [],
      "source": [
        "#@title Display 3D model\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "\n",
        "PLDDT_BANDS = [(50, '#FF7D45'),\n",
        "               (70, '#FFDB13'),\n",
        "               (90, '#65CBF3'),\n",
        "               (100, '#0053D6')]\n",
        "\n",
        "def set_b_to_plddtbands_bio(structure):\n",
        "\n",
        "    plddt_lims = np.array([_[0] for _ in PLDDT_BANDS])\n",
        "\n",
        "    for chain in structure:\n",
        "        for resi in chain:\n",
        "            for atm in resi:\n",
        "                atm.set_bfactor(float(np.argmax(plddt_lims>atm.get_bfactor())))\n",
        "\n",
        "# -----------------------------------------------------------------------------\n",
        "\n",
        "def set_3dmol_styles(\n",
        "    view,\n",
        "    viewer,\n",
        "    chain_ids=1,\n",
        "    color=[\"lDDT\", \"rainbow\", \"chain\"][0],\n",
        "    show_sidechains=False,\n",
        "    show_mainchains=False,):\n",
        "\n",
        "    if color == \"lDDT\":\n",
        "        color_map = {i: band[1] for i, band in enumerate(PLDDT_BANDS)}\n",
        "        view.setStyle(\n",
        "            {\n",
        "                \"cartoon\": {\n",
        "                    \"colorscheme\": {\"prop\": \"b\",'map': color_map}\n",
        "                }\n",
        "            }, viewer=viewer,)\n",
        "    elif color == \"chain\":\n",
        "        for cid, color in zip(\n",
        "            chain_ids, [\"lime\",\"cyan\",\"magenta\", \\\n",
        "             \"yellow\",\"salmon\",\"white\",\\\n",
        "             \"blue\",\"orange\",\"black\",\\\n",
        "             \"green\",\"gray\",]* 2,):\n",
        "            view.setStyle({\"chain\": cid},\n",
        "                          {\"cartoon\": {\"color\": color}}, viewer=viewer)\n",
        "\n",
        "# -----------------------------------------------------------------------------\n",
        "\n",
        "def bio2pdbstring(structure):\n",
        "    pio=PDBIO()\n",
        "    pio.set_structure(structure)\n",
        "    fo = io.StringIO()\n",
        "    pio.save(fo)\n",
        "    fo.seek(0)\n",
        "    return fo.read()\n",
        "\n",
        "# -----------------------------------------------------------------------------\n",
        "# -----------------------------------------------------------------------------\n",
        "import re\n",
        "import py3Dmol\n",
        "parser = PDBParser()\n",
        "\n",
        "top_model_fn = jobpath=Path('/', 'content', jobfolder, 'input', 'ranked_0.pdb')\n",
        "structure = parser.get_structure(\"AF\", top_model_fn)[0]\n",
        "parser = PDBParser(QUIET=True)\n",
        "template_structure = None\n",
        "if use_templates:\n",
        "  try:\n",
        "    template_structure = parser.get_structure(\"template\", template_files[0])[0]\n",
        "  except:\n",
        "    parser = MMCIFParser(QUIET=True)\n",
        "    template_structure = parser.get_structure(\"template\", template_files[0])[0]\n",
        "\n",
        "\n",
        "\n",
        "chain_ids = [_.id for _ in structure]\n",
        "\n",
        "view = py3Dmol.view(\n",
        "            width=1000,\n",
        "            height=500,\n",
        "            viewergrid=(1, 2),)\n",
        "\n",
        "viewer = (0, 0)\n",
        "view.addModel(bio2pdbstring(structure), \"pdb\", viewer=viewer)\n",
        "view.zoomTo(viewer=viewer)\n",
        "set_3dmol_styles( view, viewer, chain_ids=chain_ids, color=\"chain\")\n",
        "\n",
        "if use_templates:\n",
        "  view.addModel(bio2pdbstring(template_structure), \"pdb\", viewer=viewer)\n",
        "\n",
        "  view.setStyle({\"model\": -1},\n",
        "                {\"cartoon\":\n",
        "                 {\"color\": 'black',\n",
        "                  'style':'trace',\n",
        "                  'ribbon':True,\n",
        "                  'opacity':0.7}},\n",
        "              viewer=viewer)\n",
        "\n",
        "viewer = (0, 1)\n",
        "set_b_to_plddtbands_bio(structure)\n",
        "view.addModel(bio2pdbstring(structure), \"pdb\", viewer=viewer)\n",
        "view.zoomTo(viewer=viewer)\n",
        "set_3dmol_styles(view, viewer, chain_ids=chain_ids, color=\"lDDT\")\n",
        "\n",
        "if use_templates:\n",
        "  view.addModel(bio2pdbstring(template_structure), \"pdb\", viewer=viewer)\n",
        "\n",
        "  view.setStyle({\"model\": -1},\n",
        "                {\"cartoon\":\n",
        "                 {\"color\": 'black',\n",
        "                  'style':'trace',\n",
        "                  'ribbon':True,\n",
        "                  'opacity':0.7}},\n",
        "              viewer=viewer)\n",
        "\n",
        "# generate CA position dict\n",
        "CA_atom_coords = {}\n",
        "for chain in structure:\n",
        "  for res in chain:\n",
        "    CA_atom_coords[f'{res._id[1]}{chain.id}'] = list(res['CA'].coord)\n",
        "\n",
        "contact_template = r\"\\* (?P<rname1>\\w+?)/(?P<ch1>\\w+?)/(?P<res1>\\w+?)\\s+(?P<rname2>\\w+?)/(?P<ch2>\\w+?)/(?P<res2>\\w+?)\\s+(?P<pbty>[\\d\\.]*?)\"\n",
        "with open(Path('/', 'content', jobfolder, 'contacts.txt'), 'r') as ifile:\n",
        "  for contact_str in ifile:\n",
        "        m = re.match(contact_template, contact_str)\n",
        "        if not m: continue\n",
        "        xyz_from = dict([ (l,float(c)) for l,c in zip(('x','y','z'),CA_atom_coords[f\"{m.group('res1')}{m.group('ch1')}\"])])\n",
        "        xyz_to = dict([ (l,float(c)) for l,c in zip(('x','y','z'),CA_atom_coords[f\"{m.group('res2')}{m.group('ch2')}\"])])\n",
        "        view.addLine({\"color\":'red',\"start\":xyz_from,\"end\":xyz_to});\n",
        "\n",
        "view.show()\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7QqbpgDj-BJi"
      },
      "outputs": [],
      "source": [
        "#@title Pack and download results, and terminate the session (if requested)\n",
        "\n",
        "from google.colab import files\n",
        "import time\n",
        "\n",
        "if clean_outputs:\n",
        "  jobpath=Path('/', 'content', jobfolder)\n",
        "  %shell rm {jobpath}/input/features.pkl\n",
        "  %shell rm {jobpath}/input/result_model*.pkl\n",
        "  %shell rm -r {jobpath}/input/msas\n",
        "  %shell rm -r {jobpath}/msa\n",
        "\n",
        "# pack and download\n",
        "results_zip = f\"{jobfolder}.zip\"\n",
        "os.system(f\"zip -r {results_zip} {jobfolder}\")\n",
        "files.download(results_zip)\n",
        "\n",
        "\n",
        "if disconnet_when_complete:\n",
        "  time.sleep(300)\n",
        "  from google.colab import runtime\n",
        "  runtime.unassign()\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}