pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e578ba7a-6cc5-4cb7-9bd8-d63dbaba38af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import getpass\n",
    "import subprocess\n",
    "import time\n",
    "import importlib\n",
    "from shutil import copy2\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "### Path to this cloned GitHub repo:\n",
    "SCRIPT_DIR = os.path.dirname(__file__)  # edit this to the GitHub repo path. Throws an error by default.\n",
    "assert os.path.exists(SCRIPT_DIR)\n",
    "sys.path.append(SCRIPT_DIR+\"/scripts/utils\")\n",
    "import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d25115bb",
   "metadata": {},
   "source": [
    "# De novo heme binding protein design pipeline using RFdiffusionAA\n",
    "#### And other ligand binders too, I guess\n",
    "Indrek Kalvet, PhD (Institute for Protein Design, University of Washington), ikalvet@uw.edu\n",
    "\n",
    "As implemented in the publication \"<i>Generalized Biomolecular Modeling and Design with RoseTTAFold All-Atom\"\n",
    "\n",
    "Link:\n",
    "\n",
    "This notebook illustrates the design of heme-binding proteins, starting from minimal information (heme + substrate + CYS motif). It should work with minor modifications also for any other ligand.\n",
    "\n",
    "The pipeline consists of 7 steps:<br>\n",
    "\n",
    "    0) The protein backbones are generated with RFdiffusionAA\n",
    "    1) Sequence is designed with proteinMPNN (without the ligand)\n",
    "    2) Structures are predicted with AlphaFold2\n",
    "    3) Ligand binding site is designed with LigandMPNN/FastRelax, or Rosetta FastDesign\n",
    "    4) Sequences surrounding the ligand pocket are diversified with LigandMPNN\n",
    "    5) Final designed sequences are predicted with AlphaFold2\n",
    "    6) Alphafold2-predicted models are relaxed with the ligand and analyzed\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89b27018",
   "metadata": {},
   "source": [
    "## Paths of important Python scripts and programs\n",
    "\n",
    "#### RFdiffusionAA:\n",
    "Download RFdiffusionAA from here: https://github.com/baker-laboratory/rf_diffusion_all_atom\n",
    "\n",
    "#### RFjoint inpainting (proteininpainting)\n",
    "(Optional) Download RFjoint Inpainting here: https://github.com/RosettaCommons/RFDesign\n",
    "\n",
    "Inpainting is used to further resample/diversify diffusion outputs, and it may also increase AF2 success rates.\n",
    "<br>\n",
    "\n",
    "#### AlphaFold2 and LigandMPNN\n",
    "Other dependencies (ligandMPNN and AlphaFold2) are downloaded as submodules with this GitHub repository.\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "<i>After these evironments are set up and repositories are downloaded please provide paths to the inference scripts in the cell below:</i>\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "#### Python or Apptainer image\n",
    "Please define in the cell below any Python executables or Apptainer image paths that are able to run the different scripts.<br>\n",
    "\n",
    "This pipeline is tested to run based on two different Conda environments: `diffusion` and `mlfold`.<br>\n",
    "Both of these can be set up based on the yml files provided by this repository.<br>\n",
    "`mlfold` is used for AlphaFold2, and `diffusion` is used for everything else.\n",
    "\n",
    "<br>\n",
    "\n",
    "Alternatively, if your system is using Apptainers, you can set them up based on the same yml files, or do it separately. For RFdiffusionAA you can download this Apptainer image:\n",
    "http://files.ipd.uw.edu/pub/RF-All-Atom/containers/rf_se3_diffusion.sif\n",
    "\n",
    "If you are setting the environments or containers up without the provided YML files, then the minimum requirements for the different types are:<br>\n",
    "`af2` Python needs to include jax=0.4.25 and jaxlib=0.4.23<br>\n",
    "`proteinMPNN` Python needs to include pytorch and prody<br>\n",
    "`general` Python needs to include pytorch, pyrosetta, prody<br>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b49d9e31",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_script = \"PATH/TO/rf_diffusion_all_atom/run_inference.py\"  # edit this\n",
    "inpaint_script = \"PATH/TO/RFDesign/inpainting/inpaint.py\"  # edit this if needed\n",
    "proteinMPNN_script = f\"{SCRIPT_DIR}/lib/LigandMPNN/run.py\"  # from submodule\n",
    "AF2_script = f\"{SCRIPT_DIR}/scripts/af2/af2.py\"  # from submodule\n",
    "\n",
    "\n",
    "\n",
    "### Python and/or Apptainer executables needed for running the jobs\n",
    "### Please provide paths to executables that are able to run the different tasks.\n",
    "### They can all be the same if you have an environment with all of the ncessary Python modules in one\n",
    "\n",
    "# If your added Apptainer does not execute scripts directly,\n",
    "# try adding 'apptainer run' or 'apptainer run --nv' (for GPU) in front of the command\n",
    "\n",
    "CONDAPATH = \"/home/ikalvet/miniconda3\"   # edit this depending on where your Conda environments live\n",
    "PYTHON = {\"diffusion\": f\"{CONDAPATH}/envs/diffusion/bin/python\",\n",
    "          \"af2\": f\"{CONDAPATH}/envs/mlfold/bin/python\",\n",
    "          \"proteinMPNN\": f\"{CONDAPATH}/envs/diffusion/bin/python\",\n",
    "          \"general\": f\"{CONDAPATH}/envs/diffusion/bin/python\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74d4990c-bff7-44dd-b675-c0379d9a45af",
   "metadata": {},
   "source": [
    "## Project description and paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59cfb5a8-5f64-49f4-a26f-0da5f09b4120",
   "metadata": {},
   "outputs": [],
   "source": [
    "username = getpass.getuser()  # your username on the running system\n",
    "EMAIL = f\"{username}@uw.edu\"  # edit based on your organization. For Slurm job notifications.\n",
    "\n",
    "PROJECT = \"example_Heme_diffusion\"\n",
    "\n",
    "### Path where the jobs will be run and outputs dumped\n",
    "WDIR = \"/where/my/pipeline/outputs/live\"\n",
    "\n",
    "if not os.path.exists(WDIR):\n",
    "    os.makedirs(WDIR, exist_ok=True)\n",
    "\n",
    "print(f\"Working directory: {WDIR}\")\n",
    "\n",
    "USE_GPU_for_AF2 = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "204f8e5c-0235-478b-a123-cc896fa9734d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ligand information\n",
    "params = [f\"{SCRIPT_DIR}/theozyme/HBA/HBA.params\"]  # Rosetta params file(s)\n",
    "LIGAND = \"HBA\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5fca3ef-ba68-4896-8a98-dcc881832174",
   "metadata": {},
   "source": [
    "## 0: Setting up diffusion run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b9faacb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using example PDB file with ligand HBA and protein 7o2g backbone.\n",
    "## Note: the repository also contains additional HBA conformers with 7o2g and P450 motifs\n",
    "## in the same directory as a ZIP file.\n",
    "\n",
    "diffusion_inputs = glob.glob(f\"{SCRIPT_DIR}/input/*.pdb\")\n",
    "print(f\"Found {len(diffusion_inputs)} PDB files\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01405838-201f-4e1e-940d-ea1d05fdaba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up general settings for diffusion\n",
    "\n",
    "DIFFUSION_DIR = f\"{WDIR}/0_diffusion\"\n",
    "if not os.path.exists(DIFFUSION_DIR):\n",
    "    os.makedirs(DIFFUSION_DIR, exist_ok=False)\n",
    "\n",
    "os.chdir(DIFFUSION_DIR)\n",
    "\n",
    "N_designs = 5\n",
    "T_steps = 200\n",
    "\n",
    "## Edit this config based on motif residues, etc...\n",
    "config = f\"\"\"\n",
    "defaults:\n",
    "  - aa\n",
    "\n",
    "diffuser:\n",
    "  T: {T_steps}\n",
    "\n",
    "inference:\n",
    "  num_designs: {N_designs}\n",
    "  model_runner: NRBStyleSelfCond\n",
    "  ligand: '{LIGAND}'\n",
    "\n",
    "model:\n",
    "  freeze_track_motif: True\n",
    "\n",
    "contigmap:\n",
    "  contigs: [\"30-110,A15-15,30-110\"]\n",
    "  inpaint_str: null\n",
    "  length: \"100-140\"\n",
    "\n",
    "potentials:\n",
    "  guiding_potentials: [\"type:ligand_ncontacts,weight:1\"] \n",
    "  guide_scale: 2\n",
    "  guide_decay: cubic\n",
    "\"\"\"\n",
    "\n",
    "estimated_time = 3.5 * T_steps * N_designs  # assuming 3.5 seconds per timestep on A4000 GPU\n",
    "\n",
    "print(f\"Estimated time to produce {N_designs} designs = {estimated_time/60:.0f} minutes\")\n",
    "with open(\"config.yaml\", \"w\") as file:\n",
    "    file.write(config)\n",
    "print(f\"Wrote config file to {os.path.realpath('config.yaml')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8a89617-a8a7-45d5-abd5-dc85fa32bc17",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up diffusion commands based on the input PDB file(s)\n",
    "## Diffusion jobs are run in separate directories for each input PDB\n",
    "\n",
    "commands_diffusion = []\n",
    "cmds_filename = \"commands_diffusion\"\n",
    "diffusion_rundirs = []\n",
    "with open(cmds_filename, \"w\") as file:\n",
    "    for p in diffusion_inputs:\n",
    "        pdbname = os.path.basename(p).replace(\".pdb\", \"\")\n",
    "        os.makedirs(pdbname, exist_ok=True)\n",
    "        cmd = f\"cd {pdbname} ; {PYTHON['diffusion']} {diffusion_script} --config-dir=../ \"\\\n",
    "              f\"--config-name=config.yaml inference.input_pdb={p} \"\\\n",
    "              f\"inference.output_prefix='./out/{pdbname}_dif' > output.log ; cd ..\\n\"\n",
    "        commands_diffusion.append(cmd)\n",
    "        diffusion_rundirs.append(pdbname)\n",
    "        file.write(cmd)\n",
    "\n",
    "print(f\"An example diffusion command that was generated:\\n   {cmd}\")\n",
    "\n",
    "\n",
    "## Creating a Slurm submit script\n",
    "## adjust time depending on number of designs and available hardware\n",
    "submit_script = \"submit_diffusion.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"diffusion_example\", gpu=True, gres=\"gpu:a4000:1\",\n",
    "                                 mem=\"8g\", N_cores=2, time=\"1:00:00\", email=EMAIL,\n",
    "                                 array=len(commands_diffusion), array_commandfile=cmds_filename)\n",
    "\n",
    "print(f\"Writing diffusion submission script to {submit_script}\")\n",
    "print(f\"{len(commands_diffusion)} diffusion jobs to run\")\n",
    "\n",
    "if not os.path.exists(DIFFUSION_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86ed7eb2-c045-49a0-9264-292a6ef4dc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with diffusion and happy with the outputs then mark it as done\n",
    "DIFFUSION_DIR = f\"{WDIR}/0_diffusion\"\n",
    "os.chdir(DIFFUSION_DIR)\n",
    "\n",
    "if not os.path.exists(DIFFUSION_DIR+\"/.done\"):\n",
    "    with open(f\"{DIFFUSION_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a474a099-60e1-493d-85f2-12b8808a0b11",
   "metadata": {},
   "source": [
    "### Analyzing diffusion outputs\n",
    "The purpose of this step is to identify diffused backbones that meet certain quality criteria. These scaffolds should be relatively globular (measured by radius of gyration (rog), and longest helix). They should not have clashes between the ligand and the backbone, the ligand should not be too exposed (measured by relative SASA). The termini should not be too close to the ligand (term_mindist), and the backbone should not be too loopy. In the example below we are also looking for backbones that leave some part of the ligand more exposed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06c85384-baa4-41c2-bc21-42587839a44d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Analyzing diffusion outputs for clashes, ligand burial and scaffold quality\n",
    "## If it's running too slowly consider increasing --nproc\n",
    "\n",
    "analysis_script = f\"{SCRIPT_DIR}/scripts/diffusion_analysis/process_diffusion_outputs.py\"\n",
    "\n",
    "diffusion_outputs = []\n",
    "for d in diffusion_rundirs:\n",
    "    diffusion_outputs += glob.glob(f\"{d}/out/*.pdb\")\n",
    "\n",
    "# By default I don't use the --analyze flag. As a result the backbones are filtered as the script runs.\n",
    "# You can set --analyze to True to calculate all scores for all backbones.\n",
    "# This will slow the analysis down, but you can then filter the backbones separately afterwards.\n",
    "dif_analysis_cmd_dict = {\"--pdb\": \" \".join(diffusion_outputs),\n",
    "                        \"--ref\": f\"{SCRIPT_DIR}/input/*.pdb\",\n",
    "                        \"--params\": \" \".join(params),\n",
    "                        \"--term_limit\": \"15.0\",\n",
    "                        \"--SASA_limit\": \"0.3\",  # Highest allowed relative SASA of ligand\n",
    "                        \"--loop_limit\": \"0.4\",  # Fraction of backbone that can be loopy\n",
    "                        \"--ref_catres\": \"A15\",  # Position of CYS in diffusion input\n",
    "                        \"--rethread\": True,\n",
    "                        \"--fix\": True,\n",
    "                        \"--exclude_clash_atoms\": \"O1 O2 O3 O4 C5 C10\",  # Ligand atoms excluded from clashchecking because they are flexible\n",
    "                        \"--ligand_exposed_atoms\": \"C45 C46 C47\",  # Ligand atoms that need to be more exposed\n",
    "                        \"--exposed_atom_SASA\": \"10.0\",  # minimum absolute SASA for exposed ligand atoms\n",
    "                        \"--longest_helix\": \"30\",\n",
    "                        \"--rog\": \"30.0\",\n",
    "                        \"--partial\": None,\n",
    "                        \"--outdir\": None,\n",
    "                        \"--traj\": \"5/30\",  # Also random 5 models are taken from the last 30 steps of the diffusion trajectory\n",
    "                        \"--trb\": None,\n",
    "                        \"--analyze\": False,\n",
    "                        \"--nproc\": \"1\"}\n",
    "\n",
    "analysis_command = f\"{PYTHON['general']} {analysis_script}\"\n",
    "for k, val in dif_analysis_cmd_dict.items():\n",
    "    if val is not None:\n",
    "        if isinstance(val, list):\n",
    "            analysis_command += f\" {k}\"\n",
    "            analysis_command += \" \" + \" \".join(val)\n",
    "        elif isinstance(val, bool):\n",
    "            if val == True:\n",
    "                analysis_command += f\" {k}\"\n",
    "        else:\n",
    "            analysis_command += f\" {k} {val}\"\n",
    "        print(k, val)\n",
    "\n",
    "if len(diffusion_outputs) < 100:\n",
    "    ## Analyzing locally\n",
    "    p = subprocess.Popen(analysis_command, shell=True)\n",
    "    (output, err) = p.communicate()\n",
    "else:\n",
    "    ## Too many structures to analyze.\n",
    "    ## Running the analysis as a SLURM job.\n",
    "    submit_script = \"submit_diffusion_analysis.sh\"\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"diffusion_analysis\",\n",
    "                                     mem=\"8g\", N_cores=dif_analysis_cmd_dict[\"--nproc\"], time=\"0:20:00\", email=EMAIL,\n",
    "                                     command=analysis_command, outfile_name=\"output_analysis\")\n",
    "\n",
    "diffused_backbones_good = glob.glob(f\"{DIFFUSION_DIR}/filtered_structures/*.pdb\")\n",
    "\n",
    "dif_analysis_df = pd.read_csv(f\"{DIFFUSION_DIR}/diffusion_analysis.sc\", header=0, sep=\"\\s+\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c10b4314",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Visualizing the distributions of diffusion analysis metrics\n",
    "## Plotting design scores\n",
    "plt.figure(figsize=(12, 12))\n",
    "for i,k in enumerate(dif_analysis_df.keys()):\n",
    "    if k in [\"description\"]:\n",
    "        continue\n",
    "    plt.subplot(4, 3, i+1)\n",
    "    plt.hist(dif_analysis_df[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d6136b4",
   "metadata": {},
   "source": [
    "It is highly advised that you manually inspect the filtered diffusion outputs before continuing with the pipeline.\n",
    "While the filters attempt to pick out the most offending designs then nothing beats your own intuition and judgement."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9561c720",
   "metadata": {},
   "source": [
    "If you would like to perform RFjoint Inpainting on the diffusion outputs, please go to the [inpainting section](#inpainting)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f93a2f55",
   "metadata": {},
   "source": [
    "## 1: Running ProteinMPNN on diffused backbones\n",
    "\n",
    "We are first trying to just design a sequence on the backbone, without considering the ligand.\n",
    "The goal is to first find backbones that fold well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7235734c-e319-48cd-80e2-f6c5e0cacfae",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffused_backbones_good = glob.glob(f\"{DIFFUSION_DIR}/filtered_structures/*.pdb\")\n",
    "assert len(diffused_backbones_good) > 0, \"No good backbones found!\"\n",
    "\n",
    "os.chdir(WDIR)\n",
    "\n",
    "MPNN_DIR = f\"{WDIR}/1_proteinmpnn\"\n",
    "os.makedirs(MPNN_DIR, exist_ok=True)\n",
    "os.chdir(MPNN_DIR)\n",
    "\n",
    "### Parsing diffusion output TRB files to extract fixed motif residues\n",
    "## These residues will not be redesigned with proteinMPNN\n",
    "mask_json_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/make_maskdict_from_trb.py --out masked_pos.jsonl --trb\"\n",
    "for d in diffused_backbones_good:\n",
    "    mask_json_cmd += \" \" + d.replace(\".pdb\", \".trb\")\n",
    "p = subprocess.Popen(mask_json_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n",
    "\n",
    "assert os.path.exists(\"masked_pos.jsonl\"), \"Failed to create masked positions JSONL file\"\n",
    "\n",
    "\n",
    "### Setting up proteinMPNN run commands\n",
    "## We're doing design with 3 temperatures, and 5 sequences each.\n",
    "## This usually gives decent success with designable backbones.\n",
    "## For more complicated cases consider doing >100 sequences.\n",
    "\n",
    "MPNN_temperatures = [0.1, 0.2, 0.3]\n",
    "MPNN_outputs_per_temperature = 5\n",
    "MPNN_omit_AAs = \"CM\"\n",
    "\n",
    "commands_mpnn = []\n",
    "cmds_filename_mpnn = \"commands_mpnn\"\n",
    "with open(cmds_filename_mpnn, \"w\") as file:\n",
    "    for T in MPNN_temperatures:\n",
    "        for f in diffused_backbones_good:\n",
    "            commands_mpnn.append(f\"{PYTHON['proteinMPNN']} {proteinMPNN_script} \"\n",
    "                                 f\"--model_type protein_mpnn --ligand_mpnn_use_atom_context 0 \"\n",
    "                                 \"--fixed_residues_multi masked_pos.jsonl --out_folder ./ \"\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 {SCRIPT_DIR}/lib/LigandMPNN/model_params/proteinmpnn_v_48_020.pt\\n\")\n",
    "            file.write(commands_mpnn[-1])\n",
    "\n",
    "print(\"Example MPNN command:\")\n",
    "print(commands_mpnn[-1])\n",
    "\n",
    "### Running proteinMPNN with Slurm.\n",
    "### Grouping jobs with 10 commands per one array job.\n",
    "\n",
    "submit_script = \"submit_mpnn.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"1_proteinmpnn\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"0:15:00\", email=EMAIL, array=len(commands_mpnn),\n",
    "                                 array_commandfile=cmds_filename_mpnn, group=10)\n",
    "\n",
    "if not os.path.exists(MPNN_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f793273",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with diffusion and happy with the outputs then mark it as done\n",
    "MPNN_DIR = f\"{WDIR}/1_proteinmpnn\"\n",
    "os.chdir(MPNN_DIR)\n",
    "\n",
    "if not os.path.exists(MPNN_DIR+\"/.done\"):\n",
    "    with open(f\"{MPNN_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97333709",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## 2: Running AlphaFold2\n",
    "Performing AF2 single sequence predictions.<br>\n",
    "By default only using AF2 model 4 with 3 recycles. For more complicated folds 10+ recycles might be necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "065029b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(WDIR)\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/2_af2\"\n",
    "os.makedirs(AF2_DIR, exist_ok=True)\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "### First collecting MPNN outputs and creating FASTA files for AF2 input\n",
    "mpnn_fasta = utils.parse_fasta_files(glob.glob(f\"{MPNN_DIR}/seqs/*.fa\"))\n",
    "mpnn_fasta = {k: seq.strip() for k, seq in mpnn_fasta.items() if \"model_path\" not in k}  # excluding the diffused poly-A sequence\n",
    "# Giving sequences unique names based on input PDB name, temperature, and sequence identifier\n",
    "mpnn_fasta = {k.split(\",\")[0]+\"_\"+k.split(\",\")[2].replace(\" T=\", \"T\")+\"_0_\"+k.split(\",\")[1].replace(\" id=\", \"\"): seq for k, seq in mpnn_fasta.items()}\n",
    "\n",
    "print(f\"A total on {len(mpnn_fasta)} sequences will be predicted.\")\n",
    "\n",
    "## Splitting the MPNN sequences based on length\n",
    "## and grouping them in smaller batches for each AF2 job\n",
    "## Use group size of >40 when running on GPU. Also depends on how many sequences and resources you have.\n",
    "\n",
    "SEQUENCES_PER_AF2_JOB = 5  # CPU\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    SEQUENCES_PER_AF2_JOB = 100  # GPU\n",
    "mpnn_fasta_split = utils.split_fasta_based_on_length(mpnn_fasta, SEQUENCES_PER_AF2_JOB, write_files=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf52a6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up AlphaFold2 run\n",
    "\n",
    "AF2_recycles = 3\n",
    "AF2_models = \"4\"  # add other models to this string if needed, i.e. \"3 4 5\"\n",
    "\n",
    "commands_af2 = []\n",
    "cmds_filename_af2 = \"commands_af2\"\n",
    "with open(cmds_filename_af2, \"w\") as file:\n",
    "    for ff in glob.glob(\"*.fasta\"):\n",
    "        commands_af2.append(f\"{PYTHON['af2']} {AF2_script} \"\n",
    "                             f\"--af-nrecycles {AF2_recycles} --af-models {AF2_models} \"\n",
    "                             f\"--fasta {ff} --scorefile {ff.replace('.fasta', '.csv')}\\n\")\n",
    "        file.write(commands_af2[-1])\n",
    "\n",
    "print(\"Example AF2 command:\")\n",
    "print(commands_af2[-1])\n",
    "\n",
    "### Running AF2 with Slurm.\n",
    "### Running jobs on the CPU. It takes ~10 minutes per sequence\n",
    "### \n",
    "\n",
    "submit_script = \"submit_af2.sh\"\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"2_af2\", mem=\"6g\", \n",
    "                                     N_cores=2, gpu=True, gres=\"gpu:rtx2080:1\", time=\"00:12:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "else:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"2_af2\", mem=\"6g\", \n",
    "                                     N_cores=4, time=\"01:00:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a88b24e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with diffusion and happy with the outputs then mark it as done\n",
    "AF2_DIR = f\"{WDIR}/2_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    with open(f\"{AF2_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "242e9def",
   "metadata": {},
   "source": [
    "### Analyzing AlphaFold2 predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f38048",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combining all CSV scorefiles into one\n",
    "AF2_DIR = f\"{WDIR}/2_af2\"\n",
    "DIFFUSION_DIR = f\"{WDIR}/0_diffusion\"\n",
    "\n",
    "os.system(\"head -n 1 $(ls *aa*.csv | shuf -n 1) > scores.csv ; for f in *aa*.csv ; do tail -n +2 ${f} >> scores.csv ; done\")\n",
    "assert os.path.exists(\"scores.csv\"), \"Could not combine scorefiles\"\n",
    "\n",
    "### Calculating the RMSDs of AF2 predictions relative to the diffusion outputs\n",
    "### Catalytic residue sidechain RMSDs are calculated in the reference PDB has REMARK 666 line present\n",
    "\n",
    "analysis_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/analyze_af2.py --scorefile scores.csv \"\\\n",
    "               f\"--ref_path {DIFFUSION_DIR}/filtered_structures/ --mpnn --params {' '.join(params)}\"\n",
    "\n",
    "\n",
    "if len(glob.glob(f\"{AF2_DIR}/*.pdb\")) < 100:\n",
    "    ## Analyzing locally\n",
    "    p = subprocess.Popen(analysis_cmd, shell=True)\n",
    "    (output, err) = p.communicate()\n",
    "else:\n",
    "    ## Running as a Slurm job\n",
    "    submit_script = \"submit_af2_analysis.sh\"\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"af2_analysis\",\n",
    "                                     mem=\"16g\", N_cores=8, time=\"0:20:00\", email=EMAIL,\n",
    "                                     command=analysis_cmd, outfile_name=\"output_analysis\")\n",
    "\n",
    "    p = subprocess.Popen([\"sbatch\", submit_script])\n",
    "    (output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45c18f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores_af2 = pd.read_csv(\"scores.sc\", sep=\"\\s+\", header=0)\n",
    "\n",
    "### Filtering AF2 scores based on lddt and rmsd\n",
    "# Define your desired cutoffs here:\n",
    "AF2_filters = {\"lDDT\": [85.0, \">=\"],\n",
    "               \"rmsd\": [1.5, \"<=\"],\n",
    "               \"rmsd_SR1\": [1.0, \"<=\"]}  # 1st catalytic residue sc-rmsd\n",
    "\n",
    "scores_af2_filtered = utils.filter_scores(scores_af2, AF2_filters)\n",
    "utils.dump_scorefile(scores_af2_filtered, \"filtered_scores.sc\")\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 3))\n",
    "for i,k in enumerate(AF2_filters):\n",
    "    plt.subplot(1, 3, i+1)\n",
    "    plt.hist(scores_af2[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "    \n",
    "utils.plot_score_pairs(scores_af2, \"lDDT\", \"rmsd\", AF2_filters[\"lDDT\"][0], AF2_filters[\"rmsd\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aeb0952",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Copying good predictions to a separate directory\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if len(scores_af2_filtered) > 0:\n",
    "    os.makedirs(\"good\", exist_ok=True)\n",
    "    good_af2_models = [row[\"Output_PDB\"]+\".pdb\" for idx,row in scores_af2_filtered.iterrows()]\n",
    "    for pdb in good_af2_models:\n",
    "        copy2(pdb, f\"good/{pdb}\")\n",
    "    good_af2_models = glob.glob(f\"{AF2_DIR}/good/*.pdb\")\n",
    "else:\n",
    "    sys.exit(\"No good models to continue this pipeline with\")\n",
    "\n",
    "os.chdir(f\"{AF2_DIR}/good\")\n",
    "\n",
    "\n",
    "### Aligning the ligand back into the AF2 predictions.\n",
    "### This is done by aligning the AF2 model to diffusion output and copying over the ligand using PyRosetta.\n",
    "### --fix_catres option will readjust the rotamer and tautomer of \n",
    "### any catalytic residue to be the same as in the reference model.\n",
    "\n",
    "align_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/place_ligand_after_af2.py \"\\\n",
    "            f\"--outdir with_heme2 --params {' '.join(params)} --fix_catres \"\\\n",
    "            f\"--pdb {' '.join(good_af2_models)} \"\\\n",
    "            f\"--ref {' '.join(glob.glob(DIFFUSION_DIR+'/filtered_structures/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(align_cmd, shell=True)\n",
    "(output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cddf39ab",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "aac54ca5",
   "metadata": {},
   "source": [
    "## 3.1: Performing binding site design with ligandMPNN / FastRelax\n",
    "(see 3.2 down below for design with [Rosetta FastDesign](#Rosetta_design))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e730b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting up design directory and commands\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/3.1_design_pocket_ligandMPNN\"\n",
    "os.makedirs(DESIGN_DIR_ligMPNN, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_ligMPNN)\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/2_af2\"\n",
    "os.makedirs(DESIGN_DIR_ligMPNN+\"/logs\", exist_ok=True)\n",
    "\n",
    "### Performing 5 design iterations on each input structure\n",
    "NSTRUCT = 10\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_design = []\n",
    "cmds_filename_des = \"commands_design\"\n",
    "with open(cmds_filename_des, \"w\") as file:\n",
    "    for pdb in glob.glob(f\"{AF2_DIR}/good/with_heme/*.pdb\"):\n",
    "        commands_design.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/heme_pocket_ligMPNN.py \"\n",
    "                             f\"--pdb {pdb} --nstruct {NSTRUCT} \"\n",
    "                             f\"--scoring {SCRIPT_DIR}/scripts/design/scoring/heme_scoring.py \"\n",
    "                             f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(pdb).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_design[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_design[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_design.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"3.1_design_pocket_ligMPNN\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"3:00:00\", email=EMAIL, array=len(commands_design),\n",
    "                                 array_commandfile=cmds_filename_des)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_ligMPNN+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43c2821f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with design and happy with the outputs then mark it as done\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/3.1_design_pocket_ligandMPNN\"\n",
    "os.chdir(DESIGN_DIR_ligMPNN)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_ligMPNN+\"/.done\"):\n",
    "    with open(f\"{DESIGN_DIR_ligMPNN}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef0d7f43",
   "metadata": {},
   "source": [
    "### Analyzing ligMPNN / FastRelax design results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2364e89",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Analyzing ligMPNN / FastRelax designs\n",
    "if not os.path.exists(\"scorefile.txt\"):\n",
    "    sys.exit(\"Design job failed, or no successful outputs were produced.\")\n",
    "\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>=']}\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    plt.subplot(3, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{DESIGN_DIR_ligMPNN}/good\")\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2c82f60",
   "metadata": {},
   "source": [
    "## 4.1 Performing ligandMPNN redesign on the 2nd layer residues\n",
    "\n",
    "Resampling residues that are not in the pocket, but also not very far from the pocket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f470af9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb')) == 0:\n",
    "    sys.exit(\"No designs to run 2nd MPNN on.\")\n",
    "\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.1_2nd_mpnn\"\n",
    "os.makedirs(DESIGN_DIR_2nd_mpnn, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_2nd_mpnn)\n",
    "\n",
    "### Making a JSON file specifiying designable positions for each structure.\n",
    "### Will also make non-pocket ALA positions as designable.\n",
    "### This is to fix any surface ALA-patches that previous MPNN may have introduced.\n",
    "\n",
    "make_json_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/setup_ligand_mpnn_2nd_layer.py \"\\\n",
    "                f\"--params {' '.join(params)} --ligand {LIGAND} --output_path parsed_pdbs_lig.jsonl \"\\\n",
    "                 \"--output_path masked_pos.jsonl --dist_bb 6.0 --dist_sc 5.0 \"\\\n",
    "                f\"--pdb {' '.join(glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(make_json_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n",
    "\n",
    "\n",
    "### Setting up ligandMPNN run commands\n",
    "## We're doing design with 2 temperatures (more conservative than before), and 5 sequences each.\n",
    "\n",
    "MPNN_temperatures = [0.1, 0.2]\n",
    "MPNN_outputs_per_temperature = 5\n",
    "MPNN_omit_AAs = \"CM\"\n",
    "\n",
    "commands_mpnn = []\n",
    "cmds_filename_mpnn = \"commands_mpnn\"\n",
    "with open(cmds_filename_mpnn, \"w\") as file:\n",
    "    for T in MPNN_temperatures:\n",
    "        for f in glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb'):\n",
    "            commands_mpnn.append(f\"{PYTHON['proteinMPNN']} {proteinMPNN_script} \"\n",
    "                                 f\"--model_type ligand_mpnn --ligand_mpnn_use_atom_context 1 \"\n",
    "                                 \"--fixed_residues_multi masked_pos.jsonl --out_folder ./ \"\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 {SCRIPT_DIR}/lib/LigandMPNN/model_params/ligandmpnn_v_32_010_25.pt\\n\")\n",
    "            file.write(commands_mpnn[-1])\n",
    "\n",
    "print(\"Example MPNN command:\")\n",
    "print(commands_mpnn[-1])\n",
    "\n",
    "### Running ligandMPNN with Slurm.\n",
    "### Grouping jobs with 10 commands per one array job.\n",
    "\n",
    "submit_script = \"submit_mpnn.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"4.1_2nd_mpnn\", mem=\"8g\", \n",
    "                                 N_cores=1, time=\"0:15:00\", email=EMAIL, array=len(commands_mpnn),\n",
    "                                 array_commandfile=cmds_filename_mpnn, group=10)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffd97aea",
   "metadata": {},
   "source": [
    "## 5.1 AlphaFold2 predictions on the 2nd MPNN run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15e96fe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.1_2nd_mpnn\"\n",
    "assert os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"), \"2nd MPNN has not been performed!\"\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/5.1_2nd_af2\"\n",
    "os.makedirs(AF2_DIR, exist_ok=True)\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "### First collecting MPNN outputs and creating FASTA files for AF2 input\n",
    "mpnn_fasta = utils.parse_fasta_files(glob.glob(f\"{DESIGN_DIR_2nd_mpnn}/seqs/*.fa\"))\n",
    "# Giving sequences unique names based on input PDB name, temperature, and sequence identifier\n",
    "_mpnn_fasta = {}\n",
    "for k, seq in mpnn_fasta.items():\n",
    "    if \"model_path\" in k:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_native\"] = seq.strip()\n",
    "    else:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_\"+k.split(\",\")[2].replace(\" T=\", \"T\")+\"_0_\"+k.split(\",\")[1].replace(\" id=\", \"\")] = seq.strip()\n",
    "mpnn_fasta = {k:v for k,v in _mpnn_fasta.items()}\n",
    "\n",
    "print(f\"A total on {len(mpnn_fasta)} sequences will be predicted.\")\n",
    "\n",
    "## Splitting the MPNN sequences based on length\n",
    "## and grouping them in smaller batches for each AF2 job\n",
    "## Use group size of >40 when running on GPU. Also depends on how many sequences and resources you have.\n",
    "SEQUENCES_PER_AF2_JOB = 5  # CPU\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    SEQUENCES_PER_AF2_JOB = 100  # GPU\n",
    "mpnn_fasta_split = utils.split_fasta_based_on_length(mpnn_fasta, SEQUENCES_PER_AF2_JOB, write_files=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d082ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up AlphaFold2 run\n",
    "\n",
    "AF2_recycles = 3\n",
    "AF2_models = \"4\"  # add other models to this string if needed\n",
    "\n",
    "commands_af2 = []\n",
    "cmds_filename_af2 = \"commands_af2\"\n",
    "with open(cmds_filename_af2, \"w\") as file:\n",
    "    for ff in glob.glob(\"*.fasta\"):\n",
    "        commands_af2.append(f\"{PYTHON['af2']} {AF2_script} \"\n",
    "                             f\"--af-nrecycles {AF2_recycles} --af-models {AF2_models} \"\n",
    "                             f\"--fasta {ff} --scorefile {ff.replace('.fasta', '.csv')}\\n\")\n",
    "        file.write(commands_af2[-1])\n",
    "\n",
    "print(\"Example AF2 command:\")\n",
    "print(commands_af2[-1])\n",
    "\n",
    "### Running AF2 with Slurm.\n",
    "### Running jobs on the CPU. It takes ~10 minutes per sequence\n",
    "\n",
    "submit_script = \"submit_af2.sh\"\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.1_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=2, gpu=True, gres=\"gpu:rtx2080:1\", time=\"00:12:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "else:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.1_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=4, time=\"01:00:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98095a31",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with AF2 and happy with the outputs then mark it as done\n",
    "AF2_DIR = f\"{WDIR}/5.1_2nd_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/3.1_design_pocket_ligandMPNN\"\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    with open(f\"{AF2_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c770263c",
   "metadata": {},
   "source": [
    "### Analyzing AlphaFold2 predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2fd05f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combining all CSV scorefiles into one\n",
    "os.system(\"head -n 1 $(ls *aa*.csv | shuf -n 1) > scores.csv ; for f in *aa*.csv ; do tail -n +2 ${f} >> scores.csv ; done\")\n",
    "assert os.path.exists(\"scores.csv\"), \"Could not combine scorefiles\"\n",
    "\n",
    "### Calculating the RMSDs of AF2 predictions relative to the diffusion outputs\n",
    "### Catalytic residue sidechain RMSDs are calculated in the reference PDB has REMARK 666 line present\n",
    "\n",
    "analysis_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/analyze_af2.py --scorefile scores.csv \"\\\n",
    "               f\"--ref_path {DESIGN_DIR_ligMPNN}/good/ --mpnn --params {' '.join(params)}\"\n",
    "\n",
    "p = subprocess.Popen(analysis_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8512c33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualizing and filtering AF2 results\n",
    "AF2_DIR = f\"{WDIR}/5.1_2nd_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "scores_af2 = pd.read_csv(\"scores.sc\", sep=\"\\s+\", header=0)\n",
    "\n",
    "### Filtering AF2 scores based on lddt and rmsd\n",
    "# Define your desired cutoffs here:\n",
    "AF2_filters = {\"lDDT\": [85.0, \">=\"],\n",
    "               \"rmsd\": [1.2, \"<=\"],\n",
    "               \"rmsd_SR1\": [1.0, \"<=\"]}  # 1st catalytic residue sc-rmsd\n",
    "\n",
    "scores_af2_filtered = utils.filter_scores(scores_af2, AF2_filters)\n",
    "utils.dump_scorefile(scores_af2_filtered, \"filtered_scores.sc\")\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 3))\n",
    "for i,k in enumerate(AF2_filters):\n",
    "    plt.subplot(1, 3, i+1)\n",
    "    plt.hist(scores_af2[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "utils.plot_score_pairs(scores_af2, \"lDDT\", \"rmsd\", AF2_filters[\"lDDT\"][0], AF2_filters[\"rmsd\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20de13ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Copying good predictions to a separate directory\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if len(scores_af2_filtered) > 0:\n",
    "    os.makedirs(\"good\", exist_ok=True)\n",
    "    good_af2_models = [row[\"Output_PDB\"]+\".pdb\" for idx,row in scores_af2_filtered.iterrows()]\n",
    "    for pdb in good_af2_models:\n",
    "        copy2(pdb, f\"good/{pdb}\")\n",
    "    good_af2_models = glob.glob(f\"{AF2_DIR}/good/*.pdb\")\n",
    "else:\n",
    "    sys.exit(\"No good models to continue this pipeline with\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bc2ddec",
   "metadata": {},
   "source": [
    "## 6.1: Final FastRelax with the ligand\n",
    "Relaxing good AF2 models together with the ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "268bceef",
   "metadata": {},
   "outputs": [],
   "source": [
    "AF2_DIR = f\"{WDIR}/5.1_2nd_af2\"\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/3.1_design_pocket_ligandMPNN\"\n",
    "assert len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")) > 0, \"No good AF2 models to relax with\"\n",
    "\n",
    "os.chdir(WDIR)\n",
    "RELAX_DIR = f\"{WDIR}/6.1_final_relax\"\n",
    "os.makedirs(RELAX_DIR, exist_ok=True)\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "os.makedirs(RELAX_DIR+\"/logs\", exist_ok=True)\n",
    "\n",
    "## First matching up the AF2 output filenames of step 5 with pocket design filenames from step 3\n",
    "ref_and_model_pairs = []\n",
    "for r in glob.glob(f\"{DESIGN_DIR_ligMPNN}/good/*.pdb\"):\n",
    "    for pdbfile in glob.glob(f\"{AF2_DIR}/good/*.pdb\"):\n",
    "        if os.path.basename(r).replace(\".pdb\", \"_\") in pdbfile:\n",
    "            ref_and_model_pairs.append((r, pdbfile))\n",
    "\n",
    "assert len(ref_and_model_pairs) == len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")), \"Was not able to match all models with reference structures\"\n",
    "\n",
    "\n",
    "## Generating commands for relax jobs\n",
    "### Performing 1 relax iteration on each input structure\n",
    "NSTRUCT = 1\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_relax = []\n",
    "cmds_filename_rlx = \"commands_relax\"\n",
    "with open(cmds_filename_rlx, \"w\") as file:\n",
    "    for r_m in ref_and_model_pairs:\n",
    "        commands_relax.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/align_add_ligand_relax.py \"\n",
    "                              f\"--outdir ./ --ligand {LIGAND} --ref_pdb {r_m[0]} \"\n",
    "                              f\"--pdb {r_m[1]} --nstruct {NSTRUCT} \"\n",
    "                              f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(r_m[1]).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_relax[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_relax[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_relax.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"6.1_final_relax\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"0:30:00\", email=EMAIL, array=len(commands_relax),\n",
    "                                 array_commandfile=cmds_filename_rlx)\n",
    "\n",
    "if not os.path.exists(RELAX_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "366ce521",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with final relax and happy with the outputs then mark it as done\n",
    "RELAX_DIR = f\"{WDIR}/6.1_final_relax\"\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "if not os.path.exists(RELAX_DIR+\"/.done\"):\n",
    "    with open(f\"{RELAX_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d848dbbf",
   "metadata": {},
   "source": [
    "### Analyzing final relaxed structures\n",
    "Filtering them based on the same metrics as was used for the initial design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3744d0a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Analyzing Rosetta designs\n",
    "RELAX_DIR = f\"{WDIR}/6.1_final_relax\"\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>='],\n",
    " 'rmsd_CA_rlx_in': [1.0, \"<=\"]}  # rmsd_CA_rlx_in is rmsd between relaxed structure and AF2 prediction\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting relax scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    if k not in scores.keys():\n",
    "        continue\n",
    "    plt.subplot(4, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{RELAX_DIR}/good\", exist_ok=True)\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e448ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(filtered_scores) > 0:\n",
    "    print(f\"CONGRATULATIONS! You have successfully designed {len(filtered_scores)} proteins against ligand {LIGAND}\")\n",
    "    print(\"You can find the design models in the directory:\\n\"\n",
    "          f\"    {RELAX_DIR}/good\")\n",
    "    print(\"\\nIt is advised you manually inspect them before ordering.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad026b16",
   "metadata": {},
   "source": [
    "## \n",
    "## --------------------------------------------------------------------------------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02646039",
   "metadata": {},
   "source": [
    "<i>The following part of the pipeline branches off after the first AlphaFold2 step (2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37afe717",
   "metadata": {},
   "source": [
    "## 3.2: Performing binding site design with Rosetta FastDesign\n",
    "<a id='Rosetta_design'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d77879a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting up design directory and commands\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_ROSETTA = f\"{WDIR}/3.2_design_pocket_Rosetta\"\n",
    "os.makedirs(DESIGN_DIR_ROSETTA, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_ROSETTA)\n",
    "AF2_DIR = f\"{WDIR}/2_af2\"\n",
    "\n",
    "os.makedirs(DESIGN_DIR_ROSETTA+\"/logs\", exist_ok=True)\n",
    "\n",
    "### Performing 5 design iterations on each input structure\n",
    "NSTRUCT = 5\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_design = []\n",
    "cmds_filename_des = \"commands_design\"\n",
    "with open(cmds_filename_des, \"w\") as file:\n",
    "    for pdb in glob.glob(f\"{AF2_DIR}/good/with_heme/*.pdb\"):\n",
    "        commands_design.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/heme_pocket_FastDesign.py \"\n",
    "                             f\"--pdb {pdb} --norelax --nstruct {NSTRUCT} \"\n",
    "                             f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(pdb).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_design[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_design[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_design.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"3.2_design_pocket_Rosetta\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"3:00:00\", email=EMAIL, array=len(commands_design),\n",
    "                                 array_commandfile=cmds_filename_des)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_ROSETTA+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85082bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with design and happy with the outputs then mark it as done\n",
    "DESIGN_DIR_ROSETTA = f\"{WDIR}/3.2_design_pocket_Rosetta\"\n",
    "os.chdir(DESIGN_DIR_ROSETTA)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_ROSETTA+\"/.done\"):\n",
    "    with open(f\"{DESIGN_DIR_ROSETTA}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8512980",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Analyzing Rosetta designs\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>=']}\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    plt.subplot(3, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{DESIGN_DIR_ROSETTA}/good\", exist_ok=True)\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d049a3a8",
   "metadata": {},
   "source": [
    "## 4.2 Performing ligandMPNN redesign on the 2nd layer residues\n",
    "\n",
    "Resampling residues that are not in the pocket, but also not very far from the pocket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af9fe184",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(glob.glob(DESIGN_DIR_ROSETTA+'/good/*.pdb')) == 0:\n",
    "    sys.exit(\"No designs to run 2nd MPNN on.\")\n",
    "\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.2_2nd_mpnn\"\n",
    "os.makedirs(DESIGN_DIR_2nd_mpnn, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_2nd_mpnn)\n",
    "\n",
    "### Making a JSON file specifiying designable positions for each structure.\n",
    "### Will also make non-pocket ALA positions as designable.\n",
    "### This is to fix any surface ALA-patches that previous MPNN may have introduced.\n",
    "\n",
    "make_json_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/setup_ligand_mpnn_2nd_layer.py \"\\\n",
    "                f\"--params {' '.join(params)} --ligand {LIGAND} --output_path parsed_pdbs_lig.jsonl \"\\\n",
    "                 \"--output_path masked_pos.jsonl --dist_bb 6.0 --dist_sc 5.0 \"\\\n",
    "                f\"--pdb {' '.join(glob.glob(DESIGN_DIR_ROSETTA+'/good/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(make_json_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n",
    "\n",
    "\n",
    "### Setting up ligandMPNN run commands\n",
    "## We're doing design with 2 temperatures (more conservative than before), and 5 sequences each.\n",
    "\n",
    "MPNN_temperatures = [0.1, 0.2]\n",
    "MPNN_outputs_per_temperature = 5\n",
    "MPNN_omit_AAs = \"CM\"\n",
    "\n",
    "commands_mpnn = []\n",
    "cmds_filename_mpnn = \"commands_mpnn\"\n",
    "with open(cmds_filename_mpnn, \"w\") as file:\n",
    "    for T in MPNN_temperatures:\n",
    "        for f in glob.glob(DESIGN_DIR_ROSETTA+'/good/*.pdb'):\n",
    "            commands_mpnn.append(f\"{PYTHON['proteinMPNN']} {proteinMPNN_script} \"\n",
    "                                 f\"--model_type ligand_mpnn --ligand_mpnn_use_atom_context 1 \"\n",
    "                                 \"--fixed_residues_multi masked_pos.jsonl --out_folder ./ \"\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 {SCRIPT_DIR}/lib/LigandMPNN/model_params/ligandmpnn_v_32_010_25.pt\\n\")\n",
    "            file.write(commands_mpnn[-1])\n",
    "\n",
    "print(\"Example MPNN command:\")\n",
    "print(commands_mpnn[-1])\n",
    "\n",
    "\n",
    "### Running ligandMPNN with Slurm.\n",
    "### Grouping jobs with 10 commands per one array job.\n",
    "submit_script = \"submit_mpnn.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"4.2_2nd_mpnn\", mem=\"8g\", \n",
    "                                 N_cores=1, time=\"0:15:00\", email=EMAIL, array=len(commands_mpnn),\n",
    "                                 array_commandfile=cmds_filename_mpnn, group=10)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d4592e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with design and happy with the outputs then mark it as done\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.2_2nd_mpnn\"\n",
    "os.chdir(DESIGN_DIR_2nd_mpnn)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"):\n",
    "    with open(f\"{DESIGN_DIR_2nd_mpnn}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b40798b",
   "metadata": {},
   "source": [
    "## 5.2 AlphaFold2 predictions on the 2nd MPNN run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97ed25d",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.2_2nd_mpnn\"\n",
    "assert os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"), \"2nd MPNN has not been performed!\"\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/5.2_2nd_af2\"\n",
    "os.makedirs(AF2_DIR, exist_ok=True)\n",
    "os.chdir(AF2_DIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/4.2_2nd_mpnn\"\n",
    "\n",
    "### First collecting MPNN outputs and creating FASTA files for AF2 input\n",
    "mpnn_fasta = utils.parse_fasta_files(glob.glob(f\"{DESIGN_DIR_2nd_mpnn}/seqs/*.fa\"))\n",
    "# Giving sequences unique names based on input PDB name, temperature, and sequence identifier\n",
    "_mpnn_fasta = {}\n",
    "for k, seq in mpnn_fasta.items():\n",
    "    if \"model_path\" in k:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_native\"] = seq.strip()\n",
    "    else:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_\"+k.split(\",\")[2].replace(\" T=\", \"T\")+\"_0_\"+k.split(\",\")[1].replace(\" id=\", \"\")] = seq.strip()\n",
    "mpnn_fasta = {k:v for k,v in _mpnn_fasta.items()}\n",
    "\n",
    "print(f\"A total on {len(mpnn_fasta)} sequences will be predicted.\")\n",
    "\n",
    "## Splitting the MPNN sequences based on length\n",
    "## and grouping them in smaller batches for each AF2 job\n",
    "## Use group size of >40 when running on GPU. Also depends on how many sequences and resources you have.\n",
    "\n",
    "SEQUENCES_PER_AF2_JOB = 5  # CPU\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    SEQUENCES_PER_AF2_JOB = 100  # GPUSEQUENCES_PER_AF2_JOBSEQUENCES_PER_AF2_JOBSEQUENCES_PER_AF2_JOB\n",
    "mpnn_fasta_split = utils.split_fasta_based_on_length(mpnn_fasta, SEQUENCES_PER_AF2_JOB, write_files=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e01021a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up AlphaFold2 run\n",
    "\n",
    "AF2_recycles = 3\n",
    "AF2_models = \"4\"  # add other models to this string if needed\n",
    "\n",
    "commands_af2 = []\n",
    "cmds_filename_af2 = \"commands_af2\"\n",
    "with open(cmds_filename_af2, \"w\") as file:\n",
    "    for ff in glob.glob(\"*.fasta\"):\n",
    "        commands_af2.append(f\"{PYTHON['af2']} {AF2_script} \"\n",
    "                             f\"--af-nrecycles {AF2_recycles} --af-models {AF2_models} \"\n",
    "                             f\"--fasta {ff} --scorefile {ff.replace('.fasta', '.csv')}\\n\")\n",
    "        file.write(commands_af2[-1])\n",
    "\n",
    "print(\"Example AF2 command:\")\n",
    "print(commands_af2[-1])\n",
    "\n",
    "\n",
    "### Running AF2 with Slurm.\n",
    "\n",
    "submit_script = \"submit_af2.sh\"\n",
    "\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.2_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=2, gpu=True, gres=\"gpu:rtx2080:1\", time=\"00:12:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "else:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.2_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=4, time=\"01:00:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25adc9d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with AF2 and happy with the outputs then mark it as done\n",
    "AF2_DIR = f\"{WDIR}/5.2_2nd_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    with open(f\"{AF2_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08bcb115",
   "metadata": {},
   "source": [
    "### Analyzing AlphaFold2 predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b116b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combining all CSV scorefiles into one\n",
    "DESIGN_DIR_ROSETTA = f\"{WDIR}/3.2_design_pocket_Rosetta\"\n",
    "\n",
    "os.system(\"head -n 1 $(ls *aa*.csv | shuf -n 1) > scores.csv ; for f in *aa*.csv ; do tail -n +2 ${f} >> scores.csv ; done\")\n",
    "assert os.path.exists(\"scores.csv\"), \"Could not combine scorefiles\"\n",
    "\n",
    "### Calculating the RMSDs of AF2 predictions relative to the diffusion outputs\n",
    "### Catalytic residue sidechain RMSDs are calculated in the reference PDB has REMARK 666 line present\n",
    "\n",
    "analysis_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/analyze_af2.py --scorefile scores.csv \"\\\n",
    "               f\"--ref_path {DESIGN_DIR_ROSETTA}/good/ --mpnn --params {' '.join(params)}\"\n",
    "\n",
    "p = subprocess.Popen(analysis_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0021fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores_af2 = pd.read_csv(\"scores.sc\", sep=\"\\s+\", header=0)\n",
    "\n",
    "### Filtering AF2 scores based on lddt and rmsd\n",
    "# Define your desired cutoffs here:\n",
    "AF2_filters = {\"lDDT\": [85.0, \">=\"],\n",
    "               \"rmsd\": [1.2, \"<=\"],\n",
    "               \"rmsd_SR1\": [1.0, \"<=\"]}  # 1st catalytic residue sc-rmsd\n",
    "\n",
    "scores_af2_filtered = utils.filter_scores(scores_af2, AF2_filters)\n",
    "utils.dump_scorefile(scores_af2_filtered, \"filtered_scores.sc\")\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 3))\n",
    "for i,k in enumerate(AF2_filters):\n",
    "    plt.subplot(1, 3, i+1)\n",
    "    plt.hist(scores_af2[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "    \n",
    "utils.plot_score_pairs(scores_af2, \"lDDT\", \"rmsd\", AF2_filters[\"lDDT\"][0], AF2_filters[\"rmsd\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f27212",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "### Copying good predictions to a separate directory\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if len(scores_af2_filtered) > 0:\n",
    "    os.makedirs(\"good\", exist_ok=True)\n",
    "    good_af2_models = [row[\"Output_PDB\"]+\".pdb\" for idx,row in scores_af2_filtered.iterrows()]\n",
    "    for pdb in good_af2_models:\n",
    "        copy2(pdb, f\"good/{pdb}\")\n",
    "    good_af2_models = glob.glob(f\"{AF2_DIR}/good/*.pdb\")\n",
    "else:\n",
    "    sys.exit(\"No good models to continue this pipeline with\")\n",
    "\n",
    "os.chdir(f\"{AF2_DIR}/good\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2d775a",
   "metadata": {},
   "source": [
    "## 6.2: Final FastRelax with the ligand\n",
    "Relaxing good AF2 models together with the ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c7118b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")) > 0, \"No good AF2 models to relax with\"\n",
    "\n",
    "os.chdir(WDIR)\n",
    "RELAX_DIR = f\"{WDIR}/6.2_final_relax\"\n",
    "os.makedirs(RELAX_DIR, exist_ok=True)\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "DESIGN_DIR_ROSETTA = f\"{WDIR}/3.2_design_pocket_Rosetta\"\n",
    "AF2_DIR = f\"{WDIR}/5.2_2nd_af2\"\n",
    "os.makedirs(RELAX_DIR+\"/logs\", exist_ok=True)\n",
    "\n",
    "## First matching up the AF2 output filenames of step 5 with pocket design filenames from step 3\n",
    "ref_and_model_pairs = []\n",
    "for r in glob.glob(f\"{DESIGN_DIR_ROSETTA}/good/*.pdb\"):\n",
    "    for pdbfile in glob.glob(f\"{AF2_DIR}/good/*.pdb\"):\n",
    "        if os.path.basename(r).replace(\".pdb\", \"_\") in pdbfile:\n",
    "            ref_and_model_pairs.append((r, pdbfile))\n",
    "\n",
    "assert len(ref_and_model_pairs) == len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")), \"Was not able to match all models with reference structures\"\n",
    "\n",
    "\n",
    "## Generating commands for relax jobs\n",
    "### Performing 1 relax iteration on each input structure\n",
    "NSTRUCT = 1\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_relax = []\n",
    "cmds_filename_rlx = \"commands_relax\"\n",
    "with open(cmds_filename_rlx, \"w\") as file:\n",
    "    for r_m in ref_and_model_pairs:\n",
    "        commands_relax.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/align_add_ligand_relax.py \"\n",
    "                              f\"--outdir ./ --ligand {LIGAND} --ref_pdb {r_m[0]} \"\n",
    "                              f\"--pdb {r_m[1]} --nstruct {NSTRUCT} \"\n",
    "                              f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(r_m[1]).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_relax[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_relax[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_relax.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"6.2_final_relax\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"0:30:00\", email=EMAIL, array=len(commands_relax),\n",
    "                                 array_commandfile=cmds_filename_rlx)\n",
    "\n",
    "if not os.path.exists(RELAX_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8778db",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with final relax and happy with the outputs then mark it as done\n",
    "RELAX_DIR = f\"{WDIR}/6.2_final_relax\"\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "if not os.path.exists(RELAX_DIR+\"/.done\"):\n",
    "    with open(f\"{RELAX_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be98d5d9",
   "metadata": {},
   "source": [
    "### Analyzing final relaxed structures\n",
    "Filtering them based on the same metrics as was used for the initial design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d86fca0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Analyzing Rosetta designs\n",
    "RELAX_DIR = f\"{WDIR}/6.2_final_relax\"\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>='],\n",
    " 'rmsd_CA_rlx_in': [1.0, \"<=\"]}  # rmsd_CA_rlx_in is rmsd between relaxed structure and AF2 prediction\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting relax scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    if k not in scores.keys():\n",
    "        continue\n",
    "    plt.subplot(3, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{RELAX_DIR}/good\", exist_ok=True)\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f0c05b",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(filtered_scores) > 0:\n",
    "    print(f\"CONGRATULATIONS! You have successfully designed {len(filtered_scores)} proteins against ligand {LIGAND}\")\n",
    "    print(\"You can find the design models in the directory:\\n\"\n",
    "          f\"    {RELAX_DIR}/good\")\n",
    "    print(\"\\nIt is advised you manually inspect them before ordering.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb6b45e",
   "metadata": {},
   "source": [
    "## \n",
    "## --------------------------------------------------------------------------------------------------------------------\n",
    "<a id='inpainting'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5fd5f9e",
   "metadata": {},
   "source": [
    "## RFjoint Inpainting diversification of diffusion outputs\n",
    "This part of the pipeline branches off right after diffusion.\n",
    "\n",
    "The purpose of inpainting diversification is to resample the loop regions of diffused backbones. Any region that is close to the ligand will not be touched because inpainting model does not see the ligand. Resampling is done with variable sequence length.\n",
    "\n",
    "For some reason the AF2 success rates may go significantly up when performing inpainting on diffused backbones.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "012ef367",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Identifying loop regions in diffused backbones and setting up input JSON files for inpainting\n",
    "assert os.path.exists(f\"{DIFFUSION_DIR}/filtered_structures\"), \"No diffused backbones to run inpainting on\"\n",
    "INPAINT_DIR = f\"{WDIR}/inpainting/0_inpainting\"\n",
    "os.makedirs(INPAINT_DIR, exist_ok=True)\n",
    "os.chdir(INPAINT_DIR)\n",
    "\n",
    "\n",
    "inpaint_inputs_per_job = 2  # How many structures will go into each arrayed job\n",
    "inpaint_catalytic_residues = \"A15\"  # Catalytic residues that were used as motif during diffusion, space-separated string\n",
    "\n",
    "inpaint_setup_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/diffusion_analysis/setup_inpaint_from_diffusion.py \"\\\n",
    "               f\"--group {inpaint_inputs_per_job} --var --design_full --params {' '.join(params)} \"\\\n",
    "               f\"--ref_catres {inpaint_catalytic_residues} \"\\\n",
    "               f\"--pdb {' '.join(glob.glob(f'{DIFFUSION_DIR}/filtered_structures/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(inpaint_setup_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "888fd826",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting up inpainting run commands\n",
    "assert os.path.exists(f\"{DIFFUSION_DIR}/filtered_structures\"), \"No diffused backbones to run inpainting on\"\n",
    "\n",
    "commands_inpaint = []\n",
    "cmds_filename = \"commands_inpaint\"\n",
    "with open(cmds_filename, \"w\") as file:\n",
    "    for jsn in glob.glob(f\"{INPAINT_DIR}/*.json\"):\n",
    "        commands_inpaint.append(f\"{PYTHON['diffusion']} {inpaint_script} \"\n",
    "                             f\"--input_json {jsn} > {jsn.replace('.json', '.log')}\\n\")\n",
    "        file.write(commands_inpaint[-1])\n",
    "\n",
    "print(\"Example inpainting command:\")\n",
    "print(commands_inpaint[-1])\n",
    "\n",
    "\n",
    "### Running inpainting jobs with Slurm.\n",
    "submit_script = \"submit_inpaint.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"00_inpainting\", mem=\"6g\", gpu=True, gres=\"gpu:a4000:1\", \n",
    "                                 N_cores=2, time=\"0:30:00\", email=EMAIL, array=len(commands_inpaint),\n",
    "                                 array_commandfile=cmds_filename)\n",
    "\n",
    "if not os.path.exists(INPAINT_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()\n",
    "    print(output)\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04886d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with inpainting and happy with the outputs then mark it as done\n",
    "INPAINT_DIR = f\"{WDIR}/inpainting/0_inpainting\"\n",
    "os.chdir(INPAINT_DIR)\n",
    "\n",
    "if not os.path.exists(INPAINT_DIR+\"/.done\"):\n",
    "    with open(f\"{INPAINT_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6daa91b7",
   "metadata": {},
   "source": [
    "### Aligning the ligand into inpainting outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d14a1eb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Inpainting outputs will be aligned to diffusion outputs based on constant regions defined in the inpainting TRB file.\n",
    "### Ligand will be placed into the inpaining model, \n",
    "### and REMARK 666 lines defining any catalytic residues will be added to the PDB file.\n",
    "\n",
    "assert len(glob.glob(f\"{INPAINT_DIR}/*.pdb\")) > 0, \"No inpainted backbones to process\"\n",
    "os.makedirs(INPAINT_DIR+\"/alignment\", exist_ok=True)\n",
    "\n",
    "\n",
    "inpaint_align_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/diffusion_analysis/align_dif_inpaint_add_ligand.py \"\\\n",
    "               f\"--outdir alignment --params {' '.join(params)} \"\\\n",
    "               f\"--ref {' '.join(glob.glob(f'{DIFFUSION_DIR}/filtered_structures/*.pdb'))} \"\\\n",
    "               f\"--pdb {' '.join(glob.glob(f'{INPAINT_DIR}/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(inpaint_align_cmd, shell=True)\n",
    "(output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21f1bb51",
   "metadata": {},
   "source": [
    "## 1: Running ProteinMPNN on inpainted backbones\n",
    "\n",
    "We are first trying to just design a sequence on the backbone, without considering the ligand.\n",
    "The goal is to first find backbones that fold well.\n",
    "\n",
    "Alternatively, one could also go ahead and run the [LigandMPNN/FastRelax step](#ligmpnn_fr) already on the inpainted backbones.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "def0617c",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPAINT_DIR = f\"{WDIR}/inpainting/0_inpainting\"\n",
    "assert len(glob.glob(INPAINT_DIR+\"/alignment/*.pdb\")) > 0, \"No good backbones found!\"\n",
    "\n",
    "os.chdir(WDIR)\n",
    "\n",
    "MPNN_DIR = f\"{WDIR}/inpainting/1_proteinmpnn\"\n",
    "os.makedirs(MPNN_DIR, exist_ok=True)\n",
    "os.chdir(MPNN_DIR)\n",
    "\n",
    "### Parsing diffusion output TRB files to extract fixed motif residues\n",
    "## These residues will not be redesigned with proteinMPNN\n",
    "mask_json_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/make_maskdict_from_trb.py \"\\\n",
    "                f\"--out masked_pos.jsonl --trb {' '.join(glob.glob(INPAINT_DIR+'/*.trb'))}\"\n",
    "p = subprocess.Popen(mask_json_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n",
    "\n",
    "assert os.path.exists(\"masked_pos.jsonl\"), \"Failed to create masked positions JSONL file\"\n",
    "\n",
    "\n",
    "### Setting up proteinMPNN run commands\n",
    "## We're doing design with 3 temperatures, and 5 sequences each.\n",
    "## This usually gives decent success with designable backbones. For more complicated cases consider doing >100 sequences.\n",
    "\n",
    "MPNN_temperatures = [0.1, 0.2, 0.3]\n",
    "MPNN_outputs_per_temperature = 5\n",
    "MPNN_omit_AAs = \"CM\"\n",
    "\n",
    "commands_mpnn = []\n",
    "cmds_filename_mpnn = \"commands_mpnn\"\n",
    "with open(cmds_filename_mpnn, \"w\") as file:\n",
    "    for T in MPNN_temperatures:\n",
    "        for f in glob.glob(INPAINT_DIR+\"/alignment/*.pdb\"):\n",
    "            commands_mpnn.append(f\"{PYTHON['proteinMPNN']} {proteinMPNN_script} \"\n",
    "                                 f\"--model_type protein_mpnn --ligand_mpnn_use_atom_context 0 \"\n",
    "                                 \"--fixed_residues_multi masked_pos.jsonl --out_folder ./ \"\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 {SCRIPT_DIR}/lib/LigandMPNN/model_params/proteinmpnn_v_48_020.pt\\n\")\n",
    "            file.write(commands_mpnn[-1])\n",
    "\n",
    "            \n",
    "            \n",
    "print(f\"{len(commands_mpnn)} MPNN jobs to run\")\n",
    "print(\"Example MPNN command:\")\n",
    "print(commands_mpnn[-1])\n",
    "\n",
    "### Running proteinMPNN with Slurm.\n",
    "### Grouping jobs with 10 commands per one array job.\n",
    "\n",
    "submit_script = \"submit_mpnn.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"1_proteinmpnn\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"0:15:00\", email=EMAIL, array=len(commands_mpnn),\n",
    "                                 array_commandfile=cmds_filename_mpnn, group=10)\n",
    "\n",
    "if not os.path.exists(MPNN_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a57ab8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with MPNN and happy with the outputs then mark it as done\n",
    "MPNN_DIR = f\"{WDIR}/inpainting/1_proteinmpnn\"\n",
    "os.chdir(MPNN_DIR)\n",
    "\n",
    "if not os.path.exists(MPNN_DIR+\"/.done\"):\n",
    "    with open(f\"{MPNN_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f19d9ce",
   "metadata": {},
   "source": [
    "## 2: Running AlphaFold2\n",
    "Performing AF2 single sequence predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8e267d",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(WDIR)\n",
    "assert len(glob.glob(MPNN_DIR+\"/seqs/*.fa\")) > 0, \"No MPNN outputs to run AF2 on\"\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/inpainting/2_af2\"\n",
    "os.makedirs(AF2_DIR, exist_ok=True)\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "### First collecting MPNN outputs and creating FASTA files for AF2 input\n",
    "mpnn_fasta = utils.parse_fasta_files(glob.glob(f\"{MPNN_DIR}/seqs/*.fa\"))\n",
    "mpnn_fasta = {k: seq.strip() for k, seq in mpnn_fasta.items() if \"model_path\" not in k}  # excluding the diffused poly-A sequence\n",
    "# Giving sequences unique names based on input PDB name, temperature, and sequence identifier\n",
    "mpnn_fasta = {k.split(\",\")[0]+\"_\"+k.split(\",\")[2].replace(\" T=\", \"T\")+\"_0_\"+k.split(\",\")[1].replace(\" id=\", \"\"): seq for k, seq in mpnn_fasta.items()}\n",
    "\n",
    "print(f\"A total on {len(mpnn_fasta)} sequences will be predicted.\")\n",
    "\n",
    "## Splitting the MPNN sequences based on length\n",
    "## and grouping them in smaller batches for each AF2 job\n",
    "## Use group size of >40 when running on GPU. Also depends on how many sequences and resources you have.\n",
    "SEQUENCES_PER_AF2_JOB = 4  # CPU\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    SEQUENCES_PER_AF2_JOB = 100  # GPU\n",
    "mpnn_fasta_split = utils.split_fasta_based_on_length(mpnn_fasta, SEQUENCES_PER_AF2_JOB, write_files=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "411d5b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up AlphaFold2 run\n",
    "\n",
    "AF2_recycles = 3\n",
    "AF2_models = \"4\"  # add other models to this string if needed, i.e. \"3 4 5\"\n",
    "\n",
    "commands_af2 = []\n",
    "cmds_filename_af2 = \"commands_af2\"\n",
    "with open(cmds_filename_af2, \"w\") as file:\n",
    "    for ff in glob.glob(\"*.fasta\"):\n",
    "        commands_af2.append(f\"{PYTHON['af2']} {AF2_script} \"\n",
    "                             f\"--af-nrecycles {AF2_recycles} --af-models {AF2_models} \"\n",
    "                             f\"--fasta {ff} --scorefile {ff.replace('.fasta', '.csv')}\\n\")\n",
    "        file.write(commands_af2[-1])\n",
    "\n",
    "print(\"Example AF2 command:\")\n",
    "print(commands_af2[-1])\n",
    "\n",
    "### Running AF2 with Slurm.\n",
    "### Running jobs on the CPU. It takes ~10 minutes per sequence\n",
    "### \n",
    "\n",
    "submit_script = \"submit_af2.sh\"\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"2_af2\", mem=\"6g\", \n",
    "                                     N_cores=2, gpu=True, gres=\"gpu:rtx2080:1\", time=\"00:12:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "else:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"2_af2\", mem=\"6g\", \n",
    "                                     N_cores=4, time=\"01:00:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0c7b76",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with AF2 and happy with the outputs then mark it as done\n",
    "AF2_DIR = f\"{WDIR}/inpainting/2_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    with open(f\"{AF2_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184907b7",
   "metadata": {},
   "source": [
    "### Analyzing AlphaFold2 predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c85aeef",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPAINT_DIR = f\"{WDIR}/inpainting/0_inpainting\"\n",
    "\n",
    "# Combining all CSV scorefiles into one\n",
    "os.system(\"head -n 1 $(ls *aa*.csv | shuf -n 1) > scores.csv ; for f in *aa*.csv ; do tail -n +2 ${f} >> scores.csv ; done\")\n",
    "assert os.path.exists(\"scores.csv\"), \"Could not combine scorefiles\"\n",
    "\n",
    "### Calculating the RMSDs of AF2 predictions relative to the diffusion outputs\n",
    "### Catalytic residue sidechain RMSDs are calculated in the reference PDB has REMARK 666 line present\n",
    "\n",
    "analysis_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/analyze_af2.py --scorefile scores.csv \"\\\n",
    "               f\"--ref_path {INPAINT_DIR}/alignment/ --mpnn --params {' '.join(params)}\"\n",
    "\n",
    "p = subprocess.Popen(analysis_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146e5a18",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualizing and filtering AF2 predictions\n",
    "AF2_DIR = f\"{WDIR}/inpainting/2_af2\"\n",
    "\n",
    "scores_af2 = pd.read_csv(\"scores.sc\", sep=\"\\s+\", header=0)\n",
    "\n",
    "### Filtering AF2 scores based on lddt and rmsd\n",
    "# Define your desired cutoffs here:\n",
    "AF2_filters = {\"lDDT\": [85.0, \">=\"],\n",
    "               \"rmsd\": [1.5, \"<=\"],\n",
    "               \"rmsd_SR1\": [2.0, \"<=\"]}  # 1st catalytic residue sc-rmsd\n",
    "\n",
    "scores_af2_filtered = utils.filter_scores(scores_af2, AF2_filters)\n",
    "utils.dump_scorefile(scores_af2_filtered, \"filtered_scores.sc\")\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 3))\n",
    "for i,k in enumerate(AF2_filters):\n",
    "    plt.subplot(1, 3, i+1)\n",
    "    plt.hist(scores_af2[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "    \n",
    "utils.plot_score_pairs(scores_af2, \"lDDT\", \"rmsd\", AF2_filters[\"lDDT\"][0], AF2_filters[\"rmsd\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ac5be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Copying good predictions to a separate directory\n",
    "os.chdir(AF2_DIR)\n",
    "INPAINT_DIR = f\"{WDIR}/inpainting/0_inpainting\"\n",
    "\n",
    "if len(scores_af2_filtered) > 0:\n",
    "    os.makedirs(\"good\", exist_ok=True)\n",
    "    good_af2_models = [row[\"Output_PDB\"]+\".pdb\" for idx,row in scores_af2_filtered.iterrows()]\n",
    "    for pdb in good_af2_models:\n",
    "        copy2(pdb, f\"good/{pdb}\")\n",
    "    good_af2_models = glob.glob(f\"{AF2_DIR}/good/*.pdb\")\n",
    "else:\n",
    "    sys.exit(\"No good models to continue this pipeline with\")\n",
    "\n",
    "os.chdir(f\"{AF2_DIR}/good\")\n",
    "\n",
    "\n",
    "### Aligning the ligand back into the AF2 predictions.\n",
    "### This is done by aligning the AF2 model to diffusion output and copying over the ligand using PyRosetta.\n",
    "### --fix_catres option will re-adjust the rotamer and tautomer of \n",
    "### any catalytic residue to be the same as in the reference model.\n",
    "\n",
    "align_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/place_ligand_after_af2.py \"\\\n",
    "            f\"--outdir with_heme --params {' '.join(params)} --fix_catres \"\\\n",
    "            f\"--pdb {' '.join(good_af2_models)} \"\\\n",
    "            f\"--ref {' '.join(glob.glob(INPAINT_DIR+'/alignment/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(align_cmd, shell=True)\n",
    "(output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa9ad826",
   "metadata": {},
   "source": [
    "## 3.1: Performing binding site design with LigandMPNN / FastRelax\n",
    "<a id='ligmpnn_fr'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ad0fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting up design directory and commands\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/inpainting/3.1_design_pocket_ligandMPNN\"\n",
    "os.makedirs(DESIGN_DIR_ligMPNN, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_ligMPNN)\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/inpainting/2_af2\"  # change this if you want to run it on inpainting outputs\n",
    "os.makedirs(DESIGN_DIR_ligMPNN+\"/logs\", exist_ok=True)\n",
    "\n",
    "### Performing 5 design iterations on each input structure\n",
    "NSTRUCT = 10\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_design = []\n",
    "cmds_filename_des = \"commands_design\"\n",
    "with open(cmds_filename_des, \"w\") as file:\n",
    "    for pdb in glob.glob(f\"{AF2_DIR}/good/with_heme/*.pdb\"):\n",
    "        commands_design.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/heme_pocket_ligMPNN.py \"\n",
    "                             f\"--pdb {pdb} --nstruct {NSTRUCT} \"\n",
    "                             f\"--scoring {SCRIPT_DIR}/scripts/design/scoring/heme_scoring.py \"\n",
    "                             f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(pdb).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_design[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_design[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_design.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"3.1_design_pocket_ligMPNN\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"3:00:00\", email=EMAIL, array=len(commands_design),\n",
    "                                 array_commandfile=cmds_filename_des)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_ligMPNN+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77afc69",
   "metadata": {},
   "outputs": [],
   "source": [
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/inpainting/3.1_design_pocket_ligandMPNN\"\n",
    "os.chdir(DESIGN_DIR_ligMPNN)\n",
    "\n",
    "## If you're done with design and happy with the outputs then mark it as done\n",
    "if not os.path.exists(DESIGN_DIR_ligMPNN+\"/.done\"):\n",
    "    with open(f\"{DESIGN_DIR_ligMPNN}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcf6697",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Analyzing Rosetta designs\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>=']}\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting design scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    plt.subplot(3, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{DESIGN_DIR_ligMPNN}/good\", exist_ok=True)\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f56c82b3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8397b389",
   "metadata": {},
   "source": [
    "## 4.1 Performing ligandMPNN redesign on the 2nd layer residues\n",
    "\n",
    "Resampling residues that are not in the pocket, but also not very far from the pocket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac14eee",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if len(glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb')) == 0:\n",
    "    sys.exit(\"No designs to run 2nd MPNN on.\")\n",
    "\n",
    "os.chdir(WDIR)\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/inpainting/4.1_2nd_mpnn\"\n",
    "os.makedirs(DESIGN_DIR_2nd_mpnn, exist_ok=True)\n",
    "os.chdir(DESIGN_DIR_2nd_mpnn)\n",
    "\n",
    "### Making a JSON file specifiying designable positions for each structure.\n",
    "### Will also make non-pocket ALA positions as designable.\n",
    "### This is to fix any surface ALA-patches that previous MPNN may have introduced.\n",
    "\n",
    "make_json_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/setup_ligand_mpnn_2nd_layer.py \"\\\n",
    "                f\"--params {' '.join(params)} --ligand {LIGAND} --output_path parsed_pdbs_lig.jsonl \"\\\n",
    "                 \"--output_path masked_pos.jsonl --dist_bb 6.0 --dist_sc 5.0 \"\\\n",
    "                f\"--pdb {' '.join(glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb'))}\"\n",
    "\n",
    "p = subprocess.Popen(make_json_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n",
    "\n",
    "if not os.path.exists(\"masked_pos.jsonl\"):\n",
    "    sys.exit()\n",
    "\n",
    "### Setting up ligandMPNN run commands\n",
    "## We're doing design with 2 temperatures (more conservative than before), and 5 sequences each.\n",
    "\n",
    "MPNN_temperatures = [0.1, 0.2]\n",
    "MPNN_outputs_per_temperature = 5\n",
    "MPNN_omit_AAs = \"CM\"\n",
    "\n",
    "commands_mpnn = []\n",
    "cmds_filename_mpnn = \"commands_mpnn\"\n",
    "with open(cmds_filename_mpnn, \"w\") as file:\n",
    "    for T in MPNN_temperatures:\n",
    "        for f in glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb'):\n",
    "            commands_mpnn.append(f\"{PYTHON['proteinMPNN']} {proteinMPNN_script} \"\n",
    "                                 f\"--model_type ligand_mpnn --ligand_mpnn_use_atom_context 1 \"\n",
    "                                 \"--fixed_residues_multi masked_pos.jsonl --out_folder ./ \"\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 {SCRIPT_DIR}/lib/LigandMPNN/model_params/ligandmpnn_v_32_010_25.pt\\n\")\n",
    "            file.write(commands_mpnn[-1])\n",
    "\n",
    "print(\"Example MPNN command:\")\n",
    "print(commands_mpnn[-1])\n",
    "\n",
    "### Running ligandMPNN with Slurm.\n",
    "### Grouping jobs with 10 commands per one array job.\n",
    "\n",
    "submit_script = \"submit_mpnn.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"4.1_2nd_mpnn\", mem=\"8g\", \n",
    "                                 N_cores=1, time=\"0:15:00\", email=EMAIL, array=len(commands_mpnn),\n",
    "                                 array_commandfile=cmds_filename_mpnn, group=10)\n",
    "\n",
    "if not os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0233f278",
   "metadata": {},
   "outputs": [],
   "source": [
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/inpainting/3.1_design_pocket_ligandMPNN\"\n",
    "DESIGN_DIR_2nd_mpnn = f\"{WDIR}/inpainting/4.1_2nd_mpnn\"\n",
    "os.chdir(DESIGN_DIR_2nd_mpnn)\n",
    "if len(glob.glob(DESIGN_DIR_ligMPNN+'/good/*.pdb')) == 0:\n",
    "    sys.exit(\"No designs to run 2nd MPNN on.\")\n",
    "\n",
    "## If you're done with design and happy with the outputs then mark it as done\n",
    "if not os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"):\n",
    "    with open(f\"{DESIGN_DIR_2nd_mpnn}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d495bf96",
   "metadata": {},
   "source": [
    "## 5.1 AlphaFold2 predictions on the 2nd MPNN run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cba4381",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(WDIR)\n",
    "assert os.path.exists(DESIGN_DIR_2nd_mpnn+\"/.done\"), \"2nd MPNN has not been performed!\"\n",
    "\n",
    "AF2_DIR = f\"{WDIR}/inpainting/5.1_2nd_af2\"\n",
    "os.makedirs(AF2_DIR, exist_ok=True)\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "### First collecting MPNN outputs and creating FASTA files for AF2 input\n",
    "mpnn_fasta = utils.parse_fasta_files(glob.glob(f\"{DESIGN_DIR_2nd_mpnn}/seqs/*.fa\"))\n",
    "# Giving sequences unique names based on input PDB name, temperature, and sequence identifier\n",
    "_mpnn_fasta = {}\n",
    "for k, seq in mpnn_fasta.items():\n",
    "    if \"model_path\" in k:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_native\"] = seq.strip()\n",
    "    else:\n",
    "        _mpnn_fasta[k.split(\",\")[0]+\"_\"+k.split(\",\")[2].replace(\" T=\", \"T\")+\"_0_\"+k.split(\",\")[1].replace(\" id=\", \"\")] = seq.strip()\n",
    "mpnn_fasta = {k:v for k,v in _mpnn_fasta.items()}\n",
    "\n",
    "print(f\"A total on {len(mpnn_fasta)} sequences will be predicted.\")\n",
    "\n",
    "## Splitting the MPNN sequences based on length\n",
    "## and grouping them in smaller batches for each AF2 job\n",
    "## Use group size of >40 when running on GPU. Also depends on how many sequences and resources you have.\n",
    "SEQUENCES_PER_AF2_JOB = 5  # CPU\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    SEQUENCES_PER_AF2_JOB = 100  # GPU\n",
    "mpnn_fasta_split = utils.split_fasta_based_on_length(mpnn_fasta, SEQUENCES_PER_AF2_JOB, write_files=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b309a47d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setting up AlphaFold2 run\n",
    "\n",
    "AF2_recycles = 3\n",
    "AF2_models = \"4\"  # add other models to this string if needed\n",
    "\n",
    "commands_af2 = []\n",
    "cmds_filename_af2 = \"commands_af2\"\n",
    "with open(cmds_filename_af2, \"w\") as file:\n",
    "    for ff in glob.glob(\"*.fasta\"):\n",
    "        commands_af2.append(f\"{PYTHON['af2']} {AF2_script} \"\n",
    "                             f\"--af-nrecycles {AF2_recycles} --af-models {AF2_models} \"\n",
    "                             f\"--fasta {ff} --scorefile {ff.replace('.fasta', '.csv')}\\n\")\n",
    "        file.write(commands_af2[-1])\n",
    "\n",
    "print(\"Example AF2 command:\")\n",
    "print(commands_af2[-1])\n",
    "\n",
    "### Running AF2 with Slurm.\n",
    "### Running jobs on the CPU. It takes ~10 minutes per sequence\n",
    "\n",
    "submit_script = \"submit_af2.sh\"\n",
    "if USE_GPU_for_AF2 is True:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.1_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=2, gpu=True, gres=\"gpu:rtx2080:1\", time=\"00:12:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "else:\n",
    "    utils.create_slurm_submit_script(filename=submit_script, name=\"5.1_2nd_af2\", mem=\"6g\", \n",
    "                                     N_cores=4, time=\"01:00:00\", email=EMAIL, array=len(commands_af2),\n",
    "                                     array_commandfile=cmds_filename_af2)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0915fdf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you're done with AF2 and happy with the outputs then mark it as done\n",
    "AF2_DIR = f\"{WDIR}/inpainting/5.1_2nd_af2\"\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/inpainting/3.1_design_pocket_ligandMPNN\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if not os.path.exists(AF2_DIR+\"/.done\"):\n",
    "    with open(f\"{AF2_DIR}/.done\", \"w\") as file:\n",
    "        file.write(f\"Run user: {username}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d53a010",
   "metadata": {},
   "source": [
    "### Analyzing AlphaFold2 predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1266a22e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combining all CSV scorefiles into one\n",
    "os.system(\"head -n 1 $(ls *aa*.csv | shuf -n 1) > scores.csv ; for f in *aa*.csv ; do tail -n +2 ${f} >> scores.csv ; done\")\n",
    "assert os.path.exists(\"scores.csv\"), \"Could not combine scorefiles\"\n",
    "\n",
    "### Calculating the RMSDs of AF2 predictions relative to the diffusion outputs\n",
    "### Catalytic residue sidechain RMSDs are calculated in the reference PDB has REMARK 666 line present\n",
    "\n",
    "analysis_cmd = f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/utils/analyze_af2.py --scorefile scores.csv \"\\\n",
    "               f\"--ref_path {DESIGN_DIR_ligMPNN}/good/ --mpnn --params {' '.join(params)}\"\n",
    "\n",
    "p = subprocess.Popen(analysis_cmd, shell=True)\n",
    "(output, err) = p.communicate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981cc7b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualizing and filtering AF2 results\n",
    "AF2_DIR = f\"{WDIR}/inpainting/5.1_2nd_af2\"\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "scores_af2 = pd.read_csv(\"scores.sc\", sep=\"\\s+\", header=0)\n",
    "\n",
    "### Filtering AF2 scores based on lddt and rmsd\n",
    "# Define your desired cutoffs here:\n",
    "AF2_filters = {\"lDDT\": [85.0, \">=\"],\n",
    "               \"rmsd\": [1.2, \"<=\"],\n",
    "               \"rmsd_SR1\": [1.0, \"<=\"]}  # 1st catalytic residue sc-rmsd\n",
    "\n",
    "scores_af2_filtered = utils.filter_scores(scores_af2, AF2_filters)\n",
    "utils.dump_scorefile(scores_af2_filtered, \"filtered_scores.sc\")\n",
    "\n",
    "## Plotting AF2 scores\n",
    "plt.figure(figsize=(12, 3))\n",
    "for i,k in enumerate(AF2_filters):\n",
    "    plt.subplot(1, 3, i+1)\n",
    "    plt.hist(scores_af2[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "    \n",
    "utils.plot_score_pairs(scores_af2, \"lDDT\", \"rmsd\", AF2_filters[\"lDDT\"][0], AF2_filters[\"rmsd\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afad819b",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Copying good predictions to a separate directory\n",
    "os.chdir(AF2_DIR)\n",
    "\n",
    "if len(scores_af2_filtered) > 0:\n",
    "    os.makedirs(\"good\", exist_ok=True)\n",
    "    good_af2_models = [row[\"Output_PDB\"]+\".pdb\" for idx,row in scores_af2_filtered.iterrows()]\n",
    "    for pdb in good_af2_models:\n",
    "        copy2(pdb, f\"good/{pdb}\")\n",
    "    good_af2_models = glob.glob(f\"{AF2_DIR}/good/*.pdb\")\n",
    "else:\n",
    "    sys.exit(\"No good models to continue this pipeline with\")\n",
    "\n",
    "os.chdir(f\"{AF2_DIR}/good\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d059ff96",
   "metadata": {},
   "source": [
    "## 6.1: Final FastRelax with the ligand\n",
    "Relaxing good AF2 models together with the ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c7accf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "AF2_DIR = f\"{WDIR}/inpainting/5.1_2nd_af2\"\n",
    "DESIGN_DIR_ligMPNN = f\"{WDIR}/inpainting/3.1_design_pocket_ligandMPNN\"\n",
    "assert len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")) > 0, \"No good AF2 models to relax with\"\n",
    "\n",
    "os.chdir(WDIR)\n",
    "RELAX_DIR = f\"{WDIR}/inpainting/6.1_final_relax\"\n",
    "os.makedirs(RELAX_DIR, exist_ok=True)\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "\n",
    "## First matching up the AF2 output filenames of step 5 with pocket design filenames from step 3\n",
    "ref_and_model_pairs = []\n",
    "for r in glob.glob(f\"{DESIGN_DIR_ligMPNN}/good/*.pdb\"):\n",
    "    for pdbfile in glob.glob(f\"{AF2_DIR}/good/*.pdb\"):\n",
    "        if os.path.basename(r).replace(\".pdb\", \"_\") in pdbfile:\n",
    "            ref_and_model_pairs.append((r, pdbfile))\n",
    "\n",
    "assert len(ref_and_model_pairs) == len(glob.glob(f\"{AF2_DIR}/good/*.pdb\")), \"Was not able to match all models with reference structures\"\n",
    "\n",
    "\n",
    "## Generating commands for relax jobs\n",
    "### Performing 1 relax iteration on each input structure\n",
    "NSTRUCT = 1\n",
    "cstfile = f\"{SCRIPT_DIR}/theozyme/HBA/HBA_CYS_UPO.cst\"\n",
    "\n",
    "commands_relax = []\n",
    "cmds_filename_rlx = \"commands_design\"\n",
    "with open(cmds_filename_rlx, \"w\") as file:\n",
    "    for r_m in ref_and_model_pairs:\n",
    "        commands_relax.append(f\"{PYTHON['general']} {SCRIPT_DIR}/scripts/design/align_add_ligand_relax.py \"\n",
    "                              f\"--outdir ./ --ligand {LIGAND} --ref_pdb {r_m[0]}\"\n",
    "                              f\"--pdb {r_m[1]} --nstruct {NSTRUCT} \"\n",
    "                              f\"--params {' '.join(params)} --cstfile {cstfile} > logs/{os.path.basename(pdb).replace('.pdb', '.log')}\\n\")\n",
    "        file.write(commands_relax[-1])\n",
    "\n",
    "print(\"Example design command:\")\n",
    "print(commands_design[-1])\n",
    "\n",
    "\n",
    "### Running design jobs with Slurm.\n",
    "submit_script = \"submit_relax.sh\"\n",
    "utils.create_slurm_submit_script(filename=submit_script, name=\"6.1_final_relax\", mem=\"4g\", \n",
    "                                 N_cores=1, time=\"0:30:00\", email=EMAIL, array=len(commands_relax),\n",
    "                                 array_commandfile=cmds_filename_rlx)\n",
    "\n",
    "if not os.path.exists(RELAX_DIR+\"/.done\"):\n",
    "    p = subprocess.Popen(['sbatch', submit_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    (output, err) = p.communicate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19f57db8",
   "metadata": {},
   "source": [
    "### Analyzing final relaxed structures\n",
    "Filtering them based on the same metrics as was used for the initial design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b11d1f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Analyzing Rosetta designs\n",
    "RELAX_DIR = f\"{WDIR}/inpainting/6.1_final_relax\"\n",
    "os.chdir(RELAX_DIR)\n",
    "\n",
    "scores = pd.read_csv(\"scorefile.txt\", sep=\"\\s+\", header=0)\n",
    "\n",
    "filters = {'all_cst': [1.0, '<='],\n",
    " 'nlr_SR1_rms': [0.8, '<='],\n",
    " 'nlr_totrms': [1.0, '<='],\n",
    " 'L_SASA': [0.2, '<='],\n",
    " 'COO_hbond': [1.0, '='],\n",
    " 'heme_angle_wrst': [80.0, '>='],\n",
    " 'score_per_res': [0.0, '<='],\n",
    " 'corrected_ddg': [-50.0, '<='],\n",
    " 'cms_per_atom': [4.8, '>='],\n",
    " 'rmsd_CA_rlx_in': [1.0, \"<=\"]}  # rmsd_CA_rlx_in is rmsd between relaxed structure and AF2 prediction\n",
    "\n",
    "filtered_scores = utils.filter_scores(scores, filters)\n",
    "\n",
    "## Plotting relax scores\n",
    "plt.figure(figsize=(12, 9))\n",
    "for i,k in enumerate(filters):\n",
    "    if k not in scores.keys():\n",
    "        continue\n",
    "    plt.subplot(4, 3, i+1)\n",
    "    plt.hist(scores[k])\n",
    "    plt.title(k)\n",
    "    plt.xlabel(k)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "### Copying good designs over to a new directory\n",
    "if len(filtered_scores) > 0:\n",
    "    os.makedirs(f\"{RELAX_DIR}/good\", exist_ok=True)\n",
    "    for idx, row in filtered_scores.iterrows():\n",
    "        copy2(row[\"description\"]+\".pdb\", \"good/\"+row[\"description\"]+\".pdb\")\n",
    "else:\n",
    "    print(\"No good designs created, bummer...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fba174b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(filtered_scores) > 0:\n",
    "    print(f\"CONGRATULATIONS! You have successfully designed {len(filtered_scores)} proteins against ligand {LIGAND}\")\n",
    "    print(\"You can find the design models in the directory:\\n\"\n",
    "          f\"    {RELAX_DIR}/good\")\n",
    "    print(\"\\nIt is advised you manually inspect them before ordering.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "780be3c9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}