02_rerun_best.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rerun interface design on the best docks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boilerplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/home/pleung/projects/peptide_binders/r0/peptide_binders\n",
      "dig85\n"
     ]
    }
   ],
   "source": [
    "%load_ext lab_black\n",
    "# python internal\n",
    "import collections\n",
    "import copy\n",
    "import gc\n",
    "from glob import glob\n",
    "import h5py\n",
    "import itertools\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import socket\n",
    "import shutil\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "# conda/pip\n",
    "import dask\n",
    "import graphviz\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy\n",
    "import seaborn as sns\n",
    "from tqdm import tqdm\n",
    "\n",
    "# special packages on the DIGS\n",
    "import py3Dmol\n",
    "import pymol\n",
    "import pyrosetta\n",
    "\n",
    "# notebook magic\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "print(os.getcwd())\n",
    "print(socket.gethostname())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make a function to do one sided PPI design\n",
    "use `tutorial_design_v5.xml` as an example\n",
    "added CYS/CYD restriction  \n",
    "TODO CYD breaks `pyrosetta.distributed`  \n",
    "try `<Index name=\"\" resnums=\"{}\" />` instead of slices             "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyrosetta.distributed.packed_pose.core import PackedPose\n",
    "from pyrosetta.distributed import requires_init\n",
    "\n",
    "\n",
    "@requires_init\n",
    "def almost_intdes(packed_pose_in: PackedPose, **kwargs) -> PackedPose:\n",
    "    \"\"\"\n",
    "    Wrap a pyrosetta compatible version of @bcov interface design v5 xml.\n",
    "    Save scores to pose\n",
    "    \"\"\"\n",
    "    import bz2\n",
    "    import pyrosetta\n",
    "    from pyrosetta.rosetta.core.pose import Pose\n",
    "    import pyrosetta.distributed.io as io\n",
    "    from pyrosetta.distributed.tasks.rosetta_scripts import (\n",
    "        SingleoutputRosettaScriptsTask,\n",
    "    )\n",
    "\n",
    "    if packed_pose_in == None:\n",
    "        file = kwargs[\"-s\"]\n",
    "        with open(file, \"rb\") as f:\n",
    "            packed_pose_in = io.pose_from_pdbstring(bz2.decompress(f.read()).decode())\n",
    "        scores = dict(pyrosetta.distributed.cluster.get_scores_dict(file)[\"scores\"])\n",
    "    else:\n",
    "        raise RuntimeError(\"Need to supply an input\")\n",
    "\n",
    "    if \"sasa_threshold\" in kwargs.keys():\n",
    "        sasa_threshold = kwargs[\"sasa_threshold\"]\n",
    "    else:\n",
    "        sasa_threshold = 1350\n",
    "\n",
    "    xml = \"\"\"\n",
    "    <ROSETTASCRIPTS>\n",
    "        <SCOREFXNS>\n",
    "            <ScoreFunction name=\"sfxn\" weights=\"beta_nov16\" />\n",
    "            <ScoreFunction name=\"sfxn_soft\" weights=\"beta_nov16_soft\" />\n",
    "            <ScoreFunction name=\"sfxn_fa_atr\" weights=\"empty\" >\n",
    "                <Reweight scoretype=\"fa_atr\" weight=\"1\" />\n",
    "            </ScoreFunction>\n",
    "            <ScoreFunction name=\"sfxn_relax\" weights=\"beta_nov16\" >\n",
    "                <Reweight scoretype=\"arg_cation_pi\" weight=\"3\" />\n",
    "                <Reweight scoretype=\"approximate_buried_unsat_penalty\" weight=\"5\" />\n",
    "                <Set approximate_buried_unsat_penalty_burial_atomic_depth=\"3.5\" />\n",
    "                <Set approximate_buried_unsat_penalty_hbond_energy_threshold=\"-0.5\" />\n",
    "            </ScoreFunction>\n",
    "            <ScoreFunction name=\"sfxn_design\" weights=\"beta_nov16\" >\n",
    "                <Reweight scoretype=\"res_type_constraint\" weight=\"1.5\" />\n",
    "                <Reweight scoretype=\"aa_composition\" weight=\"1.0\" />\n",
    "                <Reweight scoretype=\"arg_cation_pi\" weight=\"3\" />\n",
    "                <Reweight scoretype=\"approximate_buried_unsat_penalty\" weight=\"5\" />\n",
    "                <Set approximate_buried_unsat_penalty_burial_atomic_depth=\"3.5\" />\n",
    "                <Set approximate_buried_unsat_penalty_hbond_energy_threshold=\"-0.5\" />\n",
    "                <Set approximate_buried_unsat_penalty_hbond_bonus_cross_chain=\"-1\" />\n",
    "            </ScoreFunction>\n",
    "            <ScoreFunction name=\"vdw_sol\" weights=\"empty\" >\n",
    "                <Reweight scoretype=\"fa_atr\" weight=\"1.0\" />\n",
    "                <Reweight scoretype=\"fa_rep\" weight=\"0.55\" />\n",
    "                <Reweight scoretype=\"fa_sol\" weight=\"1.0\" />\n",
    "            </ScoreFunction>\n",
    "        </SCOREFXNS>\n",
    "        <TASKOPERATIONS>\n",
    "            <SelectBySASA name=\"PR_monomer_core_sel\" mode=\"sc\" state=\"monomer\" probe_radius=\"2.2\" core_asa=\"15\" surface_asa=\"15\" core=\"0\" boundary=\"1\" surface=\"1\" verbose=\"0\" />\n",
    "        </TASKOPERATIONS>\n",
    "        <RESIDUE_SELECTORS>\n",
    "            <Chain name=\"chainA\" chains=\"A\"/>\n",
    "            <Chain name=\"chainB\" chains=\"B\"/>\n",
    "            <Neighborhood name=\"interface_chA\" selector=\"chainB\" distance=\"10.0\" />\n",
    "            <Neighborhood name=\"interface_chB\" selector=\"chainA\" distance=\"10.0\" />\n",
    "            <And name=\"AB_interface\" selectors=\"interface_chA,interface_chB\" />\n",
    "            <Not name=\"Not_interface\" selector=\"AB_interface\" />\n",
    "            <And name=\"actual_interface_chA\" selectors=\"AB_interface,chainA\" />\n",
    "            <And name=\"actual_interface_chB\" selectors=\"AB_interface,chainB\" />\n",
    "            <And name=\"chainB_not_interface\" selectors=\"Not_interface,chainB\" />\n",
    "            <Not name=\"not_chB\" selector=\"chainB\" />\n",
    "\n",
    "            <ResidueName name=\"pro_and_gly_positions\" residue_name3=\"PRO,GLY,CYS,CYD\" />\n",
    "            <ResidueName name=\"apolar\" residue_name3=\"ALA,CYS,PHE,ILE,LEU,MET,THR,PRO,VAL,TRP,TYR\" />\n",
    "            <Not name=\"polar\" selector=\"apolar\" />\n",
    "\n",
    "            <InterfaceByVector name=\"interface_by_vector\" cb_dist_cut=\"11\" nearby_atom_cut=\"5.5\" vector_angle_cut=\"75\" vector_dist_cut=\"9\" grp1_selector=\"actual_interface_chA\" grp2_selector=\"actual_interface_chB\"/>\n",
    "            <Task name=\"all_cores\" fixed=\"true\" task_operations=\"PR_monomer_core_sel\" packable=\"false\" designable=\"false\"/>\n",
    "            <And name=\"for_hydrophobic\" selectors=\"actual_interface_chA,interface_by_vector\">\n",
    "                <Not selector=\"all_cores\" />\n",
    "            </And>\n",
    "\n",
    "            <ResiduePDBInfoHasLabel name=\"HOTSPOT_res\" property=\"HOTSPOT\" />\n",
    "            <ResiduePDBInfoHasLabel name=\"DONE_res\" property=\"DONE\" />\n",
    "            \n",
    "            <Index name=\"res1\" resnums=\"1\" />\n",
    "            <Index name=\"patchdock_res\" resnums=\"{interface_hydrophobics}\" />\n",
    "\n",
    "            <Index name=\"pocket\" resnums=\"{interface_hydrophobics}\" />\n",
    "            <And name=\"target_not_pocket\" selectors=\"chainB\">\n",
    "                <Not selector=\"pocket\" />\n",
    "            </And>\n",
    "\n",
    "            <Layer name=\"surface\" select_core=\"false\" select_boundary=\"false\" select_surface=\"true\" use_sidechain_neighbors=\"true\"/>\n",
    "            <Layer name=\"boundary\" select_core=\"false\" select_boundary=\"true\" select_surface=\"false\" use_sidechain_neighbors=\"true\"/>\n",
    "            <Layer name=\"core\" select_core=\"true\" select_boundary=\"false\" select_surface=\"false\" use_sidechain_neighbors=\"true\"/>\n",
    "            <SecondaryStructure name=\"sheet\" overlap=\"0\" minH=\"3\" minE=\"2\" include_terminal_loops=\"false\" use_dssp=\"true\" ss=\"E\"/>\n",
    "            <SecondaryStructure name=\"entire_loop\" overlap=\"0\" minH=\"3\" minE=\"2\" include_terminal_loops=\"true\" use_dssp=\"true\" ss=\"L\"/>\n",
    "            <SecondaryStructure name=\"entire_helix\" overlap=\"0\" minH=\"3\" minE=\"2\" include_terminal_loops=\"false\" use_dssp=\"true\" ss=\"H\"/>\n",
    "            <And name=\"helix_cap\" selectors=\"entire_loop\">\n",
    "                <PrimarySequenceNeighborhood lower=\"1\" upper=\"0\" selector=\"entire_helix\"/>\n",
    "            </And>\n",
    "            <And name=\"helix_start\" selectors=\"entire_helix\">\n",
    "                <PrimarySequenceNeighborhood lower=\"0\" upper=\"1\" selector=\"helix_cap\"/>\n",
    "            </And>\n",
    "            <And name=\"helix\" selectors=\"entire_helix\">\n",
    "                <Not selector=\"helix_start\"/>\n",
    "            </And>\n",
    "            <And name=\"loop\" selectors=\"entire_loop\">\n",
    "                <Not selector=\"helix_cap\"/>\n",
    "            </And>\n",
    "\n",
    "        </RESIDUE_SELECTORS>\n",
    "        <TASKOPERATIONS>\n",
    "            <DesignRestrictions name=\"layer_design_F_boundary_M\">\n",
    "                <Action selector_logic=\"not_chB AND surface AND helix_start\"  aas=\"DEHKPQR\"/>\n",
    "                <Action selector_logic=\"not_chB AND surface AND helix\"        aas=\"EHKQR\"/>\n",
    "                <Action selector_logic=\"not_chB AND surface AND sheet\"        aas=\"EHKNQRST\"/>\n",
    "                <Action selector_logic=\"not_chB AND surface AND loop\"         aas=\"DEGHKNPQRST\"/>\n",
    "                <Action selector_logic=\"not_chB AND boundary AND helix_start\" aas=\"ADEFHIKLMNPQRSTVWY\"/>\n",
    "                <Action selector_logic=\"not_chB AND boundary AND helix\"       aas=\"ADEFHIKLMNQRSTVWY\"/>\n",
    "                <Action selector_logic=\"not_chB AND boundary AND sheet\"       aas=\"DEFHIKLMNQRSTVWY\"/>\n",
    "                <Action selector_logic=\"not_chB AND boundary AND loop\"        aas=\"ADEFGHIKLNPQRSTVWY\"/>\n",
    "                <Action selector_logic=\"not_chB AND core AND helix_start\"     aas=\"AFILMPVWY\"/>\n",
    "                <Action selector_logic=\"not_chB AND core AND helix\"           aas=\"AFILMVWYDENQSTH\"/>\n",
    "                <Action selector_logic=\"not_chB AND core AND sheet\"           aas=\"FILMVWYDENQSTH\"/>\n",
    "                <Action selector_logic=\"not_chB AND core AND loop\"            aas=\"AFGILPVWYDENQSTH\"/>\n",
    "                <Action selector_logic=\"not_chB AND helix_cap\"                aas=\"DNST\"/>\n",
    "            </DesignRestrictions>\n",
    "        </TASKOPERATIONS>\n",
    "        <MOVERS>\n",
    "            <AddCompositionConstraintMover name=\"3trp\" >\n",
    "                <Comp entry=\"PENALTY_DEFINITION;TYPE TRP;ABSOLUTE 0;PENALTIES 0 3;DELTA_START 0;DELTA_END 1;BEFORE_FUNCTION CONSTANT;AFTER_FUNCTION LINEAR;END_PENALTY_DEFINITION;\" />\n",
    "            </AddCompositionConstraintMover>\n",
    "            <AddCompositionConstraintMover name=\"2met\" >\n",
    "                <Comp entry=\"PENALTY_DEFINITION;TYPE MET;ABSOLUTE 0;PENALTIES 0 2;DELTA_START 0;DELTA_END 1;BEFORE_FUNCTION CONSTANT;AFTER_FUNCTION LINEAR;END_PENALTY_DEFINITION;\" />\n",
    "            </AddCompositionConstraintMover>\n",
    "            <AddCompositionConstraintMover name=\"30_percent_polar\" selector=\"for_hydrophobic\">\n",
    "                <Comp entry=\"PENALTY_DEFINITION;TYPE ASP GLU HIS LYS ASN GLN ARG SER THR TYR;FRACT_DELTA_START -0.01;FRACT_DELTA_END 0.0;PENALTIES 0.1 0 ;FRACTION 0.30;BEFORE_FUNCTION QUADRATIC;AFTER_FUNCTION CONSTANT;END_PENALTY_DEFINITION\" />\n",
    "            </AddCompositionConstraintMover>\n",
    "        </MOVERS>\n",
    "        <TASKOPERATIONS>\n",
    "            <PruneBuriedUnsats name=\"prune_buried_unsats\" allow_even_trades=\"false\" atomic_depth_cutoff=\"3.5\" minimum_hbond_energy=\"-0.5\" />\n",
    "            <ProteinProteinInterfaceUpweighter name=\"upweight_interface\" interface_weight=\"3\" />\n",
    "            <ProteinInterfaceDesign name=\"pack_long\" design_chain1=\"0\" design_chain2=\"0\" jump=\"1\" interface_distance_cutoff=\"15\"/>\n",
    "            <IncludeCurrent name=\"current\" />\n",
    "            <LimitAromaChi2 name=\"limitchi2\" chi2max=\"110\" chi2min=\"70\" include_trp=\"True\" />\n",
    "            <ExtraRotamersGeneric name=\"ex1_ex2\" ex1=\"1\" ex2aro=\"1\" />\n",
    "\n",
    "            <OperateOnResidueSubset name=\"restrict_target_not_interface\" selector=\"chainB_not_interface\">\n",
    "                <PreventRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "            <OperateOnResidueSubset name=\"restrict_to_interface\" selector=\"Not_interface\">\n",
    "                <PreventRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "            <OperateOnResidueSubset name=\"restrict_target2repacking\" selector=\"chainB\">\n",
    "                <RestrictToRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "            <OperateOnResidueSubset name=\"restrict_hotspots2repacking\" selector=\"HOTSPOT_res\">\n",
    "                <RestrictToRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "            <DisallowIfNonnative name=\"disallow_GLY\" resnum=\"0\" disallow_aas=\"G\" />\n",
    "            <DisallowIfNonnative name=\"disallow_PRO\" resnum=\"0\" disallow_aas=\"P\" />\n",
    "            <OperateOnResidueSubset name=\"restrict_PRO_GLY\" selector=\"pro_and_gly_positions\">\n",
    "                <PreventRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "\n",
    "            <SelectBySASA name=\"PR_monomer_core\" mode=\"sc\" state=\"monomer\" probe_radius=\"2.2\" core_asa=\"10\" surface_asa=\"10\" core=\"0\" boundary=\"1\" surface=\"1\" verbose=\"0\" />\n",
    "\n",
    "        </TASKOPERATIONS>\n",
    "        <RESIDUE_SELECTORS>\n",
    "            <And name=\"sheetA\" selectors=\"sheet,chainA\" />\n",
    "        </RESIDUE_SELECTORS>\n",
    "        <MOVERS>\n",
    "\n",
    "            <DeleteRegionMover name=\"delete_target_not_pocket\" residue_selector=\"target_not_pocket\" rechain=\"false\" />\n",
    "\n",
    "            <SwitchChainOrder name=\"chain1onlypre\" chain_order=\"1\" />\n",
    "            <ScoreMover name=\"scorepose\" scorefxn=\"sfxn\" verbose=\"false\" />\n",
    "            <ParsedProtocol name=\"chain1only\">\n",
    "                <Add mover=\"chain1onlypre\" />\n",
    "                <Add mover=\"scorepose\" />\n",
    "            </ParsedProtocol>\n",
    "            <TaskAwareMinMover name=\"min\" scorefxn=\"sfxn\" bb=\"0\" chi=\"1\" task_operations=\"pack_long\" />\n",
    "\n",
    "            <DeleteRegionMover name=\"delete_polar\" residue_selector=\"polar\" rechain=\"false\" />\n",
    "\n",
    "            <FastRelax name=\"relax_chain1\" scorefxn=\"sfxn\" repeats=\"1\" batch=\"false\" ramp_down_constraints=\"false\" cartesian=\"false\" bondangle=\"false\" bondlength=\"false\" min_type=\"dfpmin_armijo_nonmonotone\" task_operations=\"ex1_ex2,limitchi2\" >\n",
    "                <MoveMap name=\"MM\" >\n",
    "                    <Chain number=\"1\" chi=\"true\" bb=\"true\" />\n",
    "                </MoveMap>\n",
    "            </FastRelax>\n",
    "\n",
    "            <ParsedProtocol name=\"chain1only_relax\">\n",
    "                <Add mover=\"chain1onlypre\" />\n",
    "                <Add mover=\"relax_chain1\" />\n",
    "            </ParsedProtocol>\n",
    "\n",
    "            <StructProfileMover name=\"genProfile\" add_csts_to_pose=\"1\" consider_topN_frags=\"100\" eliminate_background=\"0\" ignore_terminal_residue=\"1\" only_loops=\"0\" burialWt=\"0\" RMSthreshold=\"0.6\" residue_selector=\"chainA\" />\n",
    "\n",
    "            <ClearConstraintsMover name=\"clear_constraints\" />\n",
    "\n",
    "            <SavePoseMover name=\"save_output\" restore_pose=\"0\" reference_name=\"pose_output\" />\n",
    "            <SavePoseMover name=\"load_output\" restore_pose=\"1\" reference_name=\"pose_output\" />\n",
    "\n",
    "            <LabelPoseFromResidueSelectorMover name=\"remove_done\" remove_property=\"DONE\"  residue_selector=\"res1\" />\n",
    "            <LabelPoseFromResidueSelectorMover name=\"add_done\" property=\"DONE\" residue_selector=\"res1\" />\n",
    "\n",
    "        </MOVERS>\n",
    "        <FILTERS>\n",
    "            <Sasa name=\"interface_buried_sasa\" confidence=\"0\" />\n",
    "            <Sasa name=\"gate_filter\" threshold=\"{sasa_threshold}\" confidence=\"1\" />\n",
    "            <Ddg name=\"ddg\"  threshold=\"-10\" jump=\"1\" repeats=\"5\" repack=\"1\" relax_mover=\"min\" confidence=\"0\" scorefxn=\"sfxn\" extreme_value_removal=\"1\" />\n",
    "\n",
    "            <ShapeComplementarity name=\"interface_sc\" verbose=\"0\" min_sc=\"0.55\" write_int_area=\"1\" write_median_dist=\"1\" jump=\"1\" confidence=\"0\"/>\n",
    "            <ShapeComplementarity name=\"interface_sc_for_fr1\" verbose=\"0\" min_sc=\"0.55\" write_int_area=\"1\" write_median_dist=\"1\" jump=\"1\" confidence=\"0\"/>\n",
    "\n",
    "            <ScoreType name=\"total_score_MBF\" scorefxn=\"sfxn\" score_type=\"total_score\" threshold=\"0\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"total_score_monomer\" mover=\"chain1only\" filter=\"total_score_MBF\" confidence=\"0\" />\n",
    "            <ResidueCount name=\"res_count_MBF\" max_residue_count=\"9999\" confidence=\"0\"/>\n",
    "            <MoveBeforeFilter name=\"res_count_monomer\" mover=\"chain1only\" filter=\"res_count_MBF\" confidence=\"0\" />\n",
    "\n",
    "            <CalculatorFilter name=\"score_per_res\" equation=\"total_score_monomer / res\" threshold=\"-3.5\" confidence=\"0\">\n",
    "                <Var name=\"total_score_monomer\" filter=\"total_score_monomer\"/>\n",
    "                <Var name=\"res\" filter=\"res_count_monomer\"/>\n",
    "            </CalculatorFilter>\n",
    "\n",
    "            <MoveBeforeFilter name=\"total_score_relax_monomer\" mover=\"chain1only_relax\" filter=\"total_score_MBF\" confidence=\"0\" />\n",
    "\n",
    "            <CalculatorFilter name=\"score_per_res_relax\" equation=\"total_score_relax_monomer / res\" threshold=\"-3.5\" confidence=\"0\">\n",
    "                <Var name=\"total_score_relax_monomer\" filter=\"total_score_relax_monomer\"/>\n",
    "                <Var name=\"res\" filter=\"res_count_monomer\"/>\n",
    "            </CalculatorFilter>\n",
    "\n",
    "            <BuriedUnsatHbonds name=\"buns_heavy_ball_1.1D\" use_reporter_behavior=\"true\" report_all_heavy_atom_unsats=\"true\" scorefxn=\"sfxn\" residue_selector=\"AB_interface\" ignore_surface_res=\"false\" print_out_info_to_pdb=\"true\" confidence=\"0\" use_ddG_style=\"true\" burial_cutoff=\"0.01\" dalphaball_sasa=\"true\" probe_radius=\"1.1\" max_hbond_energy=\"1.5\" burial_cutoff_apo=\"0.2\" />\n",
    "            <BuriedUnsatHbonds name=\"sbuns5.0_heavy_ball_1.1D\" use_reporter_behavior=\"true\" report_all_heavy_atom_unsats=\"true\" scorefxn=\"sfxn\" residue_selector=\"AB_interface\" ignore_surface_res=\"false\" print_out_info_to_pdb=\"true\" confidence=\"0\" use_ddG_style=\"true\" burial_cutoff=\"0.01\" dalphaball_sasa=\"true\" probe_radius=\"1.1\" atomic_depth_selection=\"5.0\" atomic_depth_deeper_than=\"false\" burial_cutoff_apo=\"0.2\" atomic_depth_resolution=\"0.49\" max_hbond_energy=\"1.5\"/>\n",
    "\n",
    "            <BuriedUnsatHbonds name=\"vbuns5.0_heavy_ball_1.1D\" use_reporter_behavior=\"true\" report_all_heavy_atom_unsats=\"true\" scorefxn=\"sfxn\" residue_selector=\"AB_interface\" ignore_surface_res=\"false\" print_out_info_to_pdb=\"true\" confidence=\"0\" use_ddG_style=\"true\" dalphaball_sasa=\"true\" probe_radius=\"1.1\" atomic_depth_selection=\"5.0\" burial_cutoff=\"1000\" burial_cutoff_apo=\"0.2\" atomic_depth_apo_surface=\"5.5\" atomic_depth_resolution=\"0.49\" max_hbond_energy=\"1.5\"/>\n",
    "\n",
    "            <InterfaceHydrophobicResidueContacts name=\"hydrophobic_residue_contacts\" target_selector=\"chainB\" binder_selector=\"chainA\" scorefxn=\"sfxn_soft\" confidence=\"0\"/> \n",
    "            <ContactMolecularSurface name=\"contact_molecular_surface\" distance_weight=\"0.5\" target_selector=\"chainA\" binder_selector=\"chainB\" confidence=\"0\" />\n",
    "            <ContactMolecularSurface name=\"contact_patch\" distance_weight=\"0.5\" target_selector=\"patchdock_res\" binder_selector=\"chainA\" confidence=\"0\" />\n",
    "\n",
    "            <Ddg name=\"ddg_hydrophobic_pre\"  threshold=\"-10\" jump=\"1\" repeats=\"1\" repack=\"0\" confidence=\"0\" scorefxn=\"vdw_sol\" />\n",
    "            <MoveBeforeFilter name=\"ddg_hydrophobic\" mover=\"delete_polar\" filter=\"ddg_hydrophobic_pre\" confidence=\"0\"/>\n",
    "\n",
    "            <Ddg name=\"fa_atr_pocket_pre\"  threshold=\"-15\" jump=\"1\" repeats=\"1\" repack=\"0\" confidence=\"0\" scorefxn=\"sfxn_fa_atr\" />\n",
    "            <MoveBeforeFilter name=\"fa_atr_pocket\" mover=\"delete_target_not_pocket\" filter=\"fa_atr_pocket_pre\" confidence=\"0\"/>\n",
    "\n",
    "            <SSPrediction name=\"pre_sspred_overall\" cmd=\"/software/psipred4/runpsipred_single\" use_probability=\"0\" use_svm=\"0\" threshold=\"0.85\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"sspred_overall\" mover=\"chain1only\" filter=\"pre_sspred_overall\" confidence=\"0\" />\n",
    "\n",
    "            <SSPrediction name=\"pre_mismatch_probability\" confidence=\"0\" cmd=\"/software/psipred4/runpsipred_single\" use_probability=\"1\" mismatch_probability=\"1\" use_svm=\"0\" />\n",
    "            <MoveBeforeFilter name=\"mismatch_probability\" mover=\"chain1only\" filter=\"pre_mismatch_probability\" confidence=\"0\" />\n",
    "\n",
    "            <ScorePoseSegmentFromResidueSelectorFilter name=\"has_sheet\" in_context=\"1\" residue_selector=\"sheetA\" scorefxn=\"sfxn\" confidence=\"0\" />\n",
    "            <ScorePoseSegmentFromResidueSelectorFilter name=\"was_done\" in_context=\"1\" residue_selector=\"DONE_res\" scorefxn=\"sfxn\" confidence=\"0\" />\n",
    "\n",
    "            <SSShapeComplementarity name=\"ss_sc_pre\" verbose=\"0\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"ss_sc\" mover=\"chain1only\" filter=\"ss_sc_pre\" confidence=\"0\"/>\n",
    "\n",
    "            <Time name=\"timed\"/>\n",
    "            <TrueFilter name=\"my_true_filter\" />\n",
    "        </FILTERS>\n",
    "        <MOVERS>\n",
    "            <FastDesign name=\"FastDesign\" scorefxn=\"sfxn_design\" repeats=\"1\" task_operations=\"current,limitchi2,ex1_ex2,restrict_to_interface,restrict_target2repacking,restrict_hotspots2repacking,disallow_GLY,disallow_PRO,PR_monomer_core,upweight_interface,restrict_PRO_GLY,prune_buried_unsats,layer_design_F_boundary_M\" batch=\"false\" ramp_down_constraints=\"false\" cartesian=\"false\" bondangle=\"false\" bondlength=\"false\" min_type=\"dfpmin_armijo_nonmonotone\" relaxscript=\"InterfaceDesign2019\" >\n",
    "                <MoveMap name=\"MM\" >\n",
    "                        <Chain number=\"1\" chi=\"true\" bb=\"true\" />\n",
    "                        <Chain number=\"2\" chi=\"true\" bb=\"false\" />\n",
    "                        <Jump number=\"1\" setting=\"true\" />\n",
    "                </MoveMap>\n",
    "            </FastDesign>\n",
    "            <FastRelax name=\"FastRelax\" scorefxn=\"sfxn_relax\" repeats=\"1\" batch=\"false\" ramp_down_constraints=\"false\" cartesian=\"false\" bondangle=\"false\" bondlength=\"false\" min_type=\"dfpmin_armijo_nonmonotone\" task_operations=\"current,ex1_ex2,restrict_target_not_interface,limitchi2,prune_buried_unsats\" >\n",
    "                <MoveMap name=\"MM\" >\n",
    "                    <Chain number=\"1\" chi=\"true\" bb=\"true\" />\n",
    "                    <Chain number=\"2\" chi=\"true\" bb=\"false\" />\n",
    "                    <Jump number=\"1\" setting=\"true\" />\n",
    "                </MoveMap>\n",
    "            </FastRelax>\n",
    "            <ParsedProtocol name=\"go_go_go\" mode=\"sequence\">\n",
    "                <Add filter=\"gate_filter\" />\n",
    "                Add mover=\"genProfile\" />\n",
    "                <Add mover=\"3trp\" />\n",
    "                <Add mover=\"2met\" />\n",
    "                <Add mover=\"30_percent_polar\" />\n",
    "                <Add mover=\"FastDesign\" />\n",
    "                <Add mover=\"save_output\" />\n",
    "                <Add mover=\"FastDesign\" />\n",
    "                <Add mover=\"clear_constraints\" />\n",
    "                <Add mover=\"FastRelax\" />\n",
    "                <Add mover=\"add_done\" />\n",
    "                <Add mover=\"save_output\" />\n",
    "            </ParsedProtocol>\n",
    "            <GenericMonteCarlo name=\"try_go_go_go\" mover_name=\"go_go_go\" filter_name=\"my_true_filter\" trials=\"1\" />\n",
    "        </MOVERS>\n",
    "        <APPLY_TO_POSE>\n",
    "        </APPLY_TO_POSE>\n",
    "        <PROTOCOLS>\n",
    "            <Add filter=\"timed\" />\n",
    "            <Add mover=\"remove_done\" />\n",
    "            <Add mover=\"save_output\" />\n",
    "            <Add mover=\"try_go_go_go\" />\n",
    "            <Add mover=\"load_output\" />\n",
    "            <Add filter_name=\"interface_buried_sasa\" />   \n",
    "            <Add filter_name=\"ddg\" />                     \n",
    "            <Add filter=\"ddg_hydrophobic\" />              \n",
    "            <Add filter_name=\"interface_sc\" />            \n",
    "            <Add filter_name=\"score_per_res\" />          \n",
    "            Add filter_name=\"score_per_res_relax\" />    \n",
    "            <Add filter=\"vbuns5.0_heavy_ball_1.1D\" />    \n",
    "            <Add filter=\"sbuns5.0_heavy_ball_1.1D\" />    \n",
    "            <Add filter=\"buns_heavy_ball_1.1D\" />    \n",
    "            <Add filter=\"hydrophobic_residue_contacts\" />\n",
    "            <Add filter=\"fa_atr_pocket\" />               \n",
    "            Add filter=\"mismatch_probability\" />        \n",
    "            Add filter=\"sspred_overall\" />              \n",
    "            <Add filter=\"contact_molecular_surface\" />\n",
    "            <Add filter=\"contact_patch\" />   \n",
    "            <Add filter=\"ss_sc\" />   \n",
    "            <Add filter=\"has_sheet\" />\n",
    "            <Add filter=\"was_done\" />\n",
    "            <Add filter=\"timed\" />\n",
    "        </PROTOCOLS>\n",
    "        <OUTPUT />\n",
    "    </ROSETTASCRIPTS>\n",
    "    \"\"\".format(\n",
    "        interface_hydrophobics=scores[\"interface_hydrophobics\"],\n",
    "        sasa_threshold=sasa_threshold,\n",
    "    )\n",
    "    intdes = SingleoutputRosettaScriptsTask(xml)\n",
    "    designed_ppose = intdes(packed_pose_in.pose.clone())\n",
    "    new_scores = dict(designed_ppose.pose.scores)\n",
    "    scores.update(new_scores)\n",
    "    score_xml = \"\"\"\n",
    "    <ROSETTASCRIPTS>\n",
    "        <SCOREFXNS>\n",
    "                <ScoreFunction name=\"sfxn\" weights=\"beta_nov16\" />\n",
    "        </SCOREFXNS>\n",
    "        <MOVERS>\n",
    "            <SwitchChainOrder name=\"chain1only\" chain_order=\"1\" />\n",
    "            <RollMover name=\"move_chainA_far_away\" chain=\"1\" min_angle=\"0\" max_angle=\"0\" axis=\"x\" >\n",
    "                <translate x=\"1000\" y=\"1000\" z=\"1000\" /> \n",
    "            </RollMover>\n",
    "        </MOVERS>\n",
    "        <RESIDUE_SELECTORS>\n",
    "            <Chain name=\"chainA\" chains=\"A\"/>\n",
    "            <Chain name=\"chainB\" chains=\"B\"/>\n",
    "            <Neighborhood name=\"interface_chA\" selector=\"chainB\" distance=\"10.0\" />\n",
    "            <Neighborhood name=\"interface_chB\" selector=\"chainA\" distance=\"10.0\" />\n",
    "            <And name=\"AB_interface\" selectors=\"interface_chA,interface_chB\" />\n",
    "            <Not name=\"Not_interface\" selector=\"AB_interface\" />\n",
    "            <And name=\"actual_interface_chA\" selectors=\"AB_interface,chainA\" />\n",
    "            <And name=\"actual_interface_chB\" selectors=\"AB_interface,chainB\" />\n",
    "            <ResidueName name=\"apolar\" residue_name3=\"ALA,CYS,PHE,ILE,LEU,MET,THR,PRO,VAL,TRP,TYR\" />\n",
    "            <And name=\"apolar_A\" selectors=\"apolar,actual_interface_chA\" />\n",
    "            <And name=\"apolar_B\" selectors=\"apolar,actual_interface_chB\" />\n",
    "            <Index name=\"patchdock_res\" resnums=\"{interface_hydrophobics}\" />\n",
    "            <True name=\"true_sel\" />\n",
    "        </RESIDUE_SELECTORS>\n",
    "        <SIMPLE_METRICS>\n",
    "            <SapScoreMetric name=\"sap_score\" score_selector=\"chainA\" />\n",
    "            <SapScoreMetric name=\"sap_score_target\" score_selector=\"chainB\" />\n",
    "            <SapScoreMetric name=\"binder_blocked_sap\" score_selector=\"chainA\" sap_calculate_selector=\"chainA\" sasa_selector=\"true_sel\" />\n",
    "            <SapScoreMetric name=\"target_blocked_sap\" score_selector=\"chainB\" sap_calculate_selector=\"chainB\" sasa_selector=\"true_sel\" />\n",
    "        </SIMPLE_METRICS>\n",
    "        <TASKOPERATIONS>\n",
    "            <ProteinInterfaceDesign name=\"pack_long\" design_chain1=\"0\" design_chain2=\"0\" jump=\"1\" interface_distance_cutoff=\"15\"/>\n",
    "        </TASKOPERATIONS>\n",
    "        <MOVERS>\n",
    "            <TaskAwareMinMover name=\"min\" scorefxn=\"sfxn\" bb=\"0\" chi=\"1\" task_operations=\"pack_long\" />\n",
    "        </MOVERS>\n",
    "        <FILTERS>\n",
    "            <ContactMolecularSurface name=\"contact_patch\" distance_weight=\"0.5\" target_selector=\"patchdock_res\" binder_selector=\"chainA\" confidence=\"0\" />\n",
    "            <Ddg name=\"ddg\"  threshold=\"-10\" jump=\"1\" repeats=\"5\" repack=\"1\" relax_mover=\"min\" confidence=\"0\" scorefxn=\"sfxn\" extreme_value_removal=\"1\" />\n",
    "            <ResidueCount name=\"res_count_all\" max_residue_count=\"9999\" confidence=\"0\"/>\n",
    "            <BuriedSurfaceArea name=\"buried_npsa_FAMILYVW\" select_only_FAMILYVW=\"True\" atom_mode=\"hydrophobic_atoms\" confidence=\"0.0\" />\n",
    "            <BuriedSurfaceArea name=\"buried_npsa\" select_only_FAMILYVW=\"False\" atom_mode=\"hydrophobic_atoms\" confidence=\"0.0\" />\n",
    "            <ExposedHydrophobics name=\"exposed_hydrophobics\"  confidence=\"0\" />\n",
    "            <CalculatorFilter name=\"buried_npsa_per_res\" equation=\"total_score / res\" threshold=\"-3.2\" confidence=\"0\">\n",
    "                <Var name=\"total_score\" filter=\"buried_npsa\"/>\n",
    "                <Var name=\"res\" filter=\"res_count_all\"/>\n",
    "            </CalculatorFilter>\n",
    "            <CalculatorFilter name=\"buried_npsa_FAMILYVW_per_res\" equation=\"total_score / res\" threshold=\"-3.2\" confidence=\"0\">\n",
    "                <Var name=\"total_score\" filter=\"buried_npsa_FAMILYVW\"/>\n",
    "                <Var name=\"res\" filter=\"res_count_all\"/>\n",
    "            </CalculatorFilter>\n",
    "            <MoveBeforeFilter name=\"buried_npsa_FAMILYVW_monomer\" mover=\"chain1only\" filter=\"buried_npsa_FAMILYVW\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"buried_npsa_monomer\" mover=\"chain1only\" filter=\"buried_npsa\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"exposed_hydrophobics_monomer\" mover=\"chain1only\" filter=\"exposed_hydrophobics\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"buried_npsa_per_res_monomer\" mover=\"chain1only\" filter=\"buried_npsa_per_res\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"buried_npsa_FAMILYVW_per_res_monomer\" mover=\"chain1only\" filter=\"buried_npsa_FAMILYVW_per_res\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"buried_npsa_FAMILYVW_apo\" mover=\"move_chainA_far_away\" filter=\"buried_npsa_FAMILYVW\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"buried_npsa_apo\" mover=\"move_chainA_far_away\" filter=\"buried_npsa\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"exposed_hydrophobics_apo\" mover=\"move_chainA_far_away\" filter=\"exposed_hydrophobics\" confidence=\"0\" />\n",
    "            <CalculatorFilter   name=\"delta_buried_npsa_FAMILYVW\" equation=\"complex - apo\" threshold=\"-3.2\" confidence=\"0\">\n",
    "                <Var name=\"complex\" filter=\"buried_npsa_FAMILYVW\"/>\n",
    "                <Var name=\"apo\"     filter=\"buried_npsa_FAMILYVW_apo\"/>\n",
    "            </CalculatorFilter>\n",
    "            <CalculatorFilter   name=\"delta_buried_npsa\" equation=\"complex - apo\" threshold=\"-3.2\" confidence=\"0\">\n",
    "                <Var name=\"complex\" filter=\"buried_npsa\"/>\n",
    "                <Var name=\"apo\"     filter=\"buried_npsa_apo\"/>\n",
    "            </CalculatorFilter>\n",
    "            <CalculatorFilter   name=\"delta_exposed_hydrophobics\" equation=\"complex - apo\" threshold=\"-3.2\" confidence=\"0\">\n",
    "                <Var name=\"complex\" filter=\"exposed_hydrophobics\"/>\n",
    "                <Var name=\"apo\"     filter=\"exposed_hydrophobics_apo\"/>\n",
    "            </CalculatorFilter>\n",
    "            <ContactMolecularSurface name=\"contact_molecular_surface_ap_target\" distance_weight=\"0.5\" target_selector=\"apolar_B\" binder_selector=\"chainA\" confidence=\"0\" />\n",
    "            <ContactMolecularSurface name=\"contact_molec_sq5_ap_target\" distance_weight=\"0.5\" target_selector=\"apolar_B\" binder_selector=\"chainA\" confidence=\"0\" near_squared_size=\"5\" />\n",
    "            <ContactMolecularSurface name=\"contact_molecular_surface_apap_target\" distance_weight=\"0.5\" target_selector=\"apolar_B\" binder_selector=\"chainA\" confidence=\"0\" apolar_target=\"true\" />\n",
    "            <ContactMolecularSurface name=\"contact_molec_sq5_apap_target\" distance_weight=\"0.5\" target_selector=\"apolar_B\" binder_selector=\"chainA\" confidence=\"0\" near_squared_size=\"5\" apolar_target=\"true\" />\n",
    "            <ContactMolecularSurface name=\"contact_molecular_surface_ap_binder\" distance_weight=\"0.5\" target_selector=\"apolar_A\" binder_selector=\"chainB\" confidence=\"0\" />\n",
    "            <ContactMolecularSurface name=\"contact_molec_sq5_ap_binder\" distance_weight=\"0.5\" target_selector=\"apolar_A\" binder_selector=\"chainB\" confidence=\"0\" near_squared_size=\"5\" />\n",
    "            <ContactMolecularSurface name=\"contact_molecular_surface_apap_binder\" distance_weight=\"0.5\" target_selector=\"apolar_A\" binder_selector=\"chainB\" confidence=\"0\" apolar_target=\"true\" />\n",
    "            <ContactMolecularSurface name=\"contact_molec_sq5_apap_binder\" distance_weight=\"0.5\" target_selector=\"apolar_A\" binder_selector=\"chainB\" confidence=\"0\" near_squared_size=\"5\" apolar_target=\"true\" />\n",
    "            <worst9mer name=\"pre_wnm_all\" rmsd_lookup_threshold=\"0.4\" confidence=\"0\" />\n",
    "            <worst9mer name=\"pre_wnm_hlx\" rmsd_lookup_threshold=\"0.4\" confidence=\"0\" only_helices=\"true\" />\n",
    "            <MoveBeforeFilter name=\"wnm_all\" mover=\"chain1only\" filter=\"pre_wnm_all\" confidence=\"0\" />\n",
    "            <MoveBeforeFilter name=\"wnm_hlx\" mover=\"chain1only\" filter=\"pre_wnm_hlx\" confidence=\"0\" />\n",
    "        </FILTERS>\n",
    "        <PROTOCOLS>\n",
    "            <Add filter=\"contact_molecular_surface_ap_target\" />\n",
    "            <Add filter=\"contact_molec_sq5_ap_target\" />\n",
    "            <Add filter=\"contact_molecular_surface_apap_target\" />\n",
    "            <Add filter=\"contact_molec_sq5_apap_target\" />\n",
    "            <Add filter=\"contact_patch\"/>\n",
    "            <Add filter=\"ddg\"/>\n",
    "            <Add filter=\"wnm_all\"/>\n",
    "            <Add filter=\"wnm_hlx\"/>\n",
    "            <Add metrics=\"sap_score\" />\n",
    "            <Add metrics=\"sap_score_target\" />\n",
    "            <Add metrics=\"binder_blocked_sap\" />\n",
    "            <Add metrics=\"target_blocked_sap\" />\n",
    "        </PROTOCOLS>\n",
    "        <OUTPUT />\n",
    "    </ROSETTASCRIPTS>\n",
    "    \"\"\".format(\n",
    "        interface_hydrophobics=scores[\"interface_hydrophobics\"],\n",
    "    )\n",
    "    scorer = SingleoutputRosettaScriptsTask(score_xml)\n",
    "    scored_ppose = scorer(designed_ppose.pose.clone())\n",
    "    final_scores = dict(scored_ppose.pose.scores)\n",
    "    scores.update(final_scores)\n",
    "    scores[\"file_in\"] = file\n",
    "    scored_pose = io.to_pose(scored_ppose)\n",
    "    for key, value in scores.items():\n",
    "        pyrosetta.rosetta.core.pose.setPoseExtraScore(scored_pose, key, str(value))\n",
    "    final_ppose = io.to_packed(scored_pose)\n",
    "    return final_ppose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We're going to do a huge run here, more than it makes sense to use dask for\n",
    "Will submit everything to `backfill`\n",
    "Make `SLURM` array tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run the following commands\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_0.cmds | wc -l) 02_rerun_0.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_1.cmds | wc -l) 02_rerun_1.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_2.cmds | wc -l) 02_rerun_2.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_3.cmds | wc -l) 02_rerun_3.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_4.cmds | wc -l) 02_rerun_4.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_5.cmds | wc -l) 02_rerun_5.sh\n",
      "sbatch -a 1-$(cat 02_rerun_tasks_6.cmds | wc -l) 02_rerun_6.sh\n"
     ]
    }
   ],
   "source": [
    "import os, stat, subprocess\n",
    "\n",
    "\n",
    "def create_tasks(selected):\n",
    "    with open(selected, \"r\") as f:\n",
    "        for file in f:\n",
    "            tasks = {}\n",
    "            path = file.rstrip()\n",
    "            tasks[\"-s\"] = path\n",
    "            tasks[\"-sasa_threshold\"] = 0\n",
    "            yield tasks\n",
    "\n",
    "\n",
    "def file_len(file):\n",
    "    \"\"\"https://stackoverflow.com/questions/845058/how-to-get-line-count-of-a-large-file-cheaply-in-python\"\"\"\n",
    "    with open(file) as f:\n",
    "        for i, l in enumerate(f):\n",
    "            pass\n",
    "    return i + 1\n",
    "\n",
    "\n",
    "targets = os.path.join(os.getcwd(), \"01_filter/rerun.list\")\n",
    "target_scale = int(100000 / file_len(targets))\n",
    "almost_intdes = os.path.join(os.getcwd(), \"almost_intdes.py\")\n",
    "\n",
    "jid = \"{SLURM_JOB_ID%;*}\"\n",
    "sid = \"{SLURM_ARRAY_TASK_ID}p\"\n",
    "\n",
    "for i in range(7):\n",
    "    tasklist = f\"02_rerun_tasks_{i}.cmds\"\n",
    "    run_sh = \"\"\"#!/usr/bin/env bash \\n#SBATCH -J almost_intdes \\n#SBATCH -e /mnt/home/pleung/logs/slurm_logs/almost_intdes-%J.err \\n#SBATCH -o /mnt/home/pleung/logs/slurm_logs/almost_intdes-%J.out \\n#SBATCH -p {queue} \\n#SBATCH --mem=4G \\n\\nJOB_ID=${jid} \\nCMD=$(sed -n \"${sid}\" {tasklist}) \\necho \"${c}\" | bash\"\"\".format(\n",
    "        queue=\"backfill\", jid=jid, sid=sid, tasklist=tasklist, c=\"{CMD}\"\n",
    "    )\n",
    "    shell = f\"02_rerun_{i}.sh\"\n",
    "    with open(shell, \"w+\") as f:\n",
    "        print(run_sh, file=f)\n",
    "    st = os.stat(shell)\n",
    "    os.chmod(shell, st.st_mode | stat.S_IEXEC)\n",
    "    with open(f\"02_rerun_tasks_{i}.cmds\", \"w+\") as f:\n",
    "        for nstruct in range(0, target_scale):\n",
    "            outpath = os.path.join(os.getcwd(), f\"02_reruns_{i}\")\n",
    "            full_outpath = os.path.join(os.getcwd(), outpath, f\"{nstruct}\")\n",
    "            for tasks in create_tasks(targets):\n",
    "                args_ = \" \".join([\" \".join([k, str(v)]) for k, v in tasks.items()])\n",
    "                cmd = f\"mkdir -p {full_outpath}; cd {full_outpath}; {almost_intdes} {args_}\"\n",
    "                print(cmd, file=f)\n",
    "\n",
    "# Let's go\n",
    "print(\"Run the following commands\")\n",
    "for i in range(7):\n",
    "    print(f\"sbatch -a 1-$(cat 02_rerun_tasks_{i}.cmds | wc -l) 02_rerun_{i}.sh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now need to retrieve the JSONs containing scores for the interface design runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94154 runs from 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94154/94154 [00:19<00:00, 4837.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94164 runs from 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94164/94164 [00:19<00:00, 4860.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94160 runs from 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94160/94160 [00:19<00:00, 4918.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94160 runs from 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94160/94160 [00:19<00:00, 4876.49it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94161 runs from 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94161/94161 [00:19<00:00, 4858.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94132 runs from 5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94132/94132 [00:19<00:00, 4922.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 94160 runs from 6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 94160/94160 [00:18<00:00, 5212.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 47s, sys: 40.5 s, total: 3min 27s\n",
      "Wall time: 4min 44s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import json, time\n",
    "\n",
    "\n",
    "output_path = os.path.join(os.getcwd(), \"02_filter\")\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "targets = range(0,7)\n",
    "for target in targets:\n",
    "    output_paths = glob(os.path.join(os.getcwd(), f\"02_reruns_{target}/*/*.json\"))\n",
    "    print(f\"reading {len(output_paths)} runs from {target}\")\n",
    "    time.sleep(0.5)\n",
    "    scores = {}\n",
    "    for path in tqdm(output_paths):\n",
    "        pdb = path.replace(\"json\", \"pdb.bz2\")\n",
    "        key = f\"{pdb}\"\n",
    "        with open(path, \"r\") as f:\n",
    "            try:\n",
    "                values = json.load(f)\n",
    "            except:\n",
    "                print(path)\n",
    "        scores[key] = values\n",
    "    scores_df = pd.DataFrame(scores)\n",
    "    scores_df = scores_df.T\n",
    "    output_file = os.path.join(output_path, f\"{target}_scores.json\")\n",
    "    scores_df.to_json(output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reread dumped scorefiles\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [02:58<00:00, 25.52s/it]\n"
     ]
    }
   ],
   "source": [
    "targets = range(0, 7)\n",
    "rerun_dfs = {}\n",
    "output_path = os.path.join(os.getcwd(), \"02_filter\")\n",
    "for target in tqdm(targets):\n",
    "    scores = os.path.join(output_path, f\"{target}_scores.json\")\n",
    "    rerun_dfs[(str(target))] = pd.read_json(scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pickle dict of dfs split by target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def row2state(row):\n",
    "    state = (\n",
    "        row[\"parent\"]\n",
    "        + \"_p_\"\n",
    "        + str(int(row[\"pivot_helix\"]))\n",
    "        + \"_s_\"\n",
    "        + str(int(row[\"shift\"]))\n",
    "        + \"_d_\"\n",
    "        + str(int(row[\"docked_helix\"]))\n",
    "    )\n",
    "    return state\n",
    "\n",
    "\n",
    "for _, df in rerun_dfs.items():\n",
    "    df[\"state\"] = df.apply(row2state, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"rerun_dfs.pkl\"), \"wb\") as handle:\n",
    "    pickle.dump(rerun_dfs, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.67 s, sys: 1.24 s, total: 2.91 s\n",
      "Wall time: 4.13 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import pickle\n",
    "\n",
    "\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"rerun_dfs.pkl\"), \"rb\") as handle:\n",
    "    rerun_dfs = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mv: cannot stat '02*.cmds': No such file or directory\n",
      "mv: cannot stat '02*.sh': No such file or directory\n"
     ]
    }
   ],
   "source": [
    "!mv 02*.cmds 02_filter; mv 02*.sh 02_filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Collect and combine designs from reruns and original intdes runs\n",
    "First, filter the reruns and save as a pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "189\n",
      "44165\n",
      "50\n",
      "{'THR3H_4', 'THR11', 'KH_PEP12_PLP', 'THR13'}\n",
      "1\n",
      "190\n",
      "44184\n",
      "52\n",
      "{'THR3H_4', 'THR11'}\n",
      "2\n",
      "190\n",
      "44112\n",
      "51\n",
      "{'THR3H_4', 'THR11', 'KH_PEP12_PLP'}\n",
      "3\n",
      "191\n",
      "43966\n",
      "51\n",
      "{'THR3H_4', 'THR11', 'KH_PEP12_PLP'}\n",
      "4\n",
      "190\n",
      "44288\n",
      "51\n",
      "{'THR3H_4', 'THR11', 'KH_PEP12_PLP'}\n",
      "5\n",
      "189\n",
      "44083\n",
      "52\n",
      "{'KH_PEP12_PLP', 'THR11'}\n",
      "6\n",
      "190\n",
      "43923\n",
      "52\n",
      "{'KH_PEP12_PLP', 'THR11'}\n"
     ]
    }
   ],
   "source": [
    "filtered_dfs = {}\n",
    "output_path = os.path.join(os.getcwd(), \"02_filter\")\n",
    "for target, df in rerun_dfs.items():\n",
    "    filtered = df[df[\"ddg\"] < -40]\n",
    "    filtered = filtered[filtered[\"interface_buried_sasa\"] > 1400]\n",
    "    filtered = filtered[filtered[\"score_per_res\"] < -3]\n",
    "    filtered = filtered[filtered[\"wnm_all\"] < 0.8]\n",
    "    filtered = filtered[filtered[\"wnm_hlx\"] < 0.15]\n",
    "    print(target)\n",
    "    print(len(set(filtered.state.values)))\n",
    "    print(len(filtered))\n",
    "    print(len(set(filtered.parent.values)))\n",
    "    print(set(df.parent.values) - set(filtered.parent.values))\n",
    "    filtered_dfs[target] = filtered\n",
    "\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"filtered_dfs.pkl\"), \"wb\") as handle:\n",
    "    pickle.dump(filtered_dfs, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reload pickle and concatenate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14\n"
     ]
    }
   ],
   "source": [
    "with open(os.path.join(os.getcwd(), \"01_filter\", \"filtered_dfs.pkl\"), \"rb\") as handle:\n",
    "    filtered_dfs.update(pickle.load(handle))\n",
    "print(len(filtered_dfs))\n",
    "merged_df = pd.concat(list(filtered_dfs.values()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Moved selected decoys to a new dir and generate new merged dict with updated paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "447611it [10:43, 696.00it/s]\n"
     ]
    }
   ],
   "source": [
    "output_path = os.path.join(os.getcwd(), \"02_filter\", \"decoys\")\n",
    "subdirs = []\n",
    "i = 0\n",
    "new_scores = {}\n",
    "for index, row in tqdm(merged_df.iterrows()):\n",
    "\n",
    "    subdir = int(i / 1000)\n",
    "    file = index.split(\"/\")[-1]\n",
    "    dest_path = os.path.join(output_path, str(subdir))\n",
    "    os.makedirs(dest_path, exist_ok=True)\n",
    "    full_path = os.path.join(dest_path, file)\n",
    "    shutil.move(index, full_path)\n",
    "    new_scores[full_path] = dict(row)\n",
    "    i += 1\n",
    "\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"merged_scores.pkl\"), \"wb\") as handle:\n",
    "    pickle.dump(new_scores, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter slightly more stringently"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "198\n",
      "181414\n",
      "53\n",
      "{'THR11'}\n"
     ]
    }
   ],
   "source": [
    "merged_df = pd.DataFrame(new_scores).T\n",
    "\n",
    "output_path = os.path.join(os.getcwd(), \"02_filter\")\n",
    "filtered = merged_df[merged_df[\"vbuns5.0_heavy_ball_1.1D\"] < 1]\n",
    "filtered = filtered[filtered[\"binder_blocked_sap\"] > 12]\n",
    "\n",
    "print(len(set(filtered.state.values)))\n",
    "print(len(filtered))\n",
    "print(len(set(filtered.parent.values)))\n",
    "print(set(merged_df.parent.values) - set(filtered.parent.values))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dump jsons split by target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/7 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "glucagon 51922\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 1/7 [00:06<00:41,  6.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gip 82179\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▊       | 2/7 [00:16<00:43,  8.68s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "glp 89377\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 3/7 [00:28<00:39,  9.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pth 59432\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 4/7 [00:36<00:27,  9.12s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "apoe 58634\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 71%|███████▏  | 5/7 [00:43<00:17,  8.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "covstem 9375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 6/7 [00:45<00:06,  6.12s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "neuropeptideY 96692\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:57<00:00,  8.21s/it]\n"
     ]
    }
   ],
   "source": [
    "for target in tqdm(set(merged_df.target_name.values)):\n",
    "    sub_df = merged_df[merged_df[\"target_name\"] == target]\n",
    "    print(target, len(sub_df))\n",
    "    output_file = os.path.join(os.getcwd(), \"02_filter\", f\"{target}_merged.json\")\n",
    "    sub_df.to_json(output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add refiltered dfs to dict under target names and dump as pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:01<00:00,  4.20it/s]\n"
     ]
    }
   ],
   "source": [
    "refiltered_dfs = {}\n",
    "for target in tqdm(set(merged_df.target_name.values)):\n",
    "    sub_df = merged_df[merged_df[\"target_name\"] == target]\n",
    "    refiltered_dfs[target] = sub_df\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"refiltered_dfs.pkl\"), \"wb\") as handle:\n",
    "    pickle.dump(refiltered_dfs, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup out of date scorefiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm 02_filter/*_scores.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reload refiltered dfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [04:02<00:00, 34.70s/it]\n"
     ]
    }
   ],
   "source": [
    "targets = [\"apoe\", \"covstem\", \"gip\", \"glp\", \"glucagon\", \"neuropeptideY\", \"pth\"]\n",
    "refiltered_dfs = {}\n",
    "for target in tqdm(targets):\n",
    "    input_file = os.path.join(os.getcwd(), \"02_filter\", f\"{target}_merged.json\")\n",
    "    refiltered_dfs[target] = pd.read_json(input_file, dtype=float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure out how many sequences could be ordered on a chip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "apoe\n",
      "1865\n",
      "covstem\n",
      "913\n",
      "gip\n",
      "8697\n",
      "glp\n",
      "11965\n",
      "glucagon\n",
      "6239\n",
      "neuropeptideY\n",
      "13983\n",
      "pth\n",
      "10763\n",
      "54425\n"
     ]
    }
   ],
   "source": [
    "total = 0\n",
    "cutoff = 175\n",
    "for t, df in refiltered_dfs.items():\n",
    "    print(t)\n",
    "    print(len(df[df[\"chB_start\"] < cutoff]))\n",
    "    total += len(df[df[\"chB_start\"] < cutoff])\n",
    "\n",
    "print(total)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make a final list of every design that passes filters for further detailing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "447611\n"
     ]
    }
   ],
   "source": [
    "to_detail = []\n",
    "for t, df in refiltered_dfs.items():\n",
    "    for decoy in df.index:\n",
    "        to_detail.append(decoy)\n",
    "\n",
    "print(len(to_detail))\n",
    "with open(os.path.join(os.getcwd(), \"02_filter\", \"to_detail.list\"), \"w\") as f:\n",
    "    for decoy in to_detail:\n",
    "        print(decoy, file=f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete directories of the designs that didn't pass filters for intdes runs and reruns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rm: cannot remove '02_reruns_*': No such file or directory\n"
     ]
    }
   ],
   "source": [
    "!rm -r 01_intdes_runs_*; rm -r 02_reruns_*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unused blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import pyrosetta\n",
    "from pyrosetta.distributed import cluster\n",
    "import pyrosetta.distributed.io as io\n",
    "\n",
    "flags = \"\"\"\n",
    "-out:level 300\n",
    "-corrections::beta_nov16 true\n",
    "-holes:dalphaball /home/bcov/ppi/tutorial_build/main/source/external/DAlpahBall/DAlphaBall.gcc\n",
    "-indexed_structure_store:fragment_store /home/bcov/sc/scaffold_comparison/data/ss_grouped_vall_all.h5\n",
    "\"\"\"\n",
    "pyrosetta.distributed.init(\" \".join(flags.replace(\"\\n\\t\", \" \").split()))\n",
    "t = almost_intdes(\n",
    "    None,\n",
    "    **{\n",
    "        \"-s\": \"/mnt/home/pleung/projects/peptide_binders/r0/peptide_binders/00_thread_targets/decoys/0000/2021.07.12.10.35.38.485279_00431f82398d4a56b6d2e29532614d17.pdb.bz2\",\n",
    "        \"sasa_threshold\": 0\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t.pose.dump_pdb(\"test.pdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "phil (3.8.2)",
   "language": "python",
   "name": "phil"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}