<jupyter_start><jupyter_text>AlphaFold development suite  InstructionsA. SETUP:  Run cells 1-3 to set up or hit Runtime/Run all to run everything (5 min.)B. DEVELOPMENT CYCLE:   1a. Either edit files in the alphafold directory, or   1b. Edit files in github and hit "Load current alphafold_working from github"  2. Run cell B below to run Alphafold (2 min.)<jupyter_code>#@markdown 1. Set up imports and load dependencies...this is the slow step
import os, sys
import os.path
import re
import hashlib
from pathlib import Path
from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from google.colab import files
import shutil
from string import ascii_uppercase

! echo "Installing biopython ..."
!  pip -q install biopython dm-haiku ml-collections py3Dmol

! echo "Downloading model parameters..."
!    rm -rf params
!    mkdir params
!    curl -fsSL https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar  | tar x -C params

!echo "Downloading jq curl zlib1g gawk..."
!    apt-get -qq -y update 2>&1 1>/dev/null
!    apt-get -qq -y install jq curl zlib1g gawk 2>&1 1>/dev/null

! echo "Setting up conda..."
!    wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
!    bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null
!    rm Miniconda3-latest-Linux-x86_64.sh

! echo "Setting up template search methods..."

! conda install -y -q -c conda-forge -c bioconda kalign3=3.2.2 hhsuite=3.3.0 python=3.7 2>&1 1>/dev/null

! echo "Installing openmm..."
! conda install -qy -c omnia openmm 2>&1 1>/dev/null

! echo "Installing Mock..."
!  pip install mock
 
! echo "Setting paths to site-packages and dist-packages..."
sys.path.append("/usr/local/lib/python3.7/site-packages")
sys.path.append("/usr/local/lib/python3.7/dist-packages")
sys.path.append("/usr/local/lib/python3.7/site-packages/simtk/")

target_all_atom_positions = None # Initialize
target_prot = None # initialize


! echo "Done with loading dependencies"
#@markdown 2. Load current alphafold_working from phenix-project github.  
#@markdown  This is quick...
#@markdown cycle back here after making a change in github alphafold_working
! echo "installing alphafold from https://github.com/phenix-project/af_development.git in /contents/alphafold"

! rm -rf af_development
! rm -rf alphafold
! rm -f colabfold.py


! git clone https://github.com/phenix-project/af_development.git --quiet
! (cd af_development; git checkout --quiet)
! mv af_development/alphafold alphafold
! mv af_development/colabfold.py .
! mv af_development/run_alphafold.py .
! mv alphafold/run_alphafold_test.py .



!    # remove "END" from PDBs, otherwise biopython complains
!    sed -i "s/pdb_lines.append('END')//" /content/alphafold/common/protein.py
!    sed -i "s/pdb_lines.append('ENDMDL')//" /content/alphafold/common/protein.py


! echo "Ready with alphafold in /content/alphafold"
! ls -ltr /content/alphafold

! echo "clearing out /tmp/absl_testing/"
! rm -rf /tmp/absl_testing/
! mkdir /tmp/absl_testing/

! echo " Clearing python caches ..."
for x in list(sys.modules.keys(  )) + list(globals()):
  for key in ['alphafold','protein', 'Alphafold', 'Protein', 'colabfold']:
    if x.find(key)>-1:
      if x in list(sys.modules.keys()):
        
        del(sys.modules[x])
      if x in list(globals().keys()):
      
        del globals()[x]
        assert not x in list(globals().keys())
        break


if not os.environ['PYTHONPATH'].find(":/opt/conda/bin")>-1:
  os.environ['PYTHONPATH']+=":/opt/conda/bin"
  os.environ['PYTHONPATH']+=":/opt/conda/lib/python3.7/site-packages"
  os.environ['PYTHONPATH']+=":/usr/local/lib/python3.7/dist-packages"
! echo "`grep Version run_alphafold_test.py|grep -v Apache`"
! echo "Done loading current version."
#@markdown 3a. Run ProteinTest (optional, local code)

TEST_DATA_DIR = 'alphafold/common/testdata/'

from absl.testing import parameterized
from absl.testing import absltest
from alphafold.common import protein
import numpy as np
class ProteinTest(parameterized.TestCase):

  def _check_shapes(self, prot, num_res):
    """Check that the processed shapes are correct."""
    num_atoms = residue_constants.atom_type_num
    self.assertEqual((num_res, num_atoms, 3), prot.atom_positions.shape)
    self.assertEqual((num_res,), prot.aatype.shape)
    self.assertEqual((num_res, num_atoms), prot.atom_mask.shape)
    self.assertEqual((num_res,), prot.residue_index.shape)
    self.assertEqual((num_res, num_atoms), prot.b_factors.shape)
    print("Finished _check_shapes")

  @parameterized.parameters(('2rbg.pdb', 'A', 282),
                            ('2rbg.pdb', 'B', 282))
  def test_from_pdb_str(self, pdb_file, chain_id, num_res):
    pdb_file = os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
                            pdb_file)
    with open(pdb_file) as f:
      pdb_string = f.read()
    prot = protein.from_pdb_string(pdb_string, chain_id)
    self._check_shapes(prot, num_res)
    print("Total residues: %s" %(prot.aatype.shape))
    self.assertGreaterEqual(prot.aatype.min(), 0)
    # Allow equal since unknown restypes have index equal to restype_num.
    self.assertLessEqual(prot.aatype.max(), residue_constants.restype_num)
    print("Finished test_from_pdb_str")

  def test_to_pdb(self):
    with open(
        os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
                     '2rbg.pdb')) as f:
      pdb_string = f.read()
    prot = protein.from_pdb_string(pdb_string, chain_id='A')
    pdb_string_reconstr = protein.to_pdb(prot)
    prot_reconstr = protein.from_pdb_string(pdb_string_reconstr)
    print("Total residues: %s" %(prot.aatype.shape))

    np.testing.assert_array_equal(prot_reconstr.aatype, prot.aatype)
    np.testing.assert_array_almost_equal(
        prot_reconstr.atom_positions, prot.atom_positions)
    np.testing.assert_array_almost_equal(
        prot_reconstr.atom_mask, prot.atom_mask)
    np.testing.assert_array_equal(
        prot_reconstr.residue_index, prot.residue_index)
    np.testing.assert_array_almost_equal(
        prot_reconstr.b_factors, prot.b_factors)
    print("Finished test_to_pdb")

  def test_ideal_atom_mask(self):
    with open(
        os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
                     '2rbg.pdb')) as f:
      pdb_string = f.read()
    prot = protein.from_pdb_string(pdb_string, chain_id='A')
    print("Total residues: %s" %(prot.aatype.shape))
    
    ideal_mask = protein.ideal_atom_mask(prot)
    non_ideal_residues = set([102] + list(range(127, 285)))
    for i, (res, atom_mask) in enumerate(
        zip(prot.residue_index, prot.atom_mask)):
      if res in non_ideal_residues:
        self.assertFalse(np.all(atom_mask == ideal_mask[i]), msg=f'{res}')
      else:
        self.assertTrue(np.all(atom_mask == ideal_mask[i]), msg=f'{res}')
    print("Finished test_ideal_atom_mask")
    
