heme_rfdiffusion_aa.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/engelberger/all_atom_binder_diffusion/blob/dev/heme_rfdiffusion_aa.ipynb?hl=en\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJsYHgSER7xN"
      },
      "source": [
        "#**RFdiffusion aa**\n",
        "RFdiffusion aa is a method for structure generation, with or without conditional information (a motif, target etc). It can perform a whole range of protein design challenges as we have outlined in the RFdiffusion [manuscript](https://www.science.org/doi/10.1126/science.adl2528).\n",
        "\n",
        "**<font color=\"red\">NOTE:</font>** This notebook is in development, we are still working on adding all the options from the manuscript above.\n",
        "\n",
        "For **instructions**, see end of Notebook.\n",
        "\n",
        "\n",
        "\n",
        "This is a modified version of Sergey's notebook by Felipe Engelberger, see [original version](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/rf/examples/diffusion_ori.ipynb) of this notebook (from 31Mar2023).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "ffOpPfUqHXpb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Starting: Installing aria2 for faster downloads\n",
            "Completed: Installing aria2 for faster downloads\n",
            "Starting: Installing Python dependencies\n",
            "Completed: Installing Python dependencies\n",
            "Starting: Download ColabFold Utils\n",
            "Completed: Download ColabFold Utils\n",
            "Starting: Installing DGL\n",
            "Completed: Installing DGL\n",
            "Starting: Installing SE3 Transformer\n",
            "Completed: Installing SE3 Transformer\n",
            "Starting: Downloading and setting up 'ananas'\n",
            "Completed: Downloading and setting up 'ananas'\n",
            "Environment setup complete.\n"
          ]
        }
      ],
      "source": [
        "# @title COLAB ONLY setup **RFdiffusion All Atom** (~5m) { display-mode: \"form\" }\n",
        "\n",
        "\n",
        "try:\n",
        "    from IPython import get_ipython\n",
        "\n",
        "    if \"google.colab\" in str(get_ipython()):\n",
        "        %%time\n",
        "        IN_COLAB = True\n",
        "        \n",
        "    else:\n",
        "        IN_COLAB = False\n",
        "except Exception:\n",
        "    IN_COLAB = False\n",
        "\n",
        "    \n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "import time\n",
        "import sys\n",
        "\n",
        "# Function to detect if running on Google Colab\n",
        "def is_colab():\n",
        "    return \"COLAB_GPU\" in os.environ\n",
        "\n",
        "\n",
        "# Base directory setup\n",
        "if is_colab():\n",
        "    BASE_DIR = \"/content\"\n",
        "else:\n",
        "    # For local setup, adjust this path as per your local environment\n",
        "    # By default we will assume you are in the devcontainer path\n",
        "    BASE_DIR = \"/workspaces/all_atom_binder_diffusion\"\n",
        "\n",
        "    # For setups outside the devcontainer, you may need to adjust this path\n",
        "    # BASE_DIR = os.path.expanduser(\"~\")\n",
        "\n",
        "# Adjust paths based on the environment\n",
        "PARAMS_DIR = os.path.join(BASE_DIR, \"params\")\n",
        "RF_DIFFUSION_DIR = os.path.join(BASE_DIR, \"RFdiffusion\")\n",
        "RF_ALL_ATOM_DIR = os.path.join(BASE_DIR, \"rf_diffusion_all_atom\")\n",
        "LIGAND_MPNN_DIR = os.path.join(BASE_DIR, \"LigandMPNN\")\n",
        "HEME_PROTOCOL_DIR = os.path.join(BASE_DIR, \"heme_binder_diffusion\")\n",
        "if 'RFdiffusion' not in sys.path:\n",
        "  os.environ[\"DGLBACKEND\"] = \"pytorch\"\n",
        "  sys.path.append('RFdiffusion')\n",
        "\n",
        "# Ensure the params directory exists\n",
        "os.makedirs(PARAMS_DIR, exist_ok=True)\n",
        "\n",
        "\n",
        "def run_command(command, progress_message, wait=True):\n",
        "    \"\"\"\n",
        "    Run a system command with a progress message.\n",
        "    If wait is False, the command is executed in the background.\n",
        "    \"\"\"\n",
        "    print(f\"Starting: {progress_message}\")\n",
        "    process = subprocess.Popen(\n",
        "        command,\n",
        "        shell=True,\n",
        "        stdout=subprocess.PIPE,\n",
        "        stderr=subprocess.PIPE,\n",
        "        universal_newlines=True,\n",
        "    )\n",
        "    if wait:\n",
        "        stdout, stderr = process.communicate()\n",
        "        # If the return code is an error, print the stderr but do not raise an exception\n",
        "        if process.returncode != 0:\n",
        "            print(f\"Error during {progress_message}: {stderr}\")\n",
        "            raise subprocess.CalledProcessError(process.returncode, command)\n",
        "\n",
        "        print(f\"Completed: {progress_message}\")\n",
        "    return process\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def setup_environment_colab():\n",
        "    # Install aria2 if not already installed (for faster downloads)\n",
        "    run_command(\"apt-get install -y aria2\", \"Installing aria2 for faster downloads\")\n",
        "\n",
        "    # If parameters are already downloaded, skip the download process\n",
        "    if not os.path.isfile(os.path.join(PARAMS_DIR, \"done.txt\")):\n",
        "        print(\"Downloading parameters and models...\")\n",
        "\n",
        "        # Start downloading parameters and models in the background\n",
        "        download_process = run_command(\n",
        "            f\"cd {PARAMS_DIR} && aria2c -q -x 16 https://files.ipd.uw.edu/krypton/schedules.zip && \\\n",
        "            aria2c -q -x 16 http://files.ipd.uw.edu/pub/RF-All-Atom/weights/RFDiffusionAA_paper_weights.pt && \\\n",
        "            aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar && \\\n",
        "            aria2c -q -x 16 https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_020.pt && \\\n",
        "            aria2c -q -x 16 https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_010_25.pt && \\\n",
        "            aria2c -q -x 16 https://files.ipd.uw.edu/pub/ligandmpnn/per_residue_label_membrane_mpnn_v_48_020.pt && \\\n",
        "            aria2c -q -x 16 https://files.ipd.uw.edu/pub/ligandmpnn/global_label_membrane_mpnn_v_48_020.pt && \\\n",
        "            aria2c -q -x 16 https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_020.pt && \\\n",
        "            tar -xf alphafold_params_2022-12-06.tar && \\\n",
        "            touch done.txt\",\n",
        "            \"Downloading and extracting parameters\", wait=False)\n",
        "\n",
        "\n",
        "\n",
        "    # Install Open Babel if not already installed\n",
        "    if not os.path.isfile(\"/usr/bin/obabel\"):\n",
        "        run_command(\n",
        "            \"apt-get install -y openbabel libopenbabel-dev && ln -s /usr/include/openbabel3 /usr/local/include/openbabel3\",\n",
        "            \"Installing Open Babel and its development files\",\n",
        "        )\n",
        "\n",
        "    # Install SWIG\n",
        "    if not os.path.isfile(\"/usr/bin/swig\"):\n",
        "        run_command(\n",
        "            \"apt-get remove -y swig && apt-get install -y swig3.0 && ln -sf /usr/bin/swig3.0 /usr/bin/swig\",\n",
        "            \"Installing SWIG\",\n",
        "        )\n",
        "\n",
        "    # Install Python dependencies\n",
        "    run_command(\n",
        "        \"pip install jedi omegaconf hydra-core icecream pyrsistent assertpy deepdiff fire git+https://github.com/sokrypton/ColabDesign.git@gamma py3Dmol openbabel prody pyparsing==3.1.1\",\n",
        "        \"Installing Python dependencies\",\n",
        "    )\n",
        "\n",
        "    run_command(\n",
        "        \"wget https://raw.githubusercontent.com/sokrypton/ColabFold/main/colabfold/colabfold.py -O colabfold_utils.py\",\n",
        "        \"Download ColabFold Utils\",\n",
        "    )\n",
        "    # Clone RFdiffusion repository\n",
        "    if not os.path.isdir(RF_DIFFUSION_DIR):\n",
        "        run_command(\n",
        "            f\"git clone --branch max https://github.com/engelberger/RFdiffusion.git {RF_DIFFUSION_DIR}\",\n",
        "            \"Cloning RFdiffusion repository\",\n",
        "        )\n",
        "\n",
        "    # Clone LigandMPNN repository\n",
        "    if not os.path.isdir(LIGAND_MPNN_DIR):\n",
        "        run_command(\n",
        "            f\"git clone https://github.com/ikalvet/LigandMPNN.git {LIGAND_MPNN_DIR}\",\n",
        "            \"Cloning LigandMPNN repository\",\n",
        "        )\n",
        "\n",
        "    #git clone\n",
        "    # Clone Binder Diffusion Protocol repository\n",
        "    if not os.path.isdir(HEME_PROTOCOL_DIR):\n",
        "        run_command(\n",
        "            f\"git clone --branch colab_march_2024 https://github.com/engelberger/heme_binder_diffusion.git {HEME_PROTOCOL_DIR}\",\n",
        "            \"Cloning AA Binder Protocol repository\",\n",
        "        )\n",
        "    # Clone RFdiffusion all atom repository\n",
        "    if not os.path.isdir(RF_ALL_ATOM_DIR):\n",
        "        run_command(\n",
        "            f\"git clone --recurse-submodules --branch colab_march_2024 https://github.com/engelberger/rf_diffusion_all_atom.git {RF_ALL_ATOM_DIR}\",\n",
        "            \"Cloning RFdiffusion all atom repository\",\n",
        "        )\n",
        "\n",
        "    # Install DGL\n",
        "    run_command(\n",
        "        \"pip install dgl -f https://data.dgl.ai/wheels/cu121/repo.html\",\n",
        "        \"Installing DGL\",\n",
        "    )\n",
        "\n",
        "\n",
        "    # Install SE3 Transformer\n",
        "    run_command(\n",
        "        f\"cd {os.path.join(RF_DIFFUSION_DIR, 'env/SE3Transformer')} && pip install -q --no-cache-dir -r requirements.txt && pip install -q .\",\n",
        "        \"Installing SE3 Transformer\",\n",
        "    )\n",
        "\n",
        "    # Download and set execute permissions for 'ananas'\n",
        "    run_command(\n",
        "        f\"wget -qnc https://files.ipd.uw.edu/krypton/ananas -P {BASE_DIR} && chmod +x {os.path.join(BASE_DIR, 'ananas')}\",\n",
        "        \"Downloading and setting up 'ananas'\",\n",
        "    )\n",
        "    \n",
        "    # If the download process is running in the background, wait for it to complete\n",
        "    if not os.path.isfile(os.path.join(PARAMS_DIR, \"done.txt\")):\n",
        "        download_process.communicate()\n",
        "   \n",
        "    print(\"Environment setup complete.\")\n",
        "\n",
        "\n",
        "# Call the setup function\n",
        "setup_environment_colab()\n",
        "%pip install -q -U --no-dependencies chex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting pyrosetta-installer\n",
            "  Downloading pyrosetta_installer-0.1.2-py3-none-any.whl (3.9 kB)\n",
            "Installing collected packages: pyrosetta-installer\n",
            "Successfully installed pyrosetta-installer-0.1.2\n",
            "Installing PyRosetta:\n",
            " os: ubuntu\n",
            " type: Release\n",
            " Rosetta C++ extras: \n",
            " mirror: https://west.rosettacommons.org/pyrosetta/release/release\n",
            " extra packages: numpy\n",
            "\n",
            "PyRosetta wheel url: https://:@west.rosettacommons.org/pyrosetta/release/release/PyRosetta4.Release.python310.ubuntu.wheel/pyrosetta-2024.42+release.3366cf78a3-cp310-cp310-linux_x86_64.whl\n"
          ]
        }
      ],
      "source": [
        "%pip install pyrosetta-installer \n",
        "import pyrosetta_installer\n",
        "pyrosetta_installer.install_pyrosetta()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "QZM_Q39_zwN1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Please upload your PDB file:\n",
            "Output base directory: /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif\n",
            "Final output path: /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2\n",
            "Running command: python ./rf_diffusion_all_atom/run_inference.py --config-name=config --config-dir=/workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2 diffuser.T=25\n",
            "/workspaces/all_atom_binder_diffusion/rf_diffusion_all_atom/util.py:216: UserWarning: Using torch.cross without specifying the dim arg is deprecated.\n",
            "Please either pass the dim explicitly or simply use torch.linalg.cross.\n",
            "The default value of dim will change to agree with that of linalg.cross in a future release. (Triggered internally at ../aten/src/ATen/native/Cross.cpp:63.)\n",
            "Z = torch.cross(Xn,Yn)\n",
            "/usr/local/lib/python3.10/dist-packages/hydra/_internal/defaults_list.py:251: UserWarning: In 'config': Defaults list is missing `_self_`. See https://hydra.cc/docs/1.2/upgrades/1.0_to_1.1/default_composition_order for more information\n",
            "warnings.warn(msg, UserWarning)\n",
            "[2024-11-07 09:07:13,280][inference.model_runners][INFO] - Reading checkpoint from ./params/RFDiffusionAA_paper_weights.pt\n",
            "loading ./params/RFDiffusionAA_paper_weights.pt\n",
            "loaded ./params/RFDiffusionAA_paper_weights.pt\n",
            "OVERRIDING: You are changing diffuser.T from the value this model was trained with.\n",
            "[2024-11-07 09:07:19,796][inference.model_runners][INFO] - Loading checkpoint.\n",
            "[2024-11-07 09:07:19,948][diffusion][INFO] - No IGSO3 cache found at /workspaces/all_atom_binder_diffusion/rf_diffusion_all_atom/cached_schedules/T_25_omega_1000_min_sigma_0_02_min_b_1_5_max_b_2_5_schedule_linear.pkl.\n",
            "[2024-11-07 09:07:19,948][diffusion][INFO] - Calculating IGSO3.\n",
            "[2024-11-07 09:08:33,280][__main__][INFO] - Making design /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2/sample_0\n",
            "[2024-11-07 09:08:33,280][__main__][INFO] - making design 0 of 0:1\n",
            "[2024-11-07 09:08:33,280][inference.model_runners][INFO] - Using contig: ['30-110', 'A15-15', '30-110']\n",
            "With this beta schedule (linear schedule, beta_0 = 0.08, beta_T = 0.56), alpha_bar_T = 3.613896478782408e-05\n",
            "Calculating chi_beta_T dictionary...\n",
            "Done calculating chi_beta_T dictionaries. They are now cached.\n",
            "Done calculating chi_beta_T, chi_alphas_T, and chi_abars_T dictionaries.\n",
            "[2024-11-07 09:09:11,817][inference.model_runners][INFO] - 09:09:11: Timestep 25, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:15,579][inference.model_runners][INFO] - 09:09:15: Timestep 24, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:19,330][inference.model_runners][INFO] - 09:09:19: Timestep 23, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:23,084][inference.model_runners][INFO] - 09:09:23: Timestep 22, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:26,848][inference.model_runners][INFO] - 09:09:26: Timestep 21, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:30,617][inference.model_runners][INFO] - 09:09:30: Timestep 20, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:34,367][inference.model_runners][INFO] - 09:09:34: Timestep 19, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:38,154][inference.model_runners][INFO] - 09:09:38: Timestep 18, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:41,970][inference.model_runners][INFO] - 09:09:41: Timestep 17, current sequence: ????????????????????????????????????????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:45,775][inference.model_runners][INFO] - 09:09:45: Timestep 16, current sequence: ????????????????????????????????S???????????????????????????????????????????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:49,620][inference.model_runners][INFO] - 09:09:49: Timestep 15, current sequence: ???E???????????????????G??????????S????????????????????????????????????L????????????????????????????????aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:53,467][inference.model_runners][INFO] - 09:09:53: Timestep 14, current sequence: ???EL?????E???????????SG????S???????????????????????E??T??????????????GQ??G??E?G????G????????????L???L??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:09:57,308][inference.model_runners][INFO] - 09:09:57: Timestep 13, current sequence: ???E?V???LE?????L???GQSG???????GG????????????A??????Q?ST?L????????????GQ?GG??G?G????G??????????LLL???K??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:01,161][inference.model_runners][INFO] - 09:10:01: Timestep 12, current sequence: ?????L???L??????L?????QK???????QS?TG????????????L???GGS???????????????GQ?GG??G?G?D??G???????????LL???G??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:05,016][inference.model_runners][INFO] - 09:10:05: Timestep 11, current sequence: ???A?LQ???????L??E???LSK???????G??T????S????????????G?S?TL?Q??????????QS?GS??G?G?Q????L?????????LL???Q??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:08,865][inference.model_runners][INFO] - 09:10:08: Timestep 10, current sequence: ???A??Q??????????E???L?K?T??????G?TS?????????A??????DLS??A????????????GSLGS??G?G?Q?????Q???E????LL???E??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:12,699][inference.model_runners][INFO] - 09:10:12: Timestep 9, current sequence: ??????????????L??E???S?G??LV????G?TSE??S?????A??????DQQ?TA???????????ASA?GS??G?G?Q?AE??Q???E????LA??LQQ?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:16,572][inference.model_runners][INFO] - 09:10:16: Timestep 8, current sequence: ??????????????L?LE???S?K??L?????G?TGE?Q???????????????Q?TA?Q???????L?ADA?GS?EE?G?E?AE??Q???A???LL?A??AQ?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:20,430][inference.model_runners][INFO] - 09:10:20: Timestep 7, current sequence: ????????L?????L??E???S?G???????????SA?QS?????A??????S?S?TA???????????LSALGS?EE?G?Q?A?A?AL??AA?KLL?E??AQ?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:24,249][inference.model_runners][INFO] - 09:10:24: Timestep 6, current sequence: ????????L?????L??E???LSG??????L????GG?QS?K???A??A???SGQ?TA????????A??AS??SSLE??G?Q?A?AAA???AA?K??LE???S?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:28,086][inference.model_runners][INFO] - 09:10:28: Timestep 5, current sequence: ????????L??????L?E??AL?G?????ALAQ??GA?AS?K???A??AA??G?Q??AAQ??????AALAS??SSAEE?G?Q?A?AAAS?AE??K??AE?LAS?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:31,905][inference.model_runners][INFO] - 09:10:31: Timestep 4, current sequence: ???A???????A??S??EA??SQG?????L?KQ??GQ?QS?A???A??VS??S?S??AA???????AA??S??ASASE?G?D?E?A?AA?AAA?AA?AE??A??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:35,729][inference.model_runners][INFO] - 09:10:35: Timestep 3, current sequence: ???A?????????????A???SQG???AA??SS?KGQ?SS?AA??AA?????S?S????S???A??SA??S?AAS????G?Q?AA??AS?ASAAASAAE??A??aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:39,569][inference.model_runners][INFO] - 09:10:39: Timestep 2, current sequence: ???A??SA???A?A???S???SQG???GA??AS??GQ?AS??S??AA?A???A?S??A?A??AA??SA??S?AAA?SA?K?QAAS?AAS?SSA?AS?A???AQ?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:43,388][inference.model_runners][INFO] - 09:10:43: Timestep 1, current sequence: ???A??S??????S???AA??SSG???AA??SS??GASAA??S??SA?AA??S?S????A??A???S???SAAAS?S??K?A?EA??SS?AA??S??SA??SQ?aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
            "[2024-11-07 09:10:43,389][__main__][INFO] - Finished design in 2.17 minutes\n",
            "[2024-11-07 09:10:44,252][__main__][INFO] - design : /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2/sample_0.pdb\n",
            "[2024-11-07 09:10:44,252][__main__][INFO] - Xt traj: /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2/traj/sample_0_Xt-1_traj.pdb\n",
            "[2024-11-07 09:10:44,252][__main__][INFO] - X0 traj: /workspaces/all_atom_binder_diffusion/output/ligand_protein_motif/sample_2/traj/sample_0_pX0_traj.pdb\n"
          ]
        }
      ],
      "source": [
        "\n",
        "import os\n",
        "import requests\n",
        "import random\n",
        "import string\n",
        "import yaml\n",
        "import subprocess\n",
        "\n",
        "import shutil\n",
        "import glob\n",
        "\n",
        "\n",
        "# Function to detect if running on Google Colab\n",
        "def is_colab():\n",
        "    return \"COLAB_GPU\" in os.environ\n",
        "\n",
        "\n",
        "# Base directory setup\n",
        "BASE_DIR = \"/content\" if is_colab() else \"/workspaces/all_atom_binder_diffusion\"\n",
        "INPUT_DIR = os.path.join(BASE_DIR, \"input\")\n",
        "OUTPUT_DIR = os.path.join(BASE_DIR, \"output\")\n",
        "\n",
        "# Ensure the input and output directories exist\n",
        "os.makedirs(INPUT_DIR, exist_ok=True)\n",
        "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "\n",
        "def download_pdb(pdb_code, output_dir=INPUT_DIR):\n",
        "    \"\"\"\n",
        "    Download a PDB file given a PDB code.\n",
        "    \"\"\"\n",
        "    url = f\"https://files.rcsb.org/download/{pdb_code}.pdb\"\n",
        "    response = requests.get(url)\n",
        "    if response.status_code == 200:\n",
        "        pdb_path = os.path.join(output_dir, f\"{pdb_code}.pdb\")\n",
        "        with open(pdb_path, \"w\") as file:\n",
        "            file.write(response.text)\n",
        "        return pdb_path\n",
        "    else:\n",
        "        raise ValueError(f\"Failed to download PDB file for {pdb_code}\")\n",
        "\n",
        "\n",
        "def handle_pdb_input(pdb_input_type, pdb_code=None, output_dir=INPUT_DIR):\n",
        "    \"\"\"\n",
        "    Handle PDB input by either uploading a file or downloading it using a PDB code.\n",
        "    \"\"\"\n",
        "    if pdb_input_type == \"upload\":\n",
        "        if is_colab():\n",
        "            from google.colab import files\n",
        "\n",
        "            uploaded = files.upload()\n",
        "            pdb_filename = next(iter(uploaded))\n",
        "            pdb_path = os.path.join(output_dir, pdb_filename)\n",
        "            with open(pdb_path, \"wb\") as file:\n",
        "                file.write(uploaded[pdb_filename])\n",
        "            return pdb_path\n",
        "        else:\n",
        "            #raise EnvironmentError(\"File upload is only supported on Google Colab.\")\n",
        "            pdb_path='/workspaces/all_atom_binder_diffusion/heme_binder_diffusion/input/7o2g_HBA.pdb'\n",
        "            return pdb_path\n",
        "    elif pdb_input_type == \"pdb_code\":\n",
        "        return download_pdb(pdb_code, output_dir)\n",
        "    elif pdb_input_type == \"colab_test\":\n",
        "        pdb_path='/content/heme_binder_diffusion/input/7o2g_HBA.pdb'\n",
        "    else:\n",
        "        raise ValueError(\"Invalid PDB input type\")\n",
        "\n",
        "\n",
        "def run_rfdiffusion_all_atom(\n",
        "    config,\n",
        "    output_subfolder=\"ligand_protein_motif\",\n",
        "    output_prefix=\"sample\",\n",
        "    show_last_n_lines=5,\n",
        "    save_stdout=True,\n",
        "):\n",
        "    \"\"\"\n",
        "    Wrapper function to run rfdiffusion all atom with specified options, using a YAML configuration file.\n",
        "    The configuration is passed as a dictionary.\n",
        "    \"\"\"\n",
        "    # Generate a base output directory name without duplicating parts of the path\n",
        "    base_output_path = os.path.join(OUTPUT_DIR, output_subfolder)\n",
        "    print(f\"Output base directory: {base_output_path}\")\n",
        "\n",
        "    # Initialize counter to generate a unique output directory\n",
        "    counter = 0\n",
        "    unique_output_path = f\"{base_output_path}/{output_prefix}_{counter}\"\n",
        "    while os.path.exists(unique_output_path):\n",
        "        counter += 1\n",
        "        unique_output_path = f\"{base_output_path}/{output_prefix}_{counter}\"\n",
        "\n",
        "    final_output_path = unique_output_path\n",
        "    print(f\"Final output path: {final_output_path}\")\n",
        "    # Ensure the final output directory exists\n",
        "    os.makedirs(final_output_path, exist_ok=True)\n",
        "\n",
        "    # Update the output_prefix in the config with the actual output path\n",
        "    config[\"inference\"][\"output_prefix\"] = os.path.join(\n",
        "        final_output_path, output_prefix\n",
        "    )\n",
        "\n",
        "    # Write the configuration to a YAML file inside the correct output directory\n",
        "    config_filename = \"config.yaml\"  # Configuration file name\n",
        "    config_file_path = os.path.join(\n",
        "        final_output_path, config_filename\n",
        "    )  # Full path to the configuration file\n",
        "    with open(config_file_path, \"w\") as file:\n",
        "        yaml.dump(config, file)\n",
        "\n",
        "    # Correct the command to run the inference script with the YAML config file\n",
        "    cmd = [\n",
        "        \"python\",\n",
        "        \"./rf_diffusion_all_atom/run_inference.py\",\n",
        "        f\"--config-name={config_filename[:-5]}\",  # Remove the '.yaml' extension\n",
        "        f\"--config-dir={final_output_path}\",\n",
        "        f\"diffuser.T={config['diffuser']['T']}\",  # I do not know why this need to be added again if its already in the config.yaml\n",
        "    ]\n",
        "\n",
        "    # Print the command to the console\n",
        "    print(f\"Running command: {' '.join(cmd)}\")\n",
        "\n",
        "    # Initialize a list to keep track of the output lines\n",
        "    output_lines = []\n",
        "\n",
        "    # Use subprocess.Popen to run the command and capture stdout in real-time\n",
        "    process = subprocess.Popen(\n",
        "        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True\n",
        "    )\n",
        "\n",
        "    # Correct the path for the log file to ensure it's saved in the final_output_path\n",
        "    if save_stdout:\n",
        "        log_file_path = os.path.join(\n",
        "            final_output_path, \"run_inference.log\"\n",
        "        )  # Corrected path\n",
        "        log_file = open(log_file_path, \"w\")\n",
        "\n",
        "    # Periodically check for new output\n",
        "    while True:\n",
        "        output = process.stdout.readline()\n",
        "        if output == \"\" and process.poll() is not None:\n",
        "            break\n",
        "        if output:\n",
        "            output_lines.append(output.strip())\n",
        "            # Save to log file if required\n",
        "            if save_stdout:\n",
        "                log_file.write(output)\n",
        "\n",
        "            # Display the last N lines if required\n",
        "            if show_last_n_lines > 0:\n",
        "                display_lines = output_lines[-show_last_n_lines:]\n",
        "                print(\"\\n\".join(display_lines))\n",
        "\n",
        "        # time.sleep(1)  # Adjust the sleep time as needed\n",
        "\n",
        "    # Ensure the process has finished and close the log file if it was opened\n",
        "    process.poll()\n",
        "    if save_stdout:\n",
        "        log_file.close()\n",
        "\n",
        "    return final_output_path\n",
        "\n",
        "\n",
        "# @title ### Small molecule binder design with protein motif { display-mode: \"form\" }\n",
        "\n",
        "# Interface for specifying PDB input\n",
        "pdb_input_type = \"upload\"  # @param [\"upload\", \"pdb_code\", \"manual_path\", \"colab_test\"]\n",
        "pdb_code = \"7o2g\"  # @param {type:\"string\"}\n",
        "\n",
        "if pdb_input_type == \"pdb_code\":\n",
        "    input_pdb = handle_pdb_input(pdb_input_type, pdb_code)\n",
        "elif pdb_input_type == \"upload\":\n",
        "    print(\"Please upload your PDB file:\")\n",
        "    # Assuming running in Colab\n",
        "    input_pdb = handle_pdb_input(pdb_input_type)\n",
        "\n",
        "elif pdb_input_type == \"colab_test\":\n",
        "    input_pdb = handle_pdb_input(pdb_input_type)\n",
        "elif pdb_input_type == \"manual_path\":\n",
        "    input_pdb = pdb_code  # Assuming the user has manually specified the path\n",
        "\n",
        "contigs = \"30-110,A15-15,30-110\"  # @param {type:\"string\"}\n",
        "contig_length = \"100-140\"  # @param {type:\"string\"}\n",
        "ligand = \"HBA\"  # @param {type:\"string\"}\n",
        "num_designs = 1  # @param {type:\"integer\"}\n",
        "design_startnum = 0  # @param {type:\"integer\"}\n",
        "output_prefix = \"sample\"  # @param {type:\"string\"}\n",
        "output_subfolder = \"ligand_protein_motif\"  # @param {type:\"string\"}\n",
        "\n",
        "deterministic = True  # @param {type:\"boolean\"}\n",
        "T = 25  # @param {type:\"integer\"}\n",
        "\n",
        "# Split contigs string into list\n",
        "contigs_list = contigs.split(\",\")\n",
        "\n",
        "# Convert contigs list to string format for YAML\n",
        "contigs_yaml = [f\"{contig}\" for contig in contigs_list]\n",
        "contig_length = contig_length if contig_length else None\n",
        "\n",
        "config = {\n",
        "    \"inference\": {\n",
        "        \"deterministic\": deterministic,\n",
        "        \"input_pdb\": input_pdb,\n",
        "        \"ligand\": ligand,\n",
        "        \"num_designs\": num_designs,\n",
        "        \"design_startnum\": design_startnum,\n",
        "        \"ckpt_path\": \"./params/RFDiffusionAA_paper_weights.pt\",\n",
        "        \"model_runner\": \"NRBStyleSelfCond\",\n",
        "    },\n",
        "    \"diffuser\": {\"T\": T},\n",
        "    \"contigmap\": {\n",
        "        \"contigs\": contigs_yaml,\n",
        "        \"length\": contig_length,\n",
        "        #\"inpaint_str\": \"null\",\n",
        "    },\n",
        "    \"model\": {\"freeze_track_motif\": \"True\"},\n",
        "    \"defaults\": [\"aa\"],\n",
        "}\n",
        "\n",
        "guiding_potentials = True  # @param {type:\"boolean\"}\n",
        "# If guiding potentials is true:\n",
        "if guiding_potentials:\n",
        "    guide_scale = 2  # @param {type:\"integer\"}\n",
        "    guide_decay = \"cubic\"  # @param [\"cubic\", \"linear\", \"exponential\"]\n",
        "\n",
        "    config[\"potentials\"] = {\n",
        "        \"guiding_potentials\": [f\"type:ligand_ncontacts,weight:1\"],\n",
        "        \"guide_scale\": guide_scale,\n",
        "        \"guide_decay\": guide_decay,\n",
        "    }\n",
        "\n",
        "# Specify additional options for the run function\n",
        "show_last_n_lines = 1  # Show only the last 1 lines of stdout to avoid cluttering the notebook, above 1 is not working at the moment\n",
        "save_stdout = True  # Save the stdout as a logfile in the output folder\n",
        "\n",
        "# Call the run_rfdiffusion_all_atom function with the specified configuration and options\n",
        "final_output_path = run_rfdiffusion_all_atom(\n",
        "    config,\n",
        "    output_subfolder=output_subfolder,\n",
        "    output_prefix=output_prefix,\n",
        "    show_last_n_lines=show_last_n_lines,\n",
        "    save_stdout=save_stdout,\n",
        ")\n",
        "\n",
        "# @markdown After running the diffusion function, you can zip the last job's output for download:\n",
        "\n",
        "# @markdown Run the following cell to zip and download the last job's output.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WcQD7Mk3en3i"
      },
      "outputs": [],
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "animate = \"interactive\" #@param [\"none\", \"movie\", \"interactive\"]\n",
        "color = \"chain\" #@param [\"rainbow\", \"chain\", \"plddt\"]\n",
        "denoise = True\n",
        "dpi = 100 #@param [\"100\", \"200\", \"400\"] {type:\"raw\"}\n",
        "\n",
        "from colabdesign.shared.plot import pymol_color_list\n",
        "from colabdesign.rf.utils import get_ca, get_Ls, make_animation\n",
        "from string import ascii_uppercase, ascii_lowercase\n",
        "import os\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display, HTML\n",
        "import py3Dmol\n",
        "\n",
        "alphabet_list = list(ascii_uppercase + ascii_lowercase)\n",
        "\n",
        "\n",
        "# Construct the base output directory\n",
        "base_output_dir = os.path.join(OUTPUT_DIR, output_subfolder)\n",
        "\n",
        "def find_latest_output_dir(base_dir, prefix):\n",
        "    \"\"\"Find the latest output directory based on the prefix.\"\"\"\n",
        "    dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d)) and d.startswith(prefix)]\n",
        "    if not dirs:\n",
        "        raise FileNotFoundError(f\"No output directories found with prefix '{prefix}' in '{base_dir}'\")\n",
        "    latest_dir = sorted(dirs, key=lambda x: int(x.split('_')[-1]))[-1]\n",
        "    return os.path.join(base_dir, latest_dir)\n",
        "\n",
        "def plot_pdb(num=0):\n",
        "    # Find the latest output directory\n",
        "    latest_output_dir = find_latest_output_dir(base_output_dir, output_prefix)\n",
        "\n",
        "    # Construct the path to the PDB file\n",
        "    pdb_path = os.path.join(latest_output_dir, f\"{output_prefix}_{num}.pdb\")\n",
        "\n",
        "    # Load the PDB file\n",
        "    pdb_str = open(pdb_path, 'r').read()\n",
        "\n",
        "    # Initialize the 3Dmol.js viewer\n",
        "    view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "    view.addModel(pdb_str, 'pdb')\n",
        "\n",
        "    # Apply color scheme\n",
        "    if color == \"rainbow\":\n",
        "        view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "    elif color == \"chain\":\n",
        "        # Example: Apply color by chain\n",
        "        for n, chain, c in zip(range(len(contigs)), alphabet_list, pymol_color_list):\n",
        "            view.setStyle({'chain': chain}, {'cartoon': {'color': c}})\n",
        "            # If chain == B the visualization should be atoms\n",
        "            if chain == \"B\":\n",
        "                view.setStyle({'chain': chain}, {'stick': {}})\n",
        "    else:\n",
        "        # Example: Apply a custom color scheme\n",
        "        view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':0.5,'max':0.9}}})\n",
        "\n",
        "    # Zoom to fit and display the viewer\n",
        "    view.zoomTo()\n",
        "    view.show()\n",
        "\n",
        "\n",
        "\n",
        "if num_designs > 1:\n",
        "  output = widgets.Output()\n",
        "  def on_change(change):\n",
        "    if change['name'] == 'value':\n",
        "      with output:\n",
        "        output.clear_output(wait=True)\n",
        "        plot_pdb(change['new'])\n",
        "  dropdown = widgets.Dropdown(\n",
        "      options=[(f'{k}',k) for k in range(num_designs)],\n",
        "      value=0, description='design:',\n",
        "  )\n",
        "  dropdown.observe(on_change)\n",
        "  display(widgets.VBox([dropdown, output]))\n",
        "  with output:\n",
        "    plot_pdb(dropdown.value)\n",
        "else:\n",
        "  plot_pdb()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9MJYrPbC9KwA"
      },
      "outputs": [],
      "source": [
        "#@markdown # ProteinMPNN Design Step\n",
        "#@markdown After generating designs with rfdiffusion all atom, the next step is to perform sequence design with proteinMPNN. This section sets up and runs proteinMPNN on the generated structures.\n",
        "\n",
        "import glob\n",
        "import os\n",
        "import subprocess\n",
        "\n",
        "# Assuming final_output_path is defined from the previous rfdiffusion step\n",
        "DIFFUSION_DIR = final_output_path\n",
        "\n",
        "#@markdown ### Working and Script Directories\n",
        "#@markdown We will use for google colab $HEME_PROTOCOL_DIR path defined in the first cell\n",
        "#BASE_DIR\n",
        "MPNN_DIR = os.path.join(BASE_DIR, \"LigandMPNN\")\n",
        "os.makedirs(MPNN_DIR, exist_ok=True)\n",
        "#os.chdir(MPNN_DIR)\n",
        "\n",
        "#@markdown ### Python Environment\n",
        "#@markdown Specify the Python interpreter to use. If you're using a specific environment, adjust the path accordingly.\n",
        "PYTHON_GENERAL = \"python\" #@param {type:\"string\"}\n",
        "PYTHON_PROTEINMPNN = \"python\" #@param {type:\"string\"}\n",
        "\n",
        "# Assuming final_output_path is defined from the previous rfdiffusion step\n",
        "DIFFUSION_DIR = final_output_path\n",
        "\n",
        "#@markdown ### Note on Structure Filtering\n",
        "#@markdown According to Indrek Kalvet, PhD (Institute for Protein Design, University of Washington, ikalvet@uw.edu), there are a number of metrics that could be used to filter and select the best structures from the diffusion step, typically involving PyRosetta. However, due to the complexity of installing and running PyRosetta in a Google Colab environment, we will proceed with designing all diffused structures without this filtering step in this notebook.\n",
        "\n",
        "# Find all diffused backbone structures from the diffusion step\n",
        "# Note: Skipping the filtered_structures step and using all diffused structures\n",
        "diffused_backbones_good = glob.glob(f\"{DIFFUSION_DIR}/*.pdb\")\n",
        "assert len(diffused_backbones_good) > 0, \"No diffused backbones found!\"\n",
        "\n",
        "#@markdown ## Masked Positions from TRB Files\n",
        "#@markdown Extract fixed motif residues from TRB files. These residues will not be redesigned with proteinMPNN.\n",
        "mask_json_cmd = f\"{PYTHON_GENERAL} {HEME_PROTOCOL_DIR}/scripts/design/make_maskdict_from_trb.py --out {DIFFUSION_DIR}/masked_pos.jsonl --trb\"\n",
        "for d in diffused_backbones_good:\n",
        "    mask_json_cmd += \" \" + d.replace(\".pdb\", \".trb\")\n",
        "\n",
        "print(mask_json_cmd)\n",
        "subprocess.run(mask_json_cmd, shell=True, check=True)\n",
        "\n",
        "#assert os.path.exists(f\"{BASE_DIR}/all_atom_binder_diffusion/output/ligand_protein_motif/sample_0/masked_pos.jsonl\"), \"Failed to create masked positions JSONL file\"\n",
        "\n",
        "#@markdown ## ProteinMPNN Settings\n",
        "#@markdown Configure the settings for running proteinMPNN designs.\n",
        "MPNN_temperatures = [0.1]\n",
        "MPNN_outputs_per_temperature = 5 #@param {type:\"integer\"}\n",
        "MPNN_omit_AAs = \"CM\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Running proteinMPNN designs with the specified settings.\n",
        "for T in MPNN_temperatures:\n",
        "    for f in diffused_backbones_good:\n",
        "        cmd_mpnn = (f\"{PYTHON_PROTEINMPNN} {MPNN_DIR}/run.py \"\n",
        "                    f\"--model_type protein_mpnn --ligand_mpnn_use_atom_context 0 \"\n",
        "                    f\"--fixed_residues_multi {DIFFUSION_DIR}/masked_pos.jsonl --out_folder {DIFFUSION_DIR} \"\n",
        "                    f\"--number_of_batches {MPNN_outputs_per_temperature} --temperature {T} \"\n",
        "                    f\"--omit_AA {MPNN_omit_AAs} --pdb_path {f} \"\n",
        "                    f\"--checkpoint_protein_mpnn {BASE_DIR}/params/proteinmpnn_v_48_020.pt\")\n",
        "        print(cmd_mpnn)\n",
        "        subprocess.run(cmd_mpnn, shell=True, check=True)\n",
        "\n",
        "print(\"ProteinMPNN design steps completed.\")\n",
        "\n",
        "# Load the designed sequences and print the pandas dataframe\n",
        "import pandas as pd\n",
        "import json\n",
        "\n",
        "# Parse the fasta file and load the header as multiple columns divided by the \",\" and the sequence as the sequence column\n",
        "seqs = []\n",
        "with open(f\"{DIFFUSION_DIR}/seqs/sample_0.fa\") as f:\n",
        "    for line in f:\n",
        "        if line.startswith(\">\"):\n",
        "            header = line.strip()\n",
        "            # Remove the leading '>'\n",
        "            header = header[1:]\n",
        "            # Which will result in multiple columns in the dataframe\n",
        "        else:\n",
        "            seqs.append({\"header\": header, \"sequence\": line.strip()})\n",
        "\n",
        "df = pd.DataFrame(seqs)\n",
        "df = df.iloc[1:]\n",
        "df = df['header'].str.split(',', expand=True).join(df['sequence'])\n",
        "df.columns = ['experiment_id', 'id_number', 'temperature', 'seed', 'overall_confidence', 'ligand_confidence', 'seq_rec', 'sequence']\n",
        "df['id_number'] = df['id_number'].str.replace('id=', '')\n",
        "df['temperature'] = df['temperature'].str.replace('T=', '')\n",
        "df['seed'] = df['seed'].str.replace('seed=', '')\n",
        "df['overall_confidence'] = df['overall_confidence'].str.replace('overall_confidence=', '')\n",
        "df['ligand_confidence'] = df['ligand_confidence'].str.replace('ligand_confidence=', '')\n",
        "df['seq_rec'] = df['seq_rec'].str.replace('seq_rec=', '')\n",
        "# Print the length of the sequences\n",
        "df['sequence_length'] = df['sequence'].apply(len)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vHvVOhaxzyDU"
      },
      "source": [
        "## Running AF2 on designed sequences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DmkKNd6oz41Q"
      },
      "outputs": [],
      "source": [
        "#@markdown # ColabDesign Setup\n",
        "\n",
        "# Standard library imports\n",
        "import argparse\n",
        "import gc\n",
        "import os\n",
        "import re\n",
        "import subprocess\n",
        "import sys\n",
        "import tempfile\n",
        "import time\n",
        "\n",
        "# Third-party imports\n",
        "from IPython.display import HTML\n",
        "from IPython.core import ultratb\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import plotly.graph_objects as go\n",
        "import plotly.io as pio\n",
        "import yaml\n",
        "\n",
        "# Local application/library specific imports\n",
        "from colabdesign import mk_af_model, clear_mem\n",
        "from colabdesign.af.contrib import predict\n",
        "from colabdesign.af.contrib.cyclic import add_cyclic_offset\n",
        "from colabdesign.shared.plot import plot_pseudo_3D, pymol_cmap\n",
        "from colabdesign.shared.protein import _np_rmsd, _np_kabsch\n",
        "from typing import List, Optional, Tuple\n",
        "\n",
        "# Exception hook setup\n",
        "sys.excepthook = ultratb.FormattedTB(color_scheme=\"Linux\", call_pdb=False)\n",
        "\n",
        "#pio.renderers.default = \"vscode\"\n",
        "\n",
        "class SetupAlphaFoldColabDesign:\n",
        "    def __init__(\n",
        "        self,\n",
        "        unified_memory,\n",
        "        parentPath,\n",
        "        setupPath,\n",
        "        python_colab=\"/usr/bin/python3.10\",\n",
        "        colabdesign_path=\"/usr/local/lib/python3.10/dist-packages/colabdesign\",\n",
        "    ):\n",
        "        self.unified_memory = unified_memory\n",
        "        self.parentPath = parentPath\n",
        "        self.setupPath = setupPath\n",
        "        self.python_colab = python_colab\n",
        "        self.colabdesign_path = colabdesign_path\n",
        "        self.ENV = (\n",
        "            {\"TF_FORCE_UNIFIED_MEMORY\": \"1\", \"XLA_PYTHON_CLIENT_MEM_FRACTION\": \"4.0\"}\n",
        "            if unified_memory\n",
        "            else {}\n",
        "        )\n",
        "\n",
        "    def setup(self):\n",
        "        for k, v in self.ENV.items():\n",
        "            os.environ[k] = v\n",
        "\n",
        "        os.makedirs(self.setupPath, exist_ok=True)\n",
        "\n",
        "        if os.path.isdir(os.path.join(self.setupPath, \"params\")):\n",
        "            print(\"Setup path is present.\")\n",
        "        else:\n",
        "            print(\"Setup path is not present. Installing ColabDesign...\")\n",
        "            self.install_colab_design()\n",
        "\n",
        "        if os.path.isdir(os.path.join(self.setupPath, \"hhsuite\")):\n",
        "            print(\"HHsuite is present.\")\n",
        "        else:\n",
        "            print(\"HHsuite is not present. Installing HHsuite...\")\n",
        "            self.install_hhsuite()\n",
        "\n",
        "        print(\"mmseqs2 is imported.\")\n",
        "\n",
        "    def install_colab_design(self):\n",
        "        def run_command(command):\n",
        "            process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)\n",
        "            output, error = process.communicate()\n",
        "\n",
        "            if error:\n",
        "                raise Exception(f\"Error occurred while executing command: {error}\")\n",
        "\n",
        "            return output\n",
        "\n",
        "        params_path = os.path.join(self.setupPath, \"params\")\n",
        "        # Make params_path absolute\n",
        "        params_path = os.path.abspath(params_path)\n",
        "\n",
        "        if not os.path.exists(params_path):\n",
        "            os.makedirs(params_path)\n",
        "            print(f\"Created params directory at {params_path}\")\n",
        "\n",
        "        print(\"Installing ColabDesign...\")\n",
        "\n",
        "        commands = [\n",
        "            f\"apt-get install aria2 -qq\",\n",
        "            f\"cd {self.setupPath} && aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\",\n",
        "            f\"tar -xf {self.setupPath}/alphafold_params_2022-12-06.tar -C {params_path}\",\n",
        "            f\"touch {os.path.join(params_path, 'done.txt')}\",\n",
        "            # f\"cd {self.setupPath} && rm alphafold_params_2022-12-06.tar\",\n",
        "        ]\n",
        "\n",
        "        for command in commands:\n",
        "            run_command(command)\n",
        "\n",
        "        print(\"Installing Python dependencies...\")\n",
        "\n",
        "        run_command(\n",
        "            f\"{self.python_colab} -m pip -q install git+https://github.com/sokrypton/ColabDesign.git@gamma\"\n",
        "        )\n",
        "        run_command(\n",
        "            f\"ln -s {self.colabdesign_path} {os.path.join(self.setupPath, 'colabdesign')}\"\n",
        "        )\n",
        "        run_command(\n",
        "            f\"wget https://raw.githubusercontent.com/sokrypton/ColabFold/main/colabfold/colabfold.py -O {os.path.join(self.setupPath, 'colabfold_utils.py')}\"\n",
        "        )\n",
        "\n",
        "    def install_hhsuite(self):\n",
        "        os.makedirs(os.path.join(self.setupPath, \"hhsuite\"), exist_ok=True)\n",
        "        os.system(\n",
        "            f\"curl -fsSL https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz | tar xz -C {os.path.join(self.setupPath, 'hhsuite')}\"\n",
        "        )\n",
        "\n",
        "        if \"hhsuite\" not in os.environ[\"PATH\"]:\n",
        "            os.environ[\n",
        "                \"PATH\"\n",
        "            ] += f\":{os.path.join(self.setupPath, 'hhsuite/bin')}: {os.path.join(self.setupPath, 'hhsuite/scripts')}\"\n",
        "\n",
        "class ColabDesignUtils:\n",
        "    def __init__(self, setupPath):\n",
        "        self.setupPath = setupPath\n",
        "\n",
        "    def run_hhalign(\n",
        "        self, query_sequence, target_sequence, query_a3m=None, target_a3m=None\n",
        "    ):\n",
        "        with tempfile.NamedTemporaryFile() as tmp_query, tempfile.NamedTemporaryFile() as tmp_target, tempfile.NamedTemporaryFile() as tmp_alignment:\n",
        "            if query_a3m is None:\n",
        "                tmp_query.write(f\">Q\\n{query_sequence}\\n\".encode())\n",
        "                tmp_query.flush()\n",
        "                query_a3m = tmp_query.name\n",
        "            if target_a3m is None:\n",
        "                tmp_target.write(f\">T\\n{target_sequence}\\n\".encode())\n",
        "                tmp_target.flush()\n",
        "                target_a3m = tmp_target.name\n",
        "            os.system(\n",
        "                f\"hhalign -hide_cons -i {query_a3m} -t {target_a3m} -o {tmp_alignment.name}\"\n",
        "            )\n",
        "            X, start_indices = predict.parse_hhalign_output(tmp_alignment.name)\n",
        "        return X, start_indices\n",
        "\n",
        "    def run_do_not_align(self, query_sequence, target_sequence, **arg):\n",
        "        return [query_sequence, target_sequence], [0, 0]\n",
        "\n",
        "    def run_hhfilter(self, input, output, id=90, qid=10):\n",
        "        if \"hhsuite\" not in os.environ[\"PATH\"]:\n",
        "            os.environ[\n",
        "                \"PATH\"\n",
        "            ] += f\":{os.path.join(self.setupPath, 'hhsuite/bin')}: {os.path.join(self.setupPath, 'hhsuite/scripts')}\"\n",
        "\n",
        "        os.system(f\"hhfilter -id {id} -qid {qid} -i {input} -o {output}\")\n",
        "\n",
        "    @jax.jit\n",
        "    def get_coevolution(self, X):\n",
        "        \"\"\"given one-hot encoded MSA, return contacts\"\"\"\n",
        "        Y = jax.nn.one_hot(X, 22)\n",
        "        N, L, A = Y.shape\n",
        "        Y_flat = Y.reshape(N, -1)\n",
        "        c = jnp.cov(Y_flat.T)\n",
        "        shrink = 4.5 / jnp.sqrt(N) * jnp.eye(c.shape[0])\n",
        "        ic = jnp.linalg.inv(c + shrink)\n",
        "        ic_diag = jnp.diag(ic)\n",
        "        pcc = ic / jnp.sqrt(ic_diag[:, None] * ic_diag[None, :])\n",
        "        raw = jnp.sqrt(jnp.square(pcc.reshape(L, A, L, A)[:, :20, :, :20]).sum((1, 3)))\n",
        "        i = jnp.arange(L)\n",
        "        raw = raw.at[i, i].set(0)\n",
        "        ap = raw.sum(0, keepdims=True) * raw.sum(1, keepdims=True) / raw.sum()\n",
        "        return (raw - ap).at[i, i].set(0)\n",
        "\n",
        "    def plot_3D(aux, Ls, file_name, show=False):\n",
        "        plt.figure(figsize=(10, 5))\n",
        "        xyz = aux[\"atom_positions\"][:, 1]\n",
        "        xyz = xyz @ _np_kabsch(xyz, xyz, return_v=True, use_jax=False)\n",
        "        ax = plt.subplot(1, 2, 1)\n",
        "        if len(Ls) > 1:\n",
        "            plt.title(\"chain\")\n",
        "            c = np.concatenate([[n] * L for n, L in enumerate(Ls)])\n",
        "            plot_pseudo_3D(xyz=xyz, c=c, cmap=pymol_cmap, cmin=0, cmax=39, Ls=Ls, ax=ax)\n",
        "        else:\n",
        "            plt.title(\"length\")\n",
        "            plot_pseudo_3D(xyz=xyz, Ls=Ls, ax=ax)\n",
        "        plt.axis(False)\n",
        "        ax = plt.subplot(1, 2, 2)\n",
        "        plt.title(\"plddt\")\n",
        "        plot_pseudo_3D(xyz=xyz, c=aux[\"plddt\"], cmin=0.5, cmax=0.9, Ls=Ls, ax=ax)\n",
        "        plt.axis(False)\n",
        "        plt.savefig(file_name, dpi=200, bbox_inches=\"tight\")\n",
        "        plt.show() if show else plt.close()\n",
        "\n",
        "class PrepInputs:\n",
        "    def __init__(\n",
        "        self,\n",
        "        sequence,\n",
        "        jobname,\n",
        "        copies,\n",
        "        msa_method,\n",
        "        custom_a3m_path,\n",
        "        pair_mode,\n",
        "        cov,\n",
        "        id,\n",
        "        qid,\n",
        "        do_not_filter,\n",
        "        template_mode,\n",
        "        pdb,\n",
        "        chain,\n",
        "        rm_template_seq,\n",
        "        propagate_to_copies,\n",
        "        do_not_align,\n",
        "        setupPath,\n",
        "        parentPath,\n",
        "        overwrite,\n",
        "    ):\n",
        "        self.sequence = sequence\n",
        "        self.jobname = jobname\n",
        "        self.copies = copies\n",
        "        self.msa_method = msa_method\n",
        "        self.custom_a3m_path = custom_a3m_path\n",
        "        self.pair_mode = pair_mode\n",
        "        self.cov = cov\n",
        "        self.id = id\n",
        "        self.qid = qid\n",
        "        self.do_not_filter = do_not_filter\n",
        "        self.template_mode = template_mode\n",
        "        self.pdb = pdb\n",
        "        self.chain = chain\n",
        "        self.rm_template_seq = rm_template_seq\n",
        "        self.propagate_to_copies = propagate_to_copies\n",
        "        self.do_not_align = do_not_align\n",
        "        self.rm_sidechain = rm_template_seq\n",
        "        self.rm_sequence = rm_template_seq\n",
        "        self.setupPath = setupPath\n",
        "        self.parentPath = parentPath\n",
        "        self.overwrite = overwrite\n",
        "\n",
        "    def filter_options(self):\n",
        "        self.sequence = self.sequence.upper()\n",
        "        self.sequence = re.sub(\"[^A-Z:/()]\", \"\", self.sequence.upper())\n",
        "        self.sequence = re.sub(\"\\(\", \":(\", self.sequence)\n",
        "        self.sequence = re.sub(\"\\)\", \"):\", self.sequence)\n",
        "        self.sequence = re.sub(\":+\", \":\", self.sequence)\n",
        "        self.sequence = re.sub(\"/+\", \"/\", self.sequence)\n",
        "        self.sequence = re.sub(\"^[:/]+\", \"\", self.sequence)\n",
        "        self.sequence = re.sub(\"[:/]+$\", \"\", self.sequence)\n",
        "        self.jobname = re.sub(r\"\\W+\", \"\", self.jobname)\n",
        "\n",
        "    def process_sequence(self):\n",
        "        sequences = self.sequence.split(\":\")\n",
        "        self.u_sequences = predict.get_unique_sequences(sequences)\n",
        "        self.u_cyclic = [x.startswith(\"(\") for x in self.u_sequences]\n",
        "        self.u_sub_lengths = [[len(y) for y in x.split(\"/\")] for x in self.u_sequences]\n",
        "        self.u_sequences = [\n",
        "            x.replace(\"(\", \"\").replace(\")\", \"\").replace(\"/\", \"\")\n",
        "            for x in self.u_sequences\n",
        "        ]\n",
        "        if len(sequences) > len(self.u_sequences):\n",
        "            print(\"WARNING: use copies to define homooligomers\")\n",
        "        self.u_lengths = [len(x) for x in self.u_sequences]\n",
        "        sub_seq = \"\".join(self.u_sequences)\n",
        "        seq = sub_seq * self.copies\n",
        "\n",
        "        self.jobname = f\"{self.jobname}_{predict.get_hash(seq)[:5]}\"\n",
        "\n",
        "        def check(folder):\n",
        "            return os.path.exists(f\"{self.parentPath}/{folder}\")\n",
        "\n",
        "        if check(self.jobname):\n",
        "            n = 0\n",
        "            while check(f\"{self.jobname}_{n}\"):\n",
        "                n += 1\n",
        "            # If the jobname already exists, print a warning, if overwrite is True do not change the jobname, if overwrite is False change the jobname\n",
        "            if self.overwrite:\n",
        "                print(\n",
        "                    f\"WARNING: {self.jobname} already exists. Using the same jobname. If you want to run the job with a different jobname, set overwrite to False. If you did not change other parameters your files will be overwritten.\"\n",
        "                )\n",
        "            else:\n",
        "                print(\n",
        "                    f\"WARNING: {self.jobname} already exists. Changing jobname to {self.jobname}_{n}\"\n",
        "                )\n",
        "                self.jobname = f\"{self.jobname}_{n}\"\n",
        "\n",
        "        print(\"jobname\", self.jobname)\n",
        "        print(f\"length={self.u_lengths} copies={self.copies}\")\n",
        "\n",
        "        self.input_opts = {\n",
        "            \"sequence\": self.u_sequences,\n",
        "            \"copies\": self.copies,\n",
        "            \"msa_method\": self.msa_method,\n",
        "            \"pair_mode\": self.pair_mode,\n",
        "            \"do_not_filter\": self.do_not_filter,\n",
        "            \"cov\": self.cov,\n",
        "            \"id\": self.id,\n",
        "            \"template_mode\": self.template_mode,\n",
        "            \"propagate_to_copies\": self.propagate_to_copies,\n",
        "        }\n",
        "\n",
        "    def get_msa(self):\n",
        "        def run_mmseqs2_wrapper(*args, **kwargs):\n",
        "            kwargs[\"user_agent\"] = \"colabdesign/gamma\"\n",
        "            return run_mmseqs2(*args, **kwargs)\n",
        "\n",
        "        os.makedirs(f\"{self.parentPath}/{self.jobname}\", exist_ok=True)\n",
        "\n",
        "        utils = ColabDesignUtils(self.setupPath)\n",
        "        # Import colabfold_utils.py from the setupPath\n",
        "        sys.path.append(self.setupPath)\n",
        "        from colabfold_utils import run_mmseqs2\n",
        "\n",
        "        # Create \"in\" folder path variable\n",
        "        input_path = f\"{self.parentPath}/{self.jobname}/in\"\n",
        "        os.makedirs(input_path, exist_ok=True)\n",
        "        self.Ls = [len(x) for x in self.u_sequences]\n",
        "        if self.msa_method == \"mmseqs2\":\n",
        "            # If the msa.a3m file is already present in the input\n",
        "            # folder then skip this step and print a message to the user\n",
        "            # Assuming 'input_folder' is the directory where you want to check the file\n",
        "            if os.path.isfile(os.path.join(input_path, \"msa.a3m\")):\n",
        "                print(\"msa.a3m file is already present, loading from file.\")\n",
        "                print(\n",
        "                    \"please check that this a3m file contains the sequencens you expect\"\n",
        "                )\n",
        "                print(\n",
        "                    \"this behaviour was implmented for HPC Slurm cluster usage and may not be the best for most users\"\n",
        "                )\n",
        "                self.msa, self.deletion_matrix = predict.parse_a3m(\n",
        "                    f\"{self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "                )\n",
        "                self.msa_path = f\"{self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "                self.job_path = f\"{self.parentPath}/{self.jobname}\"\n",
        "            else:\n",
        "                # The steps to be performed if the file is not present goes here\n",
        "                self.msa, self.deletion_matrix = predict.get_msa(\n",
        "                    self.u_sequences,\n",
        "                    input_path,  # Here\n",
        "                    mode=self.pair_mode,\n",
        "                    cov=self.cov,\n",
        "                    id=self.id,\n",
        "                    qid=self.qid,\n",
        "                    max_msa=4096,\n",
        "                    do_not_filter=self.do_not_filter,\n",
        "                    mmseqs2_fn=run_mmseqs2_wrapper,\n",
        "                    hhfilter_fn=utils.run_hhfilter,\n",
        "                )\n",
        "                print(f\"{self.parentPath}/{self.jobname}/in/msa.a3m\")\n",
        "                self.msa_path = f\"{self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "                self.job_path = f\"{self.parentPath}/{self.jobname}\"\n",
        "        # Else if the MSA method is single sequence\n",
        "        elif self.msa_method == \"single_sequence\":\n",
        "            with open(f\"{self.parentPath}/{self.jobname}/in/msa.a3m\", \"w\") as a3m:\n",
        "                a3m.write(f\">{self.jobname}\\n{self.sub_seq}\\n\")\n",
        "            self.msa, self.deletion_matrix = predict.parse_a3m(\n",
        "                f\"{self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "            )\n",
        "        # If the MSA method is custom_X\n",
        "        else:\n",
        "            msa_format = self.msa_method.split(\"_\")[1]\n",
        "            print(f\"MSA mode: {self.msa_method}\")\n",
        "            # TODO : Add support for Google Colab!\n",
        "            google_colab = False\n",
        "            local_run = True\n",
        "            # If google colab is used, the file is stored in the google colab environment\n",
        "            if google_colab:\n",
        "                msa_format = self.msa_method.split(\"_\")[1]\n",
        "                print(f\"upload {self.msa_method}\")\n",
        "                from google.colab import files\n",
        "\n",
        "                msa_dict = files.upload()\n",
        "                lines = []\n",
        "                for k, v in msa_dict.items():\n",
        "                    lines += v.decode().splitlines()\n",
        "            # Here we handle the custom MSA case\n",
        "            if local_run:\n",
        "                print(f\"Reading MSA from {self.custom_a3m_path}\")\n",
        "                # Check if the path is valid, in other words, if the file exists\n",
        "                if not os.path.isfile(self.custom_a3m_path):\n",
        "                    raise ValueError(\n",
        "                        f\"Invalid path: {self.custom_a3m_path}. The file does not exist.\"\n",
        "                    )\n",
        "                with open(self.custom_a3m_path, \"r\") as file:\n",
        "                    lines = file.read().splitlines()\n",
        "            input_lines = []\n",
        "            for line in lines:\n",
        "                line = line.replace(\"\\x00\", \"\")\n",
        "                if len(line) > 0 and not line.startswith(\"#\"):\n",
        "                    input_lines.append(line)\n",
        "            # The following is to avoid errors when running parallel jobs in the cluster\n",
        "            # Only write the file if it does not exist\n",
        "            if not os.path.isfile(\n",
        "                f\"{self.parentPath}/{self.jobname}/in/msa.{msa_format}\"\n",
        "            ):\n",
        "                with open(\n",
        "                    f\"{self.parentPath}/{self.jobname}/in/msa.{msa_format}\", \"w\"\n",
        "                ) as msa:\n",
        "                    msa.write(\"\\n\".join(input_lines))\n",
        "            if msa_format != \"a3m\":\n",
        "                os.system(\n",
        "                    f\"perl hhsuite/scripts/reformat.pl {msa_format} a3m {self.parentPath}/{self.jobname}/in/msa.{msa_format} {self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "                )\n",
        "            # If the user prefers to skip filtering\n",
        "            if self.do_not_filter:\n",
        "                # Print that we are not filtering the MSA and the relevant parameters\n",
        "                print(f\"WARNING: not filtering MSA. Using 0 cov, 0 qid and 100 id\")\n",
        "                print(f\"{self.parentPath}/{self.jobname}/in/msa.a3m\")\n",
        "                if \"hhsuite\" not in os.environ[\"PATH\"]:\n",
        "                    os.environ[\n",
        "                        \"PATH\"\n",
        "                    ] += f\":{os.path.join(self.setupPath, 'hhsuite/bin')}: {os.path.join(self.setupPath, 'hhsuite/scripts')}\"\n",
        "                # Only run hhfilter if the MSA is not already present in the in folder\n",
        "                # to avoid unnecesarilly running hhfilter when running multiple jobs in the cluster\n",
        "                if not os.path.isfile(\n",
        "                    f\"{self.parentPath}/{self.jobname}/in/msa.filt.a3m\"\n",
        "                ):\n",
        "                    os.system(\n",
        "                        f\"hhfilter -qid 0 -id 100 -cov 0 -i {self.parentPath}/{self.jobname}/in/msa.a3m -o {self.parentPath}/{self.jobname}/in/msa.filt.a3m\"\n",
        "                    )\n",
        "            # Else we proceed to filter the MSA with default HHFilter\n",
        "            else:\n",
        "                if \"hhsuite\" not in os.environ[\"PATH\"]:\n",
        "                    os.environ[\n",
        "                        \"PATH\"\n",
        "                    ] += f\":{os.path.join(self.setupPath, 'hhsuite/bin')}: {os.path.join(self.setupPath, 'hhsuite/scripts')}\"\n",
        "                # Only run hhfilter if the MSA is not already present in the in folder\n",
        "                # to avoid unnecesarilly running hhfilter when running multiple jobs in the cluster\n",
        "                if not os.path.isfile(\n",
        "                    f\"{self.parentPath}/{self.jobname}/in/msa.filt.a3m\"\n",
        "                ):\n",
        "                    os.system(\n",
        "                        f\"hhfilter -qid {self.qid} -id {self.id} -cov {self.cov} -i {self.parentPath}/{self.jobname}/in/msa.a3m -o {self.parentPath}/{self.jobname}/in/msa.filt.a3m\"\n",
        "                    )\n",
        "            self.msa, self.deletion_matrix = predict.parse_a3m(\n",
        "                f\"{self.parentPath}/{self.jobname}/in/msa.filt.a3m\"\n",
        "            )\n",
        "\n",
        "        if len(self.msa) > 1:\n",
        "            predict.plot_msa(self.msa, self.Ls)\n",
        "            plt.savefig(\n",
        "                f\"{self.parentPath}/{self.jobname}/in/msa_feats.png\",\n",
        "                dpi=200,\n",
        "                bbox_inches=\"tight\",\n",
        "            )\n",
        "            # If this is being run in a juptyer notebook, show the image\n",
        "            if \"ipykernel\" in sys.modules:\n",
        "                plt.show()\n",
        "\n",
        "    def use_templates(self):\n",
        "        self.use_templates = self.template_mode in [\n",
        "            \"mmseqs2\",\n",
        "            \"custom\",\n",
        "        ]  # Here we define the template mode which is either mmseqs2 or custom\n",
        "        if self.use_templates:\n",
        "            print(\"aligning template\")\n",
        "            template_msa = f\"{self.parentPath}/{self.jobname}/in/msa.a3m\"\n",
        "            if self.template_mode == \"mmseqs2\":\n",
        "                predict.get_msa(\n",
        "                    self.u_sequences,\n",
        "                    self.jobname,\n",
        "                    mode=\"unpaired\",\n",
        "                    mmseqs2_fn=lambda *x: run_mmseqs2(\n",
        "                        *x, user_agent=\"colabdesign/gamma\"\n",
        "                    ),\n",
        "                    do_not_filter=True,\n",
        "                    do_not_return=True,\n",
        "                    output_a3m=f\"{self.parentPath}/{self.jobname}/in/msa_tmp.a3m\",\n",
        "                )\n",
        "                template_msa = f\"{self.parentPath}/{self.jobname}/in/msa_tmp.a3m\"\n",
        "                if not self.propagate_to_copies and self.copies > 1:\n",
        "                    new_msa = []\n",
        "                    with open(template_msa, \"r\") as handle:\n",
        "                        for line in handle:\n",
        "                            if not line.startswith(\">\"):\n",
        "                                new_msa.append(line.rstrip())\n",
        "                    with open(template_msa, \"w\") as handle:\n",
        "                        for n, seq in enumerate(new_msa):\n",
        "                            handle.write(f\">{n}\\n{seq*self.copies}\\n\")\n",
        "\n",
        "                templates = {}\n",
        "                print(\"ID\\tpdb\\tcid\\tevalue\")\n",
        "                for line in open(\n",
        "                    f\"{self.parentPath}/{self.jobname}/in/msa/_env/pdb70.m8\", \"r\"\n",
        "                ):\n",
        "                    p = line.rstrip().split()\n",
        "                    M, target_id, qid, e_value = p[0], p[1], p[2], p[10]\n",
        "                    M = int(M)\n",
        "                    if M not in templates:\n",
        "                        templates[M] = []\n",
        "                    if len(templates[M]) < 4:\n",
        "                        print(f\"{int(M)}\\t{target_id}\\t{qid}\\t{e_value}\")\n",
        "                        templates[M].append(target_id)\n",
        "                if len(templates) == 0:\n",
        "                    use_templates = False\n",
        "                    print(\"ERROR: no templates found...\")\n",
        "                else:\n",
        "                    Ms = sorted(list(templates.keys()))\n",
        "                    pdbs, chains = [], []\n",
        "                    for M in Ms:\n",
        "                        for n, target_id in enumerate(templates[M]):\n",
        "                            pdb_id, chain_id = target_id.split(\"_\")\n",
        "                            if len(pdbs) < n + 1:\n",
        "                                pdbs.append([])\n",
        "                                chains.append([])\n",
        "                            pdbs[n].append(pdb_id)\n",
        "                            chains[n].append(chain_id)\n",
        "                    print(pdbs)\n",
        "            else:\n",
        "                pdbs, chains = [self.pdb], [self.chain]\n",
        "\n",
        "        if self.use_templates:\n",
        "            self.input_opts.update({\"pdbs\": pdbs, \"chains\": chains})\n",
        "            self.batches = []\n",
        "            for pdb, chain in zip(pdbs, chains):\n",
        "                query_seq = \"\".join(self.u_sequences)\n",
        "                batch = predict.get_template_feats(\n",
        "                    pdb,\n",
        "                    chain,\n",
        "                    query_seq=query_seq,\n",
        "                    query_a3m=template_msa,\n",
        "                    copies=self.copies,\n",
        "                    propagate_to_copies=self.propagate_to_copies,\n",
        "                    use_seq=not self.rm_sequence,\n",
        "                    get_pdb_fn=self.get_pdb,\n",
        "                    align_fn=(\n",
        "                        self.run_do_not_align if self.do_not_align else self.run_hhalign\n",
        "                    ),\n",
        "                )\n",
        "                self.batches.append(batch)\n",
        "\n",
        "            plt.figure(figsize=(3 * len(self.batches), 3))\n",
        "            for n, batch in enumerate(self.batches):\n",
        "                plt.subplot(1, len(self.batches), n + 1)\n",
        "                plt.title(f\"template features {n+1}\")\n",
        "                dgram = batch[\"dgram\"].argmax(-1).astype(float)\n",
        "                dgram[batch[\"dgram\"].sum(-1) == 0] = np.nan\n",
        "                Ln = dgram.shape[0]\n",
        "                plt.imshow(dgram, extent=(0, Ln, Ln, 0))\n",
        "                predict.plot_ticks(self.Ls * self.copies)\n",
        "            plt.savefig(\n",
        "                f\"{self.parentPath}/{self.jobname}/in/template_feats.png\",\n",
        "                dpi=200,\n",
        "                bbox_inches=\"tight\",\n",
        "            )\n",
        "            plt.show()\n",
        "        else:\n",
        "            self.batches = [None]\n",
        "\n",
        "        print(\"GC\", gc.collect())\n",
        "\n",
        "class PrepModel:\n",
        "    def __init__(\n",
        "        self,\n",
        "        model_type,\n",
        "        rank_by,\n",
        "        debug,\n",
        "        use_initial_guess,\n",
        "        num_msa,\n",
        "        num_extra_msa,\n",
        "        use_cluster_profile,\n",
        "        u_lengths,\n",
        "        copies,\n",
        "        use_templates,\n",
        "        batches,\n",
        "        msa,\n",
        "        deletion_matrix,\n",
        "        u_sub_lengths,\n",
        "        u_cyclic,\n",
        "        setupPath,\n",
        "    ):\n",
        "        self.model_type = model_type\n",
        "        self.rank_by = rank_by\n",
        "        self.debug = debug\n",
        "        self.use_initial_guess = use_initial_guess\n",
        "        self.num_msa = num_msa\n",
        "        self.num_extra_msa = num_extra_msa\n",
        "        self.use_cluster_profile = use_cluster_profile\n",
        "        self.u_lengths = u_lengths\n",
        "        self.copies = copies\n",
        "        self.use_templates = use_templates\n",
        "        self.batches = batches\n",
        "        self.msa = msa\n",
        "        self.deletion_matrix = deletion_matrix\n",
        "        self.u_sub_lengths = u_sub_lengths\n",
        "        self.u_cyclic = u_cyclic\n",
        "        self.setupPath = setupPath\n",
        "\n",
        "    def model_options(self):\n",
        "        if self.model_type == \"monomer (ptm)\":\n",
        "            use_multimer = False\n",
        "            pseudo_multimer = False\n",
        "        elif self.model_type == \"multimer (v3)\":\n",
        "            use_multimer = True\n",
        "            pseudo_multimer = False\n",
        "        elif self.model_type == \"pseudo_multimer (v3)\":\n",
        "            use_multimer = True\n",
        "            pseudo_multimer = True\n",
        "        elif len(self.u_lengths) > 1 or self.copies > 1:\n",
        "            use_multimer = True\n",
        "            pseudo_multimer = False\n",
        "        else:\n",
        "            use_multimer = False\n",
        "            pseudo_multimer = False\n",
        "\n",
        "        if self.rank_by == \"auto\":\n",
        "            self.rank_by = (\n",
        "                \"multi\" if (len(self.u_lengths) > 1 or self.copies > 1) else \"plddt\"\n",
        "            )\n",
        "\n",
        "        self.model_opts = {\n",
        "            \"num_msa\": self.num_msa,\n",
        "            \"num_extra_msa\": self.num_extra_msa,\n",
        "            \"num_templates\": len(self.batches),\n",
        "            \"use_cluster_profile\": self.use_cluster_profile,\n",
        "            \"use_multimer\": use_multimer,\n",
        "            \"pseudo_multimer\": pseudo_multimer,\n",
        "            \"use_templates\": self.use_templates,\n",
        "            \"use_batch_as_template\": False,\n",
        "            \"use_dgram\": True,\n",
        "            \"protocol\": \"hallucination\",\n",
        "            \"best_metric\": self.rank_by,\n",
        "            \"optimize_seq\": False,\n",
        "            \"debug\": self.debug,\n",
        "            \"clear_prev\": False,\n",
        "        }\n",
        "\n",
        "    def initialize_model(self):\n",
        "        if \"af\" in dir():\n",
        "            if self.model_opts != model_opts_:\n",
        "                if (\n",
        "                    self.model_opts[\"use_multimer\"] == self.af._args[\"use_multimer\"]\n",
        "                    and self.model_opts[\"use_templates\"]\n",
        "                    == self.af._args[\"use_templates\"]\n",
        "                ):\n",
        "                    old_params = dict(zip(self.af._model_names, self.af._model_params))\n",
        "                else:\n",
        "                    print(\"loading alphafold params\")\n",
        "                    old_params = {}\n",
        "                    clear_mem()\n",
        "                self.af = mk_af_model(\n",
        "                    old_params=old_params, use_mlm=True, **self.model_opts\n",
        "                )\n",
        "                model_opts_ = predict.copy_dict(self.model_opts)\n",
        "        else:\n",
        "            print(\"loading alphafold params 1\")\n",
        "            self.af = mk_af_model(\n",
        "                use_mlm=True, data_dir=f\"{self.setupPath}\", **self.model_opts\n",
        "            )\n",
        "            model_opts_ = predict.copy_dict(self.model_opts)\n",
        "\n",
        "    def prep_inputs(self):\n",
        "        self.af.prep_inputs(self.u_lengths, copies=self.copies, seed=0)\n",
        "        self.print_key = [\"plddt\", \"ptm\"]\n",
        "        if len(self.af._lengths) > 1:\n",
        "            self.print_key += [\"i_ptm\", \"multi\"]\n",
        "        self.af.set_opt(\"con\", cutoff=8.0)\n",
        "\n",
        "    def set_templates(self):\n",
        "        if self.use_templates:\n",
        "            self.af.set_opt(use_initial_guess=self.use_initial_guess)\n",
        "            for n, batch in enumerate(self.batches):\n",
        "                self.af.set_template(batch=batch, n=n)\n",
        "            self.af.set_opt(\n",
        "                \"template\",\n",
        "                rm_sc=self.rm_sidechain,\n",
        "                rm_seq=self.rm_sequence,\n",
        "                rm_ic=self.rm_interchain,\n",
        "            )\n",
        "\n",
        "    def set_msa(self):\n",
        "        self.af.set_msa(self.msa, self.deletion_matrix)\n",
        "\n",
        "    def set_chainbreaks(self):\n",
        "        L_prev = 0\n",
        "        for n, l in enumerate(self.u_sub_lengths * self.copies):\n",
        "            for L_i in l[:-1]:\n",
        "                self.af._inputs[\"residue_index\"][L_prev + L_i :] += 32\n",
        "                L_prev += L_i\n",
        "            L_prev += l[-1]\n",
        "\n",
        "    def set_cyclic_constraints(self):\n",
        "        i_cyclic = [n for n, c in enumerate(self.u_cyclic * self.copies) if c]\n",
        "        if len(i_cyclic) > 0:\n",
        "            add_cyclic_offset(self.af, i_cyclic)\n",
        "\n",
        "class RunAlphaFold:\n",
        "    def __init__(\n",
        "        self,\n",
        "        jobname,\n",
        "        model,\n",
        "        num_recycles,\n",
        "        recycle_early_stop_tolerance,\n",
        "        select_best_across_recycles,\n",
        "        use_mlm,\n",
        "        use_dropout,\n",
        "        seed,\n",
        "        num_seeds,\n",
        "        show_images,\n",
        "        use_initial_guess,\n",
        "        af,\n",
        "        copies,\n",
        "        print_key,\n",
        "        rank_by,\n",
        "        Ls,\n",
        "        parentPath,\n",
        "    ):\n",
        "        self.jobname = jobname\n",
        "        self.model = model\n",
        "        self.num_recycles = num_recycles\n",
        "        self.recycle_early_stop_tolerance = recycle_early_stop_tolerance\n",
        "        self.select_best_across_recycles = select_best_across_recycles\n",
        "        self.use_mlm = use_mlm\n",
        "        self.use_dropout = use_dropout\n",
        "        self.seed = seed\n",
        "        self.num_seeds = num_seeds\n",
        "        self.show_images = show_images\n",
        "        self.use_initial_guess = use_initial_guess\n",
        "        self.af = af\n",
        "        self.copies = copies\n",
        "        self.print_key = print_key\n",
        "        self.rank_by = rank_by\n",
        "        self.Ls = Ls\n",
        "        self.parentPath = parentPath\n",
        "\n",
        "    def run(self):\n",
        "        run_opts = {\n",
        "            \"seed\": self.seed,\n",
        "            \"use_mlm\": self.use_mlm,\n",
        "            \"use_dropout\": self.use_dropout,\n",
        "            \"num_recycles\": self.num_recycles,\n",
        "            \"model\": self.model,\n",
        "            \"use_initial_guess\": self.use_initial_guess,\n",
        "            \"select_best_across_recycles\": self.select_best_across_recycles,\n",
        "            \"recycle_early_stop_tolerance\": self.recycle_early_stop_tolerance,\n",
        "        }\n",
        "\n",
        "        # decide which models to use\n",
        "        if self.model == \"all\":\n",
        "            models = self.af._model_names\n",
        "        else:\n",
        "            models = [self.af._model_names[int(self.model) - 1]]\n",
        "\n",
        "        # set options\n",
        "        self.af.set_opt(\"mlm\", replace_fraction=0.15 if self.use_mlm else 0.0)\n",
        "\n",
        "        pdb_path = f\"{self.parentPath}/{self.jobname}/out\"\n",
        "        os.makedirs(pdb_path, exist_ok=True)\n",
        "        # Make  a figs, pdbs, and npz folder\n",
        "        os.makedirs(f\"{pdb_path}/figs\", exist_ok=True)\n",
        "        os.makedirs(f\"{pdb_path}/pdbs\", exist_ok=True)\n",
        "        os.makedirs(f\"{pdb_path}/npz\", exist_ok=True)\n",
        "\n",
        "        # keep track of results\n",
        "        info = []\n",
        "        self.af._tmp = {\n",
        "            \"traj\": {\"seq\": [], \"xyz\": [], \"plddt\": [], \"pae\": []},\n",
        "            \"log\": [],\n",
        "            \"best\": {},\n",
        "        }\n",
        "\n",
        "        # run\n",
        "        print(\"running prediction\")\n",
        "        with open(f\"{self.parentPath}/{self.jobname}/log.txt\", \"w\") as handle:\n",
        "            # go through all seeds\n",
        "            seeds = list(range(self.seed, self.seed + self.num_seeds))\n",
        "            for seed in seeds:\n",
        "                self.af.set_seed(seed)\n",
        "                # go through all models\n",
        "                for model in models:\n",
        "                    recycle = 0\n",
        "                    self.af._inputs.pop(\"prev\", None)\n",
        "                    stop_recycle = False\n",
        "                    prev_pos = None\n",
        "                    # go through all recycles\n",
        "                    while recycle < self.num_recycles + 1:\n",
        "                        print_str = (\n",
        "                            f\"seed={str(seed).zfill(3)} model={model} recycle={recycle}\"\n",
        "                        )\n",
        "                        self.af.predict(\n",
        "                            dropout=self.use_dropout, models=[model], verbose=False\n",
        "                        )\n",
        "\n",
        "                        # set previous inputs\n",
        "                        self.af._inputs[\"prev\"] = self.af.aux[\"prev\"]\n",
        "\n",
        "                        # save results\n",
        "                        if len(self.af._lengths) > 1:\n",
        "                            self.af.aux[\"log\"][\"multi\"] = (\n",
        "                                0.8 * self.af.aux[\"log\"][\"i_ptm\"]\n",
        "                                + 0.2 * self.af.aux[\"log\"][\"ptm\"]\n",
        "                            )\n",
        "\n",
        "                        \n",
        "                        self.af.save_current_pdb(\n",
        "                            f\"{pdb_path}/pdbs/{self.jobname}_{model}_r{recycle}_seed_{str(seed).zfill(3)}.pdb\"\n",
        "                        )\n",
        "\n",
        "                        # print metrics\n",
        "                        for k in self.print_key:\n",
        "                            print_str += f\" {k}={self.af.aux['log'][k]:.3f}\"\n",
        "\n",
        "                        # early stop check\n",
        "                        current_pos = self.af.aux[\"atom_positions\"][:, 1]\n",
        "                        if recycle > 0:\n",
        "                            rmsd_tol = _np_rmsd(prev_pos, current_pos, use_jax=False)\n",
        "                            if rmsd_tol < self.recycle_early_stop_tolerance:\n",
        "                                stop_recycle = True\n",
        "                            print_str += f\" rmsd_tol={rmsd_tol:.3f}\"\n",
        "                        prev_pos = current_pos\n",
        "                        # print metrics\n",
        "                        # print(print_str)\n",
        "                        handle.write(f\"{print_str}\\n\")\n",
        "\n",
        "                        tag = f\"{model}_r{recycle}_seed_{str(seed).zfill(3)}\"\n",
        "                        if self.select_best_across_recycles:\n",
        "                            info.append(\n",
        "                                [tag, print_str, self.af.aux[\"log\"][self.rank_by]]\n",
        "                            )\n",
        "                            self.af._save_results(\n",
        "                                save_best=True,\n",
        "                                best_metric=self.rank_by,\n",
        "                                metric_higher_better=True,\n",
        "                                verbose=False,\n",
        "                            )\n",
        "                            self.af._k += 1\n",
        "\n",
        "                        recycle += 1\n",
        "                        if stop_recycle:\n",
        "                            break\n",
        "\n",
        "                    # Check if the 'select_best_across_recycles' attribute is set to False\n",
        "                    if not self.select_best_across_recycles:\n",
        "                        # If it is False, append the tag, print string, and the log of the rank_by attribute to the info list\n",
        "                        info.append([tag, print_str, self.af.aux[\"log\"][self.rank_by]])\n",
        "\n",
        "                        # Call the '_save_results' method of the 'af' object to save the best results\n",
        "                        # 'save_best' is set to True to indicate that the best results should be saved\n",
        "                        # 'best_metric' is set to the 'rank_by' attribute to specify the metric to rank the results by\n",
        "                        # 'metric_higher_better' is set to True to indicate that higher values of the metric are better\n",
        "                        # 'verbose' is set to False to prevent the method from printing additional information\n",
        "                        self.af._save_results(\n",
        "                            save_best=True,\n",
        "                            best_metric=self.rank_by,\n",
        "                            metric_higher_better=True,\n",
        "                            verbose=False,\n",
        "                        )\n",
        "\n",
        "                        # Increment the '_k' attribute of the 'af' object by 1\n",
        "                        # after finishing all recycles for the current model\n",
        "                        self.af._k += 1\n",
        "\n",
        "                    # save current results for each model(after n recycles)\n",
        "                    ColabDesignUtils.plot_3D(\n",
        "                        aux=self.af.aux,\n",
        "                        Ls=self.Ls * self.copies,\n",
        "                        file_name=f\"{pdb_path}/figs/{self.jobname}_{model}_seed_{str(seed).zfill(3)}.pdf\",\n",
        "                        show=self.show_images,\n",
        "                    )\n",
        "                    predict.plot_confidence(\n",
        "                        self.af.aux[\"plddt\"] * 100,\n",
        "                        self.af.aux[\"pae\"],\n",
        "                        self.Ls * self.copies,\n",
        "                    )\n",
        "                    plt.savefig(\n",
        "                        f\"{pdb_path}/figs/{self.jobname}_{model}_seed_{str(seed).zfill(3)}.png\",\n",
        "                        dpi=200,\n",
        "                        bbox_inches=\"tight\",\n",
        "                    )\n",
        "                    plt.close()\n",
        "\n",
        "        # save best results\n",
        "        rank = np.argsort([x[2] for x in info])[::-1][:5]\n",
        "        print(f\"best_tag={info[rank[0]][0]} {info[rank[0]][1]}\")\n",
        "\n",
        "        aux_best = self.af._tmp[\"best\"][\"aux\"]\n",
        "        # Save the best pdb file\n",
        "        \n",
        "        self.af.save_pdb(\n",
        "            f\"{pdb_path}/pdbs/{self.jobname}_best_{info[rank[0]][0]}.pdb\"\n",
        "        )\n",
        "        np.savez_compressed(\n",
        "            f\"{pdb_path}/npz/{self.jobname}_best_{info[rank[0]][0]}.npz\",\n",
        "            plddt=aux_best[\"plddt\"].astype(np.float16),\n",
        "            pae=aux_best[\"pae\"].astype(np.float16),\n",
        "            tag=np.array(info[rank[0]][0]),\n",
        "            metrics=np.array(info[rank[0]][1]),\n",
        "        )\n",
        "        np.savez_compressed(\n",
        "            f\"{pdb_path}/npz/{self.jobname}_all_{info[rank[0]][0]}.npz\",\n",
        "            plddt=np.array(self.af._tmp[\"traj\"][\"plddt\"], dtype=np.float16),\n",
        "            pae=np.array(self.af._tmp[\"traj\"][\"pae\"], dtype=np.float16),\n",
        "            tag=np.array([x[0] for x in info]),\n",
        "            metrics=np.array([x[1] for x in info]),\n",
        "        )\n",
        "\n",
        "        # If this is being run in a juptyer notebook, show the image\n",
        "        if \"ipykernel\" in sys.modules:\n",
        "            ColabDesignUtils.plot_3D(\n",
        "                aux_best, self.Ls * self.copies, f\"{pdb_path}/figs/best.pdf\", show=False\n",
        "            )\n",
        "            predict.plot_confidence(\n",
        "                aux_best[\"plddt\"] * 100, aux_best[\"pae\"], self.Ls * self.copies\n",
        "            )\n",
        "            plt.savefig(f\"{pdb_path}/figs/best.png\", dpi=200, bbox_inches=\"tight\")\n",
        "            plt.close()\n",
        "\n",
        "        # garbage collection\n",
        "        print(\"GC\", gc.collect())\n",
        "\n",
        "class DefaultPipeline:\n",
        "    \"\"\"\n",
        "    This class is used to validate and manage parameters for the DefaultPipeline.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    SetupAlphaFoldColabDesign : SetupAlphaFoldColabDesign\n",
        "\n",
        "        unified_memory : bool\n",
        "            If True, use unified memory.\n",
        "\n",
        "        parentPath : str\n",
        "            The path to the parent directory.\n",
        "\n",
        "    PrepInputs : PrepInputs\n",
        "\n",
        "        sequence : str\n",
        "            The sequence to be processed.\n",
        "\n",
        "        jobname : str\n",
        "            The name of the job.\n",
        "\n",
        "        copies : int\n",
        "            The number of copies.\n",
        "\n",
        "        ## MSA retrieval options\n",
        "\n",
        "        ### MMseqs2 options\n",
        "        msa_method : str\n",
        "            The method used for multiple sequence alignment. Options are \"mmseqs2\",\"single_sequence\", \"custom_fas\", \"custom_a3m\", \"custom_sto\".\n",
        "\n",
        "        pair_mode : str\n",
        "            The pairing mode. Options are \"unpaired_paired\",\"paired\",\"unpaired\".\n",
        "\n",
        "\n",
        "        ### HHfilter options\n",
        "            https://sarata.com/manpages/hhfilter.1.html\n",
        "\n",
        "        cov : int\n",
        "            [0,100]  minimum coverage with query (%) (def=0)\n",
        "\n",
        "        id : int\n",
        "            [0,100]  maximum pairwise sequence identity (%) (def=90)\n",
        "\n",
        "        qid : int\n",
        "            [0,100]  minimum sequence identity with query (%) (def=0)\n",
        "\n",
        "        do_not_filter : bool\n",
        "            If True, do not filter.\n",
        "\n",
        "        ## Template options\n",
        "\n",
        "        template_mode : str\n",
        "            The template mode. Options are \"none\", \"mmseqs2\", \"custom\".\n",
        "\n",
        "        pdb : str\n",
        "            The pdb.\n",
        "\n",
        "        chain : str\n",
        "            The chain.\n",
        "\n",
        "        rm_template_seq : bool\n",
        "            If True, remove template sequence.\n",
        "\n",
        "        propagate_to_copies : bool\n",
        "            If True, propagate to copies.\n",
        "\n",
        "        do_not_align : bool\n",
        "            If True, do not align.\n",
        "\n",
        "    PrepModel : PrepModel\n",
        "    ## AF2 Model preparation options\n",
        "\n",
        "        model_type : str\n",
        "            The model type. Options are \"monomer (ptm)\", \"pseudo_multimer (v3)\", \"multimer (v3)\", \"auto\".\n",
        "\n",
        "        rank_by : str\n",
        "            The rank by. Options are \"auto\", \"plddt\", \"ptm\".\n",
        "\n",
        "        debug : bool\n",
        "            If True, debug.\n",
        "\n",
        "        use_initial_guess : bool\n",
        "            If True, use initial guess.\n",
        "    ## AF2 MSA options\n",
        "\n",
        "        num_msa : str\n",
        "            The number of msa. Options are \"1\",\"2\",\"4\",\"8\",\"16\",\"32\", \"64\", \"128\", \"256\", \"512\".\n",
        "\n",
        "        num_extra_msa : str\n",
        "            The number of extra msa. Options are \"1\",\"2\",\"4\",\"8\",\"16\",\"32\", \"64\", \"128\", \"256\", \"512\", \"1024\",\"2048\",\"4096\".\n",
        "\n",
        "        use_cluster_profile : bool\n",
        "            If True, use cluster profile.\n",
        "\n",
        "    RunAlphaFold : RunAlphaFold\n",
        "    ## AF2 Model run options\n",
        "        model : str\n",
        "            The model. Options are \"1\", \"2\", \"3\", \"4\", \"5\", \"all\".\n",
        "\n",
        "        num_recycles : int\n",
        "            The number of recycles. Options are \"0\", \"1\", \"2\", \"3\", \"6\", \"12\", \"24\".\n",
        "\n",
        "        recycle_early_stop_tolerance : float\n",
        "            The recycle early stop tolerance. Options are \"0.0\", \"0.5\", \"1.0\".\n",
        "\n",
        "        select_best_across_recycles : bool\n",
        "            If True, select best across recycles.\n",
        "\n",
        "    ## AF2 sthochastic options\n",
        "        use_mlm : bool\n",
        "            If True, use mlm.\n",
        "\n",
        "        use_dropout : bool\n",
        "            If True, use dropout.\n",
        "\n",
        "        seed : int\n",
        "            The seed.\n",
        "\n",
        "        num_seeds : int\n",
        "            The number of seeds. Options are \"1\", \"2\", \"4\", \"8\", \"16\", \"32\", \"64\", \"128\".\n",
        "    ## Plotting options\n",
        "        show_images : bool\n",
        "            If True, show images.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, params=None, yaml_file=None):\n",
        "        if yaml_file is not None:\n",
        "            with open(yaml_file, \"r\") as f:\n",
        "                params = yaml.safe_load(f)\n",
        "\n",
        "        self.params = params or {}\n",
        "        self.required_parameters = [\n",
        "            \"unified_memory\",\n",
        "            \"parentPath\",\n",
        "            \"setupPath\",\n",
        "            \"sequence\",\n",
        "            \"jobname\",\n",
        "            \"copies\",\n",
        "            \"msa_method\",\n",
        "            \"custom_a3m_path\",\n",
        "            \"pair_mode\",\n",
        "            \"cov\",\n",
        "            \"id\",\n",
        "            \"qid\",\n",
        "            \"do_not_filter\",\n",
        "            \"template_mode\",\n",
        "            \"pdb\",\n",
        "            \"chain\",\n",
        "            \"rm_template_seq\",\n",
        "            \"propagate_to_copies\",\n",
        "            \"do_not_align\",\n",
        "            \"model_type\",\n",
        "            \"rank_by\",\n",
        "            \"debug\",\n",
        "            \"use_initial_guess\",\n",
        "            \"num_msa\",\n",
        "            \"num_extra_msa\",\n",
        "            \"use_cluster_profile\",\n",
        "            \"model\",\n",
        "            \"num_recycles\",\n",
        "            \"recycle_early_stop_tolerance\",\n",
        "            \"select_best_across_recycles\",\n",
        "            \"use_mlm\",\n",
        "            \"use_dropout\",\n",
        "            \"seed\",\n",
        "            \"num_seeds\",\n",
        "            \"show_images\",\n",
        "        ]\n",
        "        self.param_types = {\n",
        "            \"unified_memory\": bool,\n",
        "            \"parentPath\": str,\n",
        "            \"setupPath\": str,\n",
        "            \"sequence\": str,\n",
        "            \"jobname\": str,\n",
        "            \"copies\": int,\n",
        "            \"msa_method\": str,\n",
        "            \"custom_a3m_path\": str,\n",
        "            \"pair_mode\": str,\n",
        "            \"cov\": int,\n",
        "            \"id\": int,\n",
        "            \"qid\": int,\n",
        "            \"do_not_filter\": bool,\n",
        "            \"template_mode\": str,\n",
        "            \"pdb\": str,\n",
        "            \"chain\": str,\n",
        "            \"rm_template_seq\": bool,\n",
        "            \"propagate_to_copies\": bool,\n",
        "            \"do_not_align\": bool,\n",
        "            \"model_type\": str,\n",
        "            \"rank_by\": str,\n",
        "            \"debug\": bool,\n",
        "            \"use_initial_guess\": bool,\n",
        "            \"num_msa\": int,\n",
        "            \"num_extra_msa\": int,\n",
        "            \"use_cluster_profile\": bool,\n",
        "            \"model\": str,\n",
        "            \"num_recycles\": int,\n",
        "            \"recycle_early_stop_tolerance\": float,\n",
        "            \"select_best_across_recycles\": bool,\n",
        "            \"use_mlm\": bool,\n",
        "            \"use_dropout\": bool,\n",
        "            \"seed\": int,\n",
        "            \"num_seeds\": int,\n",
        "            \"show_images\": bool,\n",
        "        }\n",
        "\n",
        "        self.param_ranges = {\n",
        "            \"copies\": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],\n",
        "            \"msa_method\": [\n",
        "                \"mmseqs2\",\n",
        "                \"single_sequence\",\n",
        "                \"custom_fas\",\n",
        "                \"custom_a3m\",\n",
        "                \"custom_sto\",\n",
        "            ],\n",
        "            \"pair_mode\": [\"unpaired_paired\", \"paired\", \"unpaired\"],\n",
        "            \"cov\": [0, 25, 50, 75, 90, 99],\n",
        "            \"id\": [90, 100],\n",
        "            \"qid\": [0, 10, 15, 20, 30],\n",
        "            \"template_mode\": [\"none\", \"mmseqs2\", \"custom\"],\n",
        "            \"model_type\": [\n",
        "                \"monomer (ptm)\",\n",
        "                \"pseudo_multimer (v3)\",\n",
        "                \"multimer (v3)\",\n",
        "                \"auto\",\n",
        "            ],\n",
        "            \"rank_by\": [\"auto\", \"plddt\", \"ptm\"],\n",
        "            \"num_msa\": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],\n",
        "            \"num_extra_msa\": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],\n",
        "            \"model\": [\"1\", \"2\", \"3\", \"4\", \"5\", \"all\"],\n",
        "            \"recycle_early_stop_tolerance\": [0.0, 0.5, 1.0],\n",
        "            \"color\": [\"pLDDT\", \"chain\", \"rainbow\"],\n",
        "        }\n",
        "        self._validate_parameters()\n",
        "\n",
        "    def _validate_parameters(self):\n",
        "        for param in self.required_parameters:\n",
        "            if param not in self.params:\n",
        "                raise ValueError(f\"{param} is a required parameter\")\n",
        "\n",
        "            # Check parameter type\n",
        "            if not isinstance(self.params[param], self.param_types[param]):\n",
        "                raise ValueError(\n",
        "                    f\"Parameter {param} should be of type {self.param_types[param]}\"\n",
        "                )\n",
        "\n",
        "            # Check parameter range if applicable\n",
        "            if (\n",
        "                param in self.param_ranges\n",
        "                and self.params[param] not in self.param_ranges[param]\n",
        "            ):\n",
        "                raise ValueError(\n",
        "                    f\"Parameter {param} should be one of {self.param_ranges[param]}\"\n",
        "                )\n",
        "\n",
        "    def __getattr__(self, attr):\n",
        "        if attr in self.params:\n",
        "            return self.params[attr]\n",
        "        raise AttributeError(f\"Attribute {attr} not found\")\n",
        "\n",
        "    def _save_config(self):\n",
        "        config = {\n",
        "            \"unified_memory\": self.unified_memory,\n",
        "            \"parentPath\": self.parentPath,\n",
        "            \"setupPath\": self.setupPath,\n",
        "            \"sequence\": self.sequence,\n",
        "            \"jobname\": self.jobname,\n",
        "            \"copies\": self.copies,\n",
        "            \"msa_method\": self.msa_method,\n",
        "            \"custom_a3m_path\": self.custom_a3m_path,\n",
        "            \"pair_mode\": self.pair_mode,\n",
        "            \"cov\": self.cov,\n",
        "            \"id\": self.id,\n",
        "            \"qid\": self.qid,\n",
        "            \"do_not_filter\": self.do_not_filter,\n",
        "            \"template_mode\": self.template_mode,\n",
        "            \"pdb\": self.pdb,\n",
        "            \"chain\": self.chain,\n",
        "            \"rm_template_seq\": self.rm_template_seq,\n",
        "            \"propagate_to_copies\": self.propagate_to_copies,\n",
        "            \"do_not_align\": self.do_not_align,\n",
        "            \"model_type\": self.model_type,\n",
        "            \"rank_by\": self.rank_by,\n",
        "            \"debug\": self.debug,\n",
        "            \"use_initial_guess\": self.use_initial_guess,\n",
        "            \"num_msa\": self.num_msa,\n",
        "            \"num_extra_msa\": self.num_extra_msa,\n",
        "            \"use_cluster_profile\": self.use_cluster_profile,\n",
        "            \"model\": self.model,\n",
        "            \"num_recycles\": self.num_recycles,\n",
        "            \"recycle_early_stop_tolerance\": self.recycle_early_stop_tolerance,\n",
        "            \"select_best_across_recycles\": self.select_best_across_recycles,\n",
        "            \"use_mlm\": self.use_mlm,\n",
        "            \"use_dropout\": self.use_dropout,\n",
        "            \"seed\": self.seed,\n",
        "            \"num_seeds\": self.num_seeds,\n",
        "            \"show_images\": self.show_images,\n",
        "        }\n",
        "\n",
        "    def run(self):\n",
        "        predictor = SetupAlphaFoldColabDesign(self.unified_memory, self.parentPath)\n",
        "        predictor.setup()\n",
        "        prep_inputs = PrepInputs(\n",
        "            self.sequence,\n",
        "            self.jobname,\n",
        "            self.copies,\n",
        "            self.msa_method,\n",
        "            self.custom_a3m_path,\n",
        "            self.pair_mode,\n",
        "            self.cov,\n",
        "            self.id,\n",
        "            self.qid,\n",
        "            self.do_not_filter,\n",
        "            self.template_mode,\n",
        "            self.pdb,\n",
        "            self.chain,\n",
        "            self.rm_template_seq,\n",
        "            self.propagate_to_copies,\n",
        "            self.do_not_align,\n",
        "            self.setupPath,\n",
        "            self.parentPath,\n",
        "        )\n",
        "        prep_inputs.filter_options()\n",
        "        prep_inputs.process_sequence()\n",
        "        prep_inputs.get_msa()\n",
        "        prep_inputs.use_templates()\n",
        "        self._save_config()\n",
        "        prep_model = PrepModel(\n",
        "            self.model_type,\n",
        "            self.rank_by,\n",
        "            self.debug,\n",
        "            self.use_initial_guess,\n",
        "            self.num_msa,\n",
        "            self.num_extra_msa,\n",
        "            self.use_cluster_profile,\n",
        "            prep_inputs.u_lengths,\n",
        "            self.copies,\n",
        "            prep_inputs.use_templates,\n",
        "            prep_inputs.batches,\n",
        "            prep_inputs.msa,\n",
        "            prep_inputs.deletion_matrix,\n",
        "            prep_inputs.u_sub_lengths,\n",
        "            prep_inputs.u_cyclic,\n",
        "            self.setupPath,\n",
        "        )\n",
        "        prep_model.model_options()\n",
        "        prep_model.initialize_model()\n",
        "        prep_model.prep_inputs()\n",
        "        prep_model.set_templates()\n",
        "        prep_model.set_msa()\n",
        "        prep_model.set_chainbreaks()\n",
        "        prep_model.set_cyclic_constraints()\n",
        "\n",
        "        run_alphafold = RunAlphaFold(\n",
        "            prep_inputs.jobname,\n",
        "            self.model,\n",
        "            self.num_recycles,\n",
        "            self.recycle_early_stop_tolerance,\n",
        "            self.select_best_across_recycles,\n",
        "            self.use_mlm,\n",
        "            self.use_dropout,\n",
        "            self.seed,\n",
        "            self.num_seeds,\n",
        "            self.show_images,\n",
        "            self.use_initial_guess,\n",
        "            prep_model.af,\n",
        "            copies=prep_inputs.input_opts[\"copies\"],\n",
        "            print_key=prep_model.print_key,\n",
        "            rank_by=prep_model.rank_by,\n",
        "            Ls=prep_inputs.Ls,\n",
        "            parentPath=self.parentPath,\n",
        "        )\n",
        "        run_alphafold.run()\n",
        "setup_dict = {\n",
        "            \"unified_memory\": False,\n",
        "            \"parentPath\": f\"{final_output_path}/af2\",\n",
        "            \"setupPath\": f\"{BASE_DIR}/params\",\n",
        "        }\n",
        "setup = SetupAlphaFoldColabDesign(setup_dict[\"unified_memory\"],setup_dict[\"parentPath\"],setup_dict[\"setupPath\"])\n",
        "setup.setup()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NdMTA9GH8Who"
      },
      "outputs": [],
      "source": [
        "#@markdown # ColabDesign Prediction\n",
        "# Iterate over the subset and print the ID and the sequence\n",
        "for index, row in df.iterrows():\n",
        "    print(row['experiment_id'], row['sequence'])\n",
        "    # Print the length of the sequence\n",
        "    print(len(row['sequence']))\n",
        "    sequence = str(row['sequence'])\n",
        "    template_mode = \"none\"\n",
        "    pdb = \"\"\n",
        "    chain = \"\"\n",
        "    cols= []\n",
        "    copies= 1\n",
        "    cov= 75\n",
        "    setupPath = f\"{BASE_DIR}/params\"\n",
        "    custom_a3m_path = \"\"\n",
        "    debug= True\n",
        "    do_not_align= False\n",
        "    do_not_filter= True\n",
        "    id= 90\n",
        "    jobname= str(f\"{row['experiment_id']}_{row['id_number']}\")\n",
        "    model= \"1\"\n",
        "    model_type= \"monomer (ptm)\"\n",
        "    msa_method= \"single_sequence\"\n",
        "    num_extra_msa= 1024\n",
        "    num_msa= 512\n",
        "    num_recycles= 2\n",
        "    num_seeds= 1\n",
        "    overwrite= True # WARNING: This will overwrite the previous results\n",
        "    pair_mode= \"unpaired_paired\"\n",
        "    parentPath = f\"{final_output_path}/af2\"\n",
        "    propagate_to_copies= True\n",
        "    qid= 0\n",
        "    rank_by= \"auto\"\n",
        "    recycle_early_stop_tolerance= 0.0\n",
        "    rm_template_seq= False\n",
        "    seed= 0\n",
        "    select_best_across_recycles= False\n",
        "    sequence= sequence\n",
        "    show_images= False\n",
        "    unified_memory= True\n",
        "    use_cluster_profile= True\n",
        "    use_dropout= False\n",
        "    use_initial_guess= False\n",
        "    use_mlm= False\n",
        "    cols_range = []\n",
        "\n",
        "\n",
        "    params = {\n",
        "        \"sequence\": sequence,\n",
        "        \"chain\": chain,\n",
        "        \"jobname\": jobname,\n",
        "        \"cols\": cols,\n",
        "        \"copies\": copies,\n",
        "        \"cov\": cov,\n",
        "        \"setupPath\": setupPath,\n",
        "        \"custom_a3m_path\": custom_a3m_path,\n",
        "        \"debug\": debug,\n",
        "        \"do_not_align\": do_not_align,\n",
        "        \"do_not_filter\": do_not_filter,\n",
        "        \"id\": id,\n",
        "        \"model\": model,\n",
        "        \"model_type\": model_type,\n",
        "        \"msa_method\": msa_method,\n",
        "        \"num_extra_msa\": num_extra_msa,\n",
        "        \"num_msa\": num_msa,\n",
        "        \"num_recycles\": num_recycles,\n",
        "        \"num_seeds\": num_seeds,\n",
        "        \"overwrite\": overwrite,\n",
        "        \"pair_mode\": pair_mode,\n",
        "        \"parentPath\": parentPath,\n",
        "        \"pdb\": pdb,\n",
        "        \"propagate_to_copies\": propagate_to_copies,\n",
        "        \"qid\": qid,\n",
        "        \"rank_by\": rank_by,\n",
        "        \"recycle_early_stop_tolerance\": recycle_early_stop_tolerance,\n",
        "        \"rm_template_seq\": rm_template_seq,\n",
        "        \"seed\": seed,\n",
        "        \"select_best_across_recycles\": select_best_across_recycles,\n",
        "        \"show_images\": show_images,\n",
        "        \"template_mode\": template_mode,\n",
        "        \"unified_memory\": unified_memory,\n",
        "        \"use_cluster_profile\": use_cluster_profile,\n",
        "        \"use_dropout\": use_dropout,\n",
        "        \"use_initial_guess\": use_initial_guess,\n",
        "        \"use_mlm\": use_mlm,\n",
        "        \"cols_range\": cols_range,\n",
        "        \"show_figures\": True,\n",
        "        \"pdb\": pdb,\n",
        "    }\n",
        "    # Run the maskin pipeline with the yaml_dict\n",
        "    default_pipeline = DefaultPipeline(params=params)\n",
        "    default_pipeline.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Analysis of designed sequences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@markdown # AF2 Prediction Analysis\n",
        "\n",
        "import os\n",
        "import glob\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "# Assuming final_output_path is defined from the previous rfdiffusion step\n",
        "DIFFUSION_DIR = final_output_path\n",
        "\n",
        "# Initialize an empty list to store the scores\n",
        "all_scores = []\n",
        "\n",
        "# Iterate over the subset and print the ID and the sequence\n",
        "for index, row in df.iterrows():\n",
        "    print(row['experiment_id'], row['sequence'])\n",
        "    # Print the length of the sequence\n",
        "    print(len(row['sequence']))\n",
        "    \n",
        "    # Find the NPZ file with the unique hash\n",
        "    npz_files = glob.glob(f\"{final_output_path}/af2/{row['experiment_id']}_{row['id_number']}_*/out/npz/{row['experiment_id']}_{row['id_number']}_*_best_model_1_ptm_r2_seed_000.npz\".replace(\" \", \"\"))\n",
        "    print(f\"{final_output_path}/af2/{row['experiment_id']}_{row['id_number']}_*/out/npz/{row['experiment_id']}_{row['id_number']}_*_best_model_1_ptm_r2_seed_000.npz\")\n",
        "    print(npz_files)\n",
        "    if npz_files:\n",
        "        npz_file = npz_files[0]  # Assuming there's only one matching NPZ file\n",
        "        with np.load(npz_file) as data:\n",
        "            print(list(data.keys()))\n",
        "            print(\"Mean pLDDT:\", np.mean(data['plddt']))\n",
        "            print(\"Mean PAE:\", np.mean(data['pae']))\n",
        "            \n",
        "            # Create a dictionary to store the scores\n",
        "            scores = {\n",
        "                'experiment_id': row['experiment_id'],\n",
        "                'id_number': row['id_number'],\n",
        "                'mean_plddt': np.mean(data['plddt']),\n",
        "                'mean_pae': np.mean(data['pae'])\n",
        "            }\n",
        "            all_scores.append(scores)\n",
        "    else:\n",
        "        print(f\"No NPZ file found for {row['experiment_id']}_{row['id_number']}\")\n",
        "\n",
        "# Create a DataFrame from the scores\n",
        "scores_df = pd.DataFrame(all_scores)\n",
        "\n",
        "# Save the combined scores to a CSV file\n",
        "scores_df.to_csv(f\"{final_output_path}/af2/combined_scores.csv\", index=False)\n",
        "\n",
        "# Sort the scores by the mean_plddt and show it as interactive table in google colab\n",
        "scores_df.sort_values(by='mean_plddt', ascending=False).reset_index(drop=True).style\n",
        "scores_df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PhJ2xFUqRSt1"
      },
      "outputs": [],
      "source": [
        "#@title Package and download results\n",
        "#@markdown If you are having issues downloading the result archive,\n",
        "#@markdown try disabling your adblocker and run this cell again.\n",
        "#@markdown  If that fails click on the little folder icon to the\n",
        "#@markdown  left, navigate to file: `name.result.zip`,\n",
        "#@markdown  right-click and select \\\"Download\\\"\n",
        "#@markdown (see [screenshot](https://pbs.twimg.com/media/E6wRW2lWUAEOuoe?format=jpg&name=small)).\n",
        "import shutil\n",
        "\n",
        "def zip_last_job(final_output_path, counter):\n",
        "    \"\"\"\n",
        "    Zip the last job's output directory for download.\n",
        "    \"\"\"\n",
        "    base_output_path = os.path.join(OUTPUT_DIR, final_output_path)\n",
        "\n",
        "    if final_output_path:\n",
        "        output_path = f\"{base_output_path}\"\n",
        "        shutil.make_archive(output_path, 'zip', output_path)\n",
        "        return f\"{output_path}.zip\"\n",
        "    else:\n",
        "        print(\"No output directory found.\")\n",
        "        return None\n",
        "\n",
        "\n",
        "# Assuming you're in a Jupyter notebook cell\n",
        "from google.colab import files\n",
        "\n",
        "# This cell should be run after the diffusion function to zip and download the output\n",
        "zip_path = zip_last_job(final_output_path=final_output_path, counter=0)\n",
        "\n",
        "if zip_path and is_colab():\n",
        "    from google.colab import files\n",
        "    files.download(zip_path)\n",
        "else:\n",
        "    print(\"Zip file path:\", zip_path)\n",
        "    print(\"Note: Automatic download is only supported in Google Colab.\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}