alphafold_ft_colab_pipeline_v1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/phbradley/alphafold_finetune/blob/main/alphafold_ft_colab_pipeline_v1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# alphafold_finetune colab\n",
        "\n",
        "This draft colab notebook has examples of fine-tuning and binder prediction. Hopefully it will give an idea of one route to installing the software, but note that the colab machines already have the GPU stuff (CUDA/CUDNN) set up, which simplifies things. Starting from scratch may require additional GPU-specific installation depending on the machine.\n",
        "\n",
        "The plan for the future is to add a \"forms\" interface that will let users run structure/binding predictions for peptide-MHC targets starting from the allele and peptide information.  \n",
        "\n",
        "This notebook is based on the AlphaFold colab notebook https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb -- many thanks to the AlphaFold developers for creating and sharing their code and related content.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uC1dKAwk2eyl"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Start by running the 3 cells below to set up the code and download data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "woIxeCPygt7K"
      },
      "outputs": [],
      "source": [
        "# Set environment variables before running any other code.\n",
        "import os\n",
        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '4.0'\n",
        "\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",
        "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\n",
        "      pbar.update(6)\n",
        "\n",
        "      # Install py3dmol.\n",
        "      %shell pip install py3dmol\n",
        "      pbar.update(2)\n",
        "\n",
        "      # Install OpenMM and pdbfixer.\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",
        "      pbar.update(12)\n",
        "\n",
        "      PATH=%env PATH\n",
        "      %env PATH=/opt/conda/bin:{PATH}\n",
        "\n",
        "      # phil changing to python 3.8 from 3.10 for compat with local versions...\n",
        "      %shell conda install -qy conda==4.13.0 python=3.8\n",
        "\n",
        "      #%shell conda install -qy conda==4.13.0 \\\n",
        "      #    && conda install -qy -c conda-forge \\\n",
        "      #      python=3.8\n",
        "      pbar.update(80)\n",
        "\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "#print(captured)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VzJ5iMjTtoZw"
      },
      "outputs": [],
      "source": [
        "GIT_REPO = 'https://github.com/phbradley/alphafold_finetune'\n",
        "\n",
        "PARAMS_URLS = ['https://www.dropbox.com/s/e3uz9mwxkmmv35z/params_model_2_ptm.npz',\n",
        "]\n",
        "\n",
        "PARAMS_DIR = './alphafold_params/params'\n",
        "\n",
        "try:\n",
        "  with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "      %shell rm -rf alphafold_finetune\n",
        "      %shell git clone --branch main {GIT_REPO} alphafold_finetune\n",
        "      pbar.update(20)\n",
        "      # Install the required versions of all dependencies.\n",
        "      %shell pip3 install -r ./alphafold_finetune/requirements_colab_python38_v2.txt\n",
        "      pbar.update(50)\n",
        "\n",
        "      # Load parameters\n",
        "      %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "      for URL in PARAMS_URLS:\n",
        "        PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(URL))\n",
        "        %shell wget -O \"{PARAMS_PATH}\" \"{URL}\"\n",
        "      pbar.update(15)\n",
        "\n",
        "      # download alphafold fine-tune dataset\n",
        "      %shell wget https://files.ipd.uw.edu/pub/alphafold_finetune_motmaen_pnas_2023/datasets_alphafold_finetune_v2_2023-02-20.tgz\n",
        "      %shell tar -xzf datasets_alphafold_finetune_v2_2023-02-20.tgz\n",
        "      %shell mv datasets_alphafold_finetune alphafold_finetune/\n",
        "      pbar.update(15)\n",
        "\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "#print(captured)\n",
        "\n",
        "import jax\n",
        "if jax.local_devices()[0].platform == 'tpu':\n",
        "  raise RuntimeError('Colab TPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "elif jax.local_devices()[0].platform == 'cpu':\n",
        "  raise RuntimeError('Colab CPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "else:\n",
        "  print(f'Running with {jax.local_devices()[0].device_kind} GPU')\n",
        "\n",
        "# Make sure everything we need is on the path.\n",
        "import sys\n",
        "sys.path.append('/opt/conda/lib/python3.8/site-packages')\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "cd alphafold_finetune/"
      ],
      "metadata": {
        "id": "lM9cAcKxiklC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## This command fine-tunes AlphaFold's parameters for peptide-MHC binding prediction."
      ],
      "metadata": {
        "id": "t3Ugyj0fwJBZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%shell python run_finetuning.py \\\n",
        "    --data_dir /content/alphafold_params/ \\\n",
        "    --outprefix testrun1 \\\n",
        "    --binder_intercepts 0.80367635 --binder_intercepts 0.43373787 \\\n",
        "    --freeze_binder  \\\n",
        "    --train_dataset datasets_alphafold_finetune/pmhc_finetune/combo_1and2_train.tsv \\\n",
        "    --valid_dataset datasets_alphafold_finetune/pmhc_finetune/combo_1and2_valid.tsv\n"
      ],
      "metadata": {
        "id": "z7Fh5GS2F2Ag"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## This command makes structure and binding predictions for a set of 10mer targets using the fine-tuned parameters from the paper."
      ],
      "metadata": {
        "id": "hkgT4PWowqS9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%shell python run_prediction.py --targets examples/pmhc_hcv_polg_10mers/targets.tsv \\\n",
        "    --outfile_prefix polg_test2 --model_names model_2_ptm_ft \\\n",
        "    --model_params_files datasets_alphafold_finetune/params/mixed_mhc_pae_run6_af_mhc_params_20640.pkl \\\n",
        "    --ignore_identities"
      ],
      "metadata": {
        "id": "ECWIiE6QF18i"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "V100",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}