t = ProteinTest()

t.test_to_pdb()
t.test_ideal_atom_mask()


#@markdown 3c. test lddt (optional, code in repository)


from alphafold.model.all_atom_test import AllAtomTest
aat = AllAtomTest()

aat.test_frame_aligned_point_error_perfect_on_global_transform_rot_174_trans_1()

#@markdown 3d. Run run_alphafold_test (optional, code in repository)
from run_alphafold_test import RunAlphafoldTest
r= RunAlphafoldTest()
r.test_end_to_end()
#@markdown 3e. RunAlphafoldTest (optional, local code)
import os

from absl.testing import absltest
from absl.testing import parameterized
from run_alphafold import predict_structure
from unittest import mock
import numpy as np
import run_alphafold
# Internal import (7716).


class RunAlphafoldTest(parameterized.TestCase):

  def test_end_to_end(self):

    data_pipeline_mock = mock.Mock()
    model_runner_mock = mock.Mock()
    amber_relaxer_mock = mock.Mock()

    data_pipeline_mock.process.return_value = {}
    model_runner_mock.process_features.return_value = {
        'aatype': np.zeros((12, 10), dtype=np.int32),
        'residue_index': np.tile(np.arange(10, dtype=np.int32)[None], (12, 1)),
    }
    model_runner_mock.predict.return_value = {
        'structure_module': {
            'final_atom_positions': np.zeros((10, 37, 3)),
            'final_atom_mask': np.ones((10, 37)),
        },
        'predicted_lddt': {
            'logits': np.ones((10, 50)),
        },
        'plddt': np.ones(10) * 42,
        'ptm': np.array(0.),
        'aligned_confidence_probs': np.zeros((10, 10, 50)),
        'predicted_aligned_error': np.zeros((10, 10)),
        'max_predicted_aligned_error': np.array(0.),
    }
    amber_relaxer_mock.process.return_value = ('RELAXED', None, None)

    fasta_path = os.path.join(absltest.get_default_test_tmpdir(),
                              'target.fasta')
    with open(fasta_path, 'wt') as f:
      f.write('>A\nAAAAAAAAAAAAA')
    fasta_name = 'test'

    out_dir = absltest.get_default_test_tmpdir()

    run_alphafold.predict_structure(
        fasta_path=fasta_path,
        fasta_name=fasta_name,
        output_dir_base=out_dir,
        data_pipeline=data_pipeline_mock,
        model_runners={'model1': model_runner_mock},
        amber_relaxer=amber_relaxer_mock,
        benchmark=False,
        random_seed=0)

    base_output_files = os.listdir(out_dir)
    self.assertIn('target.fasta', base_output_files)
    self.assertIn('test', base_output_files)


    target_output_files = os.listdir(os.path.join(out_dir, 'test'))
    self.assertCountEqual(
        ['features.pkl', 'msas', 'ranked_0.pdb', 'ranking_debug.json',
         'relaxed_model1.pdb', 'result_model1.pkl', 'timings.json',
         'unrelaxed_model1.pdb'], target_output_files)

    # Check that pLDDT is set in the B-factor column.
    with open(os.path.join(out_dir, 'test', 'unrelaxed_model1.pdb')) as f:
      for line in f:
        if line.startswith('ATOM'):
          self.assertEqual(line[61:66], '42.00')
    print("End-to-end completed successfully")      

r= RunAlphafoldTest()
r.test_end_to_end()

#@markdown Upload files to /content/alphafold/model/ .
#@markdown Check box to select and upload files.
upload_files = False #@param {type:"boolean"}
if upload_files:
  uploaded = files.upload()
  for filename,contents in uploaded.items():
   if filename == 'protein.py':
     filepath = Path("/content/alphafold/common",filename)
   else:
     filepath = Path("/content/alphafold/model",filename)
   with filepath.open("w") as fh:
        fh.write(contents.decode("UTF-8"))
   print("Uploaded to %s" %(filepath))
#@markdown Optional code goes here. Type in anything, check box, and run cell
run_optional_code = False #@param {type:"boolean"}

! # use ! at start of line to indicate bash command
#   use any python command as-is
#   comment out commands you don't want to run
! # mv /content/folding.py /content/alphafold/model/
! # mv /content/modules.py /content/alphafold/model/

# USER INPUT SECTION

# IMPORTS, STANDARD PARAMETERS AND METHODS

import os, sys
import os.path
import re
import hashlib
from pathlib import Path
from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from google.colab import files
import shutil
from string import ascii_uppercase

# Local methods

def add_hash(x,y):
  return x+"_"+hashlib.sha1(y.encode()).hexdigest()[:5]

def clear_directories(all_dirs):

  for d in all_dirs:
    if d.exists():
      shutil.rmtree(d)
    d.mkdir(parents=True)


def clean_query(query_sequence):
  query_sequence = "".join(query_sequence.split())
  query_sequence = re.sub(r'[^a-zA-Z]','', query_sequence).upper()
  return query_sequence

def clean_jobname(jobname):
  jobname = "".join(jobname.split())
  jobname = re.sub(r'\W+', '', jobname)
  if len(jobname.split("_")) == 1:
    jobname = add_hash(jobname, query_sequence)
  return jobname

def save_sequence(jobname, query_sequence):
  # save sequence as text file
  filename = f"{jobname}.fasta"
  with open(filename, "w") as text_file:
    text_file.write(">1\n%s" % query_sequence)
  print("Saved sequence in %s: %s" %(filename, query_sequence))

def upload_target():
  target_dict = {}
  from alphafold.common import protein
  with redirect_stdout(StringIO()) as out:
    uploaded = files.upload()
    for filename,contents in uploaded.items():
      if not str(filename).endswith(".cif"):
        continue
      filepath = Path(cif_dir,filename)
      text_as_mmcif = contents.decode("UTF-8")
      text_as_bio_pdb = protein.mmcif_text_to_BioPDB(text_as_mmcif)
      text_as_pdb = protein.BioPDB_to_pdb_text(text_as_bio_pdb)
      target_dict[filename] = {}
      target_dict[filename]['prot'] = prot = protein.from_pdb_string(text_as_pdb)
      target_dict[filename]['mmcif_text'] = text_as_mmcif

  print("Targets uploaded: %s" %(target_dict.keys())) 
  if len( list(target_dict.keys())) < 1:
    print("\n*** WARNING: no templates uploaded...Please use only .cif files ***\n")
  return target_dict

def upload_msa():
  a3m_file = f"{jobname}.a3m"
  print("\nPlease select MSA file...")
  sys.stdout.flush()
  with redirect_stdout(StringIO()) as out:
    uploaded = files.upload()
  for filename,contents in uploaded.items():
      a3m_lines=contents.decode("UTF-8")
      break
  print("MSA file %s with %s lines uploaded " %(
            filename, len(a3m_lines.splitlines())))
  sys.stdout.flush()
  return a3m_lines

def get_jobnames_sequences_from_file(
    upload_manual_templates = None, cif_dir = None):
  from io import StringIO
  from google.colab import files
  print("Upload file with one jobname, a space and one sequence on each line")

  uploaded = files.upload()
  s = StringIO()
  query_sequences = []
  jobnames = []
  cif_filename_dict = {}
  for filename,contents in uploaded.items():
    print(contents.decode("UTF-8"), file = s)
    text = s.getvalue()
    for line in text.splitlines():
      spl = line.split()
      if len(spl) < 2:
        pass # empty line
      else: # usual
        jobname = spl[0]
        query_sequence = "".join(spl[1:])
        jobname = clean_jobname(jobname)
        query_sequence = clean_query(query_sequence)

        if jobname in jobnames:
          pass # already there
        else:
          query_sequences.append(query_sequence)
          jobnames.append(jobname)
          if upload_manual_templates:
            print("\nPlease upload CIF template for %s" %(jobname))
            sys.stdout.flush()
            cif_filename_dict[jobname] = upload_templates(cif_dir)
  return jobnames, query_sequences, cif_filename_dict

# Set working directory
os.chdir("/content/")

# Clear out directories
parent_dir = Path("/content/manual_templates")
cif_dir = Path(parent_dir,"mmcif")

# GET INPUTS

#@title A. Enter sequence and jobname. Load with <i><b>Run</b></i> button to left.
#@markdown <b><i><font color=green>Protein sequence and job name</font></i></b>

query_sequence = 'PSRIPLTLSEIEDLRRKGFNQTEIAELYGVTRQAVSWHKKTYGGRLTTRQIVQQNWPWDTRKPHDKSKAFQRLRDHGEYM RVGSFRTMSEDKKKRLLSWWKMYEGMAGGGFRYVPRDISDDDLLIRVNEHTQLTAEGELLWSWPDDIEELL' #@param {type:"string"}
jobname = 'protein_box_with_no_msa' #@param {type:"string"}

templates_to_use = "None"  

upload_manual_templates = False 
include_templates_from_pdb = False

disable_jit = False #@param {type:"boolean"}
number_of_ensembles = 1 #@param {type:"integer"}
af_iterations = 3 #@param {type:"integer"}
upload_target_pdb_file = True #@param {type:"boolean"}
use_target_pdb_as_template = True #@param {type:"boolean"}
use_target_pdb_as_recycle = False #@param {type:"boolean"}
use_custom_msa = False #@param {type:"boolean"}
skip_all_msa = False #@param {type:"boolean"}

if skip_all_msa:
  use_custom_msa = False

if upload_manual_templates:
  print("Templates will be uploaded")
if include_templates_from_pdb:
  print("Templates from the PDB will be included")
upload_file_with_jobname_space_sequence_lines = False 
maximum_templates_from_pdb = 20 
clear_saved_sequences_and_jobnames = True 

# Initialize query_sequences so we can loop through input
if clear_saved_sequences_and_jobnames or (
     not locals().get('query_sequences', None)):
  query_sequences = []
  jobnames = []
  cif_filename_dict = {}
  clear_directories([parent_dir,cif_dir])
del locals()['clear_saved_sequences_and_jobnames'] # so it updates

target_dict = {} # initialize

if upload_file_with_jobname_space_sequence_lines:
  del locals()['upload_file_with_jobname_space_sequence_lines'] # so it updates
  jobnames, query_sequences, cif_filename_dict = \
    get_jobnames_sequences_from_file(
        upload_manual_templates = upload_manual_templates,
        cif_dir = cif_dir)
else: # usual
  jobname = clean_jobname(jobname)
  query_sequence = clean_query(query_sequence)
  if query_sequence and not jobname:
    print("Please enter a job name and rerun")
    raise AssertionError("Please enter a job name and rerun")

  if jobname and not query_sequence:
    print("Please enter a query_sequence and rerun")
    raise AssertionError("Please enter a query_sequence rerun")
  
  # Add sequence and jobname if new
  if (jobname and query_sequence) and (
       not query_sequence in query_sequences) and (
       not jobname in jobnames):
      query_sequences.append(query_sequence)
      jobnames.append(jobname)
      if upload_target_pdb_file:
        print("\nPlease upload template for %s (must match exactly in sequence)" %(jobname))
        sys.stdout.flush()
        target_dict = upload_target()
        cif_filename_dict[jobname] = list(target_dict.keys())
  
        target_dict['jobname'] = list(target_dict.keys())[0]
  
        assert len(jobnames) == 1 # need different code if multiple


# Save sequence
for i in range(len(query_sequences)):
  # save the sequence as a file with name jobname.fasta
  save_sequence(jobnames[i], query_sequences[i])
  
print("\nCurrent jobs, sequences, and templates:")

for qs,jn in zip(query_sequences, jobnames):
  template_list = []
  for t in cif_filename_dict.get(jn,[]):
    template_list.append(os.path.split(str(t))[-1])
  print(jn, qs, template_list)

  target_dict['use_target_pdb_as_template'] = False
  target_dict['use_target_pdb_as_recycle'] = False
if use_target_pdb_as_template:
  print("Using target_pdb as a template")
  target_dict['use_target_pdb_as_template'] = True
if use_target_pdb_as_recycle:
  print("Using target_pdb as if it were recycled model")
  target_dict['use_target_pdb_as_recycle'] = True


sys.stdout.flush()  # seems to overwrite otherwise


if not query_sequences:
  print("Please supply a query sequence and run again")
  raise AssertionError("Need a query sequence")

# STANDARD PARAMETERS AND METHODS

#standard values of parameters
msa_mode = "MMseqs2 (UniRef+Environmental)" 
num_models = 1 
homooligomer = 1
use_msa = (not skip_all_msa)
use_env = True

use_amber = False 
use_templates = True


#@title B. Create AlphaFold models with the <b><i>Run</i></b> button to the left

! echo "Clearing and re-importing python modules and tmp directories..."
! rm -rf /tmp/absl_testing/
! mkdir /tmp/absl_testing/

for x in list(sys.modules.keys(  )) + list(globals()):
  for key in ['alphafold','protein', 'Alphafold', 'Protein','haiku',"layer_norm","base","colabfold","control_flow","check_tree_and_avals"]:
    if x.find(key)>-1:
      if x in list(sys.modules.keys()):
        
        del(sys.modules[x])
      if x in list(globals().keys()):
      
        del globals()[x]
        assert not x in list(globals().keys())
        break
    

if not os.environ['PYTHONPATH'].find(":/opt/conda/bin")>-1:
  os.environ['PYTHONPATH']+=":/opt/conda/bin"
  os.environ['PYTHONPATH']+=":/opt/conda/lib/python3.7/site-packages"
  os.environ['PYTHONPATH']+=":/usr/local/lib/python3.7/dist-packages"
! echo "VERSION:  `grep Version run_alphafold_test.py|grep -v Apache`"


from contextlib import redirect_stderr, redirect_stdout
from dataclasses import dataclass, replace
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio import SeqIO


print("Setting up methods...", end = "")
import_alphafold_items = True
# setup the model
if import_alphafold_items:

  # hiding warning messages
  import warnings
  from absl import logging
  import os
  import tensorflow as tf
  warnings.filterwarnings('ignore')
  logging.set_verbosity("error")
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
  tf.get_logger().setLevel('ERROR')

  import sys
  import numpy as np
  import pickle
  from alphafold.common import protein
  from alphafold.data import pipeline
  from alphafold.data import templates
  from alphafold.model import data
  from alphafold.model import config
  from alphafold.model import model
  from alphafold.data.tools import hhsearch
  import colabfold as cf

  # plotting libraries
  import py3Dmol
  import matplotlib.pyplot as plt
  import ipywidgets
  from ipywidgets import interact, fixed, GridspecLayout, Output



from alphafold.data import mmcif_parsing
from alphafold.data.templates import (_get_pdb_id_and_chain,
                                      _process_single_hit,
                                      _assess_hhsearch_hit,
                                      _build_query_to_hit_index_mapping,
                                      _extract_template_features,
                                      SingleHitResult,
                                      TEMPLATE_FEATURES)

def mk_mock_template(query_sequence):
  # since alphafold's model requires a template input
  # we create a blank example w/ zero input, confidence -1
  ln = len(query_sequence)
  output_templates_sequence = "-"*ln
  output_confidence_scores = np.full(ln,-1)
  templates_all_atom_positions = np.zeros((ln, templates.residue_constants.atom_type_num, 3))
  templates_all_atom_masks = np.zeros((ln, templates.residue_constants.atom_type_num))
  templates_aatype = templates.residue_constants.sequence_to_onehot(output_templates_sequence,
                                                                    templates.residue_constants.HHBLITS_AA_TO_ID)
  template_features = {'template_all_atom_positions': templates_all_atom_positions[None],
                       'template_all_atom_masks': templates_all_atom_masks[None],
                       'template_sequence': [f'none'.encode()],
                       'template_aatype': np.array(templates_aatype)[None],
                       'template_confidence_scores': output_confidence_scores[None],
                       'template_domain_names': [f'none'.encode()],
                       'template_release_date': [f'none'.encode()]}
  return template_features

def mk_template(a3m_lines, template_paths):
  template_featurizer = templates.TemplateHitFeaturizer(
      mmcif_dir=template_paths,
      max_template_date="2100-01-01",
      max_hits=20,
      kalign_binary_path="kalign",
      release_dates_path=None,
      obsolete_pdbs_path=None)

  hhsearch_pdb70_runner = hhsearch.HHSearch(binary_path="hhsearch", databases=[f"{template_paths}/pdb70"])

  hhsearch_result = hhsearch_pdb70_runner.query(a3m_lines)
  hhsearch_hits = pipeline.parsers.parse_hhr(hhsearch_result)
  templates_result = template_featurizer.get_templates(query_sequence=query_sequence,
                                                       query_pdb_code=None,
                                                       query_release_date=None,
                                                       hits=hhsearch_hits)
  return templates_result.features

def set_bfactor(pdb_filename, bfac, idx_res, chains):
  I = open(pdb_filename,"r").readlines()
  O = open(pdb_filename,"w")
  for line in I:
    if line[0:6] == "ATOM  ":
      seq_id = int(line[22:26].strip()) - 1
      seq_id = np.where(idx_res == seq_id)[0][0]
      O.write(f"{line[:21]}{chains[seq_id]}{line[22:60]}{bfac[seq_id]:6.2f}{line[66:]}")
  O.close()

def predict_structure(prefix, feature_dict, Ls, model_params, 
  use_model,
  model_runner_1,
  model_runner_3,
  do_relax=False, random_seed=0):  
  """Predicts structure using AlphaFold for the given sequence."""

  # Minkyung's code
  # add big enough number to residue index to indicate chain breaks
  idx_res = feature_dict['residue_index']
  L_prev = 0
  # Ls: number of residues in each chain
  for L_i in Ls[:-1]:
      idx_res[L_prev+L_i:] += 200
      L_prev += L_i  
  chains = list("".join([ascii_uppercase[n]*L for n,L in enumerate(Ls)]))
  feature_dict['residue_index'] = idx_res

  # Run the models.
  plddts,paes = [],[]
  unrelaxed_pdb_lines = []
  relaxed_pdb_lines = []

  for model_name, params in model_params.items():
    if model_name in use_model:
      print(f"running {model_name}")
      # swap params to avoid recompiling
      # note: models 1,2 have diff number of params compared to models 3,4,5
      if any(str(m) in model_name for m in [1,2]): model_runner = model_runner_1
      if any(str(m) in model_name for m in [3,4,5]): model_runner = model_runner_3
      model_runner.params = params
      
      processed_feature_dict = model_runner.process_features(feature_dict, random_seed=random_seed)
      prediction_result = model_runner.predict(processed_feature_dict)
      unrelaxed_protein = protein.from_prediction(processed_feature_dict,prediction_result)
      unrelaxed_pdb_lines.append(protein.to_pdb(unrelaxed_protein))
      plddts.append(prediction_result['plddt'])
      paes.append(prediction_result['predicted_aligned_error'])


  # rerank models based on predicted lddt
  lddt_rank = np.mean(plddts,-1).argsort()[::-1]
  out = {}
  
  for n,r in enumerate(lddt_rank):
    print(f"model_{n+1} {np.mean(plddts[r])}")

    unrelaxed_pdb_path = f'{prefix}_unrelaxed_model_{n+1}.pdb'    
    with open(unrelaxed_pdb_path, 'w') as f: f.write(unrelaxed_pdb_lines[r])
    set_bfactor(unrelaxed_pdb_path, plddts[r], idx_res, chains)


    out[f"model_{n+1}"] = {"plddt":plddts[r], "pae":paes[r]}
  return out



def hh_process_seq(query_seq,template_seq,hhDB_dir,db_prefix="DB"):
  """
  This is a hack to get hhsuite output strings to pass on
  to the AlphaFold template featurizer. 
  
  Note: that in the case of multiple templates, this would be faster to build one database for
  all the templates. Currently it builds a database with only one template at a time. Even 
  better would be to get an hhsuite alignment without using a database at all, just between
  pairs of sequence files. However, I have not figured out how to do this.

  Update: I think the hhsearch can be replaced completely, and we can just do a pairwise 
  alignment with biopython, or skip alignment if the seqs match. TODO
  """
  # set up directory for hhsuite DB. Place one template fasta file to be the DB contents
  if hhDB_dir.exists():
    shutil.rmtree(hhDB_dir)
  
  msa_dir = Path(hhDB_dir,"msa")
  msa_dir.mkdir(parents=True)
  template_seq_path = Path(msa_dir,"template.fasta")
  with template_seq_path.open("w") as fh:
    SeqIO.write([template_seq], fh, "fasta")

  # make hhsuite DB
  with redirect_stdout(StringIO()) as out:
    os.chdir(msa_dir)
    %shell ffindex_build -s ../DB_msa.ff{data,index} .
    os.chdir(hhDB_dir)
    %shell ffindex_apply DB_msa.ff{data,index}  -i DB_a3m.ffindex -d DB_a3m.ffdata  -- hhconsensus -M 50 -maxres 65535 -i stdin -oa3m stdout -v 0
    %shell rm DB_msa.ff{data,index}
    %shell ffindex_apply DB_a3m.ff{data,index} -i DB_hhm.ffindex -d DB_hhm.ffdata -- hhmake -i stdin -o stdout -v 0
    %shell cstranslate -f -x 0.3 -c 4 -I a3m -i DB_a3m -o DB_cs219 
    %shell sort -k3 -n -r DB_cs219.ffindex | cut -f1 > sorting.dat

    %shell ffindex_order sorting.dat DB_hhm.ff{data,index} DB_hhm_ordered.ff{data,index}
    %shell mv DB_hhm_ordered.ffindex DB_hhm.ffindex
    %shell mv DB_hhm_ordered.ffdata DB_hhm.ffdata

    %shell ffindex_order sorting.dat DB_a3m.ff{data,index} DB_a3m_ordered.ff{data,index}
    %shell mv DB_a3m_ordered.ffindex DB_a3m.ffindex
    %shell mv DB_a3m_ordered.ffdata DB_a3m.ffdata

  # run hhsearch
  hhsearch_runner = hhsearch.HHSearch(binary_path="hhsearch", databases=[hhDB_dir.as_posix()+"/"+db_prefix])
  with StringIO() as fh:
    SeqIO.write([query_seq], fh, "fasta")
    seq_fasta = fh.getvalue()
  hhsearch_result = hhsearch_runner.query(seq_fasta)

  # process hits
  hhsearch_hits = pipeline.parsers.parse_hhr(hhsearch_result)
  if len(hhsearch_hits) >0:
    hit = hhsearch_hits[0]
    hit = replace(hit,**{"name":template_seq.id})
  else:
    hit = None
  return hit

def plot_plddt_legend():
  thresh = ['plDDT:','Very low (<50)','Low (60)','OK (70)','Confident (80)','Very high (>90)']
  plt.figure(figsize=(1,0.1),dpi=100)
  ########################################
  for c in ["#FFFFFF","#FF0000","#FFFF00","#00FF00","#00FFFF","#0000FF"]:
    plt.bar(0, 0, color=c)
  plt.legend(thresh, frameon=False,
             loc='center', ncol=6,
             handletextpad=1,
             columnspacing=1,
             markerscale=0.5,)
  plt.axis(False)
  return plt

def plot_confidence(outs, model_num=1):
  model_name = f"model_{model_num}"
  plt.figure(figsize=(10,3),dpi=100)
  """Plots the legend for plDDT."""
  #########################################
  plt.subplot(1,2,1); plt.title('Predicted lDDT')
  plt.plot(outs[model_name]["plddt"])
  for n in range(homooligomer+1):
    x = n*(len(query_sequence))
    plt.plot([x,x],[0,100],color="black")
  plt.ylabel('plDDT')
  plt.xlabel('position')
  #########################################
  plt.subplot(1,2,2);plt.title('Predicted Aligned Error')
  plt.imshow(outs[model_name]["pae"], cmap="bwr",vmin=0,vmax=30)
  plt.colorbar()
  plt.xlabel('Scored residue')
  plt.ylabel('Aligned residue')
  #########################################
  return plt

def show_pdb(model_num=1, show_sidechains=False, show_mainchains=False, color="lDDT"):
  model_name = f"model_{model_num}"
  if use_amber:
    pdb_filename = f"{jobname}_relaxed_{model_name}.pdb"
  else:
    pdb_filename = f"{jobname}_unrelaxed_{model_name}.pdb"

  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)
  view.addModel(open(pdb_filename,'r').read(),'pdb')

  if color == "lDDT":
    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})
  elif color == "rainbow":
    view.setStyle({'cartoon': {'color':'spectrum'}})
  elif color == "chain":
    for n,chain,color in zip(range(homooligomer),list("ABCDEFGH"),
                     ["lime","cyan","magenta","yellow","salmon","white","blue","orange"]):
       view.setStyle({'chain':chain},{'cartoon': {'color':color}})
  if show_sidechains:
    BB = ['C','O','N']
    view.addStyle({'and':[{'resn':["GLY","PRO"],'invert':True},{'atom':BB,'invert':True}]},
                        {'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
    view.addStyle({'and':[{'resn':"GLY"},{'atom':'CA'}]},
                        {'sphere':{'colorscheme':f"WhiteCarbon",'radius':0.3}})
    view.addStyle({'and':[{'resn':"PRO"},{'atom':['C','O'],'invert':True}]},
                        {'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})  
  if show_mainchains:
    BB = ['C','O','N','CA']
    view.addStyle({'atom':BB},{'stick':{'colorscheme':f"WhiteCarbon",'radius':0.3}})

  view.zoomTo()
  return view

def run_job(query_sequence,
        jobname,
        target_dict,
        num_models,
        homooligomer,
        use_msa,
        use_env,
        use_custom_msa,
        use_amber,
        use_templates,
        include_templates_from_pdb,
        number_of_ensembles,
        af_iterations,
        disable_jit,):


  key = target_dict.get('jobname',None)
  prot = target_dict.get(key,{}).get('prot',None)
  if key is not None:
    use_target_pdb_as_recycle = target_dict['use_target_pdb_as_recycle']
    use_target_pdb_as_template = target_dict['use_target_pdb_as_template']
  else:
    use_target_pdb_as_recycle = False
    use_target_pdb_as_template = False

  if prot is not None and use_target_pdb_as_recycle:
    target_all_atom_positions = prot.atom_positions
  else:
    target_all_atom_positions = None
  if prot is not None and use_target_pdb_as_template:
    mmcif_text = target_dict[key]['mmcif_text']
    print("Using target model...",jobname, key)
  else:
    mmcif_text = ""


  #@title Get MSA and templates
  print("Getting MSA and templates...")

  template_paths = None # toss these ... get template_paths later
  template_features = mk_mock_template(query_sequence * homooligomer)

  if (not use_msa):
    a3m_lines = ">query sequence \n%s" %(query_sequence)
  elif use_custom_msa:
    a3m_lines = upload_msa()
  elif use_msa:
    a3m_lines = cf.run_mmseqs2(query_sequence, jobname, use_env)
  else:
    pass # nothing to do
 
  # File for a3m
  a3m_file = f"{jobname}.a3m"


  with open(a3m_file, "w") as text_file:
      text_file.write(a3m_lines)
 
  # parse MSA
  msa, deletion_matrix = pipeline.parsers.parse_a3m(a3m_lines)
  
  print("Done with MSA and templates")
  
  #Process templates
  print("PROCESSING TEMPLATES")
  
  os.chdir("/content/")
  
  other_cif_dir = Path("/content/%s" %(template_paths))
  parent_dir = Path("/content/manual_templates")
  cif_dir = Path(parent_dir,"mmcif")
  fasta_dir = Path(parent_dir,"fasta")
  hhDB_dir = Path(parent_dir,"hhDB")
  msa_dir = Path(hhDB_dir,"msa")
  clear_directories([fasta_dir,hhDB_dir,msa_dir])
  
  if use_target_pdb_as_template:
    cif_files = ['text_a.cif']
  else:
    cif_files = []
  print("CIF files to include:",cif_files)
  query_seq = SeqRecord(Seq(query_sequence),id="query",name="",description="")
  query_seq_path = Path(fasta_dir,"query.fasta")
  with query_seq_path.open("w") as fh:
      SeqIO.write([query_seq], fh, "fasta")
  
  shutil.copyfile(query_seq_path,Path(msa_dir,"query.fasta"))
  seqs = []
  template_hit_list = []
  
  n_used = 0
  for i,filepath in enumerate(cif_files):
    n_used += 1
    assert filepath == 'text_a.cif'
    filestr = mmcif_text

    mmcif_obj = mmcif_parsing.parse(file_id='text_a',mmcif_string=filestr)
    mmcif = mmcif_obj.mmcif_object
    if not mmcif: continue
    print("CIF file included:",i+1,str(filepath))
  
    for chain_id,template_sequence in mmcif.chain_to_seqres.items():
        template_sequence = mmcif.chain_to_seqres[chain_id]
        seq_name = filepath.upper()+"_"+chain_id
        seq = SeqRecord(Seq(template_sequence),id=seq_name,name="",description="")
        seqs.append(seq)
  
        with  Path(fasta_dir,seq.id+".fasta").open("w") as fh:
          SeqIO.write([seq], fh, "fasta")
  
        """
        At this stage, we have a template sequence.
        and a query sequence. 
        There are two options to generate template features:
          1. Write new code to manually generate template features
          2. Get an hhr alignment string, and pass that
            to the existing template featurizer. 
            
        I chose the second, implemented in hh_process_seq()
        """
        SeqIO.write([seq], sys.stdout, "fasta")
        SeqIO.write([query_seq], sys.stdout, "fasta")
        try:
          hit = hh_process_seq(query_seq,seq,hhDB_dir)
        except Exception as e:
          hit = None
        if hit is not None:
          template_hit_list.append(hit)
  
  if template_hit_list:
    #process hits into template features
    template_hit_list = [replace(hit,**{"index":i+1}) for i,hit in enumerate(template_hit_list)]
  
  if (len(manual_templates_uploaded) > 0) and upload_manual_templates and (not template_hit_list):
    # check to make sure we got something
    # need template and did not get any
      print("\n",80*"-")
      print("\nNo templates obtained...")
      print("\n",80*"-")
      raise AssertionError("Failed to read template file")
  elif use_templates and template_hit_list:
    # have new templates to work with
  
    template_features = {}
    for template_feature_name in TEMPLATE_FEATURES:
      template_features[template_feature_name] = []
  
    for i,hit in enumerate(sorted(template_hit_list, key=lambda x: x.sum_probs, reverse=True)):
      # modifications to alphafold/data/templates.py _process_single_hit
      hit_pdb_code, hit_chain_id = _get_pdb_id_and_chain(hit)
      mapping = _build_query_to_hit_index_mapping(
      hit.query, hit.hit_sequence, hit.indices_hit, hit.indices_query,
      query_sequence)
      template_sequence = hit.hit_sequence.replace('-', '')
  

      features, realign_warning = _extract_template_features(
          mmcif_object=mmcif,
          pdb_id=hit_pdb_code,
          mapping=mapping,
          template_sequence=template_sequence,
          query_sequence=query_sequence,
          template_chain_id=hit_chain_id,
          kalign_binary_path="kalign")

      features['template_sum_probs'] = [hit.sum_probs]
  
      single_hit_result = SingleHitResult(features=features, error=None, warning=None)
      for k in template_features:
        template_features[k].append(features[k])
  
    for name in template_features:
      template_features[name] = np.stack(
          template_features[name], axis=0).astype(TEMPLATE_FEATURES[name])
      
    #overwrite template data
    template_paths = cif_dir.as_posix()


    # Select only one chain from any cif file
    unique_template_hits = []
    pdb_text_list = []
    for hit in template_hit_list:
      pdb_text = hit.name.split()[0].split("_")[0]
      if not pdb_text in pdb_text_list:
        pdb_text_list.append(pdb_text)
        unique_template_hits.append(hit)
    template_hit_list = unique_template_hits
    template_hits = template_hit_list

    print("\nIncluding templates:")
    for hit in template_hit_list:
      print("\t",hit.name.split()[0])
    if len(template_hit_list) == 0:
      print("No templates found...quitting")
      raise AssertionError("No templates found...quitting")
    os.chdir("/content/")
  
    for key,value in template_features.items():
      if np.all(value==0):
        print("ERROR: Some template features are empty")
  else:  # no templates
    print("Not using any templates")
  
  print("\nPREDICTING STRUCTURE")

  
  # collect model weights
  use_model = {}
  model_params = {}
  model_runner_1 = None
  model_runner_3 = None

  for model_name in ["model_1","model_2","model_3","model_4","model_5"][:num_models]:
    use_model[model_name] = True

    if model_name not in list(model_params.keys()):
      model_params[model_name] = data.get_model_haiku_params(model_name=model_name+"_ptm", data_dir=".")
      if model_name == "model_1":
        model_config = config.model_config(model_name+"_ptm")
        model_config.model.num_recycle = af_iterations
        model_config.data.common.num_recycle = af_iterations
        print("Recycle iterations will be %s" %(model_config.data.common.num_recycle))
        model_config.data.eval.num_ensemble = number_of_ensembles
        print("Number of ensembles will be %s" %(model_config.data.eval.num_ensemble))
        if disable_jit:
          model_config.data.common.disable_jit = True
          model_config.model.global_config.disable_jit = True

        if use_target_pdb_as_recycle and target_all_atom_positions is not None:
          print("Setting all_atom",target_all_atom_positions.shape[0])
          model_config.model.global_config.target_all_atom_positions = target_all_atom_positions
          model_config.model.global_config.target_prot = prot
          print("Set target_all_atom_positions")
        model_runner_1 = model.RunModel(model_config, model_params[model_name])
        print("Done running model.RunModel to get model_runner_1")

      if model_name == "model_3":
        model_config = config.model_config(model_name+"_ptm")
        model_config.data.eval.num_ensemble = 1
        model_runner_3 = model.RunModel(model_config, model_params[model_name])


  if homooligomer == 1:
    msas = [msa]
    deletion_matrices = [deletion_matrix]
  else:
    # make multiple copies of msa for each copy
    # AAA------
    # ---AAA---
    # ------AAA
    #
    # note: if you concat the sequences (as below), it does NOT work
    # AAAAAAAAA
    msas = []
    deletion_matrices = []
    Ln = len(query_sequence)
    for o in range(homooligomer):
      L = Ln * o
      R = Ln * (homooligomer-(o+1))
      msas.append(["-"*L+seq+"-"*R for seq in msa])
      deletion_matrices.append([[0]*L+mtx+[0]*R for mtx in deletion_matrix])
  
  # gather features
  feature_dict = {
      **pipeline.make_sequence_features(sequence=query_sequence*homooligomer,
                                        description="none",
                                        num_res=len(query_sequence)*homooligomer),
      **pipeline.make_msa_features(msas=msas,deletion_matrices=deletion_matrices),
      **template_features
  }
  outs = predict_structure(jobname, feature_dict,
                           Ls=[len(query_sequence)]*homooligomer,
                           model_params=model_params, use_model=use_model,
                           model_runner_1=model_runner_1,
                           model_runner_3=model_runner_3,
                           do_relax=use_amber)
  print("DONE WITH STRUCTURE")
  
  #@title Making plots...
  
  # gather MSA info
  deduped_full_msa = list(dict.fromkeys(msa))
  msa_arr = np.array([list(seq) for seq in deduped_full_msa])
  seqid = (np.array(list(query_sequence)) == msa_arr).mean(-1)
  seqid_sort = seqid.argsort() #[::-1]
  non_gaps = (msa_arr != "-").astype(float)
  non_gaps[non_gaps == 0] = np.nan
  
  ##################################################################
  plt.figure(figsize=(14,4),dpi=100)
  ##################################################################
  plt.subplot(1,2,1); plt.title("Sequence coverage")
  plt.imshow(non_gaps[seqid_sort]*seqid[seqid_sort,None],
             interpolation='nearest', aspect='auto',
             cmap="rainbow_r", vmin=0, vmax=1, origin='lower')
  plt.plot((msa_arr != "-").sum(0), color='black')
  plt.xlim(-0.5,msa_arr.shape[1]-0.5)
  plt.ylim(-0.5,msa_arr.shape[0]-0.5)
  plt.colorbar(label="Sequence identity to query",)
  plt.xlabel("Positions")
  plt.ylabel("Sequences")
  
  ##################################################################
  plt.subplot(1,2,2); plt.title("Predicted lDDT per position")
  for model_name,value in outs.items():
    plt.plot(value["plddt"],label=model_name)
  if homooligomer > 0:
    for n in range(homooligomer+1):
      x = n*(len(query_sequence)-1)
      plt.plot([x,x],[0,100],color="black")
  plt.legend()
  plt.ylim(0,100)
  plt.ylabel("Predicted lDDT")
  plt.xlabel("Positions")
  plt.savefig(jobname+"_coverage_lDDT.png")
  ##################################################################
  plt.show()
  
  print("Predicted Alignment Error")
  ##################################################################
  plt.figure(figsize=(3*num_models,2), dpi=100)
  for n,(model_name,value) in enumerate(outs.items()):
    plt.subplot(1,num_models,n+1)
    plt.title(model_name)
    plt.imshow(value["pae"],label=model_name,cmap="bwr",vmin=0,vmax=30)
    plt.colorbar()
  plt.savefig(jobname+"_PAE.png")
  plt.show()
  ##################################################################
  #@title Displaying 3D structure... {run: "auto"}
  model_num = 1 
  color = "lDDT" 
  show_sidechains = False 
  show_mainchains = False 
  
  
  
  show_pdb(model_num,show_sidechains, show_mainchains, color).show()
  if color == "lDDT": plot_plddt_legend().show()  
  plot_confidence(outs, model_num).show()
  #@title Packaging and downloading results...
  
  #@markdown When modeling is complete .zip files with results will be downloaded automatically.
  
  citations = {
  "Mirdita2021":  """@article{Mirdita2021,
  author = {<NAME>, <NAME>, Martin},
  doi = {10.1101/2021.08.15.456425},
  journal = {bioRxiv},
  title = {{ColabFold - Making Protein folding accessible to all}},
  year = {2021},
  comment = {ColabFold including MMseqs2 MSA server}
  }""",
    "Mitchell2019": """@article{Mitchell2019,
  author = {Mitchell, <NAME> Almeida, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME> and Richardson, <NAME> and Sakharova, <NAME>, <NAME>, <NAME>, <NAME>, Olga and Lapidus, Alla and Finn, <NAME>},
  doi = {10.1093/nar/gkz1035},
  journal = {Nucleic Acids Res.},
  title = {{MGnify: the microbiome analysis resource in 2020}},
  year = {2019},
  comment = {MGnify database}
  }""",
    "Eastman2017": """@article{Eastman2017,
  author = {<NAME>, <NAME>, <NAME>. and McGibbon, <NAME>. and Zhao, Yutong and Beauchamp, <NAME>. and <NAME> and Simmonett, <NAME>. and Harrigan, <NAME>, <NAME>. and Brooks, <NAME>. and <NAME>.},
  doi = {10.1371/journal.pcbi.1005659},
  journal = {PLOS Comput. Biol.},
  number = {7},
  title = {{OpenMM 7: Rapid development of high performance algorithms for molecular dynamics}},
  volume = {13},
  year = {2017},
  comment = {Amber relaxation}
  }""",
    "Jumper2021": """@article{Jumper2021,
  author = {<NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME> and {\v{Z}}{\'{i}}dek, Augustin and Potapenko, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, Andrew and Romera-Paredes, Bernardino and Nikolov, Stanislav and Jain, Rishub and Adler, Jonas and Back, Trevor and Petersen, Stig and Reiman, David and Clancy, Ellen and Zielinski, Michal and Steinegger, Martin and Pacholska, Michalina and Berghammer, Tamas and Bodenstein, Sebastian and Silver, David and Vinyals, Oriol and Senior, <NAME>. and Kavukcuoglu, Koray and <NAME> and Hassabis, Demis},
  doi = {10.1038/s41586-021-03819-2},
  journal = {Nature},
  pmid = {34265844},
  title = {{Highly accurate protein structure prediction with AlphaFold.}},
  year = {2021},
  comment = {AlphaFold2 + BFD Database}
  }""",
    "Mirdita2019": """@article{Mirdita2019,
  author = {<NAME> <NAME> and S{\"{o}}ding, Johannes},
  doi = {10.1093/bioinformatics/bty1057},
  journal = {Bioinformatics},
  number = {16},
  pages = {2856--2858},
  pmid = {30615063},
  title = {{MMseqs2 desktop and local web server app for fast, interactive sequence searches}},
  volume = {35},
  year = {2019},
  comment = {MMseqs2 search server}
  }""",
    "Steinegger2019": """@article{Steinegger2019,
  author = {<NAME> Meier, <NAME> Mirdita, <NAME> V{\"{o}}hringer, <NAME>, <NAME>. and S{\"{o}}ding, Johannes},
  doi = {10.1186/s12859-019-3019-7},
  journal = {BMC Bioinform.},
  number = {1},
  pages = {473},
  pmid = {31521110},
  title = {{HH-suite3 for fast remote homology detection and deep protein annotation}},
  volume = {20},
  year = {2019},
  comment = {PDB70 database}
  }""",
    "Mirdita2017": """@article{Mirdita2017,
  author = {<NAME> <NAME>, <NAME>, <NAME> Martin, <NAME>. and S{\"{o}}ding, <NAME>},
  doi = {10.1093/nar/gkw1081},
  journal = {Nucleic Acids Res.},
  number = {D1},
  pages = {D170--D176},
  pmid = {27899574},
  title = {{Uniclust databases of clustered and deeply annotated protein sequences and alignments}},
  volume = {45},
  year = {2017},
  comment = {Uniclust30/UniRef30 database},
  }""",
    "Berman2003": """@misc{Berman2003,
  author = {<NAME> <NAME> <NAME>},
  booktitle = {Nat. Struct. Biol.},
  doi = {10.1038/nsb1203-980},
  number = {12},
  pages = {980},
  pmid = {14634627},
  title = {{Announcing the worldwide Protein Data Bank}},
  volume = {10},
  year = {2003},
  comment = {templates downloaded from wwPDB server}
  }""",
  }
  
  to_cite = [ "Mirdita2021", "Jumper2021" ]
  if use_msa:       to_cite += ["Mirdita2019"]
  if use_msa:       to_cite += ["Mirdita2017"]
  if use_env:       to_cite += ["Mitchell2019"]
  if use_templates: to_cite += ["Steinegger2019"]
  if use_templates: to_cite += ["Berman2003"]
  if use_amber:     to_cite += ["Eastman2017"]
  
  with open(f"{jobname}.bibtex", 'w') as writer:
    for i in to_cite:
      writer.write(citations[i])
      writer.write("\n")
  
  print(f"Found {len(to_cite)} citation{'s' if len(to_cite) > 1 else ''} for tools or databases.")
  if use_custom_msa:
    print("Don't forget to cite your custom MSA generation method.")
  
  !echo 'FILES TO PACKAGE: $a3m_file $jobname"_"*"relaxed_model_"*".pdb" $jobname"_coverage_lDDT.png" $jobname".bibtex" $jobname"_PAE.png" '
  try:
    print("zipping files...")
    !zip -FSr $jobname".result.zip" $a3m_file $jobname"_"*"relaxed_model_"*".pdb" $jobname"_coverage_lDDT.png" $jobname".bibtex" $jobname"_PAE.png"
  except Exception as e:
    print("unable to zip files")

  filename = f"{jobname}.result.zip"
  if os.path.isfile(filename):
    print("About to download %s" %(filename))
  
    try:
      print("Downloading zip file %s" %(filename))
      files.download(filename)
      print("Start of download successful (NOTE: if the download symbol does not go away it did not work. Download it manually using the folder icon to the left)")
      return filename
    except Exception as e:
      print("Unable to download zip file %s" %(filename))
      return None
  else:
    print("No .zip file %s created" %(filename))
    return None

# RUN THE JOBS HERE

for query_sequence, jobname in zip(query_sequences, jobnames):
  print("\n","****************************************","\n",
         "RUNNING JOB %s with sequence %s\n" %(
    jobname, query_sequence),
    "****************************************","\n")
  # GET TEMPLATES AND SET UP FILES



  # User input of manual templates
  manual_templates_uploaded = cif_filename_dict.get(
      jobname,[])
  if manual_templates_uploaded:
    print("Using uploaded templates %s for this run" %(
        manual_templates_uploaded))

  if 1:
    filename = run_job(query_sequence,
        jobname,
        target_dict,
        num_models,
        homooligomer,
        use_msa,
        use_env,
        use_custom_msa,
        use_amber,
        use_templates,
        include_templates_from_pdb,
        number_of_ensembles,
        af_iterations,
        disable_jit,
        )
    if filename:
      print("FINISHED JOB (%s) %s with sequence %s\n" %(
        filename, jobname, query_sequence),
        "****************************************","\n")
    else:
      print("NO RESULT FOR JOB %s with sequence %s\n" %(
    jobname, query_sequence),
    "****************************************","\n")

  if 0:
    print("FAILED: JOB %s with sequence %s\n\n%s\n" %(
    jobname, query_sequence, str(e)),
    "****************************************","\n")


print("\nDOWNLOADING FILES NOW:\n")
for query_sequence, jobname in zip(query_sequences, jobnames):
  filename = f"{jobname}.result.zip"
  if os.path.isfile(filename):
    print(filename)

print("\nALL DONE\n")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13915999125404996, 'fraction_numerical': 0.016160130394164068, 'mean_word_length': 3.697917639368755, 'pattern_counts': {'":': 27, '<': 79, '<?xml version=': 0, '>': 99, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '362224', 'license': 'MIT', 'n_tokens_mistral': 18126, 'n_tokens_neox': 17100, 'n_words': 3702, 'path': 'alphafold2/AlphaFold_testing.ipynb', 'repo_name': 'phenix-project/Colabs', 'size': '50935'}